[None][fix] Size the trtllm-gen multi-CTA KV counter buffer for beam search

[pranav-nvidia]

Description

The trtllm-gen decode kernel keeps one multi-CTA KV counter per attention head per decoder sequence. _get_multi_ctas_kv_counter_size was sized from metadata.max_num_requests alone, omitting the beam-width expansion, so beam search rejects the buffer at warmup with:

TVM_FFI_CHECK(multi_ctas_kv_counter_size >= counter_bytes)

The max(..., multi_processor_count) floor hides the missing factor at small batch, so the failure appears as a threshold rather than a universal beam-search break. With 6 decoder heads on a 148-SM part at beam 2:

max_batch_size	allocated max(6B, 148)	required 6·B·2	result
8	148	96	pass
16	148	192	fail
32	192	384	fail

Not model-specific: any model reaching num_heads · batch · beam_width > SM count under beam search hits it. Present since the function was introduced in 487330e8a03 (2026-06-09); ships in v1.3.0rc19–rc22.

The fix passes metadata.max_num_requests * metadata.beam_width at the single call site. Sizing up is safe — the kernel check is >=, and for cross-attention (K/V not beam-packed) the only cost is 4 · num_heads · batch · (beam − 1) bytes.

Test Coverage

• tests/unittest/_torch/attention/test_fmha_page_index.py — two cases: the old sizing under-allocates for a beam-expanded batch above the SM floor, and the multi-processor floor is preserved at batch 1.
• Validated by controlled A/B on B200 (single-file read-only shim, so the fix is the only delta) with a full concurrent batch in flight: max_batch_size 16 and 32 go fail → pass, 8 stays passing.

Why CI missed it: every in-tree beam test sits in the masked region — the disagg beam test uses max_beam_width=2 at low concurrency, test_cli_flow beam tests target the legacy CLI, and the Whisper beam_search case uses max_batch_size=2.

PR Checklist

• PR title follows [JIRA/NVBUG/None][type] summary
• Commit is DCO signed off
• New unit tests added
• NVBug to be filed and title updated

GitHub Bot Help

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...