feat: add dense CuTe DSL NVFP4 W4A16 mode

[zianglih]

Motivation

@HumansAnd

Extend the opt-in SGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16 mode from #35120 to serialized ModelOpt NVFP4 dense linear layers selected by --fp4-gemm-backend flashinfer_cutedsl. Dense weights remain NVFP4 while activations and outputs remain BF16.

Stack and dependency contract:

• This branch is stacked by ancestry on [FlashInfer v0.6.18] add FlashInfer CuTe DSL NVFP4 W4A16 mode #35120 at d01db73b433ce975ee055fe21499d31d72571098; [FlashInfer v0.6.18] add FlashInfer CuTe DSL NVFP4 W4A16 mode #35120 must merge first. Review the dense-only commit range in this stack comparison. After [FlashInfer v0.6.18] add FlashInfer CuTe DSL NVFP4 W4A16 mode #35120 lands, this branch will be rebased onto current main so the upstream PR contains only the dense integration. Dense commit: d7959d0ca0ddc47326ddd1a05f2fd41dbf1864b2.
• Merge gate: SGLang's pinned FlashInfer dependency and CI image must use a stable release containing the dependency chain inherited from [FlashInfer v0.6.18] add FlashInfer CuTe DSL NVFP4 W4A16 mode #35120 (#4048, #4394, and #4558) plus the dense SM100 CuTe DSL W4A16 GEMM from #4466, merged as 94f6ae8ee90fa4d99c8fe91729da57e89fdc6997. This integration intentionally has no compatibility fallback for earlier FlashInfer APIs.
• Upstream dependency status (2026-08-18): all four FlashInfer PRs above are merged. The latest stable FlashInfer release remains v0.6.17 and does not contain the full required chain, so this PR must remain gated on a later stable release.
• Validation uses FlashInfer tag nightly-v0.6.18-20260818 at 0751be928d24187bd39a0dde45ce96418bc03ddf, whose ancestry includes [Fix] use torch.cat instead of torch.concat to prevent entering the Autograd backends. #4466. The published Python wheel reports build commit 02647e81072fe23da164fdf5e179adf78e03a2ec, the immediate pre-[Fix] use torch.cat instead of torch.concat to prevent entering the Autograd backends. #4466 parent, so validation installs the requested nightly Python/cubin/cu130-JIT package set and overlays the exact tag source. This ensures the run exercises [Fix] use torch.cat instead of torch.concat to prevent entering the Autograd backends. #4466 rather than treating the release tag alone as sufficient evidence.

Modifications

• Reuse the existing ModelOptNvFp4A16LinearMethod weight-only path. When the environment variable and flashinfer_cutedsl dense backend are both selected, regular serialized non-AWQ ModelOpt NVFP4 linear layers use this method; all existing W4A4 selections remain unchanged when the opt-in is off or another dense backend is selected.
• Prepare packed checkpoint weights through FlashInfer's public prepare_bf16_fp4_weights(..., backend="cute-dsl") API, then invoke mm_bf16_fp4(..., backend="cute-dsl", out_dtype=torch.bfloat16) through a registered SGLang custom op with a fake implementation for compile/graph tracing. This folds the new kernel into the existing A16 linear backbone rather than adding a parallel wrapper.
• Preserve the loader-owned raw scales on the FlashInfer path and update prepared weights, MMA-layout scales, and global alpha through stable parameter bindings. /update_weights_from_disk can therefore reload the original checkpoint tensors and rebuild the derived runtime state without invalidating existing parameter identities when shape and dtype are unchanged.
• Extend the generic LM-head quant-method guard to recognize FlashInfer's prepared uint8 weight plus MMA-layout scale/global-alpha attributes in addition to the existing Marlin int32 representation; this does not add model-specific LM-head handling.
• Keep AWQ outside this mode. MoE behavior, FlashInfer A2A behavior, and the environment-variable contract remain owned by [FlashInfer v0.6.18] add FlashInfer CuTe DSL NVFP4 W4A16 mode #35120.
• Reuse TestServerUpdateWeightsFromDiskNVFP4W4A16CuteDSL from [FlashInfer v0.6.18] add FlashInfer CuTe DSL NVFP4 W4A16 mode #35120 rather than adding another test class. Its --fp4-gemm-backend flashinfer_cutedsl plus SGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16=1 now exercises the dense W4A16 path alongside the existing MoE path.

Accuracy Tests

Environment:

