deepseek-ai/DeepGEMM

Fork 0

mirror of https://github.com/deepseek-ai/DeepGEMM.git synced 2026-11-05 00:47:26 +00:00

[PR #416] [SM100][MegaMoE] Fuse and publish Kimi K3's BF16 shared expert #430

New issue

Open

opened 2026-10-26 21:17:41 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:17:41 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/416
Author: @gcanlin
Created: 8/24/2026
Status: 🔄 Open

Base: nv_dev ← Head: feat/kimi-k3-fp8fp4-bf16-shared

📝 Commits (5)

1b07639 feat: fuse BF16 shared expert into FP8 FP4 MegaMoE
ea00574 Add destination-scattered BF16 shared MegaMoE output
7c84df6 Fix BF16 shared ReduceScatter publication
ebef8f1 Document and deduplicate BF16 shared publication API
f6e92ef Support SP-token BF16 shared publication

📊 Changes

7 files changed (+1107 additions, -70 deletions)

View changed files

📝 csrc/apis/mega.hpp (+195 -0)
📝 csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp (+96 -33)
📝 deep_gemm/__init__.py (+4 -0)
📝 deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh (+394 -24)
📝 deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh (+8 -7)
📝 deep_gemm/mega/__init__.py (+172 -5)
📝 tests/test_mega_moe_situ.py (+238 -1)

📄 Description

Summary

This PR extends the SM100 FP8/FP4 MegaMoE path with an optional BF16 shared expert for latent-MoE models such as Kimi K3. It schedules routed FP8/FP4 experts and BF16 shared L1/L2 together, applies the routed RMSNorm contract, and can publish each TP-sharded shared partial directly to a downstream symmetric workspace.

Kimi K3's routed output is latent-width (3584), while its shared output is hidden-width (7168). The final 3584 -> 7168 projection and residual combination therefore remain in the dependent vLLM tail rather than in this kernel.

This update adds sequence-parallel token publication. The routed branch still processes SP-local tokens, while the shared branch may process a different, gathered token count. Each shared partial is written directly to the rank that owns the corresponding output token.

Implementation

Thread num_shared_tokens independently from routed num_tokens through host validation, tensor-map creation, kernel arguments, and the MegaMoE scheduler. Shared L1/L2 task generation now uses the shared token count; routed dispatch remains unchanged.
Add fp8_fp4_mega_moe_bf16_shared_sp_rs as an explicit opt-in wrapper. Existing dense shared output and hidden-dimension ReduceScatter publication APIs retain their behavior.
In SP publication mode, shared epilogue stores use:
- tokens_per_rank = num_shared_tokens / num_ranks;
- destination_rank = global_token / tokens_per_rank;
- destination_token = global_token % tokens_per_rank;
- workspace layout [generation, destination_token, source_rank, shared_hidden].
The destination receives the complete hidden vector from every TP source rank. The downstream consumer reduces the source-rank dimension locally and combines it with the routed up-projection.
Validate that the global shared token count is rank-divisible, workspace dimensions match the selected publication mode, peer pointers cover every rank, and the generation/byte-stride metadata satisfies the cross-repository ABI.
Preserve the existing hidden-sharded publication path, including the generalized vector_n - destination_rank * shard column offset that avoids underflow when a shard boundary crosses a 128-column store tile.

The new behavior is compile-time specialized and optional. Calls that do not request BF16 shared work or publication instantiate the existing kernel behavior.

Integration dependencies

The consumer is vllm-project/vllm#53556. The TP8 x PP2 sequence-parallel benchmark temporarily stacked that PR on vllm-project/vllm#54347; #54347 is not merged into #53556 and is not a DeepGEMM dependency.

Test Plan

Repository regression

tests/test_mega_moe_situ.py validates independent routed/shared token counts, fused RMSNorm, dense BF16 shared output, the existing published output, and the new SP publication API. Published BF16 values are compared bitwise with dense output.

python3 tests/test_mega_moe_situ.py

TP8 remote-publication regression

An 8-GPU B200 harness gives each TP rank different shared weights, gathers every dense partial as the reference, calls fp8_fp4_mega_moe_bf16_shared_sp_rs, and compares every destination workspace value in [local_token, source_rank, hidden] order. The test passed twice with Kimi K3 matrix widths.

End-to-end setup

The dependent vLLM integration used two 8xB200 nodes, TP8 x PP2 x EP8, VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT=1, VLLM_KIMI_K3_GEMM_RS=1, max_num_seqs=128, and max_num_batched_tokens=32768. Baseline and fused runs used identical nodes, image, model, scheduler settings, exact request lengths, four warmups, and seed=2026.

Test Results

TP8 producer microbenchmark

The benchmark uses Kimi K3 matrix widths (routed H/I=3584/3072, shared H/I per TP rank=7168/768), local routed M=8, global shared M=64, one routed expert per rank for the isolated replay, and 100 measured iterations. Reported latency is the maximum rank time.

Run	Dense shared output	SP token-owner publication	Change
1	42.487 us	42.944 us	+1.08%
2	42.368 us	43.002 us	+1.50%
Mean	42.428 us	42.973 us	+1.29%

The remote token-owner publication adds about 0.55 us to the producer kernel and replaces downstream full shared-output materialization plus SP reduction.

The original non-SP destination-scatter replay remains neutral at its measured shapes:

M	Existing shared output	Published shared output	Change
32	0.2032 ms	0.2030 ms	-0.09%
64	0.2664 ms	0.2666 ms	+0.04%

Downstream SP-sharded prefill

Workload: 8192 input / 1 output, concurrency 8, 128 requests.

Run	SHARD=1 baseline total tok/s	Fused + SP publication total tok/s	Change
1	53,543.40	54,523.92	+1.83%
2	53,084.73	54,776.67	+3.19%
Mean	53,314.06	54,650.29	+2.51%

Mean median TTFT improves by 3.55% (1230.43 -> 1186.76 ms). These end-to-end results include the dependent vLLM targeted gather and local projection tail; they are not a standalone claim for the DeepGEMM producer.

Downstream decode

For 128 input / 1024 output at concurrency 32, the fused runs measured 773.91, 799.72, and 800.06 output tok/s; the SHARD=1 baseline measured 780.46 and 783.86 tok/s. Including the first cold fused run gives +1.16%; steady-state second-run comparison gives +2.02%.

Accuracy

The original fused non-SP integration was evaluated on the complete GSM8K v3 test split with lm-eval 0.4.12, 5-shot chat-template evaluation, and greedy decoding. All 1,319 samples were retained. The SP-token extension is covered by the TP8 bitwise and serving regressions above; this full-dataset run was not repeated for the temporary #54347 benchmark stack.

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9682	±	0.0048
		strict-match	5	exact_match	↑	0.9689	±	0.0048

The routed output passed the real-checkpoint comparison, and TP8 SP-published shared fragments were bit-exact against dense output.

Compatibility

Existing FP8/FP4 MegaMoE calls are unchanged.
BF16 shared work, hidden-sharded publication, and SP token publication are independent opt-in modes.
Host checks reject incompatible token counts, workspace shapes, dtypes, peer metadata, and devices before launch.
The dependent vLLM integration retains its native fallback when the new API or supported topology is unavailable.

Contribution notes

AI assistance was used during implementation, debugging, validation, benchmark analysis, and preparation of this PR description.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/416 **Author:** [@gcanlin](https://github.com/gcanlin) **Created:** 8/24/2026 **Status:** 🔄 Open **Base:** `nv_dev` ← **Head:** `feat/kimi-k3-fp8fp4-bf16-shared` --- ### 📝 Commits (5) - [`1b07639`](https://github.com/deepseek-ai/DeepGEMM/commit/1b07639096f35b8820ec7eba86e47163090a0a3e) feat: fuse BF16 shared expert into FP8 FP4 MegaMoE - [`ea00574`](https://github.com/deepseek-ai/DeepGEMM/commit/ea0057479f9be1ed12fad5a442b4835baf2e9f5f) Add destination-scattered BF16 shared MegaMoE output - [`7c84df6`](https://github.com/deepseek-ai/DeepGEMM/commit/7c84df6b68eabc3ac4350c2df4ad218415cbcd22) Fix BF16 shared ReduceScatter publication - [`ebef8f1`](https://github.com/deepseek-ai/DeepGEMM/commit/ebef8f186de992ba86b0e8962b8e62052305d6aa) Document and deduplicate BF16 shared publication API - [`f6e92ef`](https://github.com/deepseek-ai/DeepGEMM/commit/f6e92ef24e48a85e91945f93467b14772d3b6137) Support SP-token BF16 shared publication ### 📊 Changes **7 files changed** (+1107 additions, -70 deletions) <details> <summary>View changed files</summary> 📝 `csrc/apis/mega.hpp` (+195 -0) 📝 `csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp` (+96 -33) 📝 `deep_gemm/__init__.py` (+4 -0) 📝 `deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh` (+394 -24) 📝 `deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh` (+8 -7) 📝 `deep_gemm/mega/__init__.py` (+172 -5) 📝 `tests/test_mega_moe_situ.py` (+238 -1) </details> ### 📄 Description ## Summary This PR extends the SM100 FP8/FP4 MegaMoE path with an optional BF16 shared expert for latent-MoE models such as Kimi K3. It schedules routed FP8/FP4 experts and BF16 shared L1/L2 together, applies the routed RMSNorm contract, and can publish each TP-sharded shared partial directly to a downstream symmetric workspace. Kimi K3's routed output is latent-width (3584), while its shared output is hidden-width (7168). The final `3584 -> 7168` projection and residual combination therefore remain in the dependent vLLM tail rather than in this kernel. This update adds sequence-parallel token publication. The routed branch still processes SP-local tokens, while the shared branch may process a different, gathered token count. Each shared partial is written directly to the rank that owns the corresponding output token. ## Implementation - Thread `num_shared_tokens` independently from routed `num_tokens` through host validation, tensor-map creation, kernel arguments, and the MegaMoE scheduler. Shared L1/L2 task generation now uses the shared token count; routed dispatch remains unchanged. - Add `fp8_fp4_mega_moe_bf16_shared_sp_rs` as an explicit opt-in wrapper. Existing dense shared output and hidden-dimension ReduceScatter publication APIs retain their behavior. - In SP publication mode, shared epilogue stores use: - `tokens_per_rank = num_shared_tokens / num_ranks`; - `destination_rank = global_token / tokens_per_rank`; - `destination_token = global_token % tokens_per_rank`; - workspace layout `[generation, destination_token, source_rank, shared_hidden]`. - The destination receives the complete hidden vector from every TP source rank. The downstream consumer reduces the source-rank dimension locally and combines it with the routed up-projection. - Validate that the global shared token count is rank-divisible, workspace dimensions match the selected publication mode, peer pointers cover every rank, and the generation/byte-stride metadata satisfies the cross-repository ABI. - Preserve the existing hidden-sharded publication path, including the generalized `vector_n - destination_rank * shard` column offset that avoids underflow when a shard boundary crosses a 128-column store tile. The new behavior is compile-time specialized and optional. Calls that do not request BF16 shared work or publication instantiate the existing kernel behavior. ## Integration dependencies The consumer is [vllm-project/vllm#53556](https://github.com/vllm-project/vllm/pull/53556). The TP8 x PP2 sequence-parallel benchmark temporarily stacked that PR on [vllm-project/vllm#54347](https://github.com/vllm-project/vllm/pull/54347); #54347 is not merged into #53556 and is not a DeepGEMM dependency. ## Test Plan ### Repository regression `tests/test_mega_moe_situ.py` validates independent routed/shared token counts, fused RMSNorm, dense BF16 shared output, the existing published output, and the new SP publication API. Published BF16 values are compared bitwise with dense output. ```bash python3 tests/test_mega_moe_situ.py ``` ### TP8 remote-publication regression An 8-GPU B200 harness gives each TP rank different shared weights, gathers every dense partial as the reference, calls `fp8_fp4_mega_moe_bf16_shared_sp_rs`, and compares every destination workspace value in `[local_token, source_rank, hidden]` order. The test passed twice with Kimi K3 matrix widths. ### End-to-end setup The dependent vLLM integration used two 8xB200 nodes, TP8 x PP2 x EP8, `VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT=1`, `VLLM_KIMI_K3_GEMM_RS=1`, `max_num_seqs=128`, and `max_num_batched_tokens=32768`. Baseline and fused runs used identical nodes, image, model, scheduler settings, exact request lengths, four warmups, and `seed=2026`. ## Test Results ### TP8 producer microbenchmark The benchmark uses Kimi K3 matrix widths (`routed H/I=3584/3072`, `shared H/I per TP rank=7168/768`), local routed M=8, global shared M=64, one routed expert per rank for the isolated replay, and 100 measured iterations. Reported latency is the maximum rank time. Run | Dense shared output | SP token-owner publication | Change ---: | ---: | ---: | ---: 1 | 42.487 us | 42.944 us | +1.08% 2 | 42.368 us | 43.002 us | +1.50% Mean | 42.428 us | 42.973 us | **+1.29%** The remote token-owner publication adds about `0.55 us` to the producer kernel and replaces downstream full shared-output materialization plus SP reduction. The original non-SP destination-scatter replay remains neutral at its measured shapes: M | Existing shared output | Published shared output | Change ---: | ---: | ---: | ---: 32 | 0.2032 ms | 0.2030 ms | -0.09% 64 | 0.2664 ms | 0.2666 ms | +0.04% ### Downstream SP-sharded prefill Workload: 8192 input / 1 output, concurrency 8, 128 requests. Run | SHARD=1 baseline total tok/s | Fused + SP publication total tok/s | Change ---: | ---: | ---: | ---: 1 | 53,543.40 | 54,523.92 | +1.83% 2 | 53,084.73 | 54,776.67 | +3.19% Mean | 53,314.06 | 54,650.29 | **+2.51%** Mean median TTFT improves by **3.55%** (`1230.43 -> 1186.76 ms`). These end-to-end results include the dependent vLLM targeted gather and local projection tail; they are not a standalone claim for the DeepGEMM producer. ### Downstream decode For 128 input / 1024 output at concurrency 32, the fused runs measured `773.91`, `799.72`, and `800.06 output tok/s`; the SHARD=1 baseline measured `780.46` and `783.86 tok/s`. Including the first cold fused run gives **+1.16%**; steady-state second-run comparison gives **+2.02%**. ## Accuracy The original fused non-SP integration was evaluated on the complete GSM8K v3 test split with `lm-eval 0.4.12`, 5-shot chat-template evaluation, and greedy decoding. All 1,319 samples were retained. The SP-token extension is covered by the TP8 bitwise and serving regressions above; this full-dataset run was not repeated for the temporary #54347 benchmark stack. |Tasks|Version|Filter|n-shot|Metric| |Value| |Stderr| |---|---:|---|---:|---|:---:|---:|:---:|---:| |gsm8k|3|flexible-extract|5|exact_match|↑|0.9682|±|0.0048| |||strict-match|5|exact_match|↑|0.9689|±|0.0048| The routed output passed the real-checkpoint comparison, and TP8 SP-published shared fragments were bit-exact against dense output. ## Compatibility - Existing FP8/FP4 MegaMoE calls are unchanged. - BF16 shared work, hidden-sharded publication, and SP token publication are independent opt-in modes. - Host checks reject incompatible token counts, workspace shapes, dtypes, peer metadata, and devices before launch. - The dependent vLLM integration retains its native fallback when the new API or supported topology is unavailable. ## Contribution notes AI assistance was used during implementation, debugging, validation, benchmark analysis, and preparation of this PR description. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-26 21:17:41 +00:00

bsakaguchi referenced this issue

2026-10-26 21:17:49 +00:00

[PR #431] fix: port the #343 tensormap drain to the SM120 k_grouped NT contiguous kernel #441