gqa_pack_shape_ok

Function gqa_pack_shape_ok 

Source
pub fn gqa_pack_shape_ok(
    num_q_heads: u32,
    num_kv_heads: u32,
    head_dim: u32,
) -> bool
Expand description

Whether this launch’s SHAPE is one the packed kernels can serve.

Every condition is a hard precondition of the kernel, not a preference:

  • num_kv_heads > 0 — it is the grid’s x extent and the divisor below.
  • num_q_heads == num_kv_heads * DECODE_GQA_PACK_WIDTH — the kernel recovers its heads as kv_head * PD_GQA + h, the exact inverse of the unpacked kernel’s q_head / gqa_ratio. Any other ratio reads and writes the wrong heads, silently, so it is refused rather than clamped. Stated as a MULTIPLICATION on purpose: num_q_heads / num_kv_heads truncates, so 25 q heads over 4 kv heads would pass a division test and then have the last head read past the group.
  • head_dim == DECODE_GQA_PACK_HEAD_DIM — see that constant.

The caller keeps the unpacked kernel when this is false. That is a documented route, not a fallback that hides an error: the unpacked kernel is the one every target has been serving.