Module attn_splitk 

The paged-decode attention SPLIT-K policy: how many KV splits a launch uses, as a pure function of CONFIGURATION.

The defect (#928)

crates/spark-model/.../decode/run_paged_decode.rs imported atlas_core::device::sm121::NUM_SMS — the GB10 constant, 48 — and picked

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 };

On Qwen3.8-27B at --max-batch-size 16 that is 24 * 16 = 384 >= 48, so num_splits was 1 at EVERY batch size including C=1. nsys on 1xH100 80GB HBM3 (round 13 cell T1N) has the receipt: paged_decode_attn_fp8 runs grid=(24,1,1) — 24 CTAs on 132 SMs — 231.51 us/launch for 9.93 MB of KV, i.e. 42.9 GB/s = 1.28% of HBM, and with the BF16-KV sibling the pair is 3.79 ms of a 16.69 ms C=1 decode step (22.7%). The source comment dismissing occupancy here (“attention occupancy is NOT the long-ctx bottleneck”) records a GB10 A/B on a 48-SM part.

Two things were wrong and both are fixed here: the SM count is a property of the compiled TARGET (kernels/<hw>/HARDWARE.toml [hardware] sm_count, baked as crate::TARGET_SM_COUNT), and the occupancy the rule should size for is the SINGLE-STREAM shape, which is the one that starves.

The determinism invariant — why this module is pure

The online-softmax split-merge is NON-ASSOCIATIVE. If num_splits moved with the runtime co-batched count, one sequence would traverse a different reduction tree alone than beside fifteen others and flip its temp-0 argmax — the nondeterminism split_ref_seqs was introduced to stop (tasks/determinism_investigation.md). So SplitkPolicy::Auto and SplitkPolicy::Pinned read (sm_count, num_q_heads, max_decode_seqs) and NOTHING else: the split count is fixed for the life of a serve. SplitkPolicy::Legacy is the pre-#928 rule preserved verbatim, including its dependence on split_ref_seqs, because every target but Hopper still runs it and “unchanged” has to mean unchanged.

Short contexts are handled INSIDE the kernel (kernels/hopper/common/paged_decode_splitk_hopper.cuh, PD_MIN_KV_PER_SPLIT), from each sequence’s own seq_len: the host may not branch on seq_lens, which is device memory and behind a captured CUDA graph, and a per-sequence rule stays co-batch invariant where a per-batch one would not.

Enums

SplitkPolicy

How a target picks its paged-decode split count.

Constants

MAX_DECODE_SPLITS

Hard ceiling on the split count, and the bound the split-K workspace is sized against.

SPLITK_TARGET_WAVES

Waves of CTAs the auto policy aims to put on the device at the single-stream shape.

Functions

auto_splits

clamp(ceil(WAVES * sm_count / num_q_heads), 1, MAX_DECODE_SPLITS).

legacy_splits

The pre-#928 rule, preserved verbatim for every target that still declares it. ref_seqs is split_ref_seqs(num_seqs, max_decode_seqs).

num_splits

The split count for a launch.

parse

legacy | auto | 0/off/false/no | a decimal count.

policy_from_env

resolve_policy against the process environment and this binary’s baked declaration.

resolve_policy

The target’s declaration, overridden by ATLAS_ATTN_DECODE_SPLITK.

workspace_slots

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