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

[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 the arena has to cover the WIDEST batch the decode-metadata layout will accept (DecodeMetaLayout::rows()), not the pinned max batch — a short allocation here is an out-of-bounds device write, silently.

Legacy keeps its old constant bound and its old arena: that rule picks sm_count / (num_q_heads * max(max_batch, num_seqs)), so num_seqs * num_q_heads * num_splits <= sm_count for every batch, which is exactly what sizes.rs allocated before #928.