Conversation
| static_cast<uint64_t>(current_index) * bytes_per_buffer | ||
| + static_cast<uint64_t>(sym_buffer.rank_idx) * kSharedShard | ||
| * sizeof(nv_bfloat16); | ||
| shared_column_offset = n_idx - destination_rank * kSharedShard; |
There was a problem hiding this comment.
🔴 critical: 修正跨分片 tile 的无符号偏移下溢: 当 kSharedHidden / kNumRanks 仅满足当前的 8 元素对齐、但不是 128 元素 tile 对齐时（例如 2304/8=288），一个输出 tile 会跨越目标 rank。跨界 lane 的 destination_rank 已递增，但这里的两个 uint32_t 相减会下溢，随后形成约 8 GiB 的错误地址偏移并执行越界远端写入。应直接根据 vector_n 计算分片内列偏移，或限制分片宽度必须按 BLOCK_N 对齐。
🤖 v6
There was a problem hiding this comment.
Fixed in 7c84df6: the RS path now derives the in-shard column from vector_n and removes the separate lane byte offset. This supports shards that are not multiples of BLOCK_N. I also validated the exact crossing case on 8 B200 GPUs with shared_hidden=2304 and shard=288 for all three generations; routed output and every published fragment were bit-exact.
| def _bf16_shared_worker(local_rank: int, master_port: int) -> None: | ||
| os.environ['MASTER_ADDR'] = '127.0.0.1' | ||
| os.environ['MASTER_PORT'] = str(master_port) | ||
| os.environ['WORLD_SIZE'] = '1' | ||
| os.environ['RANK'] = '0' | ||
|
|
||
| buffer = None | ||
| try: | ||
| _, _, group = init_dist(local_rank, NUM_RANKS) | ||
| generator = torch.Generator(device='cuda') | ||
| generator.manual_seed(20260823) | ||
|
|
||
| def randn(shape: Tuple[int, ...], scale: float) -> torch.Tensor: | ||
| return torch.randn( | ||
| shape, | ||
| dtype=torch.bfloat16, | ||
| device='cuda', | ||
| generator=generator).mul_(scale) | ||
|
|
||
| routed_x = randn((NUM_TOKENS, HIDDEN), 0.5) | ||
| routed_l1 = randn( | ||
| (NUM_EXPERTS, INTERMEDIATE * 2, HIDDEN), 0.04) | ||
| routed_l2 = randn( | ||
| (NUM_EXPERTS, HIDDEN, INTERMEDIATE), 0.04) | ||
| shared_x = randn((NUM_TOKENS, SHARED_HIDDEN), 0.25) | ||
| shared_l1 = randn( | ||
| (SHARED_INTERMEDIATE * 2, SHARED_HIDDEN), 0.02) | ||
| shared_l2 = randn( | ||
| (SHARED_HIDDEN, SHARED_INTERMEDIATE), 0.02) | ||
| rms_weight = randn((HIDDEN,), 0.1).add_(1.0) | ||
| rms_epsilon = 1e-6 | ||
| situ_beta = 1.25 | ||
| situ_linear_beta = 1.5 | ||
|
|
||
| routed_x_fp8, routed_x_sf = per_token_cast_to_fp8( | ||
| routed_x, | ||
| use_ue8m0=True, | ||
| gran_k=32, | ||
| use_packed_ue8m0=True) | ||
| transformed_routed_l1, transformed_routed_l2 = ( | ||
| deep_gemm.transform_weights_for_mega_moe( | ||
| _cast_weights_to_fp4(routed_l1), | ||
| _cast_weights_to_fp4(routed_l2), | ||
| activation='situ')) | ||
| transformed_shared_l1, transformed_shared_l2 = ( | ||
| deep_gemm.transform_weights_for_mega_moe( | ||
| shared_l1, shared_l2, activation='situ')) | ||
| topk_idx = ( | ||
| torch.arange(NUM_TOKENS, device='cuda', dtype=torch.long) | ||
| % NUM_EXPERTS | ||
| ).view(NUM_TOKENS, NUM_TOPK) | ||
| topk_weights = torch.ones( | ||
| (NUM_TOKENS, NUM_TOPK), device='cuda', dtype=torch.float) | ||
|
|
||
| buffer = deep_gemm.get_symm_buffer_for_mega_moe( | ||
| group, | ||
| NUM_EXPERTS, | ||
| NUM_TOKENS, | ||
| NUM_TOPK, | ||
| HIDDEN, | ||
| INTERMEDIATE, | ||
| mma_type='fp8xfp4', | ||
| activation='situ', | ||
| bf16_shared_intermediate_hidden=SHARED_INTERMEDIATE) | ||
|
|
||
| def prepare_inputs() -> None: | ||
| buffer.buffer.zero_() | ||
| buffer.x[:NUM_TOKENS].copy_(routed_x_fp8) | ||
| buffer.x_sf[:NUM_TOKENS].copy_(routed_x_sf) | ||
| buffer.topk_idx[:NUM_TOKENS].copy_(topk_idx) | ||
| buffer.topk_weights[:NUM_TOKENS].copy_(topk_weights) | ||
|
|
||
| # Reference routed output: existing routed-only kernel followed by the | ||
| # BF16-input RMSNorm contract used by Kimi K3. | ||
| prepare_inputs() | ||
| routed_unnormalized = torch.empty( | ||
| (NUM_TOKENS, HIDDEN), dtype=torch.bfloat16, device='cuda') | ||
| deep_gemm.fp8_fp4_mega_moe( | ||
| routed_unnormalized, | ||
| transformed_routed_l1, | ||
| transformed_routed_l2, | ||
| buffer, | ||
| activation='situ', | ||
| situ_beta=situ_beta, | ||
| situ_linear_beta=situ_linear_beta) | ||
| routed_fp32 = routed_unnormalized.float() | ||
| routed_reference = ( | ||
| routed_fp32 | ||
| * torch.rsqrt(routed_fp32.square().mean(-1, keepdim=True) | ||
| + rms_epsilon) | ||
| * rms_weight.float()).to(torch.bfloat16) | ||
|
|
||
| prepare_inputs() | ||
| routed_actual = torch.empty_like(routed_unnormalized) | ||
| shared_actual = torch.empty( | ||
| (NUM_TOKENS, SHARED_HIDDEN), | ||
| dtype=torch.bfloat16, | ||
| device='cuda') | ||
| deep_gemm.fp8_fp4_mega_moe_bf16_shared( | ||
| routed_actual, | ||
| shared_actual, | ||
| transformed_routed_l1, | ||
| transformed_routed_l2, | ||
| shared_x, | ||
| transformed_shared_l1, | ||
| transformed_shared_l2, | ||
| rms_weight, | ||
| rms_epsilon, | ||
| buffer, | ||
| activation='situ', | ||
| situ_beta=situ_beta, | ||
| situ_linear_beta=situ_linear_beta) | ||
| torch.cuda.synchronize() | ||
|
|
||
| gate, up = F.linear(shared_x, shared_l1).chunk(2, dim=-1) | ||
| gate_fp32, up_fp32 = gate.float(), up.float() | ||
| shared_intermediate = ( | ||
| torch.sigmoid(gate_fp32) | ||
| * (situ_beta * torch.tanh(gate_fp32 / situ_beta)) | ||
| * (situ_linear_beta | ||
| * torch.tanh(up_fp32 / situ_linear_beta))).to(torch.bfloat16) | ||
| shared_reference = F.linear( | ||
| shared_intermediate, shared_l2).to(torch.bfloat16) | ||
|
|
||
| routed_error = _relative_l2(routed_actual, routed_reference) | ||
| shared_error = _relative_l2(shared_actual, shared_reference) | ||
| if routed_error >= 2e-3: | ||
| raise AssertionError( | ||
| f'fused RMSNorm relative L2 too high: {routed_error:.6f}') | ||
| if shared_error >= 1e-2: | ||
| raise AssertionError( | ||
| f'BF16 shared expert relative L2 too high: {shared_error:.6f}') | ||
| finally: | ||
| if buffer is not None: | ||
| buffer.destroy() | ||
| if dist.is_initialized(): | ||
| dist.destroy_process_group() |
There was a problem hiding this comment.
🟡 warning: The in-repo test covers the fused BF16 shared path (fp8_fp4_mega_moe_bf16_shared) but not the RS publication path; per-fragment and consecutive-generation validation was done only in the external test plan. Please add a single-GPU deterministic test of fp8_fp4_mega_moe_bf16_shared_rs (workspace shape (3, tokens, 1, shared_hidden), flags[0] cycling 0..2, peer_ptrs = [workspace.data_ptr()], bit-exact comparison against the dense shared output) so this path is guarded in CI.
🤖 v5
There was a problem hiding this comment.
Added a deterministic single-GPU RS regression in 7c84df6. It cycles flags[0] through generations 0, 1, and 2, uses the local workspace pointer as peer_ptrs, and requires bit-exact equality with dense shared output.
| DG_HOST_ASSERT(shared_rs_flags.is_cuda()); | ||
| DG_HOST_ASSERT(shared_rs_flags.scalar_type() == torch::kInt); | ||
| DG_HOST_ASSERT(shared_rs_flags.is_contiguous() and shared_rs_flags.numel() >= 9); | ||
| DG_HOST_ASSERT(shared_rs_peer_ptrs.is_cuda()); | ||
| DG_HOST_ASSERT(shared_rs_peer_ptrs.scalar_type() == torch::kInt64); | ||
| DG_HOST_ASSERT(shared_rs_peer_ptrs.is_contiguous()); | ||
| DG_HOST_ASSERT(shared_rs_peer_ptrs.dim() == 1); | ||
| DG_HOST_ASSERT(shared_rs_peer_ptrs.numel() == static_cast<int64_t>(sym_buffer_ptrs.size())); |
There was a problem hiding this comment.
🔵 suggestion: The shared_rs_flags contract (flags[0] = current generation index, flags[2] = bytes per generation buffer, numel >= 9) is only implied by the kernel's __ldg reads and the >=9 host assert. Document this layout in a comment or docstring, since it is a cross-repo contract with the vLLM consumer.
🤖 v5
There was a problem hiding this comment.
Documented the cross-repository flags ABI in csrc/apis/mega.hpp in 7c84df6: flags[0] is the generation and flags[2] is the byte stride per generation.
| void* shared_output_base = shared_y; | ||
| uint64_t shared_output_offset = 0; | ||
| uint64_t shared_column_offset = n_idx; | ||
| if constexpr (kPublishSharedRS) { | ||
| const auto current_index = __ldg(shared_rs_flags); | ||
| const auto bytes_per_buffer = __ldg(shared_rs_flags + 2); | ||
| constexpr uint32_t kSharedShard = | ||
| kSharedHidden / kNumRanks; | ||
| const auto vector_n = | ||
| n_idx + (lane_idx % 16) * 8; | ||
| const auto destination_rank = | ||
| vector_n / kSharedShard; | ||
| shared_output_base = reinterpret_cast<void*>( | ||
| __ldg(shared_rs_peer_ptrs + destination_rank)); | ||
| shared_output_offset = | ||
| static_cast<uint64_t>(current_index) * bytes_per_buffer | ||
| + static_cast<uint64_t>(sym_buffer.rank_idx) * kSharedShard | ||
| * sizeof(nv_bfloat16); | ||
| shared_column_offset = n_idx - destination_rank * kSharedShard; | ||
| } | ||
| const auto dst_ptr = math::advance_ptr<float4>( | ||
| shared_output_base, | ||
| shared_output_offset | ||
| + static_cast<uint64_t>(dst_token_idx) * kSharedHidden * sizeof(nv_bfloat16) | ||
| + shared_column_offset * sizeof(nv_bfloat16) | ||
| + (lane_idx % 16) * sizeof(float4)); | ||
| *dst_ptr = packed; |
There was a problem hiding this comment.
🔵 suggestion: The kernel only reads the flags and never signals completion; cross-rank consumption of published fragments relies entirely on external synchronization in the downstream vLLM tail. State explicitly in the fp8_fp4_mega_moe_bf16_shared_rs docstring that ordering/visibility across ranks is the caller's responsibility.
🤖 v5
| def fp8_fp4_mega_moe_bf16_shared_rs( | ||
| y: torch.Tensor, | ||
| l1_weights: Tuple[torch.Tensor, torch.Tensor], | ||
| l2_weights: Tuple[torch.Tensor, torch.Tensor], | ||
| shared_x: torch.Tensor, | ||
| shared_l1_weights: torch.Tensor, | ||
| shared_l2_weights: torch.Tensor, | ||
| rms_weight: torch.Tensor, | ||
| rms_epsilon: float, | ||
| shared_rs_workspace: torch.Tensor, | ||
| shared_rs_flags: torch.Tensor, | ||
| shared_rs_peer_ptrs: torch.Tensor, | ||
| sym_buffer: SymmBuffer, | ||
| cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None, | ||
| recipe: Tuple[int, int, int] = (1, 1, 32), | ||
| activation: str = 'situ', | ||
| fast_math: bool = True, | ||
| situ_beta: Optional[float] = None, | ||
| situ_linear_beta: Optional[float] = None): | ||
| """Publish BF16 shared outputs into a symmetric ReduceScatter buffer.""" | ||
| assert sym_buffer.shared_bf16_l2_acts is not None, \ | ||
| 'SymmBuffer was not initialized with a BF16 shared intermediate' | ||
| _C.fp8_fp4_mega_moe_bf16_shared( | ||
| y, None, | ||
| l1_weights, l2_weights, | ||
| shared_x, sym_buffer.shared_bf16_l2_acts, | ||
| shared_l1_weights, shared_l2_weights, | ||
| rms_weight, rms_epsilon, | ||
| cumulative_local_expert_recv_stats, | ||
| sym_buffer.buffer, | ||
| sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(), | ||
| sym_buffer.num_max_tokens_per_rank, | ||
| sym_buffer.num_experts, sym_buffer.num_topk, | ||
| recipe, activation, fast_math, | ||
| situ_beta, situ_linear_beta, | ||
| shared_rs_workspace, shared_rs_flags, shared_rs_peer_ptrs | ||
| ) | ||
|
|
||
| def bf16_mega_moe(y: torch.Tensor, |
There was a problem hiding this comment.
🔵 suggestion: fp8_fp4_mega_moe_bf16_shared_rs duplicates the argument marshaling of fp8_fp4_mega_moe_bf16_shared except the shared_y/RS trio; consider a shared private helper to avoid drift. Also missing a blank line before bf16_mega_moe (PEP8 two-blank-line rule).
🤖 v5
There was a problem hiding this comment.
Addressed in ebef8f1: both public wrappers now share one private validation/C++ argument-marshaling helper, and the public definitions retain PEP 8 separation.
| static_cast<uint64_t>(current_index) * bytes_per_buffer | ||
| + static_cast<uint64_t>(sym_buffer.rank_idx) * kSharedShard | ||
| * sizeof(nv_bfloat16); | ||
| shared_column_offset = n_idx - destination_rank * kSharedShard; |
There was a problem hiding this comment.
🔴 critical: shared_column_offset 声明为 uint64_t，但右式 n_idx - destination_rank * kSharedShard 中的量都是 uint32_t，会先按 32 位无符号计算再转为 64 位。当 kSharedShard 不是 BLOCK_N(128) 的倍数时（例如 shared_hidden=1024、num_ranks=16，shard=64；或 shared_hidden=256、num_ranks=4，shard=64），一个 128 列的 N block 会跨越 rank 分片，部分 lane 的 vector_n 落入下一个 destination rank，此时 destination_rank * kSharedShard > n_idx，32 位减法下溢成巨大 uint64 偏移，最终导致 dst_ptr 越界写。建议改为基于 vector_n 计算列内偏移（vector_n - destination_rank * kSharedShard，并移除后续 (lane_idx % 16) * sizeof(float4)），或改用 int64_t 完成减法；同时可在 host/static assert 中增加 kSharedShard % BLOCK_N == 0 的约束。
🤖 v4p
| fast_math: bool = True, | ||
| situ_beta: Optional[float] = None, | ||
| situ_linear_beta: Optional[float] = None): | ||
| assert sym_buffer.shared_bf16_l2_acts is not None, \ |
There was a problem hiding this comment.
🟡 warning: 这里只检查 sym_buffer.shared_bf16_l2_acts is not None，但 SymmBuffer 允许同时传入 num_shared_experts > 0 和 bf16_shared_intermediate_hidden > 0。这种 buffer 会通过该检查，而 C++ 侧 fp8_fp4_mega_moe_bf16_shared 固定用 num_shared_experts=0 重新 slice sym_buffer（csrc/apis/mega.hpp:442），导致 routed buffer 视图与实际按 num_shared_experts > 0 分配的布局不一致，内核会读写错误的缓冲区。建议在 SymmBuffer 中记录 num_shared_experts，并在此 assert 其为 0（或让 C++ 按 buffer 的实际布局 slice）。
🤖 v4p
There was a problem hiding this comment.
Added an explicit num_shared_experts == 0 guard to both BF16 shared wrappers in 7c84df6, so a routed shared-expert symmetric layout is rejected before entering C++.
🤖 ds-review-bot Code Reviewv6发布路径对已通过现有校验的部分 shared-hidden/并行度组合会计算出下溢偏移，导致远端显存越界写入。因此该补丁不能视为正确。 v5The change adds an optional BF16 shared-expert path to the SM100 FP8/FP4 MegaMoE producer (fp8_fp4_mega_moe_bf16_shared) plus a destination-scattered publication variant (fp8_fp4_mega_moe_bf16_shared_rs) that writes each TP-sharded shared partial directly into the destination rank's symmetric workspace, removing the standalone shared-output all-reduce. Correctness verified during review: (1) publication addressing is sound — each lane stores one float4 (8 bf16), vector_n = n_idx + (lane%16)8 selects the per-lane destination rank, the in-shard column resolves to vector_n - destination_rankkSharedShard, and a static assert on (kSharedHidden/kNumRanks) % 8 == 0 guarantees no vector straddles two shards; (2) the host-side workspace contract (dim-4 BF16 tensor, >=3 generations, num_ranks x shard with size(2)size(3) == shared_hidden) matches the kernel's per-token stride kSharedHidden and source-rank offset rank_idxshard; (3) publish_shared_rs is mutually exclusive with a dense shared_y and requires the full workspace/flags/peer_ptrs trio with dtype/contiguity/device checks; (4) static asserts enforce kPublishSharedRS -> kUseBF16Shared -> kNumSharedExperts > 0 and kApplyRMSNorm only with BF16 shared mode; (5) JIT codegen, launch signature, Python bindings, and exports are consistent. Backward compatibility is preserved: new template parameters and host arguments default off, and the scheduler change only lifts SHARED_*_SHAPE/BLOCK_K into template parameters with unchanged defaults, so the existing FP8/FP4 MegaMoE path is untouched. v4pMR 为 SM100 FP8/FP4 MegaMoE 增加可选 BF16 shared expert 路径，将 Kimi K3 的 routed experts、BF16 shared expert、RMSNorm 以及 ReduceScatter 发布融合进 MegaMoE 调度，并以 publish 方式避免 shared 输出 TP all-reduce。整体实现基本遵循既有 JIT/调度/TMA 模式，现有 FP8/FP4 路径在非 BF16 shared 场景下保持兼容，测试覆盖了单卡非发布路径；但发布路径存在一个 32 位无符号下溢导致的越界写风险，且 Python 包装对 SymmBuffer 配置缺少互斥校验。 Files reviewed: 7 |
Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
7c6d78f to
7c84df6
Compare
|
Rebased this PR onto current Also addressed the latest review findings in
Validation:
I also swept the existing MegaMoE block tiers and an earlier shared-L2 scheduling variant. Neither improved the current heuristic, so those experimental changes were not retained. |
Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
2 participants
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 -> 7168projection 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
num_shared_tokensindependently from routednum_tokensthrough 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.fp8_fp4_mega_moe_bf16_shared_sp_rsas an explicit opt-in wrapper. Existing dense shared output and hidden-dimension ReduceScatter publication APIs retain their behavior.tokens_per_rank = num_shared_tokens / num_ranks;destination_rank = global_token / tokens_per_rank;destination_token = global_token % tokens_per_rank;[generation, destination_token, source_rank, shared_hidden].vector_n - destination_rank * shardcolumn 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.pyvalidates 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.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, andmax_num_batched_tokens=32768. Baseline and fused runs used identical nodes, image, model, scheduler settings, exact request lengths, four warmups, andseed=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.The remote token-owner publication adds about
0.55 usto 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:
Downstream SP-sharded prefill
Workload: 8192 input / 1 output, concurrency 8, 128 requests.
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, and800.06 output tok/s; the SHARD=1 baseline measured780.46and783.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.The routed output passed the real-checkpoint comparison, and TP8 SP-published shared fragments were bit-exact against dense output.
Compatibility
Contribution notes
AI assistance was used during implementation, debugging, validation, benchmark analysis, and preparation of this PR description.