[Perf][DSV4.1] Fuse MXFP8 wo_b GEMM with sequence-parallel reduce-scatter

[gcanlin]

Purpose

Fuse DeepSeek-V4.1's MXFP8 wo_b GEMM and sequence-parallel reduce-scatter into one SM100 CuTe DSL kernel. Opt in with VLLM_DSV41_GEMM_RS=1 on TP4 within one NVLink multicast domain; batches below 1,024 tokens retain the existing GEMM + RS path. The kernel preserves MXFP8 operands and block-32 scales, reduces BF16 partials with FP32 accumulation, and shares one workspace across layers. Activation quantization remains a separate kernel.

Duplicate-work check: #56611 is closed and concerned shifted mHC overlap; #54151 implements Kimi-K3 AG-GEMM; #56751 optimizes tiny-M DSV4.1 projections. This change targets the large-M MXFP8 output projection and its RS, using Kimi's symmetric allocation/reduction helpers without invoking its BF16 GEMM. No matching open MXFP8 GEMM-RS PR was found.

Test Plan

Deploy

Executed independently on node17 and node30, four B200 GPUs per server, with TP4 + EP4 + MegaMoE + SP. Start fresh servers for VLLM_DSV41_GEMM_RS=0 and 1, then swap baseline/candidate across nodes. Both arms use source/native build 78eaaad2c7a7f8c7864d9c9c72772e3f9f2669d0 (including merged Mega-mHC), PyTorch 2.13.0+cu130, FlashInfer 0.6.18.post1, and CUTLASS DSL 4.7.1. The dependency image is vllm/vllm-openai:nightly-dc36fcce902a63eab06c1b93a5c4a5ee178a0c56; /opt/dsv41-native overlays native libraries and DeepGEMM from the exact source-commit wheel. Other FlashInfer versions have not been validated; this kernel currently imports private FlashInfer helpers.

export PYTHONPATH=/opt/dsv41-native:/gpfs/mszn/workspace/vllm
export OMP_NUM_THREADS=1 VLLM_USE_V2_MODEL_RUNNER=1
export VLLM_DSV41_GEMM_RS=1  # 0 for baseline
uv run --no-project /opt/dsv41-gemm-rs/.venv/bin/python \
  -m vllm.entrypoints.cli.main serve \
  /gpfs/mszn/models/deepseek-ai/DeepSeek-V4.1-Flash \
  --served-model-name dsv41 --host 127.0.0.1 --port 18517 \
  --tensor-parallel-size 4 --enable-expert-parallel \
  --moe-backend deep_gemm_mega_moe \
  --language-model-only --tokenizer-mode deepseek_v41 \
  --engram-config '{"cpu_offload":true}' \
  --max-model-len 9216 --max-num-batched-tokens 8192 --max-num-seqs 1024 \
  --gpu-memory-utilization 0.90 --no-enable-prefix-caching --seed 2026

Benchmark

The commands below retain the measured flags; result paths are shortened. Startup/autotuning and eight warmup requests per case are excluded from timing.

for concurrency in 16 64 964; do
  requests=$((concurrency * 2))
  if (( requests < 128 )); then requests=128; fi
  uv run --no-project /opt/dsv41-gemm-rs/.venv/bin/python \
    -m vllm.entrypoints.cli.main bench serve \
    --backend vllm --base-url http://127.0.0.1:18517 --model dsv41 \
    --tokenizer /gpfs/mszn/models/deepseek-ai/DeepSeek-V4.1-Flash \
    --tokenizer-mode deepseek_v41 \
    --dataset-name random --random-input-len 8192 --random-output-len 1024 \
    --random-range-ratio 0 --ignore-eos --temperature 0 --seed 2026 \
    --max-concurrency "$concurrency" --request-rate inf \
    --num-prompts "$requests" --num-warmups 8 \
    --percentile-metrics ttft,tpot,itl,e2el --metric-percentiles 50,90,99 \
    --save-result --save-detailed --result-dir results \
    --result-filename "gemm-rs-${VLLM_DSV41_GEMM_RS}-c${concurrency}.json"
done

Correctness and kernel benchmark (same Python/runtime environment):

uv run --no-project /opt/dsv41-gemm-rs/.venv/bin/python \
  -m pytest tests/kernels/test_dsv41_gemm_rs.py -q
uv run --no-project /opt/dsv41-gemm-rs/.venv/bin/python \
  -m torch.distributed.run --standalone --nproc-per-node 4 \
  benchmarks/kernels/benchmark_dsv41_gemm_rs.py \
  --tokens 128 256 512 1024 2048 4096 8192 \
  --workspaces 10 --samples 40 --output results/kernel.json

Test Result

DeepSeek-V4.1-Flash, 4×B200/server, TP4/EP4/MegaMoE/SP, random 8192-input/1024-output tokens. One measured run per node/mode/concurrency: 128, 128, and 1,928 requests respectively. Aggregate throughput is total output tokens divided by total duration across both nodes; mean latencies are request-weighted.

Concurrency	Output tok/s: baseline → fused	Change	Mean TTFT (s): baseline → fused	Mean TPOT (ms): baseline → fused
16	1469.0 → 1498.0	+2.0%	0.831 → 0.763	10.08 → 9.94
64	3147.2 → 3197.8	+1.6%	2.664 → 2.464	17.60 → 17.48
964	6419.4 → 6799.7	+5.9%	36.524 → 34.376	110.43 → 103.96

Concurrency	Node17 output throughput change	Node30 output throughput change
16	+1.3%	+2.6%
64	+10.9%	-7.0%
964	+4.4%	+7.5%

C64 is inconsistent across nodes and does not establish a stable gain. These are single paired runs per node, not repeated-run confidence intervals. All 8,736 measured requests succeeded with exact 8192/1024 token lengths and zero KV preemptions. Routed experts use MegaMoE; main's shared-expert shape guard disables shared-expert fusion for this checkpoint in both arms.

The TP4 correctness test passed, covering independent dequantized FP32 references, nonuniform scales, odd-M padding, workspace reuse, changed-input CUDA graph replay, and output lifetime. The kernel benchmark passed on both nodes, including the model dispatch's small-M fallback. On node17, quantization + GEMM + RS latency at M=2048/4096/8192 decreased from 100.08/139.81/225.74 µs to 60.91/94.21/160.78 µs (39.1%/32.6%/28.8% reductions). Smaller shapes regressed in the native kernel and use the existing path below M=1024.

Model validation is limited to serving and generation smoke tests. Baseline and fused generated text differed; the long-prompt baseline also varied across runs. No task-level accuracy acceptance has been established. The feature remains off by default, and formal accuracy evaluation is pending.

All applicable staged-file pre-commit hooks passed, including Ruff, formatting, typos, mypy, SPDX, and device API checks. Submission cleanup only changes annotations, local naming, license attribution, and test/benchmark device APIs; the fusion algorithm and serving configuration are unchanged from the measured runs.

AI assistance: Codex implemented the change and ran validation. This draft is submitted for the human author's line-by-line review and independent test verification; those reviews are not claimed complete.