pub fn workspace_slots(
policy: SplitkPolicy,
sm_count: u32,
num_q_heads: u32,
max_rows: u32,
legacy_ref_seqs: u32,
) -> u32Expand 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.