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