[leejnau]

Motivation

This PR bumps FlashInfer to 0.7.3. It also brings in @trevor-m's migration of the NVLink one-sided all-to-all to flashinfer.moe_ep from #46885, because neither change can land without the other:

• The bump needs the migration. FlashInfer 0.7.3 removes flashinfer.comm.MoeAlltoAll and moe_a2a_get_workspace_size_per_rank, which the current FlashinferDispatcher uses (feat(moe_ep)!: sync NVLinkOneSidedAlltoAll with TensorRT-LLM's kernels; remove the deprecated MoeAlltoAll/MnnvlMoe flashinfer-ai/flashinfer#6047, [Feature] support add custom middlewares #6621).
• The migration needs the bump. [Chore] flashinfer: Use flashinfer moe-ep API instead of deprecated NVLinkOneSidedAllToAll API #46885 passes runtime_max_tokens_per_rank to flashinfer.moe_ep, and that argument only exists from 0.7.3. In 0.7.2 it is max_tokens_per_rank.

#46926 (0.7.2.post1) stays open for now as a fallback.

Depends on

• [Fix] [NVIDIA] deepseek-v4: Pass a caller-owned multi-CTA counter buffer to the trtllm sparse MLA kernel #48467 (DeepSeek-V4 trtllm counter buffer), now merged. Main's deepseek_v4_trtllm_backend.py patches flashinfer.mla._core._get_trtllm_gen_multi_ctas_kv_counter_buffer. 0.7.3 no longer has that name, so --dsv4-attn-backend trtllm fails at init. [Fix] [NVIDIA] deepseek-v4: Pass a caller-owned multi-CTA counter buffer to the trtllm sparse MLA kernel #48467 works on 0.7.2 and 0.7.3, so it can merge first.

Modifications

• Commits from [Chore] flashinfer: Use flashinfer moe-ep API instead of deprecated NVLinkOneSidedAllToAll API #46885, cherry-picked with @trevor-m as author:
  • Switch NVLinkOneSidedAllToAll to the flashinfer.moe_ep API.
  • Lint fix.
  • Pass runtime_max_tokens_per_rank in dispatch.
• Pass the new last_valid_index argument to flashinfer::sampling::DeviceSamplingFromProb in SGLang's JIT tree speculative-sampling kernel (sgl_kernel/speculative/sampling.cuh). perf(sampling): defer last-valid-index reduction until fallback flashinfer-ai/flashinfer#5352 (in 0.7.3) added that parameter. Without this change the kernel fails to compile, and the first EAGLE verify step crashes the scheduler. The change mirrors FlashInfer's own chain speculative-sampling kernel, and it only compiles against 0.7.3 headers, so it has to land with the bump.
• Pass the new fixed_uniform_q_len argument in the FlashInfer backend's fast_prefill_plan. FlashInfer 0.7.3's FA2 paged-prefill plan() takes 21 arguments instead of 20. The fast plan also clears _eagle_verify_fp8kv_sm90_state, as FlashInfer's own plan() now does on every call.
• MiniMax-M3.1's fused NVFP4 KV store now zeroes the payload when a block's scale underflows to zero, to match FlashInfer 0.7.3's quantizer (fix(quant): guard the fp8-rounded SF against underflow in block-scale quantization flashinfer-ai/flashinfer#6996 changed its guard from amax != 0 to sf != 0).
• Temporary workaround for [Bug] trtllm-gen MoE small-batch finalize (#6546) makes MoE output depend on batch size, breaking prefill/decode determinism flashinfer-ai/flashinfer#8831. FlashInfer 0.7.2.post1 and 0.7.3 include perf(moe): split the topK reduction in small-batch trtllm-gen MoE finalize flashinfer-ai/flashinfer#6546, which adds a small-batch MoE finalize kernel for 4 < top-k <= 64. That kernel adds the top-k expert outputs in a different order from the large-batch kernel, so a token's MoE output depends on its batch size, and prefill and decode logprobs stop matching. With --enable-deterministic-inference, the pre-routed NVFP4 trtllm MoE path now calls FlashInfer with do_finalize=False and sums the outputs with SGLang's moe_finalize_fuse_shared, which uses one order at every batch size. moe_finalize_fuse_shared gains an out= argument so the sum lands in the usual output buffer. Nothing changes outside deterministic mode. Remove this once [Bugfix] Fixing the bug when max_new_tokens is None. #8831 is fixed in a FlashInfer release.
• Bump FlashInfer from 0.7.2 to 0.7.3 in docker/Dockerfile, docker/Dockerfile.cu134, python/pyproject.toml and the engine.py version check.

0.7.3 requires nvidia-cutlass-dsl[cu13]>=4.8.0, and SGLang already pins 4.8.0.

Tests

• CI. The first run is expected to hit the cold SGLang JIT cache in the jit-kernel suite, as on earlier FlashInfer bumps.

• On GB200, [Fix] [NVIDIA] deepseek-v4: Pass a caller-owned multi-CTA counter buffer to the trtllm sparse MLA kernel #48467's DeepSeek-V4 trtllm prefill test passes on 0.7.3 with the fix, and fails without it.

• On GB200 with FlashInfer 0.7.3, test/registered/kernels/ops/speculative/test_speculative_sampling.py passes (11 passed).

• TestInklingSmallNvfp4DsparkDeterministic.test_input_output_logprobs_match (KL threshold 1e-9) on GB200, with the test's --tp changed from 4 to 2 (GB200 at TP=4 fails even on main with 0.7.2):

  Code	FlashInfer	avg KL
  main	0.7.2	0.0 (3 runs)
  main	0.7.3	5.6e-6 to 1.6e-3 (several runs)
  main	0.7.2.post1, small-batch finalize disabled	0.0
  this PR, with the workaround	0.7.3 + matching jit-cache	0.0 (2 of 2 runs)

• @trevor-m validated the moe_ep migration on 0.7.3 on 4×GB200 (see [Chore] flashinfer: Use flashinfer moe-ep API instead of deprecated NVLinkOneSidedAllToAll API #46885): GSM8K 96%, and test_flashinfer_a2a_wide_ep.py, test_flashinfer_dispatcher.py and test_muse_spark_runner_dispatch.py 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%
  [Fix] [NVIDIA] deepseek-v4: Pass a caller-owned multi-CTA counter buffer to the trtllm sparse MLA kernel #48467	0.7.2	96.74%	0%
  this PR + [Fix] [NVIDIA] deepseek-v4: Pass a caller-owned multi-CTA counter buffer to the trtllm sparse MLA kernel #48467	0.7.3	96.59%	0%

  Before the speculative-sampling fix, the 0.7.3 run scored 0%, because every request failed when the scheduler crashed on the JIT build error.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ⏳ Run #42122889789
Latest PR Test (Extra): ⏳ Run #42122889726
Latest PR Test (AMD ROCm 10): ⏳ Run #42122889643

Tip

Base CI only covers the core cases, and its workload may not exercise your change. Use /rerun-test to run the tests your change touches, or copy this prompt for an agent:

Find the registered tests this PR affects and run them with /rerun-test, following https://github.com/sgl-project/sglang/blob/main/.agents/skills/rerun-test/SKILL.md

[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.

[bobboli]

Hi @leejnau , could you rename the backend name to flashinfer_nvlink_one_sided, as stated in #46885 (comment)?

[leejnau]

Thanks @bobboli. I'd like to keep this PR to the FlashInfer 0.7.3 bump and the changes 0.7.3 requires, so that it stays easy to review and revert. As @trevor-m noted on #46885, renaming the backend everywhere touches 40+ files.

A smaller way to get the name you suggested: accept --moe-a2a-backend flashinfer_nvlink_one_sided as the documented name and keep flashinfer as a deprecated alias with a warning, normalized in one place in server args, with the docs updated. Internal identifiers can follow later. That doesn't depend on 0.7.3, so I'll open it as a separate PR right after this one. Does that work for you?

[bobboli]

accept --moe-a2a-backend flashinfer_nvlink_one_sided as the documented name and keep flashinfer as a deprecated alias with a warning

Yes, that's the best way. Feel free to do it in the follow up PR. Thanks!

[nvpohanh]

[by Claude Code] This review is done by agents. Please feel free to disagree if they do not make sense.

[nvpohanh]

[by Claude Code] Severity: functional | Confidence: High

Issue

test/manual/ep/test_flashinfer_dispatcher.py still uses the removed MoeAlltoAll API, and this migration breaks it. Instances built here share one workspace when their geometry matches, and FlashInfer 0.7.3's check_round_owner rejects a dispatch while another instance on that workspace has not called combine. Each manual test calls dispatcher.dispatch() but never combine(). Its cleanup helpers import flashinfer.comm.trtllm_moe_alltoall, catch the ImportError, and do nothing.

Example

With torchrun --nproc_per_node=4 test_flashinfer_dispatcher.py, test_dispatch_basic leaves its round open, so test_dispatch_with_empty_tokens raises "another NVLinkOneSidedAlltoAll sharing this workspace has a dispatch awaiting its combine".

Fix

End each test with dispatcher.combine(...) or dispatcher.moe_a2a.destroy(), and port _zero_moe_a2a_dispatch_payloads and tearDownClass to the new API.

[leejnau]

Thanks. This is right, and it's tracked in #48506 along with a second problem in the same test. Its global_num_tokens layout doesn't match the dispatcher's sizing rules for its plain TP4 setup.

This PR doesn't make the test worse than main. On main with 0.7.2, all 5 tests already error while the dispatcher is built (num_fused_shared_experts=None in the fallback MoeRunnerConfig, which #46885's commit here fixes). On this PR test_dispatch_basic passes and the other 4 fail as you describe.

It's a manual test that CI doesn't run, and fixing it means porting the cleanup helpers to moe_ep and adding combine() or destroy() to each test. So I'd rather keep it out of the FlashInfer bump and fix it in a follow-up under #48506.

[Fridge003]

I think it will be OK if we are not moving the affected test to CI test