[Fix] Don't free the multi-CTAs KV counter the decode graphs captured

[mmangkad]

Motivation

DeepseekSparseAttnBackend._forward_trtllm rebound the counter buffer:

self._multi_ctas_kv_counter_buffer = grow_multi_ctas_kv_counter_buffer_if_needed(...)

grow_..._if_needed returns a new torch.zeros when it grows, so the assignment frees the buffer the decode CUDA graphs captured. Later replays write through a dangling pointer.

batch_size is page_table_1.shape[0] — query rows, not requests — so a prefill batch trips it, not decode. Reasoning in decode batch sizes (max_running_requests x draft_tokens) makes TRTLLM_MLA_MAX_BATCH_SIZE=8192 look unreachable, which is why it hid.

Stock main, default flags, GLM-5.3-Flash tp4 / 4x GB300, 8192 in / 1024 out, temporary logging on all 4 ranks:

capture uses ptr=0xf0eb9d06dc00 bs=288      target-verify graph
capture uses ptr=0xf0e91ee00000 bs=48       draft-decode graph
grow fired: bs=16384 old_ptr=0xf0eb9d06dc00 new_ptr=0xf0e7c8e00000
            mode=EXTEND OLD_IS_THE_CAPTURED_ONE=True

Last capture ended 22 s before the first grow. Latent, not benign: a probe right after the grow did not land on the freed block, so today the stale writes hit memory nothing owns. Perturb the allocation pattern and it becomes a reproducible illegal access inside fmhaSm100fKernel...DynamicTokenSparse, the kernel that consumes the counter.

Modifications

_multi_ctas_kv_counter_for	oversized eager call takes a temporary; the captured allocation is never replaced
_ensure_multi_ctas_kv_counter_capacity	init_cuda_graph_state sizes for max(max_bs, max_num_tokens) before any capture; grow-only
class default = None	init_cuda_graph_state sizes it for every backend, but only the TRT-LLM branch allocates one

Capture-time growth is then unreachable, and the assert guarding it is an invariant check rather than a live path.

Accuracy Tests

Not applicable — buffer lifetime only, no numerics touched.

Benchmarking and Profiling

GLM-5.3-Flash tp4 / 4x GB300, 8192 in / 1024 out, 3 runs each, same build:

	tok/s	TTFT	TPOT
main	404.0	269.3 ms	2.21 ms
+ fix	402.4	269.9 ms	2.22 ms

TTFT is the number at risk (the oversized path now allocates per call across 11 DSA layers), so it was measured rather than assumed. All inside noise.

Checklist

• test/registered/kernel/attention/test_dsa_multi_ctas_counter.py, 7 cases
• Each fix mutated independently: 2 / 1 / 1 failures, 7 passed unmutated
• check_registered_tests.py and git diff --check exit 0

---

CI States

Latest PR Test (Base): ❌ Run #34830045556
Latest PR Test (Extra): ❌ Run #34830045213
Latest PR Test (AMD ROCm 10): ❌ Run #34830045737

[mmangkad]

/rerun-test test/registered/kernel/attention/test_dsa_multi_ctas_counter.py test/registered/e2e/models/test_glm53_flash_b200.py test/registered/e2e/models/test_dsa_glm52_nvfp4_tp_mtp.py

[github-actions]

Results for /rerun-test test/registered/kernel/attention/test_dsa_multi_ctas_counter.py test/registered/e2e/models/test_glm53_flash_b200.py test/registered/e2e/models/test_dsa_glm52_nvfp4_tp_mtp.py:

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/kernel/attention/test_dsa_multi_ctas_counter.py

🚀 4-gpu-b200 (2 tests): ❌ View workflow run

cd test/ && python3 registered/e2e/models/test_glm53_flash_b200.py
cd test/ && python3 registered/e2e/models/test_dsa_glm52_nvfp4_tp_mtp.py

[github-actions]

Results for /rerun-test test/registered/kernel/attention/test_dsa_multi_ctas_counter.py test/registered/e2e/models/test_glm53_flash_b200.py test/registered/e2e/models/test_dsa_glm52_nvfp4_tp_mtp.py:

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/kernel/attention/test_dsa_multi_ctas_counter.py

🚀 4-gpu-b200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/e2e/models/test_glm53_flash_b200.py
cd test/ && python3 registered/e2e/models/test_dsa_glm52_nvfp4_tp_mtp.py