Module attn_splitk 

2026-09-25: Paged-decode attention launch geometry: the split-K count, the split-K workspace size and the GQA-packed kernel gate, as pure functions of configuration.

The online-softmax split merge is not associative, so a split count that moved with the co-batched sequence count would change a sequence’s output bytes depending on what it is batched with.

Owner: kernels crate (paged-decode attention policy). Invariants:

• Under SplitkPolicy::Auto and SplitkPolicy::Pinned the split count depends only on (sm_count, num_q_heads) and the pin; only SplitkPolicy::Legacy reads the reference batch.
• resolve_policy is the one resolution of the target declaration and METRALE_ATTN_DECODE_SPLITK; the metrale-model-layers dispatch and the metrale-gpu-runtime workspace sizing both call it.
• gqa_pack_shape_ok admits only num_q_heads == num_kv_heads * DECODE_GQA_PACK_WIDTH at head_dim == DECODE_GQA_PACK_HEAD_DIM.

SplitkPolicy::Legacy computes, with NUM_SMS standing for the target’s sm_count:

current_ctas = num_q_heads * split_ref_seqs(num_seqs, max_decode_seqs);
num_splits   = if current_ctas >= NUM_SMS { 1 } else { NUM_SMS / current_ctas };

Short contexts are handled inside the Hopper split-K kernel, which floors each sequence’s split size at PD_MIN_KV_PER_SPLIT from that sequence’s own seq_len (kernels/hopper/common/paged_decode_splitk_hopper.cuh).

Enums

SplitkPolicy

2026-09-25: How a target picks its paged-decode split count.

Constants

DECODE_GQA_PACK_DECLARED

2026-09-25: The packed kernels’ declared state: off, for every target. METRALE_ATTN_DECODE_GQA_PACK overrides it (resolve_gqa_pack).

DECODE_GQA_PACK_HEAD_DIM

2026-09-25: The head dim the packed kernels are compiled for: both sources #define PD_HDIM 256, and gqa_pack_shape_ok refuses any other head dim.

DECODE_GQA_PACK_WIDTH

2026-09-25: Query heads one packed CTA carries: the Rust copy of the kernels’ #define PD_GQA, which sizes their q_reg/o_reg register arrays at compile time. cuda_sources_declare_the_pack_width_rust_dispatches_on fails if the two differ.

MAX_DECODE_SPLITS

2026-09-25: Ceiling on the split count under every policy but Legacy, and so on the split-K workspace, which is rows * num_q_heads * num_splits * (head_dim + 2) F32 values. Without it a model with few q heads would ask for 2 * sm_count / num_q_heads splits (264 for one head on 132 SMs). At 24 q heads on the 148-SM targets (kernels/b200, kernels/b300) auto asks for 13. A pinned count is clamped to it.

SPLITK_TARGET_WAVES

2026-09-25: Waves of CTAs SplitkPolicy::Auto aims to put on the device with one sequence in flight (see auto_splits).

Functions

auto_splits

2026-09-25: clamp(ceil(WAVES * sm_count / num_q_heads), 1, MAX_DECODE_SPLITS).

gqa_pack_ctas

2026-09-25: CTAs a packed launch puts on the device: num_kv_heads * num_seqs, a DECODE_GQA_PACK_WIDTH-fold cut from the unpacked num_q_heads * num_seqs. With 4 KV heads on gb10’s 48 SMs the packed grid covers the SMs only from num_seqs = 12.

gqa_pack_enabled

2026-09-25: Whether the packed kernels are armed in this process, resolved once. The dispatch reads it on every non-split BF16 and FP8 paged-decode launch, and a fixed answer keeps a captured CUDA graph consistent with later launches.

gqa_pack_shape_ok

2026-09-25: Whether this launch’s shape is one the packed kernels can serve. Each condition is a precondition of the kernel:

legacy_splits

2026-09-25: The SplitkPolicy::Legacy rule. The dispatch passes ref_seqs = split_ref_seqs(num_seqs, max_decode_seqs), which is max(max_decode_seqs, num_seqs).

num_splits

2026-09-25: The split count for a launch.

parse

2026-09-25: legacy | auto | 0/off/false/no | a decimal count (clamped to 1..=MAX_DECODE_SPLITS), trimmed and case-insensitive.

parse_gqa_pack

2026-09-25: 1/on/true/yes | 0/off/false/no, trimmed and case-insensitive.

policy_from_env

2026-09-25: resolve_policy against the process environment and this binary’s baked declaration. Used by the split-K workspace sizing in metrale-gpu-runtime (sizes.rs), which metrale-model-layers depends on and so cannot call the model-layers resolver.

resolve_gqa_pack

2026-09-25: declared, overridden by a parseable METRALE_ATTN_DECODE_GQA_PACK value. Pure; gqa_pack_enabled supplies the environment.

resolve_policy

2026-09-25: The target’s declaration, overridden by a parseable METRALE_ATTN_DECODE_SPLITK. An unparseable declaration is Legacy.

workspace_slots

2026-09-25: [o[head_dim], m, l] slots the split-K workspace must hold.