workspace_slots

Function workspace_slots 

Source
pub fn workspace_slots(
    policy: SplitkPolicy,
    sm_count: u32,
    num_q_heads: u32,
    max_rows: u32,
    legacy_ref_seqs: u32,
) -> u32
Expand description

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

The split-K kernel addresses ((seq * num_q_heads) + head) * num_splits + split, so max_rows must be the widest batch the decode metadata accepts (sizes.rs passes DecodeMetaLayout::rows()); a short allocation is an out-of-bounds device write with no error.

Legacy needs sm_count slots: when it splits, num_q_heads * ref_seqs < sm_count with ref_seqs >= num_seqs, so num_seqs * num_q_heads * num_splits <= sm_count; at one split the non-split kernel runs and the workspace is not touched.