• Image: lmsysorg/sglang:nightly-dev-cu13-20260818-c0b6474b
• Image manifest digest: sha256:51e576f02368480c055c7aadb67590d82b172e2392123ce4cf4cc8251b2d8caf
• Hardware: C2 devbox, 4 of 8 NVIDIA B300 SXM6 AC GPUs (SM103), driver 590.48.01; TP4/DP4/EP4 with DP attention and FlashInfer A2A
• Runtime: Python 3.12.3, PyTorch 2.13.0+cu130, CUDA 13.0, nvidia-cutlass-dsl==4.6.2
• Model: nvidia/Qwen3-30B-A3B-NVFP4
• Model snapshot: 2538ded2a4edb247b4d2b4a8ba24e44bd4c017c3
• SGLang stack base: d01db73b433ce975ee055fe21499d31d72571098; tested dense head: d7959d0ca0ddc47326ddd1a05f2fd41dbf1864b2
• FlashInfer runtime source: tag nightly-v0.6.18-20260818 at 0751be928d24187bd39a0dde45ce96418bc03ddf
• Installed release packages: flashinfer-python==0.6.18.dev20260818, flashinfer-cubin==0.6.18.dev20260818, and flashinfer-jit-cache==0.6.18.dev20260818+cu130 (cp39-abi3, manylinux_2_28_x86_64), with the five [Fix] use torch.cat instead of torch.concat to prevent entering the Autograd backends. #4466 runtime files overlaid from the exact tag source because the published Python wheel predates the merge

FlashInfer package and exact-tag runtime setup

base=/hai-workspace/sglang-w4a16-dense-validation
mkdir -p "$base"/{cache,logs,src,wheels}
cd "$base/wheels"

curl -fL --retry 3 -o flashinfer_python-0.6.18.dev20260818-py3-none-any.whl \
  https://github.com/flashinfer-ai/flashinfer/releases/download/nightly-v0.6.18-20260818/flashinfer_python-0.6.18.dev20260818-py3-none-any.whl
curl -fL --retry 3 -o flashinfer_cubin-0.6.18.dev20260818-py3-none-any.whl \
  https://github.com/flashinfer-ai/flashinfer/releases/download/nightly-v0.6.18-20260818/flashinfer_cubin-0.6.18.dev20260818-py3-none-any.whl
curl -fL --retry 3 -o flashinfer_jit_cache-0.6.18.dev20260818+cu130-cp39-abi3-manylinux_2_28_x86_64.whl \
  'https://github.com/flashinfer-ai/flashinfer/releases/download/nightly-v0.6.18-20260818/flashinfer_jit_cache-0.6.18.dev20260818%2Bcu130-cp39-abi3-manylinux_2_28_x86_64.whl'

