[leejnau]

Motivation

--dsv4-attn-backend trtllm needs a multi-CTA KV counter buffer sized by query rows. FlashInfer's default buffer is sized by request count, while the DSv4 VarSeq kernel indexes it by query row.

The backend worked around this by replacing the private flashinfer.mla._core._get_trtllm_gen_multi_ctas_kv_counter_buffer. That has two problems.

1. The patch doesn't take effect on the current FlashInfer pin.
   • In 0.7.2 and 0.7.2.post1, trtllm_batch_decode_sparse_mla_dsv4 gets its buffer through flashinfer.utils._resolve_trtllm_gen_multi_ctas_kv_counter_buffer.
   • When no buffer is passed in, that function calls _get_trtllm_gen_multi_ctas_kv_counter_buffer from flashinfer.utils's own namespace, not from flashinfer.mla._core.
   • So the replaced function is never called on this path. FlashInfer allocates a new buffer sized by request count on every launch.
2. It breaks with FlashInfer 0.7.3. flashinfer.mla._core no longer imports that name, so backend init raises AttributeError. This blocks the FlashInfer 0.7.3 bump ([Chore] deps: Bump FlashInfer to 0.7.2.post1 #46926).

The workaround's docstring says to drop it "once FlashInfer accepts a caller-owned buffer". The function has accepted one through its public multi_ctas_kv_counter_buffer argument since at least 0.7.2. The MLA, MHA and DSA trtllm backends already pass their own buffers this way.

Modifications

All changes are in python/sglang/srt/layers/attention/deepseek_v4_trtllm_backend.py.

• Remove the monkeypatch and the import of flashinfer.mla._core.
• _reserve_trtllm_query_rows (renamed from _install_persistent_trtllm_semaphores) only records the query-row capacity, computed as before from chunked prefill, max prefill tokens and running requests × draft tokens.
• _get_trtllm_kv_counter_buffer(num_heads, device) returns one zeroed persistent buffer per head count.
  • It is sized with FlashInfer's public get_trtllm_gen_multi_ctas_kv_counter_bytes(capacity_rows, num_heads, sm_count).
  • It is held in the runtime context's persistent buffers, like the existing trtllm workspace.
  • It is created outside graph capture; the assert is unchanged.
• Both trtllm_batch_decode_sparse_mla_dsv4 call sites (decode and varlen prefill) pass it as multi_ctas_kv_counter_buffer. The existing _check_trtllm_query_rows guard still applies.
• The manual Blackwell prefill test also passes the buffer.

This works with FlashInfer 0.7.2, 0.7.2.post1 and 0.7.3. The public argument and helpers exist in all three.

Tests

• New unit test test_trtllm_kv_counter_buffer_is_sized_by_query_rows_and_shared in test/registered/attention/unittests/dsv4/test_deepseek_v4.py. It checks:

  • the buffer is sized by reserved query rows, not request count;
  • it is zero-initialized;
  • it is reused across calls;
  • each head count gets its own buffer.

• On GB200 (SM100), in lmsysorg/sglang:dev (arm64) with torch 2.14.1+cu130 and nvidia-cutlass-dsl 4.8.0:

  FlashInfer	Code	test_dsv4_trtllm_prefill.py (rows 1/16/17/8192)	New unit tests
  0.7.2	main	pass	n/a
  0.7.2	this PR	pass	2 passed
  0.7.3	main	fail: AttributeError: module 'flashinfer.mla._core' has no attribute '_get_trtllm_gen_multi_ctas_kv_counter_buffer'	n/a
  0.7.3	this PR	pass	2 passed

• End-to-end on one 4×GB200 node: DeepSeek-V4-Flash with the cookbook's verified GB200 low-latency recipe (TP=4, flashinfer_mxfp4 MoE, EAGLE with 3 steps and 4 draft tokens) plus --dsv4-attn-backend trtllm, image lmsysorg/sglang:dev (arm64). Full GSM8K (1319 questions) through sgl-eval 0.1.3 at temperature 1.0 and top_p 1.0, the cookbook's recommended sampling:

  Code	FlashInfer	GSM8K	Request errors
  main	0.7.2	96.66%	0%
  this PR	0.7.2	96.74%	0%
  [Chore] [NVIDIA] deps: Bump FlashInfer to 0.7.3 and migrate the NVLink one-sided all-to-all to flashinfer.moe_ep #48468 + this PR	0.7.3	96.59%	0%

🤖 Generated with Claude Code

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[trevor-m]

LGTM

[trevor-m]

Can you do an e2e GSM8K test?

[Fridge003]

/rerun-test test/registered/attention/unittests/dsv4/test_deepseek_v4.py test/registered/e2e/models/test_deepseek_v4_flash_fp4_b200_trtllm.py

[github-actions]

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

cd test/ && python3 registered/attention/unittests/dsv4/test_deepseek_v4.py
cd test/ && python3 registered/e2e/models/test_deepseek_v4_flash_fp4_b200_trtllm.py

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

cd test/ && python3 registered/attention/unittests/dsv4/test_deepseek_v4.py

[leejnau]

@trevor-m Done. Full GSM8K (1319 questions) on 4×GB200 with DeepSeek-V4-Flash, using the cookbook's GB200 low-latency recipe (TP=4, flashinfer_mxfp4 MoE, EAGLE with 3 steps and 4 draft tokens) plus --dsv4-attn-backend trtllm. sgl-eval 0.1.3, temperature 1.0, top_p 1.0:

Code	FlashInfer	GSM8K
main	0.7.2	96.66%
this PR	0.7.2	96.74%
#48468 + this PR	0.7.3	96.59%

No request errors in any run. The 0.7.3 run also caught a third incompatibility: FlashInfer #5352 changed DeviceSamplingFromProb, which breaks SGLang's JIT speculative-sampling kernel. That fix is now in #48468.