Repository navigation
[Kernel][DSV4.1] Fuse MXFP8 wo_b GEMM with sequence-parallel reduce-scatter - #57428
Merged
Merged
Conversation
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Use torch._scaled_mm for the MXFP8 reference and require contiguous weights in the kernel wrapper. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Select compatible backends before weight loading, adapt loaded weights and scales at callsites, and warm up each selected dtype. Exercise BF16 and both FlashInfer MXFP8 layouts against torch._scaled_mm. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Fold the tensor launch path into apply and migrate the tracked benchmark to the projection-based API. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
…tter Bind the DeepSeek-V4.1 sequence-parallel wo_b projection to the SM100 BF16/MXFP8 GEMM-RS kernel (opt-in via VLLM_DSV41_GEMM_RS=1) and move that kernel out of the Kimi-K3 tree into vllm/model_executor/kernels/linear/cute_dsl/gemm_rs_ar.py so both models share one workspace, init gate, warmup, test and benchmark. - DeepseekV4Attention gains bind_gemm_rs()/_wo_b_proj(); every _o_proj routes wo_b through it and the decoder layer skips its own sp_reduce_scatter when the fused path is bound. DSpark draft layers bind as well. - GemmRsAr.apply zeroes RS output rows beyond M so callers that pass an unpadded token count see the same padding as sp_reduce_scatter. - Benchmark gains --model deepseek_v41 and --backend flashinfer_cutlass|flashinfer_cutedsl presets. Builds on gau-nernst's MXFP8 support in the Kimi-K3 GEMM-RS/AR kernel. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>
mergify
Bot
added
ci/build
deepseek
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
17 of 19 tasks
…r path The online-MXFP8 baseline projections all-reduced inside RowParallelLinear, so the reference partial was already summed and every collective candidate failed the correctness check by the TP factor. Build them with reduce_results=False like sequence-parallel layers do, label assertion failures with the candidate and M, and add a candidate that runs the model's unfused path (projection kernel + sp_reduce_scatter) so the table compares against what the fused kernel actually replaces. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>
gcanlin
marked this pull request as ready for review
September 18, 2026 05:00
gcanlin
requested review from
AndreasKaratzas,
Harry-Chen,
WoosukKwon,
khluu,
mgoin,
tlrmchlsmth,
yewentao256 and
zyongye
as code owners
September 18, 2026 05:00
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Signed-off-by: Canlin <canlinguosdu@gmail.com>
gcanlin
force-pushed
the
perf/dsv41-gemm-rs-mxfp8-k3
branch
from
September 21, 2026 04:24
940b37c to
8f272b2
Compare
Contributor
Author
Done. Thanks! |
Contributor
|
/ci run |
|
✅ Triggered Buildkite CI #90137 for commit |
Contributor
Author
|
/ci retry |
|
❌ @gcanlin, A reviewer with write access must run |
Contributor
|
/ci retry |
|
✅ Queued 1 failed job(s) for retry in Buildkite CI #90137. |
jeejeelee
reviewed
Sep 21, 2026
… flag Both models bind the same SM100 GEMM-RS kernel, so replace the per-model VLLM_KIMI_K3_GEMM_RS and VLLM_DSV41_GEMM_RS switches (both default off) with a single VLLM_ENABLE_GEMM_RS. VLLM_KIMI_K3_GEMM_AR is unchanged. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>
jeejeelee
approved these changes
Sep 22, 2026
Member
|
/ci run |
jeejeelee
enabled auto-merge (squash)
September 22, 2026 00:42
github-actions
Bot
added
the
ready
Contributor
|
/ci run |
|
❌ This PR is 8 commits behind upstream |
Contributor
|
/ci run |
|
✅ Triggered Buildkite CI #90359 for commit |
Contributor
Author
|
Test with AIPerf + Agentx datasets:
|
Contributor
|
/ci retry |
|
✅ Queued 2 failed job(s) for retry in Buildkite CI #90359. |
ywang96
disabled auto-merge
September 22, 2026 08:07
vrdn-23
added a commit
to vrdn-23/vllm
that referenced
this pull request
Oct 8, 2026
…nflicts Took the branch's vllm/envs.py wholesale: main's legacy TYPE_CHECKING block and environment_variables dict are already superseded by the pydantic *Settings models here. Then ported main's semantic delta from the 57 main-side commits that touched vllm/envs.py since 03a8d0b: vllm-project#52431 VLLM_USE_MODELSCOPE: description (bool already accepts "1") vllm-project#52216 prefix_cache_retention_interval promoted to an arg (see vllm-project#55353) vllm-project#48484 + VLLM_REPLICATE_EMBED vllm-project#52633 + VLLM_CPU_CI_ENV vllm-project#52182 - VLLM_TEST_FORCE_FP8_MARLIN vllm-project#52523 VLLM_API_KEY not a compile factor (branch already marked it) vllm-project#52998 VLLM_ALLREDUCE_USE_FLASHINFER default -> True vllm-project#53040 + VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION vllm-project#53053 + VLLM_KIMI_K3_GEMM_AR vllm-project#50501 + VLLM_XPU_INC_WNA16_BACKEND (case-sensitive Literal) vllm-project#52018 + VLLM_B12X_MOE_FP4_FORCE_A16 vllm-project#43018 + VLLM_KV_OFFLOAD_MAX_BATCH_DESCRIPTORS vllm-project#51718 VLLM_KV_CACHE_LAYOUT Literal widened to 8 layouts + description vllm-project#51896 + VLLM_MAX_MEDIA_DOWNLOAD_SIZE_MB vllm-project#53219 VLLM_ENABLE_COHERE_API: description vllm-project#53559 - VLLM_TRITON_ATTN_USE_TD shim (already absent on the branch) vllm-project#53561 VLLM_MAX_AUDIO_CLIP_FILESIZE_MB: description vllm-project#53378 VLLM_ELASTIC_EP_SCALE_UP_LAUNCH not a compile factor vllm-project#53141 VLLM_ROCM_USE_AITER_FP4_ASM_GEMM removed; restored by vllm-project#57055 vllm-project#54160 + VLLM_ENABLE_HPC_OPS vllm-project#53536 + VLLM_XPU_FORCE_N_CONTIG_WEIGHT vllm-project#53576 + VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC vllm-project#54632 + VLLM_MAX_EMBED_DECODE_BYTES (not a compile factor) vllm-project#54579 + VLLM_ENABLE_SCALE_OUT_ENDPOINTS; removed by vllm-project#55176 (net zero) vllm-project#53491 VLLM_GPU_SYNC_CHECK: description vllm-project#54371 + VLLM_PLE_CPU_OFFLOAD; reworked vllm-project#56512, dropped vllm-project#57937 (net zero) vllm-project#48247 VLLM_ROCM_USE_AITER_CUSTOM_AR: description vllm-project#55353 - VLLM_MM_HASHER_ALGORITHM (+ _lower_mm_hasher), - VLLM_PREFIX_CACHE_RETENTION_INTERVAL vllm-project#54985 - VLLM_NIXL_EP_MAX_NUM_RANKS vllm-project#56512 VLLM_PLE_CPU_OFFLOAD rework (net zero, see vllm-project#54371) vllm-project#55176 - VLLM_ENABLE_SCALE_OUT_ENDPOINTS (net zero, see vllm-project#54579) vllm-project#53940 + VLLM_ROCM_USE_AITER_MOE_SITUV2 (case-insensitive Literal) vllm-project#41074 + VLLM_SKIP_VERSION_SUFFIX (not a compile factor) vllm-project#48599 - envs.validate_environ (moved to Platform.validate_environ, which reads the branch's environment_variables shim) vllm-project#52136 pydocstyle docstring restyle of legacy helpers (no-op here) vllm-project#57212 docstring rejoin of legacy helpers (no-op here) vllm-project#56913 + FORCE_COLOR (alias, non-"0" is True, not a compile factor) vllm-project#57055 VLLM_ROCM_USE_AITER_FP4_ASM_GEMM restored (net zero, see vllm-project#53141) vllm-project#57142 shm buffer name autogen -> new _generate_shm_name() (<=30 chars) vllm-project#56685 VLLM_HUMMING_INPUT_QUANT_CONFIG: description vllm-project#57937 - VLLM_PLE_CPU_OFFLOAD (net zero, see vllm-project#54371) vllm-project#50814 + VLLM_MXFP4_EMULATION_DEQUANT_AT_LOAD vllm-project#57428 + VLLM_ENABLE_GEMM_RS, - VLLM_KIMI_K3_GEMM_RS vllm-project#51600 - VLLM_XPU_ENABLE_XPU_GRAPH vllm-project#58582 + VLLM_TRITON_JIT_WARMUP_NUM_THREADS (not a compile factor) vllm-project#55145 + VLLM_XPU_PP_MICROBATCH vllm-project#58201 VLLM_ROCM_USE_AITER_MOE_SITUV2 6-way choice, - VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 vllm-project#58919 + VLLM_MOONCAKE_CONNECTOR_TIMEOUT vllm-project#56861 + VLLM_ROCM_AITER_MLA_DCP_VERIFY (case-insensitive Literal) vllm-project#53492 + VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA vllm-project#57991 VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE default True (branch already True) + description vllm-project#58819 + VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4 (tri-state, bool(int(v))) vllm-project#50300 + VLLM_CHAT_TEMPLATE_RENDER_TIMEOUT vllm-project#54024 VLLM_CPU_INT4_W4A8: description vllm-project#59781 - VLLM_ENGINE_ITERATION_TIMEOUT_S vllm-project#57995 + VLLM_FLASHINFER_MOE_A2A_LOW_PRECISION_COMBINE vllm-project#59965 VLLM_ROCM_AITER_MLA_DCP_VERIFY default route (final "auto") Also VLLM_SSM_CONV_STATE_LAYOUT's description follows vllm-project#51718's layout rename. Parity: 311 branch fields vs 311 main runtime entries, no difference in either direction. Every var's value matches main's lambdas in an empty environment and under 15 override scenarios, except case normalization of the two lowercased ROCm choices (both callers already call .lower()). Intentionally not ported: main's new consumer tests that match env_with_choices' error text (test_invalid_backend_value_raises, test_dcp_verify_env_rejects_unknown_route, and the ROCm-gated test_aiter_moe_situv2_rejects_unknown_activation) fail on this branch: pydantic's ValidationError names the field, not the env var. Invalid values are still rejected. Tracked as a branch-wide follow-up. Other conflicts and merge fallout: - requirements/docs.txt, requirements/test/{cpu,cuda,rocm,xpu}.txt: took main's lockfiles and regenerated them with the pip-compile pre-commit hooks; they differ from main only by the branch's pydantic-settings lines. - tests/test_envs_pydantic.py: dropped the mm-hasher test, pointed the validate_environ tests at Platform.validate_environ, updated the shm-name prefix assertion. - check-test-tethering (new on main): wired test_envs_pydantic.py into the CPU params/env job; allowlisted tests/docs/test_env_vars_gen.py, which cannot run from a CI job's tests/ cwd (tests/docs shadows docs/). Tests: pytest tests/test_envs.py tests/test_envs_pydantic.py tests/docs/test_env_vars_gen.py -> 55 passed; pre-commit run --files on all changed files -> clean. AI assistance (Claude) was used for this merge resolution. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Vinay Damodaran <vrdn@hey.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Under sequence parallel, DeepSeek-V4.1's
wo_bruns as an MXFP8 row-parallel GEMM followed by a TP reduce-scatter: two kernels and one HBM round trip per layer. This PR fuses them with the SM100 GEMM-RS kernel. It builds on @gau-nernst's MXFP8 support for the Kimi-K3 GEMM-RS/AR kernel (the first five commits), moves that kernel tovllm/model_executor/kernels/linear/cute_dsl/gemm_rs_ar.pyso both models share one workspace, init gate, warmup, test and benchmark, and binds V4.1'swo_bto it behindVLLM_ENABLE_GEMM_RS=1(default off; this single flag replacesVLLM_KIMI_K3_GEMM_RSand gates the Kimi-K3 RS path too —VLLM_KIMI_K3_GEMM_ARis unchanged). The decoder layer skips its ownsp_reduce_scatterwhen bound; RS zero-fills output rows beyond M for callers that pass an unpadded token count. Requires SM100, BF16 model dtype, sequence parallel, TP in 2..16, and FlashInfer CUTLASS or CuTe-DSL MXFP8 weights (whatlinear_backend=autoselects); anything else falls back with a warning.No open PR fuses the V4.1
wo_breduce-scatter: #57257 changes the SP all-gather side and #54151 is Kimi-K3 AG-GEMM.Test Plan
pytest -v -s tests/kernels/test_gemm_rs_ar.py # 2 GPUs: bf16, flashinfer_cutlass, flashinfer_cutedsl torchrun --nproc-per-node=4 -- benchmarks/kernels/benchmark_gemm_rs_ar.py \ --model deepseek_v41 --backend flashinfer_cutedsl --m 128 192 256 384 512 1024 2048 4096 8192 16384Deploy
VLLM_ENABLE_GEMM_RS=1 vllm serve deepseek-ai/DeepSeek-V4.1-Flash --tensor-parallel-size 4 --enable-expert-parallel \ --kernel-config '{"moe_backend":"deep_gemm_mega_moe"}' --attention-backend FLASHINFER_MLA_SPARSE_DSV41 \ --kv-cache-dtype fp8 --tokenizer-mode deepseek_v41 --max-model-len 32768 --max-num-batched-tokens 16384 \ --max-num-seqs 32 --no-enable-prefix-caching --seed 2026Baseline: the same command with
VLLM_ENABLE_GEMM_RS=0(main'swo_bpath).Benchmark
vllm bench serve --backend vllm --model dsv41 --dataset-name random --random-input-len 8192 --random-output-len 1024 \ --random-range-ratio 0 --ignore-eos --temperature 0 --seed 2026 --max-concurrency 32 --num-prompts 96 --num-warmups 8GSM8K: lm-eval 0.4.12,
gsm8k5-shot chat (--apply_chat_template --fewshot_as_multiturn, thinking off), greedy, 1024 max tokens, all 1319 questions.Test Result
4xB200, TP4 + EP4 with MegaMoE (sequence parallel), FlashInfer sparse MLA, fp8 KV, no speculative decoding. Kernel test: 3 passed.
Operator level (
benchmark_gemm_rs_ar.py), TP4,wo_bshape N=5120 with K=2048 per rank, MXFP8 CuTe-DSL layout, activation quantization included in both candidates, CUDA graphs, median over 40 samples of the slowest rank. The baseline is the model's unfused path with vLLM's own collective dispatch (custom Lamport reduce-scatter for small messages, NCCL above its size cap):sp_reduce_scatter(us)The fused kernel wins from about 1024 tokens; below that vLLM's Lamport reduce-scatter is faster than the fused kernel's fixed flag/exit handshake. Without speculative decoding, decode batches (at most
max_num_seqstokens) never reach the fused path and prefill chunks are large, so serving only sees the winning region.Torch profiler, one 8192-token prefill step on rank 0: the 40 NCCL reduce-scatter launches (9.1 ms) and 40
wo_bGEMMs (2.8 ms) are replaced by 40 fused launches (6.6 ms); GPU busy time per step drops from 127.1 ms to 118.1 ms (-7.1%).Online serving,
bench servewith fixed-length random requests, mean of two runs per arm (arms off, on, on, off with a fresh server each), zero preemptions:Run-to-run spread: baseline runs differ by 0.5% (c16), 5.4% (c32) and 16% (1k/2k); fused runs by 0.1%, 1.7% and 7%. The c16 change is within drift; c32 is at its edge with every percentile improving (p99 TTFT 4760 -> 3610 ms); the 1k/2k row is not resolved by two runs.
GSM8K, same server config:
Both within one standard error; 417 of 1319 generations are byte-identical, the rest differ in BF16 partial rounding order between the fused FP32 multicast reduction and NCCL's ring.
Not measured: speculative decoding (decode batches near 192 tokens would also take the fused path) and TP8. The
offarm runs this branch with the flag unset rather than a separate main checkout; the earlier commits only touch Kimi-K3 andcute_utils, so V4.1'swo_bpath is main's.AI assistance was used (Claude Code) for the V4.1 integration, the kernel relocation and driving the benchmarks; every changed line and result was reviewed by the author.
🤖 Generated with Claude Code