sha256sum *.whl
python3 -m pip install --no-deps --force-reinstall ./*.whl

src="$base/src/flashinfer-nightly-v0.6.18-20260818"
dst=/usr/local/lib/python3.12/dist-packages
backup="$base/overlay-backup"
git clone --depth 1 --branch nightly-v0.6.18-20260818 \
  https://github.com/flashinfer-ai/flashinfer.git "$src"
test "$(git -C "$src" rev-parse HEAD)" = \
  0751be928d24187bd39a0dde45ce96418bc03ddf

files=(
  flashinfer/gemm/gemm_bf16_fp4.py
  flashinfer/gemm/gemm_bf16_fp4_cute_dsl.py
  flashinfer/gemm/kernels/cute_dsl/dense_gemm_bf16_fp4_sm100.py
  flashinfer/gemm/kernels/cute_dsl/dense_gemm_bf16_fp4_sm100_utils.py
  flashinfer/gemm/kernels/cute_dsl/dense_gemm_bf16_fp4_sm12x.py
)
for f in "${files[@]}"; do
  mkdir -p "$backup/$(dirname "$f")" "$dst/$(dirname "$f")"
  if [ -f "$dst/$f" ]; then
    cp -a "$dst/$f" "$backup/$f"
  fi
  cp -a "$src/$f" "$dst/$f"
  test "$(sha256sum "$src/$f" | cut -d' ' -f1)" = \
    "$(sha256sum "$dst/$f" | cut -d' ' -f1)"
done

sglang_src="$base/src/sglang-d7959d0ca"
git clone --filter=blob:none https://github.com/zianglih/sglang.git "$sglang_src"
git -C "$sglang_src" fetch origin \
  d7959d0ca0ddc47326ddd1a05f2fd41dbf1864b2
git -C "$sglang_src" checkout --detach \
  d7959d0ca0ddc47326ddd1a05f2fd41dbf1864b2
test "$(git -C "$sglang_src" rev-parse HEAD)" = \
  d7959d0ca0ddc47326ddd1a05f2fd41dbf1864b2

Wheel SHA-256 values, in Python/cubin/cu130-JIT order: 23ed75e3af20e704f4bbdcc7e3e1d9471620ab5a16443b3e7bc28477a2244c3c, 430eff3b305be5aabc7345b5c19b88f4dada61749e278bd801a6173728de157b, and e26092488089a0ab87b1e3360a03e3faa6a2b9b653d59eaf5f287830af5b25cb.

Command:

export PYTHONPATH=/hai-workspace/sglang-w4a16-dense-validation/src/sglang-d7959d0ca/python
export HF_HOME=/hai-workspace/.cache/huggingface
export SGLANG_CACHE_DIR=/hai-workspace/sglang-w4a16-dense-validation/cache/d7959d0ca-w4a16-dense-002/sglang
export FLASHINFER_WORKSPACE_BASE=/hai-workspace/sglang-w4a16-dense-validation/cache/d7959d0ca-w4a16-dense-002/flashinfer
export CUDA_VISIBLE_DEVICES=0,1,2,3
export SGLANG_ENABLE_ASYNC_ASSERT=0
export TOKENIZERS_PARALLELISM=false
cd /hai-workspace/sglang-w4a16-dense-validation/src/sglang-d7959d0ca/test

python3 registered/rl/test_update_weights_from_disk_blackwell.py \
  TestServerUpdateWeightsFromDiskNVFP4W4A16CuteDSL.test_parameterized_update_weights_from_disk \
  -v -f

This registered test launches deterministic inference with BF16 dense activations and outputs using --fp4-gemm-backend flashinfer_cutedsl, --moe-runner-backend flashinfer_cutedsl, and --moe-a2a-backend flashinfer. It performs both disk-reload variants (flush_cache=true and false) across weight/CUDA-graph release and resume, then requires unchanged decode text and token IDs and token logprobs within 1e-4. Decode CUDA-graph capture completed all 52 batch sizes through 512, and runtime decode used CUDA graphs.

The CuTe DSL dense runner does not emit its op key in this deterministic configuration. Dense-path provenance was therefore checked from the resolved SGLang implementation: this env/backend pair selects ModelOptNvFp4A16LinearMethod, whose prepare/apply path calls the two FlashInfer W4A16 APIs above. The runtime-overlay manifest records matching exact-tag source and installed hashes for all five #4466 files.

Unittest and timing tail (the full server log is identified below):

ok
----------------------------------------------------------------------
Ran 1 test in 298.955s

OK
real 320.67
user 64.99
sys 7.96
end_utc=2026-08-18T06:50:00Z
elapsed_seconds=321
exit_status=0

GPU-validation log: /hai-workspace/sglang-w4a16-dense-validation/logs/d7959d0ca-w4a16-dense-002.log, SHA-256 744304b54ab3d525ecc0b10a1b496e90dd6ce41d4bda9a18d0dea437c3d881c9.

Static validation:

git diff --check d01db73b433ce975ee055fe21499d31d72571098..d7959d0ca0ddc47326ddd1a05f2fd41dbf1864b2
pre-commit run --files \
  docs/docs/references/environment_variables.mdx \
  python/sglang/srt/environ.py \
  python/sglang/srt/layers/logits_processor.py \
  python/sglang/srt/layers/quantization/modelopt_quant.py

Both commands passed; the pre-commit run included the repository's selected Ruff checks, Black, isort, codespell, and Python AST checks.

Scope notes:

• The Qwen3 checkpoint covers serialized dense projections and MoE reload behavior. The generic quantized-LM-head guard is updated, but this checkpoint does not independently exercise a quantized LM head.
• The passing test proves unchanged text, token IDs, and token logprobs for both reload variants; those values are asserted by the test and are not printed. The fatal-log scan found no traceback, watchdog timeout, CUDA/NCCL error, non-finite-output, or unittest-failure marker.
• AWQ, non-CuTe-DSL dense backends, and earlier FlashInfer releases are deliberately outside this PR.

Speed Tests and Profiling

No independent speed claim is made here. This PR only routes SGLang's existing A16 linear method to FlashInfer's public dense W4A16 API; FlashInfer #4466 owns the kernel microbenchmarks. The C2 run above is a correctness and disk-reload test, not a serving benchmark.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #32223961782
Latest PR Test (Extra): ❌ Run #32223961598

[zianglih]

Folded this dense W4A16 commit into #35120 at e48dd82 now that FlashInfer 0.6.18 contains both the MoE and dense dependencies. Continuing review and CI on #35120.