[NVIDiA][MoE][FP4] Enable NVFP4 shared-expert fusion for FlashInfer TRTLLM-gen kernels

[wenscarl]

---

Summary

Wire num_fused_shared_experts through the FlashInfer TRTLLM FP4 MoE forward path so that the shared expert in nvidia/DeepSeek-R1-0528-NVFP4-v2 is fused into the same kernel call as the routed experts, eliminating a separate dense GEMM per MoE layer.

This is the FP4 counterpart of the FP8 shared-expert fusion already present in this codebase, and depends on the FlashInfer-side implementation in flashinfer-ai/flashinfer#4239.

Target model: nvidia/DeepSeek-R1-0528-NVFP4-v2 (https://huggingface.co/nvidia/DeepSeek-R1-0528-NVFP4-v2) — the only publicly available model with NVFP4-quantized MoE experts including a shared expert slot.

---

Changes

python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py

• Add _FP4_TRTLLM_HAS_FUSED_SHARED and _FP4_TRTLLM_ROUTED_HAS_FUSED_SHARED module-level flags, computed once at import via inspect.signature on trtllm_fp4_block_scale_moe and trtllm_fp4_block_scale_routed_moe
• In fused_experts_none_to_flashinfer_trtllm_fp4: inject num_fused_shared_experts into the bypassed (logits) path kwargs when _FP4_TRTLLM_HAS_FUSED_SHARED is True; inject into routed path when _FP4_TRTLLM_ROUTED_HAS_FUSED_SHARED is True (not yet available in current nightly)
• Add EP guard: raise NotImplementedError when num_fused_shared > 0 and local_num_experts < global_num_experts

python/sglang/srt/arg_groups/overrides.py

• Add _fp4_trtllm_supports_fused_shared() helper that inspects the installed FlashInfer signature
• Lift the blanket auto-disable of shared-expert fusion for flashinfer_trtllm / flashinfer_trtllm_routed when quantization == "modelopt_fp4" and the installed FlashInfer exposes the parameter. All other runners and quantizations are unchanged.

test/registered/quant/test_nvfp4_shared_expert_fusion.py (new)

• 12 dummy-weight tests, no GPU kernels executed, covering: capability flag detection, overrides gate logic (FP4 vs non-FP4, supported vs unsupported FlashInfer), EP + capability runtime guards, and kwarg pass-through for both forward paths

---

Backward compatibility

The gate is fully backward-compatible. On FlashInfer ≤ 0.6.15.post1 (without the param), both flags are False and behavior is identical to before this PR — shared-expert fusion remains auto-disabled for the FP4 path.

---

EP not supported

Expert parallelism is explicitly unsupported for this fused FP4 path, consistent with the FlashInfer maintainers' guidance. Both determine_num_fused_shared_experts (model level) and the runner-level NotImplementedError guard it independently.

---

Test plan

• 12 unit tests pass (test/registered/quant/test_nvfp4_shared_expert_fusion.py)
• Server launch with --moe-runner-backend flashinfer_trtllm --tp 8 on nvidia/DeepSeek-R1-0528-NVFP4-v2 confirmed num_fused_shared_experts=1 passed to FlashInfer bypassed kernel
  Accuracy validation

Server A — fused path:

python3 -m sglang.launch_server \
  --model-path nvidia/DeepSeek-R1-0528-NVFP4-v2 \
  --trust-remote-code \
  --disable-radix-cache \
  --kv-cache-dtype fp8_e4m3 \
  --max-running-requests 1024 \
  --chunked-prefill-size 16384 \
  --mem-fraction-static 0.8 \
  --max-prefill-tokens 16384 \
  --moe-runner-backend flashinfer_trtllm \
  --tp 8 \
  --port 30000

Server B — baseline (fusion disabled by --disable-shared-experts-fusion \):

Client — A/B accuracy check (temperature=0 for deterministic output):
curl -s http://localhost:30000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"model":"nvidia/DeepSeek-R1-0528-NVFP4-v2","messages":[{"role":"user","content":"1+1=?"}],"max_tokens":32,"temperature":0}' \
  | jq -r '.choices[0].message.content'

curl -s http://localhost:30001/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"model":"nvidia/DeepSeek-R1-0528-NVFP4-v2","messages":[{"role":"user","content":"1+1=?"}],"max_tokens":32,"temperature":0}' \
  | jq -r '.choices[0].message.content'

• Numerical equivalence check between fused and non-fused paths

---

CI States

Latest PR Test (Base): ❌ Run #30660223310
Latest PR Test (Extra): ❌ Run #30660222877

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.