@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 77871126f9

[Kernel][DSV4.1] Fuse MoE finalize into the TP all-reduce + mHC boundary (#58586)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Yongye Zhu committed Sep 26, 2026 at 12:08 UTC 77871126f9b69cff9feffff390d1171a60dbd7e2
8 files changed +1807 -309
CMakeLists.txt
-2
@@ -460,8 +460,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
460 DEFINED CMAKE_CUDA_COMPILER_VERSION AND
461 CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.0)
462
463 - list(APPEND VLLM_STABLE_EXT_SRC "csrc/libtorch_stable/all_reduce_mhc.cu")
464 -
463 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
464 cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
465 "9.0a;10.0f;10.1f;10.3f;10.7f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
csrc/libtorch_stable/all_reduce_mhc.cu deleted
-212
@@ -1,212 +0,0 @@
1 -// SPDX-License-Identifier: Apache-2.0
2 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -#include "torch_utils.h"
5 -#include "custom_all_gather_reduce_scatter.cuh"
6 -#include <cooperative_groups.h>
7 -#include <torch/csrc/stable/library.h>
8 -
9 -namespace {
10 -using Tensor = torch::stable::Tensor;
11 -using Pack = vllm::packed_t<nv_bfloat16>::P;
12 -constexpr int kHidden = 5120;
13 -constexpr int kRanks = 4;
14 -constexpr int kThreads = 128;
15 -constexpr int kCluster = 5;
16 -constexpr int kPacks = kHidden / Pack::size;
17 -
18 -__global__ __launch_bounds__(kThreads)
19 - __cluster_dims__(1, kCluster, 1) void all_reduce_mhc_kernel(
20 - const Pack* input, const Pack* residual, const float* post,
21 - const float* comb, const float* pre, const Pack* weight, Pack* output,
22 - Pack* normalized, Pack* local, Pack* multicast, uint32_t* epochs,
23 - int rank, int stage_size, float eps) {
24 -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
25 - cudaGridDependencySynchronize();
26 - const int token = blockIdx.x;
27 - const int column = blockIdx.y * kThreads + threadIdx.x;
28 - const int idx = token * kPacks + column;
29 - const int size = gridDim.x * kPacks;
30 - const int stage = epochs[0] % 3;
31 - const int dirty = vllm::mnnvl_lamport_dirty_stage(stage);
32 - const int dirty_size = epochs[2 + dirty];
33 - Pack value = input[idx];
34 - vllm::store_multimem_lamport_payload(
35 - multicast + stage * stage_size + rank * size + idx,
36 - vllm::sanitize_lamport_payload(value));
37 - // Multi-token forwards benefit from launching dependents before mHC finishes.
38 - if (gridDim.x > 1) cudaTriggerProgrammaticLaunchCompletion();
39 - vllm::lamport_cta_arrive(epochs + 1);
40 - for (int i = idx; i < dirty_size; i += size) {
41 - local[dirty * stage_size + i] = vllm::lamport_sentinel<Pack>();
42 - }
43 - Pack peers[kRanks];
44 - bool pending;
45 - do {
46 - vllm::LamportPack<Pack> packets[kRanks];
47 - #pragma unroll
48 - for (int peer = 0; peer < kRanks; ++peer) {
49 - auto ptr = local + stage * stage_size + idx + peer * size;
50 - asm volatile("ld.relaxed.sys.global.v4.u32 {%0,%1,%2,%3}, [%4];"
51 - : "=r"(packets[peer].words[0]), "=r"(packets[peer].words[1]),
52 - "=r"(packets[peer].words[2]), "=r"(packets[peer].words[3])
53 - : "l"(ptr)
54 - : "memory");
55 - }
56 - pending = false;
57 - #pragma unroll
58 - for (int peer = 0; peer < kRanks; ++peer) {
59 - pending |= vllm::is_lamport_dirty(packets[peer]);
60 - peers[peer] = packets[peer].packed;
61 - }
62 - } while (pending);
63 - auto sum = vllm::upcast(peers[0]);
64 - #pragma unroll
65 - for (int r = 1; r < kRanks; ++r) {
66 - vllm::packed_assign_add(sum, vllm::upcast(peers[r]));
67 - }
68 - // Keep the BF16 boundary between the collective and mHC post.
69 - auto reduced = vllm::upcast(vllm::downcast<Pack>(sum));
70 - decltype(sum) streams[4];
71 - #pragma unroll
72 - for (int r = 0; r < 4; ++r) {
73 - streams[r] = vllm::upcast(residual[(token * 4 + r) * kPacks + column]);
74 - }
75 - decltype(sum) collapsed{};
76 - #pragma unroll
77 - for (int c = 0; c < 4; ++c) {
78 - decltype(sum) mixed;
79 - #pragma unroll
80 - for (int e = 0; e < Pack::size; ++e) {
81 - // Match post*x followed by the four residual FMAs.
82 - float acc = __fmul_rn(post[token * 4 + c], reduced.data[e]);
83 - #pragma unroll
84 - for (int r = 0; r < 4; ++r) {
85 - acc = fmaf(comb[token * 16 + r * 4 + c], streams[r].data[e], acc);
86 - }
87 - mixed.data[e] = acc;
88 - }
89 - Pack rounded = vllm::downcast<Pack>(mixed);
90 - output[(token * 4 + c) * kPacks + column] = rounded;
91 - auto as_float = vllm::upcast(rounded);
92 - #pragma unroll
93 - for (int e = 0; e < Pack::size; ++e) {
94 - collapsed.data[e] =
95 - fmaf(pre[token * 4 + c], as_float.data[e], collapsed.data[e]);
96 - }
97 - }
98 - // RMSNorm consumes the rounded BF16 collapse in the unfused path.
99 - collapsed = vllm::upcast(vllm::downcast<Pack>(collapsed));
100 - float sqr = 0;
101 - #pragma unroll
102 - for (int e = 0; e < Pack::size; ++e) {
103 - sqr = fmaf(collapsed.data[e], collapsed.data[e], sqr);
104 - }
105 - #pragma unroll
106 - for (int delta = 16; delta; delta /= 2) {
107 - sqr += __shfl_xor_sync(0xffffffff, sqr, delta);
108 - }
109 - __shared__ float warp_sums[kCluster][kThreads / 32];
110 - auto cluster = cooperative_groups::this_cluster();
111 - const int lane = threadIdx.x % 32;
112 - if (lane < kCluster) {
113 - *cluster.map_shared_rank(&warp_sums[blockIdx.y][threadIdx.x / 32], lane) =
114 - sqr;
115 - }
116 - cluster.sync();
117 - float total = 0;
118 - #pragma unroll
119 - for (int c = 0; c < kCluster; ++c) {
120 - #pragma unroll
121 - for (int w = 0; w < kThreads / 32; ++w) total += warp_sums[c][w];
122 - }
123 - const float scale = rsqrtf(__fadd_rn(total / kHidden, eps));
124 - auto weights = vllm::upcast(weight[column]);
125 - #pragma unroll
126 - for (int e = 0; e < Pack::size; ++e) {
127 - collapsed.data[e] = collapsed.data[e] * scale * weights.data[e];
128 - }
129 - normalized[idx] = vllm::downcast<Pack>(collapsed);
130 - if (idx == 0) {
131 - while (*reinterpret_cast<volatile uint32_t*>(epochs + 1) <
132 - gridDim.x * gridDim.y) {
133 - }
134 - epochs[2 + stage] = kRanks * size;
135 - epochs[0] = vllm::mnnvl_lamport_next_stage(stage);
136 - epochs[1] = 0;
137 - }
138 - if (gridDim.x == 1) cudaTriggerProgrammaticLaunchCompletion();
139 -#else
140 - asm volatile("trap;");
141 -#endif
142 -}
143 -
144 -void all_reduce_mhc(Tensor input, Tensor residual, Tensor post, Tensor comb,
145 - Tensor pre, Tensor weight, Tensor output, Tensor normalized,
146 - int64_t local, int64_t multicast, Tensor epochs,
147 - int64_t rank, int64_t stage_bytes, double eps) {
148 - using torch::headeronly::ScalarType;
149 - const torch::stable::accelerator::DeviceGuard guard(input.get_device_index());
150 - const int64_t n = input.size(0);
151 - STD_TORCH_CHECK(input.dim() == 2 && input.size(1) == kHidden && n > 0 &&
152 - n <= 16);
153 - for (const auto& t : {input, residual, weight, output, normalized}) {
154 - STD_TORCH_CHECK(t.is_contiguous() &&
155 - t.scalar_type() == ScalarType::BFloat16);
156 - STD_TORCH_CHECK(t.get_device_index() == input.get_device_index());
157 - STD_TORCH_CHECK(reinterpret_cast<uintptr_t>(t.const_data_ptr()) % 16 == 0);
158 - }
159 - for (const auto& t : {post, comb, pre}) {
160 - STD_TORCH_CHECK(t.is_contiguous() && t.scalar_type() == ScalarType::Float);
161 - STD_TORCH_CHECK(t.get_device_index() == input.get_device_index());
162 - }
163 - STD_TORCH_CHECK(residual.numel() == n * 4 * kHidden &&
164 - output.numel() == residual.numel() &&
165 - normalized.numel() == input.numel() &&
166 - weight.numel() == kHidden);
167 - STD_TORCH_CHECK(post.numel() == n * 4 && pre.numel() == n * 4 &&
168 - comb.numel() == n * 16);
169 - STD_TORCH_CHECK(epochs.is_contiguous() &&
170 - epochs.scalar_type() == ScalarType::Int &&
171 - epochs.numel() >= 5 &&
172 - epochs.get_device_index() == input.get_device_index());
173 - STD_TORCH_CHECK(local && multicast && rank >= 0 && rank < kRanks &&
174 - stage_bytes % 16 == 0 &&
175 - stage_bytes >= n * kHidden * 2 * kRanks);
176 - STD_TORCH_CHECK(get_device_prop()->major == 10);
177 - cudaLaunchAttribute attr{};
178 - attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
179 - attr.val.programmaticStreamSerializationAllowed = 1;
180 - cudaLaunchConfig_t config{
181 - .gridDim = dim3(n, kCluster),
182 - .blockDim = dim3(kThreads),
183 - .dynamicSmemBytes = 0,
184 - .stream = get_current_cuda_stream(input.get_device_index()),
185 - .attrs = &attr,
186 - .numAttrs = 1};
187 - STD_CUDA_CHECK(cudaLaunchKernelEx(
188 - &config, all_reduce_mhc_kernel,
189 - static_cast<const Pack*>(input.const_data_ptr()),
190 - static_cast<const Pack*>(residual.const_data_ptr()),
191 - static_cast<const float*>(post.const_data_ptr()),
192 - static_cast<const float*>(comb.const_data_ptr()),
193 - static_cast<const float*>(pre.const_data_ptr()),
194 - static_cast<const Pack*>(weight.const_data_ptr()),
195 - static_cast<Pack*>(output.mutable_data_ptr()),
196 - static_cast<Pack*>(normalized.mutable_data_ptr()),
197 - reinterpret_cast<Pack*>(local), reinterpret_cast<Pack*>(multicast),
198 - static_cast<uint32_t*>(epochs.mutable_data_ptr()), static_cast<int>(rank),
199 - static_cast<int>(stage_bytes / sizeof(Pack)), static_cast<float>(eps)));
200 -}
201 -} // namespace
202 -
203 -STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, mhc) {
204 - mhc.def(
205 - "all_reduce_mhc(Tensor input, Tensor residual, Tensor post, Tensor comb, "
206 - "Tensor pre, Tensor weight, Tensor! output, Tensor! normalized, "
207 - "int local, int multicast, Tensor! epochs, int rank, int stage_bytes, "
208 - "float eps) -> ()");
209 -}
210 -STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, mhc) {
211 - mhc.impl("all_reduce_mhc", TORCH_BOX(&all_reduce_mhc));
212 -}
tests/distributed/test_custom_all_reduce.py
+120 -73
@@ -25,12 +25,43 @@ for i, v in enumerate(test_sizes):
25 test_sizes[i] -= v % 8
26
27
28 +def _bf16_ulps(a: torch.Tensor, b: torch.Tensor) -> int:
29 + """Largest distance between two BF16 tensors in units in the last place."""
30 +
31 + def ordered(t: torch.Tensor) -> torch.Tensor:
32 + bits = t.view(torch.int16).int() & 0xFFFF
33 + return torch.where(bits >= 0x8000, 0x8000 - bits, bits)
34 +
35 + return int((ordered(a) - ordered(b)).abs().max())
36 +
37 +
38 +def _unfused_all_reduce_mhc(peers, residual, post, comb, pre, weight, eps):
39 + """All-reduce, mHC post, collapse and RMSNorm in the fused kernel's FP32
40 + order, rounding to BF16 where the unfused path returns BF16."""
41 + reduced = peers[0].float()
42 + for peer in peers[1:]:
43 + reduced = reduced + peer.float()
44 + reduced = reduced.bfloat16().float()
45 + output = torch.empty_like(residual)
46 + collapse = torch.zeros_like(reduced)
47 + for target in range(4):
48 + mixed = reduced * post[:, target : target + 1]
49 + for source in range(4):
50 + mixed = torch.addcmul(
51 + mixed, residual[:, source].float(), comb[:, source, target, None]
52 + )
53 + output[:, target] = mixed.bfloat16()
54 + collapse = torch.addcmul(
55 + collapse, output[:, target].float(), pre[:, target : target + 1]
56 + )
57 + prenorm = collapse.bfloat16().float()
58 + inv_rms = torch.rsqrt(prenorm.square().mean(-1, keepdim=True) + eps)
59 + return output, (prenorm * inv_rms * weight.float()).bfloat16()
60 +
61 +
62 @ray.remote(num_gpus=1, max_calls=1)
63 def _all_reduce_mhc(monkeypatch, tp_size, pp_size, rank, distributed_init_port):
30 - from vllm.model_executor.kernels.mhc.tilelang import (
31 - mhc_post_tilelang,
32 - mhc_pre_delayed_tilelang,
33 - )
64 + from vllm.models.deepseek_v41.nvidia.ops.cute_dsl import AllReduceMHC
65
66 with monkeypatch.context() as m:
67 m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
@@ -38,98 +69,114 @@ def _all_reduce_mhc(monkeypatch, tp_size, pp_size, rank, distributed_init_port):
69 torch.accelerator.set_device_index(device)
70 init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
71 ensure_model_parallel_initialized(tp_size, pp_size)
41 - comm = get_tp_group().device_communicator.ca_comm
42 - assert comm is not None and comm.mnnvl_lamport_ag_multicast_ptr
43 - fn = torch.zeros(24, 20480, device=device)
44 - scale = torch.ones(3, device=device)
45 - base = torch.zeros(24, device=device)
46 -
47 - # Changing shapes and interleaving all-gather exercise the shared
48 - # Lamport stages, including cleanup of a larger previous payload.
49 - def run(n):
50 - torch.manual_seed(42 + rank)
51 - x = torch.randn(n, 5120, device=device, dtype=torch.bfloat16)
52 - # Packed +0/-0 pairs collide with the Lamport sentinel.
53 - x[:, :16] = 0
54 - x[:, 9:16:2] = -0.0
72 + op = AllReduceMHC(
73 + hidden_size=5120, hc_mult=4, max_num_tokens=16, top_k=6, device=device
74 + )
75 +
76 + def check_eager_and_replayed(fused, check, halved, doubled):
77 + check(fused())
78 + graph = torch.cuda.CUDAGraph()
79 + with torch.cuda.graph(graph):
80 + for _ in range(5):
81 + fused()
82 + captured = fused()
83 + for _ in range(20):
84 + graph.replay()
85 + check(captured)
86 + # Replays must pick up in-place input changes.
87 + halved.mul_(0.5)
88 + doubled.mul_(2)
89 + graph.replay()
90 + check(captured)
91 +
92 + def mhc_inputs(n):
93 torch.manual_seed(123)
94 residual = torch.randn(n, 4, 5120, device=device, dtype=torch.bfloat16)
95 post = torch.rand(n, 4, device=device)
96 comb = torch.randn(n, 4, 4, device=device) * 0.1
97 pre = torch.rand(n, 4, device=device)
98 weight = torch.randn(5120, device=device, dtype=torch.bfloat16)
61 - output = torch.empty_like(residual)
62 - normalized = torch.empty_like(x)
99 + return residual, post, comb, pre, weight
100 +
101 + def run(n):
102 + torch.manual_seed(42 + rank)
103 + x = torch.randn(n, 5120, device=device, dtype=torch.bfloat16)
104 + # Packed +0/-0 pairs collide with the Lamport sentinel.
105 + x[:, :16] = 0
106 + x[:, 9:16:2] = -0.0
107 + residual, post, comb, pre, weight = mhc_inputs(n)
108 +
109 + def fused():
110 + return op(x, residual, post, comb, pre, weight, 1e-6)
111 +
112 + def check(outputs):
113 + peers = get_tp_group().all_gather(x, dim=0).view(tp_size, n, 5120)
114 + output, normalized = _unfused_all_reduce_mhc(
115 + peers, residual, post, comb, pre, weight, 1e-6
116 + )
117 + # The mixed hc streams match bit for bit. The RMSNorm sums the
118 + # squares in another order and uses an approximate rsqrt.
119 + assert torch.equal(outputs[0], output)
120 + assert _bf16_ulps(outputs[1], normalized) <= 1
121 +
122 + check_eager_and_replayed(fused, check, x, residual)
123 +
124 + def run_finalize(n):
125 + torch.manual_seed(7 + rank)
126 + # A padded permuted GEMM2 buffer, like the MoE's.
127 + rows = n * 6 + 5
128 + gemm2 = torch.randn(rows, 5120, device=device, dtype=torch.bfloat16)
129 + permuted = torch.randperm(rows, device=device)[: n * 6].view(n, 6).int()
130 + # A route to an expert this rank does not hold.
131 + permuted[0, -1] = -1
132 + weights = torch.rand(n, 6, device=device)
133 + shared = torch.randn(n, 5120, device=device, dtype=torch.bfloat16)
134 + residual, post, comb, pre, weight = mhc_inputs(n)
135
136 def fused():
65 - torch.ops._C_custom_ar.all_reduce_mhc(
66 - x,
137 + return op.finalize(
138 + gemm2,
139 + weights,
140 + permuted,
141 + shared,
142 residual,
143 post,
144 comb,
145 pre,
146 weight,
72 - output,
73 - normalized,
74 - comm.mnnvl_lamport_ag_local_ptr,
75 - comm.mnnvl_lamport_ag_multicast_ptr,
76 - comm.mnnvl_lamport_epochs[0],
77 - rank,
78 - comm.mnnvl_buffer_size,
79 - 1e-6,
80 - )
81 -
82 - def check():
83 - gathered = comm.custom_all_gather(x)
84 - assert gathered is not None
85 - peers = gathered.view(tp_size, n, 5120).float()
86 - reduced = peers[0].clone()
87 - for peer in peers[1:]:
88 - reduced.add_(peer)
89 - expected = mhc_post_tilelang(
90 - reduced.bfloat16(), residual, post.unsqueeze(-1), comb
91 - )
92 - expected_norm = mhc_pre_delayed_tilelang(
93 - expected,
94 - fn,
95 - scale,
96 - base,
97 - 1e-6,
98 - 1e-6,
147 1e-6,
100 - 2.0,
101 - 20,
102 - pre_mix=pre,
103 - norm_weight=weight,
104 - )[2]
105 - torch.testing.assert_close(output, expected, rtol=0.008, atol=1e-6)
106 - torch.testing.assert_close(
107 - normalized, expected_norm, rtol=0.008, atol=0.008
148 )
149
110 - fused()
111 - check()
112 - graph = torch.cuda.CUDAGraph()
113 - with torch.cuda.graph(graph):
114 - for _ in range(6):
115 - fused()
116 - for _ in range(20):
117 - graph.replay()
118 - check()
119 - x.mul_(0.5)
120 - residual.mul_(2)
121 - graph.replay()
122 - check()
123 -
124 - # Cover fixed Q6 and shrinking/growing adaptive verification batches.
150 + def check(outputs):
151 + # One FP32 FMA per route in route order, the shared add, one
152 + # rounding, then the plain path: only the finalize differs.
153 + acc = torch.zeros(n, 5120, device=device)
154 + for k in range(6):
155 + valid = (permuted[:, k] >= 0).unsqueeze(-1)
156 + rows_k = gemm2[permuted[:, k].clamp_min(0).long()].float()
157 + torch.addcmul(
158 + acc,
159 + torch.where(valid, rows_k, 0.0),
160 + torch.where(valid, weights[:, k : k + 1], 0.0),
161 + out=acc,
162 + )
163 + x = (acc + shared.float()).bfloat16()
164 + expected = op(x, residual, post, comb, pre, weight, 1e-6)
165 + assert torch.equal(outputs[0], expected[0])
166 + assert torch.equal(outputs[1], expected[1])
167 +
168 + check_eager_and_replayed(fused, check, gemm2, shared)
169 +
170 + # Changing shapes cover shrinking and growing batches in one mailbox.
171 for n in (1, 6, 12, 8, 16, 3, 5, 2, 4, 1):
172 run(n)
173 + run_finalize(n)
174
175
176 @pytest.mark.skipif(
177 not current_platform.is_device_capability_family(100), reason="Requires SM100"
178 )
132 -def test_all_reduce_mhc_preserves_bf16_boundaries_and_graph_replay(monkeypatch):
179 +def test_all_reduce_mhc_matches_unfused_path(monkeypatch):
180 if torch.accelerator.device_count() < 4:
181 pytest.skip("Requires four GPUs with NVLink multicast")
182 multi_process_parallel(monkeypatch, 4, 1, _all_reduce_mhc)
vllm/models/deepseek_v41/nvidia/model.py
+76 -4
@@ -33,6 +33,14 @@ from vllm.model_executor.kernels.mhc.triton import hc_collapse_triton
33 from vllm.model_executor.layers.fused_moe import (
34 fused_moe_make_expert_params_mapping,
35 )
36 +from vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe import (
37 + TrtLlmMxfp4ExpertsModular,
38 +)
39 +from vllm.model_executor.layers.fused_moe.moe_output import (
40 + MoEOutput,
41 + UnfinalizedMoEOutput,
42 +)
43 +from vllm.model_executor.layers.fused_moe.runner.moe_runner import _unpack
44 from vllm.model_executor.layers.layernorm import RMSNorm
45 from vllm.model_executor.layers.logits_processor import LogitsProcessor
46 from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -94,6 +102,7 @@ from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask
102 from .engram import Engram, can_share_engram_tables, gather_engram_hashes
103 from .ops.mhc import (
104 MHC_OVERLAP_MAX_TOKENS,
105 + init_mhc_all_reduce,
106 mhc_pre_delayed_overlap,
107 mhc_shifted_post_pre,
108 supports_mhc_all_reduce,
@@ -135,6 +144,48 @@ class DeepseekV4MoE(DeepseekV4MoEBase):
144 image_sentinel_lo=IMAGE_SENTINEL_BASE_ID,
145 )
146
147 + def defer_finalize(self) -> None:
148 + """Leave the routed top-k reduction to the next fused all-reduce + mHC.
149 +
150 + Only the modular TRTLLM MXFP4 experts can stop after GEMM2; the fused
151 + kernel takes as many tokens as the mHC overlap path.
152 + """
153 + moe_config = self.experts.moe_config
154 + experts_cls = getattr(
155 + self.experts.routed_experts.quant_method, "experts_cls", None
156 + )
157 + if (
158 + experts_cls is TrtLlmMxfp4ExpertsModular
159 + and self.experts.routed_scaling_factor == 1.0
160 + ):
161 + moe_config.defer_moe_finalize(MHC_OVERLAP_MAX_TOKENS)
162 + if moe_config.use_deferred_moe_finalize:
163 + logger.info_once(
164 + "DSV4.1 mHC: MoE top-k finalize fused into the all-reduce for "
165 + "up to %d tokens.",
166 + MHC_OVERLAP_MAX_TOKENS,
167 + )
168 +
169 + def defers_finalize(self, num_tokens: int) -> bool:
170 + return (
171 + not self.use_mega_moe
172 + and self.experts.moe_config.should_defer_moe_finalize(num_tokens)
173 + )
174 +
175 + def forward_unfinalized(
176 + self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None
177 + ) -> MoEOutput:
178 + """``forward`` with the routed top-k reduction and all-reduce left open."""
179 + # The runner's custom op returns tensors only, so run its body directly.
180 + shared_output, routed = _unpack(
181 + self.experts._forward_impl(
182 + hidden_states, hidden_states, hidden_states, input_ids
183 + )
184 + )
185 + assert shared_output is not None
186 + assert isinstance(routed, UnfinalizedMoEOutput)
187 + return MoEOutput(routed=routed, shared_output=shared_output)
188 +
189
190 def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]:
191 """Pick the CUDA sparse-MLA attention class for the configured backend.
@@ -394,7 +445,7 @@ class DeepseekV4DecoderLayer(nn.Module):
445
446 def forward(
447 self,
397 - x: torch.Tensor,
448 + x: torch.Tensor | MoEOutput,
449 positions: torch.Tensor,
450 input_ids: torch.Tensor | None,
451 pre_mix: torch.Tensor | None = None,
@@ -407,7 +458,7 @@ class DeepseekV4DecoderLayer(nn.Module):
458 capture_previous_aux: bool = False,
459 mega_gate_metadata: MegaGateRoutingMetadata | None = None,
460 ) -> tuple[
410 - torch.Tensor,
461 + torch.Tensor | MoEOutput,
462 torch.Tensor,
463 torch.Tensor,
464 torch.Tensor,
@@ -432,6 +483,7 @@ class DeepseekV4DecoderLayer(nn.Module):
483 # sublayer's pre-mix: attention uses the pre-mix carried in (identity
484 # for the first layer), the FFN uses this layer's attention pre-mix.
485 if residual is None:
486 + assert isinstance(x, torch.Tensor)
487 if x.dim() == 2:
488 # First layer: the stream is the embedding broadcast to hc
489 # copies and the identity pre-mix selects copy 0.
@@ -471,7 +523,9 @@ class DeepseekV4DecoderLayer(nn.Module):
523 # Engram injection happens between the previous sublayer's post
524 # and this block's pre, on the full hc stream, so the mix
525 # coefficients see the injected stream. The injection also keeps
474 - # the post out of the pre-norm GEMM's fused prologue.
526 + # the post out of the pre-norm GEMM's fused prologue. The layer
527 + # before an engram layer finalizes its own MoE.
528 + assert isinstance(x, torch.Tensor)
529 if self.fuse_mhc_all_reduce:
530 x = tensor_model_parallel_all_reduce(x)
531 previous_post = mhc_post_tilelang(x, residual, post_mix, res_mix)
@@ -552,7 +606,11 @@ class DeepseekV4DecoderLayer(nn.Module):
606 stream=mhc_stream,
607 reduce_results=self.fuse_mhc_all_reduce,
608 )
555 - x = self.ffn(x, input_ids, mega_gate_metadata)
609 + if self.fuse_mhc_all_reduce and self.ffn.defers_finalize(x.shape[0]):
610 + # The next layer's first boundary finalizes it with the all-reduce.
611 + x = self.ffn.forward_unfinalized(x, input_ids)
612 + else:
613 + x = self.ffn(x, input_ids, mega_gate_metadata)
614 if mhc_stream is not None:
615 torch.cuda.current_stream().wait_stream(mhc_stream)
616 return x, residual, post_mix, res_mix, ffn_pre, previous_aux
@@ -591,6 +649,8 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
649 self.fuse_mhc_all_reduce = mhc_stream is not None and supports_mhc_all_reduce(
650 vllm_config
651 )
652 + if self.fuse_mhc_all_reduce:
653 + init_mhc_all_reduce(vllm_config)
654
655 # Reserved topk indices buffer for all Indexer layers to reuse.
656 self.topk_indices_buffer = torch.empty(
@@ -670,6 +730,16 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
730 ),
731 prefix=f"{prefix}.layers",
732 )
733 + if self.fuse_mhc_all_reduce:
734 + # A MoE's top-k finalize folds into the next layer's first mHC
735 + # boundary, which the last local layer lacks and an engram layer
736 + # replaces with its own all-reduce.
737 + local_layers = list(islice(self.layers, self.start_layer, self.end_layer))
738 + for layer, successor in zip(local_layers, local_layers[1:]):
739 + assert isinstance(layer, DeepseekV4DecoderLayer)
740 + assert isinstance(successor, DeepseekV4DecoderLayer)
741 + if successor.engram is None:
742 + layer.ffn.defer_finalize()
743
744 # The n-gram hash needs a slot-keyed rolling store of compressed ids
745 # (chunked prefill / decode lookback); key it off the first local
@@ -874,6 +944,8 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
944 if self.use_sequence_parallel:
945 previous_aux = sp_all_gather(previous_aux)[:full_num_tokens]
946 aux_hidden_by_layer[idx] = previous_aux
947 + # Without a successor boundary, the last layer finalized its own MoE.
948 + assert isinstance(hidden_states, torch.Tensor)
949 if layer is not None:
950 # The last layer has no successor to fold its post into.
951 if self.fuse_mhc_all_reduce:
vllm/models/deepseek_v41/nvidia/ops/cute_dsl/__init__.py new
+8
@@ -0,0 +1,8 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""CuTe DSL kernels for DSV4.1 mHC."""
5 +
6 +from .all_reduce_mhc import AllReduceMHC
7 +
8 +__all__ = ["AllReduceMHC"]
vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_reduce_mhc.py new
+1134
@@ -0,0 +1,1134 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# Adapted from FlashInfer's low-latency MNNVL CuTe DSL all-reduce
4 +# (flashinfer/comm/mnnvl_cutedsl/kernel_ll, flashinfer-ai/flashinfer@139af6f6).
5 +
6 +"""Lamport all-reduce fused with DSV4.1's mHC post, collapse and RMSNorm.
7 +
8 +FlashInfer's LL kernels, copied: a publish kernel (or its MoE finalize
9 +variant) multicasts each rank's contribution into a three-generation Lamport
10 +mailbox, and a PDL-chained consumer spins until every rank's fragment lands,
11 +reduces it and normalizes. The consumer's residual add and copy are replaced by
12 +the mHC post-mix of the hc residual streams and their pre-mix collapse, which
13 +the RMSNorm then normalizes. The finalize variant also reads FP32 routing
14 +weights.
15 +"""
16 +
17 +from __future__ import annotations
18 +
19 +import math
20 +
21 +import cuda.bindings.driver as cuda
22 +import cutlass
23 +import cutlass.cute as cute
24 +import torch
25 +import torch.distributed as dist
26 +import torch.distributed._symmetric_memory as symm_mem
27 +from cutlass import BFloat16, Float32, Int32, Int64, Uint32
28 +from cutlass.cute.runtime import make_fake_compact_tensor
29 +
30 +from vllm.distributed import get_tp_group
31 +
32 +from .primitives import (
33 + NEGATIVE_ZERO_BF16_BITS,
34 + QUAD_BF16,
35 + VEC_BF16,
36 + WARP_SIZE,
37 + bf16x4_to_packed_u32x2,
38 + bf16x8_to_packed_u32x4,
39 + current_cu_stream,
40 + f32_to_bf16_bits,
41 + fragment_has_negative_zero,
42 + load_global_bf16_as_f32,
43 + load_global_u32x2,
44 + load_global_u32x4,
45 + load_volatile_u32,
46 + make_fake_dynamic_compact_tensor,
47 + map_shared_to_peer,
48 + packed_u32x2_to_bf16x4,
49 + packed_u32x4_to_bf16x8,
50 + sanitize_negative_zero_u32x2,
51 + sanitize_negative_zero_u32x4,
52 + shuffle_sync_idx_u32,
53 + stmc_bf16x2,
54 + stmc_bf16x4,
55 + stmc_bf16x8,
56 + store_global_u32,
57 + store_global_u32x4,
58 + store_lamport_sentinel_u32x4,
59 + store_shared_cluster_f32,
60 + to_cute,
61 + to_cute_dynamic,
62 +)
63 +
64 +LAMPORT_GENERATIONS = 3
65 +NEXT_STAGE = 0
66 +ACTIVE_STAGE = 1
67 +
68 +
69 +@cute.jit
70 +def _group_leader_block_sum(
71 + value: Float32,
72 + warp_sums: cute.Tensor,
73 + warps: cutlass.Constexpr[int],
74 + leader_stride: cutlass.Constexpr[int],
75 +) -> Float32:
76 + lane = cute.arch.lane_idx()
77 + warp = cute.arch.warp_idx()
78 + for offset in cutlass.range_constexpr(1, WARP_SIZE):
79 + if cutlass.const_expr(offset >= leader_stride and (offset & (offset - 1)) == 0):
80 + value = value + cute.arch.shuffle_sync_bfly(
81 + value,
82 + offset=offset,
83 + mask=-1,
84 + mask_and_clamp=31,
85 + )
86 + if lane == 0:
87 + cute.arch.store((warp_sums + warp).llvm_ptr, value)
88 + cute.arch.barrier()
89 +
90 + result = Float32(0.0)
91 + if warp == 0:
92 + if lane < Int32(warps):
93 + result = cute.arch.load(
94 + (warp_sums + lane).llvm_ptr,
95 + Float32,
96 + )
97 + result = cute.arch.warp_reduction_sum(result)
98 + if lane == 0:
99 + cute.arch.store(warp_sums.llvm_ptr, result)
100 + cute.arch.barrier()
101 + return cute.arch.load(warp_sums.llvm_ptr, Float32)
102 +
103 +
104 +class _SharedOnlyPublishDeviceKernel:
105 + def __init__(
106 + self,
107 + *,
108 + hidden: int,
109 + tp: int,
110 + rank: int,
111 + capacity_m: int,
112 + elements_per_thread: int,
113 + threads: int,
114 + release_before_store: bool,
115 + enable_pdl: bool,
116 + ) -> None:
117 + if elements_per_thread not in (1, QUAD_BF16, VEC_BF16):
118 + raise ValueError("elements_per_thread must be 1, 4, or 8")
119 + if hidden <= 0 or hidden % elements_per_thread:
120 + raise ValueError("hidden must divide evenly across thread fragments")
121 + self.hidden = hidden
122 + self.tp = tp
123 + self.rank = rank
124 + self.capacity_m = capacity_m
125 + self.elements_per_thread = elements_per_thread
126 + self.threads = threads
127 + self.fragments = hidden // elements_per_thread
128 + self.ctas_per_token = math.ceil(self.fragments / threads)
129 + self.release_before_store = release_before_store
130 + self.enable_pdl = enable_pdl
131 +
132 + @cute.jit
133 + def __call__(
134 + self,
135 + shared_output: cute.Tensor,
136 + stage_state: cute.Tensor,
137 + contribution_mailbox_multicast_address: Int64,
138 + m: Int32,
139 + stream: cuda.CUstream,
140 + ) -> None:
141 + self.kernel(
142 + shared_output,
143 + stage_state,
144 + contribution_mailbox_multicast_address,
145 + ).launch(
146 + grid=(m * self.ctas_per_token, 1, 1),
147 + block=(self.threads, 1, 1),
148 + stream=stream,
149 + use_pdl=self.enable_pdl,
150 + )
151 +
152 + @cute.kernel
153 + def kernel(
154 + self,
155 + shared_output: cute.Tensor,
156 + stage_state: cute.Tensor,
157 + contribution_mailbox_multicast_address: Int64,
158 + ) -> None:
159 + block, _, _ = cute.arch.block_idx()
160 + tidx, _, _ = cute.arch.thread_idx()
161 + token = block // self.ctas_per_token
162 + cta_in_token = block % self.ctas_per_token
163 + fragment = cta_in_token * self.threads + tidx
164 +
165 + if cutlass.const_expr(self.enable_pdl):
166 + cute.arch.griddepcontrol_wait()
167 +
168 + if cutlass.const_expr(self.elements_per_thread == 1):
169 + bits = Uint32(0)
170 + if fragment < self.fragments:
171 + element = Int64(token) * self.hidden + Int64(fragment)
172 + bits = f32_to_bf16_bits(
173 + load_global_bf16_as_f32(
174 + Int64((shared_output.iterator + element).toint())
175 + )
176 + )
177 + if bits == Uint32(NEGATIVE_ZERO_BF16_BITS):
178 + bits = Uint32(0)
179 + partner_bits = shuffle_sync_idx_u32(
180 + bits,
181 + Int32(cute.arch.lane_idx() | Int32(1)),
182 + )
183 + packed_word = bits | (partner_bits << Uint32(16))
184 + elif cutlass.const_expr(self.elements_per_thread == QUAD_BF16):
185 + packed_words = cute.make_rmem_tensor(cute.make_layout((2,)), Uint32)
186 + packed_words.fill(Uint32(0))
187 + if fragment < self.fragments:
188 + element = Int64(token) * self.hidden + Int64(fragment) * QUAD_BF16
189 + pointer = cute.make_ptr(
190 + BFloat16,
191 + (shared_output.iterator + element).llvm_ptr,
192 + cute.AddressSpace.gmem,
193 + assumed_align=8,
194 + )
195 + packed_words.store(
196 + sanitize_negative_zero_u32x2(load_global_u32x2(pointer))
197 + )
198 + else:
199 + packed_words = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32)
200 + packed_words.fill(Uint32(0))
201 + if fragment < self.fragments:
202 + element = Int64(token) * self.hidden + Int64(fragment) * VEC_BF16
203 + pointer = cute.make_ptr(
204 + BFloat16,
205 + (shared_output.iterator + element).llvm_ptr,
206 + cute.AddressSpace.gmem,
207 + assumed_align=16,
208 + )
209 + packed_words.store(
210 + sanitize_negative_zero_u32x4(load_global_u32x4(pointer))
211 + )
212 +
213 + stage = load_volatile_u32(stage_state.iterator + NEXT_STAGE)
214 + if cutlass.const_expr(self.release_before_store):
215 + if block == 0 and tidx == 0:
216 + store_global_u32(
217 + stage_state.iterator + ACTIVE_STAGE,
218 + stage,
219 + )
220 + cute.arch.barrier()
221 + if cutlass.const_expr(self.enable_pdl):
222 + cute.arch.griddepcontrol_launch_dependents()
223 +
224 + if cutlass.const_expr(self.elements_per_thread == 1):
225 + if (cute.arch.lane_idx() & Int32(1)) == Int32(
226 + 0
227 + ) and fragment < self.fragments:
228 + mailbox_element = (
229 + (Int64(stage) * self.tp + self.rank) * self.capacity_m
230 + + Int64(token)
231 + ) * self.hidden + Int64(fragment)
232 + stmc_bf16x2(
233 + contribution_mailbox_multicast_address + mailbox_element * 2,
234 + packed_word,
235 + )
236 + elif cutlass.const_expr(self.elements_per_thread == QUAD_BF16):
237 + if fragment < self.fragments:
238 + mailbox_element = (
239 + (Int64(stage) * self.tp + self.rank) * self.capacity_m
240 + + Int64(token)
241 + ) * self.hidden + Int64(fragment) * QUAD_BF16
242 + stmc_bf16x4(
243 + contribution_mailbox_multicast_address + mailbox_element * 2,
244 + packed_words,
245 + )
246 + else:
247 + if fragment < self.fragments:
248 + mailbox_element = (
249 + (Int64(stage) * self.tp + self.rank) * self.capacity_m
250 + + Int64(token)
251 + ) * self.hidden + Int64(fragment) * VEC_BF16
252 + stmc_bf16x8(
253 + contribution_mailbox_multicast_address + mailbox_element * 2,
254 + packed_words,
255 + )
256 +
257 + if cutlass.const_expr(not self.release_before_store):
258 + if block == 0 and tidx == 0:
259 + store_global_u32(
260 + stage_state.iterator + ACTIVE_STAGE,
261 + stage,
262 + )
263 + cute.arch.barrier()
264 + if cutlass.const_expr(self.enable_pdl):
265 + cute.arch.griddepcontrol_launch_dependents()
266 +
267 +
268 +class _QuadFinalizePublishDeviceKernel:
269 + def __init__(
270 + self,
271 + *,
272 + hidden: int,
273 + top_k: int,
274 + tp: int,
275 + rank: int,
276 + capacity_m: int,
277 + threads: int,
278 + routed_scaling_factor: float,
279 + include_shared_expert: bool,
280 + load_shared_expert_before_pdl: bool,
281 + enable_pdl: bool,
282 + prefetch_group: int,
283 + fp32_weights: bool = False,
284 + ) -> None:
285 + if hidden <= 0 or hidden % QUAD_BF16:
286 + raise ValueError("hidden must be a positive multiple of 4")
287 + self.hidden = hidden
288 + self.top_k = top_k
289 + self.tp = tp
290 + self.rank = rank
291 + self.capacity_m = capacity_m
292 + self.threads = threads
293 + self.routed_scaling_factor = routed_scaling_factor
294 + self.fragments = hidden // QUAD_BF16
295 + self.ctas_per_token = math.ceil(self.fragments / threads)
296 + self.include_shared_expert = include_shared_expert
297 + self.load_shared_expert_before_pdl = load_shared_expert_before_pdl
298 + self.enable_pdl = enable_pdl
299 + self.prefetch_group = prefetch_group
300 + self.prefetch_groups = (top_k + prefetch_group - 1) // prefetch_group
301 + self.fp32_weights = fp32_weights
302 +
303 + def smem_size_in_bytes(self) -> int:
304 + return self.top_k * 8
305 +
306 + @cute.jit
307 + def __call__(
308 + self,
309 + routed_output: cute.Tensor,
310 + expert_weights: cute.Tensor,
311 + permuted_indices: cute.Tensor,
312 + shared_output: cute.Tensor,
313 + stage_state: cute.Tensor,
314 + contribution_mailbox_multicast_address: Int64,
315 + m: Int32,
316 + stream: cuda.CUstream,
317 + ) -> None:
318 + self.kernel(
319 + routed_output,
320 + expert_weights,
321 + permuted_indices,
322 + shared_output,
323 + stage_state,
324 + contribution_mailbox_multicast_address,
325 + ).launch(
326 + grid=(m * self.ctas_per_token, 1, 1),
327 + block=(self.threads, 1, 1),
328 + smem=self.smem_size_in_bytes(),
329 + stream=stream,
330 + use_pdl=self.enable_pdl,
331 + )
332 +
333 + @cute.kernel
334 + def kernel(
335 + self,
336 + routed_output: cute.Tensor,
337 + expert_weights: cute.Tensor,
338 + permuted_indices: cute.Tensor,
339 + shared_output: cute.Tensor,
340 + stage_state: cute.Tensor,
341 + contribution_mailbox_multicast_address: Int64,
342 + ) -> None:
343 + block, _, _ = cute.arch.block_idx()
344 + tidx, _, _ = cute.arch.thread_idx()
345 + token = block // self.ctas_per_token
346 + cta_in_token = block % self.ctas_per_token
347 + fragment = cta_in_token * self.threads + tidx
348 +
349 + smem = cutlass.utils.SmemAllocator()
350 + staged_indices = smem.allocate_array(Int32, self.top_k)
351 + staged_weights = smem.allocate_array(Float32, self.top_k)
352 + metadata_index = Int32(tidx)
353 + while metadata_index < Int32(self.top_k):
354 + element = Int64(token) * self.top_k + Int64(metadata_index)
355 + row = cute.arch.load(
356 + (permuted_indices.iterator + element).llvm_ptr,
357 + Int32,
358 + )
359 + if cutlass.const_expr(self.fp32_weights):
360 + weight = cute.arch.load(
361 + (expert_weights.iterator + element).llvm_ptr,
362 + Float32,
363 + )
364 + else:
365 + weight = load_global_bf16_as_f32(
366 + Int64((expert_weights.iterator + element).toint())
367 + )
368 + if cutlass.const_expr(self.routed_scaling_factor != 1.0):
369 + weight = weight * Float32(self.routed_scaling_factor)
370 + if row == Int32(-1):
371 + weight = Float32(0.0)
372 + cute.arch.store(
373 + (staged_indices + metadata_index).llvm_ptr,
374 + row,
375 + )
376 + cute.arch.store(
377 + (staged_weights + metadata_index).llvm_ptr,
378 + weight,
379 + )
380 + metadata_index = metadata_index + self.threads
381 + cute.arch.barrier()
382 +
383 + if cutlass.const_expr(self.include_shared_expert):
384 + shared_values = cute.make_rmem_tensor(
385 + cute.make_layout((QUAD_BF16,)), BFloat16
386 + )
387 + shared_values.fill(BFloat16(0.0))
388 + if cutlass.const_expr(self.load_shared_expert_before_pdl): # noqa: SIM102
389 + if fragment < self.fragments:
390 + shared_element = (
391 + Int64(token) * self.hidden + Int64(fragment) * QUAD_BF16
392 + )
393 + shared_pointer = cute.make_ptr(
394 + BFloat16,
395 + (shared_output.iterator + shared_element).llvm_ptr,
396 + cute.AddressSpace.gmem,
397 + assumed_align=8,
398 + )
399 + shared_values.store(
400 + packed_u32x2_to_bf16x4(load_global_u32x2(shared_pointer))
401 + )
402 +
403 + if cutlass.const_expr(self.enable_pdl):
404 + cute.arch.griddepcontrol_wait()
405 +
406 + if cutlass.const_expr( # noqa: SIM102
407 + self.include_shared_expert and not self.load_shared_expert_before_pdl
408 + ):
409 + if fragment < self.fragments:
410 + shared_element = (
411 + Int64(token) * self.hidden + Int64(fragment) * QUAD_BF16
412 + )
413 + shared_pointer = cute.make_ptr(
414 + BFloat16,
415 + (shared_output.iterator + shared_element).llvm_ptr,
416 + cute.AddressSpace.gmem,
417 + assumed_align=8,
418 + )
419 + shared_values.store(
420 + packed_u32x2_to_bf16x4(load_global_u32x2(shared_pointer))
421 + )
422 +
423 + stage = load_volatile_u32(stage_state.iterator + NEXT_STAGE)
424 + if fragment < self.fragments:
425 + accumulator = cute.make_rmem_tensor(cute.make_layout((QUAD_BF16,)), Float32)
426 + accumulator.fill(Float32(0.0))
427 + if cutlass.const_expr(self.prefetch_group == 1):
428 + for k in cutlass.range_constexpr(self.top_k):
429 + row = cute.arch.load((staged_indices + k).llvm_ptr, Int32)
430 + weight = cute.arch.load(
431 + (staged_weights + k).llvm_ptr,
432 + Float32,
433 + )
434 + if row != Int32(-1):
435 + source_element = (
436 + Int64(row) * self.hidden + Int64(fragment) * QUAD_BF16
437 + )
438 + source_pointer = cute.make_ptr(
439 + BFloat16,
440 + (routed_output.iterator + source_element).llvm_ptr,
441 + cute.AddressSpace.gmem,
442 + assumed_align=8,
443 + )
444 + accumulator.store(
445 + accumulator.load()
446 + + packed_u32x2_to_bf16x4(
447 + load_global_u32x2(source_pointer)
448 + ).to(Float32)
449 + * weight
450 + )
451 + else:
452 + inputs = cute.make_rmem_tensor(
453 + cute.make_layout((self.prefetch_group, 2)),
454 + Uint32,
455 + )
456 + inputs.fill(Uint32(0))
457 + for group in cutlass.range_constexpr(self.prefetch_groups):
458 + for item in cutlass.range_constexpr(self.prefetch_group):
459 + k = group * self.prefetch_group + item
460 + if cutlass.const_expr(k < self.top_k):
461 + row = cute.arch.load(
462 + (staged_indices + k).llvm_ptr,
463 + Int32,
464 + )
465 + for word in cutlass.range_constexpr(2):
466 + inputs[item, word] = Uint32(0)
467 + if row != Int32(-1):
468 + source_element = (
469 + Int64(row) * self.hidden
470 + + Int64(fragment) * QUAD_BF16
471 + )
472 + source_pointer = cute.make_ptr(
473 + BFloat16,
474 + (routed_output.iterator + source_element).llvm_ptr,
475 + cute.AddressSpace.gmem,
476 + assumed_align=8,
477 + )
478 + source = load_global_u32x2(source_pointer)
479 + for word in cutlass.range_constexpr(2):
480 + inputs[item, word] = source[word]
481 + for item in cutlass.range_constexpr(self.prefetch_group):
482 + k = group * self.prefetch_group + item
483 + if cutlass.const_expr(k < self.top_k):
484 + source = cute.make_rmem_tensor(
485 + cute.make_layout((2,)),
486 + Uint32,
487 + )
488 + for word in cutlass.range_constexpr(2):
489 + source[word] = inputs[item, word]
490 + accumulator.store(
491 + accumulator.load()
492 + + packed_u32x2_to_bf16x4(source.load()).to(Float32)
493 + * cute.arch.load(
494 + (staged_weights + k).llvm_ptr,
495 + Float32,
496 + )
497 + )
498 + result = accumulator.load()
499 + if cutlass.const_expr(self.include_shared_expert):
500 + result = result + shared_values.load().to(Float32)
501 + result_packed = bf16x4_to_packed_u32x2(result.to(BFloat16))
502 + packed = sanitize_negative_zero_u32x2(result_packed)
503 + mailbox_element = (
504 + (Int64(stage) * self.tp + self.rank) * self.capacity_m + Int64(token)
505 + ) * self.hidden + Int64(fragment) * QUAD_BF16
506 + stmc_bf16x4(
507 + contribution_mailbox_multicast_address + mailbox_element * 2,
508 + packed,
509 + )
510 +
511 + if block == 0 and tidx == 0:
512 + store_global_u32(
513 + stage_state.iterator + ACTIVE_STAGE,
514 + stage,
515 + )
516 + cute.arch.barrier()
517 + if cutlass.const_expr(self.enable_pdl):
518 + cute.arch.griddepcontrol_launch_dependents()
519 +
520 +
521 +class _LamportMHCDeviceKernel:
522 + """FlashInfer's ``_LamportResidualRMSNormDeviceKernel`` with mHC.
523 +
524 + The residual add and residual copy become the mHC post-mix of the hc
525 + streams, stored as the residual output, and their pre-mix collapse, which
526 + the unchanged RMSNorm normalizes.
527 + """
528 +
529 + def __init__(
530 + self,
531 + *,
532 + hidden: int,
533 + hc: int,
534 + tp: int,
535 + capacity_m: int,
536 + cluster_size: int,
537 + rank_lanes: int,
538 + threads: int,
539 + enable_pdl: bool,
540 + ) -> None:
541 + if rank_lanes not in (1, 2, 4, 8):
542 + raise ValueError("rank_lanes must be 1, 2, 4, or 8")
543 + if tp % rank_lanes:
544 + raise ValueError("tp must be divisible by rank_lanes")
545 + if threads <= 0 or threads % WARP_SIZE or threads % rank_lanes:
546 + raise ValueError("threads must be a positive warp and rank-lane multiple")
547 + if hidden <= 0 or hidden % VEC_BF16:
548 + raise ValueError("hidden must be a positive multiple of 8")
549 + self.hidden = hidden
550 + self.hc = hc
551 + self.tp = tp
552 + self.capacity_m = capacity_m
553 + self.cluster_size = cluster_size
554 + self.rank_lanes = rank_lanes
555 + self.threads = threads
556 + self.enable_pdl = enable_pdl
557 + self.fragments = hidden // VEC_BF16
558 + self.groups_per_cta = threads // rank_lanes
559 + self.fragment_stride = cluster_size * self.groups_per_cta
560 + self.trips = math.ceil(self.fragments / self.fragment_stride)
561 + self.warps = threads // WARP_SIZE
562 + self.rank_waves = tp // rank_lanes
563 +
564 + def smem_size_in_bytes(self) -> int:
565 + return (self.warps + self.cluster_size) * 4
566 +
567 + @cute.jit
568 + def __call__(
569 + self,
570 + contribution_mailbox: cute.Tensor,
571 + residual_source: cute.Tensor,
572 + post: cute.Tensor,
573 + comb: cute.Tensor,
574 + pre: cute.Tensor,
575 + gamma: cute.Tensor,
576 + residual_output: cute.Tensor,
577 + norm_output: cute.Tensor,
578 + stage_state: cute.Tensor,
579 + eps: Float32,
580 + m: Int32,
581 + stream: cuda.CUstream,
582 + ) -> None:
583 + self.kernel(
584 + contribution_mailbox,
585 + residual_source,
586 + post,
587 + comb,
588 + pre,
589 + gamma,
590 + residual_output,
591 + norm_output,
592 + stage_state,
593 + eps,
594 + ).launch(
595 + grid=(m, self.cluster_size, 1),
596 + block=(self.threads, 1, 1),
597 + cluster=(1, self.cluster_size, 1),
598 + smem=self.smem_size_in_bytes(),
599 + stream=stream,
600 + use_pdl=self.enable_pdl,
601 + )
602 +
603 + @cute.kernel
604 + def kernel(
605 + self,
606 + contribution_mailbox: cute.Tensor,
607 + residual_source: cute.Tensor,
608 + post: cute.Tensor,
609 + comb: cute.Tensor,
610 + pre: cute.Tensor,
611 + gamma: cute.Tensor,
612 + residual_output: cute.Tensor,
613 + norm_output: cute.Tensor,
614 + stage_state: cute.Tensor,
615 + eps: Float32,
616 + ) -> None:
617 + hc = self.hc
618 + tidx, _, _ = cute.arch.thread_idx()
619 + token, _, _ = cute.arch.block_idx()
620 + cluster_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
621 + rank_lane = tidx % self.rank_lanes
622 + group = tidx // self.rank_lanes
623 + base_fragment = cluster_rank * self.groups_per_cta + group
624 +
625 + prenorm_fragments = cute.make_rmem_tensor(
626 + cute.make_layout(
627 + (self.trips, VEC_BF16),
628 + stride=(VEC_BF16, 1),
629 + ),
630 + BFloat16,
631 + )
632 + prenorm_fragments.fill(BFloat16(0.0))
633 + gamma_fragments = cute.make_rmem_tensor(
634 + cute.make_layout(
635 + (self.trips, VEC_BF16),
636 + stride=(VEC_BF16, 1),
637 + ),
638 + BFloat16,
639 + )
640 + gamma_fragments.fill(BFloat16(0.0))
641 + residual_fragments = cute.make_rmem_tensor(
642 + cute.make_layout(
643 + (self.trips, hc, VEC_BF16),
644 + stride=(hc * VEC_BF16, VEC_BF16, 1),
645 + ),
646 + BFloat16,
647 + )
648 + residual_fragments.fill(BFloat16(0.0))
649 +
650 + for trip in cutlass.range_constexpr(self.trips):
651 + fragment = base_fragment + trip * self.fragment_stride
652 + if fragment < self.fragments and rank_lane == 0:
653 + gamma_element = Int64(fragment) * VEC_BF16
654 + gamma_pointer = cute.make_ptr(
655 + BFloat16,
656 + (gamma.iterator + gamma_element).llvm_ptr,
657 + cute.AddressSpace.gmem,
658 + assumed_align=16,
659 + )
660 + gamma_fragments[trip, None].store(
661 + packed_u32x4_to_bf16x8(load_global_u32x4(gamma_pointer))
662 + )
663 + for source in cutlass.range_constexpr(hc):
664 + residual_element = (
665 + Int64(token) * hc + source
666 + ) * self.hidden + gamma_element
667 + residual_pointer = cute.make_ptr(
668 + BFloat16,
669 + (residual_source.iterator + residual_element).llvm_ptr,
670 + cute.AddressSpace.gmem,
671 + assumed_align=16,
672 + )
673 + residual_fragments[trip, source, None].store(
674 + packed_u32x4_to_bf16x8(load_global_u32x4(residual_pointer))
675 + )
676 + post_mix = cute.make_rmem_tensor(cute.make_layout((hc,)), Float32)
677 + pre_mix = cute.make_rmem_tensor(cute.make_layout((hc,)), Float32)
678 + comb_mix = cute.make_rmem_tensor(
679 + cute.make_layout((hc, hc), stride=(hc, 1)), Float32
680 + )
681 + for target in cutlass.range_constexpr(hc):
682 + mix_index = Int64(token) * hc + target
683 + post_mix[target] = cute.arch.load(
684 + (post.iterator + mix_index).llvm_ptr, Float32
685 + )
686 + pre_mix[target] = cute.arch.load(
687 + (pre.iterator + mix_index).llvm_ptr, Float32
688 + )
689 + for source in cutlass.range_constexpr(hc):
690 + comb_index = (Int64(token) * hc + source) * hc + target
691 + comb_mix[source, target] = cute.arch.load(
692 + (comb.iterator + comb_index).llvm_ptr, Float32
693 + )
694 +
695 + if cutlass.const_expr(self.enable_pdl):
696 + cute.arch.griddepcontrol_wait()
697 +
698 + active_stage = load_volatile_u32(stage_state.iterator + ACTIVE_STAGE)
699 + for trip in cutlass.range_constexpr(self.trips):
700 + fragment = base_fragment + trip * self.fragment_stride
701 + lane_packed = cute.make_rmem_tensor(
702 + cute.make_layout((self.rank_waves, 4)),
703 + Uint32,
704 + )
705 + lane_packed.fill(Uint32(0))
706 + dirty = fragment < self.fragments
707 + while dirty:
708 + dirty = False
709 + for wave in cutlass.range_constexpr(self.rank_waves):
710 + source_rank = wave * self.rank_lanes + rank_lane
711 + if fragment < self.fragments:
712 + source_element = (
713 + (Int64(active_stage) * self.tp + Int64(source_rank))
714 + * self.capacity_m
715 + + Int64(token)
716 + ) * self.hidden + Int64(fragment) * VEC_BF16
717 + source_pointer = cute.make_ptr(
718 + BFloat16,
719 + (contribution_mailbox.iterator + source_element).llvm_ptr,
720 + cute.AddressSpace.gmem,
721 + assumed_align=16,
722 + )
723 + packed = load_global_u32x4(
724 + source_pointer,
725 + volatile=True,
726 + )
727 + dirty = dirty | fragment_has_negative_zero(packed)
728 + for word in cutlass.range_constexpr(4):
729 + lane_packed[wave, word] = packed[word]
730 +
731 + lane_sum = cute.make_rmem_tensor(cute.make_layout((VEC_BF16,)), Float32)
732 + lane_sum.fill(Float32(0.0))
733 + for wave in cutlass.range_constexpr(self.rank_waves):
734 + packed = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32)
735 + for word in cutlass.range_constexpr(4):
736 + packed[word] = lane_packed[wave, word]
737 + lane_sum.store(
738 + lane_sum.load() + packed_u32x4_to_bf16x8(packed.load()).to(Float32)
739 + )
740 + for offset in cutlass.range_constexpr(1, 5):
741 + if cutlass.const_expr(offset < self.rank_lanes and offset in (1, 2, 4)):
742 + for element in cutlass.range_constexpr(VEC_BF16):
743 + lane_sum[element] = lane_sum[
744 + element
745 + ] + cute.arch.shuffle_sync_bfly(
746 + lane_sum[element],
747 + offset=offset,
748 + mask=-1,
749 + mask_and_clamp=31,
750 + )
751 +
752 + if fragment < self.fragments and rank_lane == 0:
753 + # The all-reduced sublayer output is BF16, as the unfused
754 + # collective returns it, before mHC post mixes it in.
755 + reduced = lane_sum.load().to(BFloat16).to(Float32)
756 + collapse = cute.make_rmem_tensor(cute.make_layout((VEC_BF16,)), Float32)
757 + collapse.fill(Float32(0.0))
758 + for target in cutlass.range_constexpr(hc):
759 + mixed = reduced * post_mix[target]
760 + for source in cutlass.range_constexpr(hc):
761 + mixed = (
762 + mixed
763 + + residual_fragments[trip, source, None].load().to(Float32)
764 + * comb_mix[source, target]
765 + )
766 + mixed_bf16 = mixed.to(BFloat16)
767 + output_element = (Int64(token) * hc + target) * self.hidden + Int64(
768 + fragment
769 + ) * VEC_BF16
770 + store_global_u32x4(
771 + Int64((residual_output.iterator + output_element).toint()),
772 + bf16x8_to_packed_u32x4(mixed_bf16),
773 + )
774 + collapse.store(
775 + collapse.load() + mixed_bf16.to(Float32) * pre_mix[target]
776 + )
777 + prenorm_fragments[trip, None].store(collapse.load().to(BFloat16))
778 +
779 + if token == 0 and cluster_rank == 0 and tidx == 0:
780 + store_global_u32(
781 + stage_state.iterator + NEXT_STAGE,
782 + (active_stage + Uint32(1)) % Uint32(LAMPORT_GENERATIONS),
783 + )
784 + if cutlass.const_expr(self.enable_pdl):
785 + cute.arch.griddepcontrol_launch_dependents()
786 +
787 + for trip in cutlass.range_constexpr(self.trips):
788 + fragment = base_fragment + trip * self.fragment_stride
789 + for wave in cutlass.range_constexpr(self.rank_waves):
790 + source_rank = wave * self.rank_lanes + rank_lane
791 + if fragment < self.fragments:
792 + source_element = (
793 + (Int64(active_stage) * self.tp + Int64(source_rank))
794 + * self.capacity_m
795 + + Int64(token)
796 + ) * self.hidden + Int64(fragment) * VEC_BF16
797 + store_lamport_sentinel_u32x4(
798 + Int64((contribution_mailbox.iterator + source_element).toint())
799 + )
800 +
801 + thread_sum = Float32(0.0)
802 + for trip in cutlass.range_constexpr(self.trips):
803 + fragment = base_fragment + trip * self.fragment_stride
804 + if fragment < self.fragments and rank_lane == 0:
805 + values = prenorm_fragments[trip, None].load().to(Float32)
806 + thread_sum = thread_sum + (values * values).reduce(
807 + cute.ReductionOp.ADD,
808 + init_val=Float32(0.0),
809 + reduction_profile=0,
810 + )
811 +
812 + smem = cutlass.utils.SmemAllocator()
813 + warp_sums = smem.allocate_array(Float32, self.warps)
814 + cluster_sums = smem.allocate_array(Float32, self.cluster_size)
815 + cta_sum = _group_leader_block_sum(
816 + thread_sum,
817 + warp_sums,
818 + self.warps,
819 + self.rank_lanes,
820 + )
821 + if tidx < self.cluster_size:
822 + local_slot = cluster_sums + cluster_rank
823 + remote_slot = map_shared_to_peer(local_slot, Int32(tidx))
824 + store_shared_cluster_f32(remote_slot, cta_sum)
825 + cute.arch.cluster_arrive()
826 + cute.arch.cluster_wait()
827 + full_sum = Float32(0.0)
828 + for peer in cutlass.range_constexpr(self.cluster_size):
829 + full_sum = full_sum + cute.arch.load(
830 + (cluster_sums + peer).llvm_ptr,
831 + Float32,
832 + )
833 + inv_rms = cute.math.rsqrt(
834 + full_sum / Float32(self.hidden) + eps,
835 + fastmath=True,
836 + )
837 + for trip in cutlass.range_constexpr(self.trips):
838 + fragment = base_fragment + trip * self.fragment_stride
839 + if fragment < self.fragments and rank_lane == 0:
840 + gamma_values = gamma_fragments[trip, None].load().to(Float32)
841 + result = (
842 + prenorm_fragments[trip, None].load().to(Float32)
843 + * inv_rms
844 + * gamma_values
845 + ).to(BFloat16)
846 + output_element = Int64(token) * self.hidden + Int64(fragment) * VEC_BF16
847 + store_global_u32x4(
848 + Int64((norm_output.iterator + output_element).toint()),
849 + bf16x8_to_packed_u32x4(result),
850 + )
851 +
852 +
853 +class AllReduceMHC:
854 + """Lamport TP all-reduce fused with mHC post, collapse and RMSNorm.
855 +
856 + ``__call__`` reduces ``x`` across the TP group, writes the post-mixed hc
857 + streams and the normalized collapse, and leaves the next mixes to the
858 + caller. ``finalize`` does the same for an unfinalized MoE output, folding
859 + the top-k reduction and the shared-expert add into the publish. Each
860 + instance owns a Lamport mailbox; construction is collective over the TP
861 + group.
862 + """
863 +
864 + def __init__(
865 + self,
866 + *,
867 + hidden_size: int,
868 + hc_mult: int,
869 + max_num_tokens: int,
870 + top_k: int,
871 + device: torch.device,
872 + ) -> None:
873 + group = get_tp_group().device_group
874 + self.hidden_size = hidden = hidden_size
875 + self.hc_mult = hc = hc_mult
876 + self.capacity = capacity = max_num_tokens
877 + self.top_k = top_k
878 + tp = dist.get_world_size(group)
879 + rank = dist.get_rank(group)
880 + device = torch.device(device)
881 +
882 + # FlashInfer's LL_ALL_REDUCE_GB300_TP{4,8}_H5120 and
883 + # LL_FINALIZE_..._H5120_K* presets: a cluster of 5 x 128 threads covers
884 + # a 640-fragment token in one fully used trip, and the finalize stages
885 + # all top_k rows at once.
886 + with torch.accelerator.device_index(device.index):
887 + self._publish = cute.compile(
888 + _SharedOnlyPublishDeviceKernel(
889 + hidden=hidden,
890 + tp=tp,
891 + rank=rank,
892 + capacity_m=capacity,
893 + elements_per_thread=VEC_BF16,
894 + threads=128,
895 + release_before_store=False,
896 + enable_pdl=True,
897 + ),
898 + make_fake_dynamic_compact_tensor(
899 + BFloat16, alignment=16, divisibility=hidden
900 + ),
901 + make_fake_compact_tensor(Int32, (2,), assumed_align=4),
902 + Int64(0),
903 + Int32(capacity),
904 + current_cu_stream(),
905 + )
906 + self._finalize_publish = cute.compile(
907 + _QuadFinalizePublishDeviceKernel(
908 + hidden=hidden,
909 + top_k=top_k,
910 + tp=tp,
911 + rank=rank,
912 + capacity_m=capacity,
913 + threads=128,
914 + # DSV4.1's router folds the scale into the weights.
915 + routed_scaling_factor=1.0,
916 + include_shared_expert=True,
917 + load_shared_expert_before_pdl=False,
918 + enable_pdl=True,
919 + prefetch_group=top_k,
920 + fp32_weights=True,
921 + ),
922 + make_fake_dynamic_compact_tensor(
923 + BFloat16, alignment=16, divisibility=hidden
924 + ),
925 + make_fake_dynamic_compact_tensor(
926 + Float32, alignment=4, divisibility=top_k
927 + ),
928 + make_fake_dynamic_compact_tensor(
929 + Int32, alignment=4, divisibility=top_k
930 + ),
931 + make_fake_dynamic_compact_tensor(
932 + BFloat16, alignment=16, divisibility=hidden
933 + ),
934 + make_fake_compact_tensor(Int32, (2,), assumed_align=4),
935 + Int64(0),
936 + Int32(capacity),
937 + current_cu_stream(),
938 + )
939 + self._collective = cute.compile(
940 + _LamportMHCDeviceKernel(
941 + hidden=hidden,
942 + hc=hc,
943 + tp=tp,
944 + capacity_m=capacity,
945 + cluster_size=5,
946 + rank_lanes=1,
947 + threads=128,
948 + enable_pdl=True,
949 + ),
950 + make_fake_compact_tensor(
951 + BFloat16,
952 + (LAMPORT_GENERATIONS * tp * capacity * hidden,),
953 + assumed_align=16,
954 + ),
955 + make_fake_dynamic_compact_tensor(
956 + BFloat16, alignment=16, divisibility=hc * hidden
957 + ),
958 + make_fake_dynamic_compact_tensor(Float32, alignment=4, divisibility=hc),
959 + make_fake_dynamic_compact_tensor(
960 + Float32, alignment=4, divisibility=hc * hc
961 + ),
962 + make_fake_dynamic_compact_tensor(Float32, alignment=4, divisibility=hc),
963 + make_fake_compact_tensor(BFloat16, (hidden,), assumed_align=16),
964 + make_fake_dynamic_compact_tensor(
965 + BFloat16, alignment=16, divisibility=hc * hidden
966 + ),
967 + make_fake_dynamic_compact_tensor(
968 + BFloat16, alignment=16, divisibility=hidden
969 + ),
970 + make_fake_compact_tensor(Int32, (2,), assumed_align=4),
971 + Float32(0.0),
972 + Int32(capacity),
973 + current_cu_stream(),
974 + )
975 +
976 + self._mailbox = symm_mem.empty(
977 + (LAMPORT_GENERATIONS, tp, capacity, hidden),
978 + dtype=torch.bfloat16,
979 + device=device,
980 + )
981 + self._mailbox_handle = symm_mem.rendezvous(self._mailbox, group)
982 + multicast = int(self._mailbox_handle.multicast_ptr or 0)
983 + if not multicast or multicast % 16:
984 + raise RuntimeError("NVLink multicast mapping is unavailable")
985 + self._multicast = multicast
986 + # Every slot starts as the Lamport sentinel, BF16 negative zero.
987 + self._mailbox.view(torch.int16).fill_(-32768)
988 + self._stage_state = torch.zeros(2, dtype=torch.int32, device=device)
989 + torch.accelerator.synchronize()
990 + dist.barrier(group=group)
991 +
992 + def __call__(
993 + self,
994 + x: torch.Tensor,
995 + residual: torch.Tensor,
996 + post: torch.Tensor,
997 + comb: torch.Tensor,
998 + pre: torch.Tensor,
999 + norm_weight: torch.Tensor,
1000 + norm_eps: float,
1001 + ) -> tuple[torch.Tensor, torch.Tensor]:
1002 + """Return (post-mixed residual streams, normalized layer input).
1003 +
1004 + Args:
1005 + x: [M, hidden] BF16 TP-partial sublayer output.
1006 + residual: [M, hc, hidden] BF16 residual streams.
1007 + post: [M, hc(, 1)] FP32 post-mix.
1008 + comb: [M, hc, hc] FP32 residual mix, indexed [source, target].
1009 + pre: [M, hc] FP32 carried pre-mix for the collapse.
1010 + norm_weight: [hidden] BF16 RMSNorm weight.
1011 + norm_eps: RMSNorm epsilon.
1012 +
1013 + """
1014 + m = x.shape[0]
1015 + self._check(m, residual, post, comb, pre)
1016 + if x.shape != (m, self.hidden_size):
1017 + raise ValueError("unexpected x shape")
1018 + self._publish(
1019 + to_cute_dynamic(x.flatten(), 16, divisibility=self.hidden_size),
1020 + to_cute(self._stage_state, 4),
1021 + Int64(self._multicast),
1022 + Int32(m),
1023 + current_cu_stream(),
1024 + )
1025 + return self._collect(residual, post, comb, pre, norm_weight, norm_eps, m)
1026 +
1027 + def finalize(
1028 + self,
1029 + gemm2_permuted: torch.Tensor,
1030 + expert_weights: torch.Tensor,
1031 + expanded_idx_to_permuted_idx: torch.Tensor,
1032 + shared_output: torch.Tensor,
1033 + residual: torch.Tensor,
1034 + post: torch.Tensor,
1035 + comb: torch.Tensor,
1036 + pre: torch.Tensor,
1037 + norm_weight: torch.Tensor,
1038 + norm_eps: float,
1039 + ) -> tuple[torch.Tensor, torch.Tensor]:
1040 + """``__call__`` for a MoE output left unfinalized.
1041 +
1042 + Args:
1043 + gemm2_permuted: [rows, hidden] BF16 unweighted GEMM2 output in the
1044 + MoE's permuted order.
1045 + expert_weights: [M, top_k] FP32 routing weights, already scaled.
1046 + expanded_idx_to_permuted_idx: [M, top_k] int32 row of each route,
1047 + -1 for an expert this rank does not hold.
1048 + shared_output: [M, hidden] BF16 TP-partial shared-expert output.
1049 + residual: [M, hc, hidden] BF16 residual streams.
1050 + post: [M, hc(, 1)] FP32 post-mix.
1051 + comb: [M, hc, hc] FP32 residual mix, indexed [source, target].
1052 + pre: [M, hc] FP32 carried pre-mix for the collapse.
1053 + norm_weight: [hidden] BF16 RMSNorm weight.
1054 + norm_eps: RMSNorm epsilon.
1055 +
1056 + """
1057 + m = shared_output.shape[0]
1058 + hidden, top_k = self.hidden_size, self.top_k
1059 + self._check(m, residual, post, comb, pre)
1060 + if gemm2_permuted.dim() != 2 or gemm2_permuted.shape[1] != hidden:
1061 + raise ValueError("gemm2_permuted must be [rows, hidden]")
1062 + if shared_output.shape != (m, hidden):
1063 + raise ValueError("unexpected shared_output shape")
1064 + if expert_weights.dtype != torch.float32 or expert_weights.numel() != m * top_k:
1065 + raise ValueError("expert_weights must be FP32 [M, top_k]")
1066 + if (
1067 + expanded_idx_to_permuted_idx.dtype != torch.int32
1068 + or expanded_idx_to_permuted_idx.numel() != m * top_k
1069 + ):
1070 + raise ValueError("expanded_idx_to_permuted_idx must be int32 [M, top_k]")
1071 + self._finalize_publish(
1072 + to_cute_dynamic(gemm2_permuted.flatten(), 16, divisibility=hidden),
1073 + to_cute_dynamic(expert_weights.flatten(), 4, divisibility=top_k),
1074 + to_cute_dynamic(
1075 + expanded_idx_to_permuted_idx.flatten(), 4, divisibility=top_k
1076 + ),
1077 + to_cute_dynamic(shared_output.flatten(), 16, divisibility=hidden),
1078 + to_cute(self._stage_state, 4),
1079 + Int64(self._multicast),
1080 + Int32(m),
1081 + current_cu_stream(),
1082 + )
1083 + return self._collect(residual, post, comb, pre, norm_weight, norm_eps, m)
1084 +
1085 + def _check(
1086 + self,
1087 + m: int,
1088 + residual: torch.Tensor,
1089 + post: torch.Tensor,
1090 + comb: torch.Tensor,
1091 + pre: torch.Tensor,
1092 + ) -> None:
1093 + hc = self.hc_mult
1094 + if not 1 <= m <= self.capacity:
1095 + raise ValueError(f"M={m} is outside [1, {self.capacity}]")
1096 + if residual.shape != (m, hc, self.hidden_size):
1097 + raise ValueError("unexpected residual shape")
1098 + if (
1099 + post.numel() != m * hc
1100 + or pre.numel() != m * hc
1101 + or comb.numel() != m * hc * hc
1102 + ):
1103 + raise ValueError("unexpected mix shapes")
1104 +
1105 + def _collect(
1106 + self,
1107 + residual: torch.Tensor,
1108 + post: torch.Tensor,
1109 + comb: torch.Tensor,
1110 + pre: torch.Tensor,
1111 + norm_weight: torch.Tensor,
1112 + norm_eps: float,
1113 + m: int,
1114 + ) -> tuple[torch.Tensor, torch.Tensor]:
1115 + hidden, hc = self.hidden_size, self.hc_mult
1116 + residual_output = torch.empty_like(residual)
1117 + layer_input = torch.empty(
1118 + (m, hidden), dtype=torch.bfloat16, device=residual.device
1119 + )
1120 + self._collective(
1121 + to_cute(self._mailbox.flatten(), 16),
1122 + to_cute_dynamic(residual.flatten(), 16, divisibility=hc * hidden),
1123 + to_cute_dynamic(post.flatten(), 4, divisibility=hc),
1124 + to_cute_dynamic(comb.flatten(), 4, divisibility=hc * hc),
1125 + to_cute_dynamic(pre.flatten(), 4, divisibility=hc),
1126 + to_cute(norm_weight, 16),
1127 + to_cute_dynamic(residual_output.flatten(), 16, divisibility=hc * hidden),
1128 + to_cute_dynamic(layer_input.flatten(), 16, divisibility=hidden),
1129 + to_cute(self._stage_state, 4),
1130 + Float32(norm_eps),
1131 + Int32(m),
1132 + current_cu_stream(),
1133 + )
1134 + return residual_output, layer_input
vllm/models/deepseek_v41/nvidia/ops/cute_dsl/primitives.py new
+396
@@ -0,0 +1,396 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# Adapted from FlashInfer's flashinfer/comm/mnnvl_cutedsl/cute_dsl_primitives.py
4 +# and runtime.py (flashinfer-ai/flashinfer@139af6f6).
5 +
6 +"""PTX building blocks and launch helpers for the DSV4.1 LL all-reduce."""
7 +
8 +import cuda.bindings.driver as cuda
9 +import cutlass
10 +import cutlass.cute as cute
11 +import torch
12 +from cutlass import Float32, Int32, Int64, Uint16, Uint32
13 +from cutlass._mlir import ir
14 +from cutlass._mlir.dialects import llvm, vector
15 +from cutlass.cute.runtime import from_dlpack, make_fake_compact_tensor
16 +from cutlass.cutlass_dsl import T, dsl_user_op
17 +
18 +WARP_SIZE = 32
19 +VEC_BF16 = 8
20 +QUAD_BF16 = 4
21 +NEGATIVE_ZERO_BF16_BITS = 0x8000
22 +NEGATIVE_ZERO_BF16_PAIR = 0x80008000
23 +
24 +
25 +def _asm(result, operands, text, constraints, *, side_effects, loc, ip):
26 + return llvm.inline_asm(
27 + result,
28 + operands,
29 + text,
30 + constraints,
31 + has_side_effects=side_effects,
32 + is_align_stack=False,
33 + asm_dialect=llvm.AsmDialect.AD_ATT,
34 + loc=loc,
35 + ip=ip,
36 + )
37 +
38 +
39 +@dsl_user_op
40 +def load_global_u32x4(
41 + pointer: cute.Pointer,
42 + *,
43 + volatile: cutlass.Constexpr[bool] = False,
44 + loc=None,
45 + ip=None,
46 +):
47 + address = pointer.toint(loc=loc, ip=ip)
48 + opcode = "ld.volatile.global.v4.u32" if volatile else "ld.global.v4.u32"
49 + loaded = _asm(
50 + llvm.StructType.get_literal([T.i32()] * 4),
51 + [address.ir_value(loc=loc, ip=ip)],
52 + f"{opcode} {{$0, $1, $2, $3}}, [$4];",
53 + "=r,=r,=r,=r,l",
54 + side_effects=volatile,
55 + loc=loc,
56 + ip=ip,
57 + )
58 + packed = vector.from_elements(
59 + ir.VectorType.get([4], T.i32(), loc=loc),
60 + [
61 + llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip)
62 + for index in range(4)
63 + ],
64 + loc=loc,
65 + ip=ip,
66 + )
67 + return cute.TensorSSA(packed, 4, Uint32)
68 +
69 +
70 +@dsl_user_op
71 +def load_global_u32x2(pointer: cute.Pointer, *, loc=None, ip=None):
72 + address = pointer.toint(loc=loc, ip=ip)
73 + loaded = _asm(
74 + llvm.StructType.get_literal([T.i32()] * 2),
75 + [address.ir_value(loc=loc, ip=ip)],
76 + "ld.global.v2.u32 {$0, $1}, [$2];",
77 + "=r,=r,l",
78 + side_effects=False,
79 + loc=loc,
80 + ip=ip,
81 + )
82 + packed = vector.from_elements(
83 + ir.VectorType.get([2], T.i32(), loc=loc),
84 + [
85 + llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip)
86 + for index in range(2)
87 + ],
88 + loc=loc,
89 + ip=ip,
90 + )
91 + return cute.TensorSSA(packed, 2, Uint32)
92 +
93 +
94 +@dsl_user_op
95 +def load_global_bf16_as_f32(address: Int64, *, loc=None, ip=None) -> Float32:
96 + return Float32(
97 + _asm(
98 + T.f32(),
99 + [address.ir_value(loc=loc, ip=ip)],
100 + "{\n\t.reg .b16 bits;\n\tld.global.b16 bits, [$1];\n\t"
101 + "cvt.f32.bf16 $0, bits;\n\t}",
102 + "=f,l",
103 + side_effects=False,
104 + loc=loc,
105 + ip=ip,
106 + )
107 + )
108 +
109 +
110 +@dsl_user_op
111 +def store_global_u32x4(address: Int64, packed, *, loc=None, ip=None) -> None:
112 + words = [packed[index].ir_value(loc=loc, ip=ip) for index in range(4)]
113 + _asm(
114 + None,
115 + [address.ir_value(loc=loc, ip=ip), *words],
116 + "st.global.v4.u32 [$0], {$1, $2, $3, $4};",
117 + "l,r,r,r,r",
118 + side_effects=True,
119 + loc=loc,
120 + ip=ip,
121 + )
122 +
123 +
124 +@dsl_user_op
125 +def store_lamport_sentinel_u32x4(address: Int64, *, loc=None, ip=None) -> None:
126 + sentinel = Uint32(NEGATIVE_ZERO_BF16_PAIR).ir_value(loc=loc, ip=ip)
127 + _asm(
128 + None,
129 + [address.ir_value(loc=loc, ip=ip), sentinel, sentinel, sentinel, sentinel],
130 + "st.global.v4.u32 [$0], {$1, $2, $3, $4};",
131 + "l,r,r,r,r",
132 + side_effects=True,
133 + loc=loc,
134 + ip=ip,
135 + )
136 +
137 +
138 +@dsl_user_op
139 +def load_volatile_u32(pointer: cute.Pointer, *, loc=None, ip=None) -> Uint32:
140 + address = pointer.toint(loc=loc, ip=ip)
141 + return Uint32(
142 + _asm(
143 + T.i32(),
144 + [address.ir_value(loc=loc, ip=ip)],
145 + "ld.volatile.global.u32 $0, [$1];",
146 + "=r,l",
147 + side_effects=True,
148 + loc=loc,
149 + ip=ip,
150 + )
151 + )
152 +
153 +
154 +@dsl_user_op
155 +def store_global_u32(
156 + pointer: cute.Pointer, value: Uint32, *, loc=None, ip=None
157 +) -> None:
158 + address = pointer.toint(loc=loc, ip=ip)
159 + _asm(
160 + None,
161 + [address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)],
162 + "st.global.u32 [$0], $1;",
163 + "l,r",
164 + side_effects=True,
165 + loc=loc,
166 + ip=ip,
167 + )
168 +
169 +
170 +@dsl_user_op
171 +def packed_u32x4_to_bf16x8(packed, *, loc=None, ip=None):
172 + values = llvm.bitcast(
173 + ir.VectorType.get([VEC_BF16], cutlass.BFloat16.mlir_type, loc=loc),
174 + packed.ir_value(loc=loc, ip=ip),
175 + loc=loc,
176 + ip=ip,
177 + )
178 + return cute.TensorSSA(values, VEC_BF16, cutlass.BFloat16)
179 +
180 +
181 +@dsl_user_op
182 +def packed_u32x2_to_bf16x4(packed, *, loc=None, ip=None):
183 + values = llvm.bitcast(
184 + ir.VectorType.get([QUAD_BF16], cutlass.BFloat16.mlir_type, loc=loc),
185 + packed.ir_value(loc=loc, ip=ip),
186 + loc=loc,
187 + ip=ip,
188 + )
189 + return cute.TensorSSA(values, QUAD_BF16, cutlass.BFloat16)
190 +
191 +
192 +@dsl_user_op
193 +def bf16x4_to_packed_u32x2(values, *, loc=None, ip=None):
194 + packed = llvm.bitcast(
195 + ir.VectorType.get([2], T.i32(), loc=loc),
196 + values.ir_value(loc=loc, ip=ip),
197 + loc=loc,
198 + ip=ip,
199 + )
200 + return cute.TensorSSA(packed, 2, Uint32)
201 +
202 +
203 +@dsl_user_op
204 +def bf16x8_to_packed_u32x4(values, *, loc=None, ip=None):
205 + packed = llvm.bitcast(
206 + ir.VectorType.get([4], T.i32(), loc=loc),
207 + values.ir_value(loc=loc, ip=ip),
208 + loc=loc,
209 + ip=ip,
210 + )
211 + return cute.TensorSSA(packed, 4, Uint32)
212 +
213 +
214 +@cute.jit
215 +def sanitize_negative_zero_u32(word: Uint32) -> Uint32:
216 + low = Uint16(word & Uint32(0xFFFF))
217 + high = Uint16(word >> Uint32(16))
218 + if low == Uint16(NEGATIVE_ZERO_BF16_BITS):
219 + word = word & Uint32(0xFFFF0000)
220 + if high == Uint16(NEGATIVE_ZERO_BF16_BITS):
221 + word = word & Uint32(0x0000FFFF)
222 + return word
223 +
224 +
225 +@cute.jit
226 +def sanitize_negative_zero_u32x4(packed):
227 + sanitized = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32)
228 + for index in cutlass.range_constexpr(4):
229 + sanitized[index] = sanitize_negative_zero_u32(packed[index])
230 + return sanitized.load()
231 +
232 +
233 +@cute.jit
234 +def sanitize_negative_zero_u32x2(packed):
235 + sanitized = cute.make_rmem_tensor(cute.make_layout((2,)), Uint32)
236 + for index in cutlass.range_constexpr(2):
237 + sanitized[index] = sanitize_negative_zero_u32(packed[index])
238 + return sanitized.load()
239 +
240 +
241 +@cute.jit
242 +def fragment_has_negative_zero(packed):
243 + dirty = False
244 + for index in cutlass.range_constexpr(4):
245 + word = packed[index]
246 + dirty = (
247 + dirty
248 + | (Uint16(word & Uint32(0xFFFF)) == Uint16(NEGATIVE_ZERO_BF16_BITS))
249 + | (Uint16(word >> Uint32(16)) == Uint16(NEGATIVE_ZERO_BF16_BITS))
250 + )
251 + return dirty
252 +
253 +
254 +@dsl_user_op
255 +def map_shared_to_peer(
256 + smem_pointer: cute.Pointer, peer_rank: Int32, *, loc=None, ip=None
257 +) -> Int32:
258 + address = smem_pointer.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip)
259 + return Int32(
260 + _asm(
261 + T.i32(),
262 + [address, peer_rank.ir_value(loc=loc, ip=ip)],
263 + "mapa.shared::cluster.u32 $0, $1, $2;",
264 + "=r,r,r",
265 + side_effects=False,
266 + loc=loc,
267 + ip=ip,
268 + )
269 + )
270 +
271 +
272 +@dsl_user_op
273 +def store_shared_cluster_f32(
274 + remote_address: Int32, value: Float32, *, loc=None, ip=None
275 +) -> None:
276 + _asm(
277 + None,
278 + [remote_address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)],
279 + "st.shared::cluster.f32 [$0], $1;",
280 + "r,f",
281 + side_effects=True,
282 + loc=loc,
283 + ip=ip,
284 + )
285 +
286 +
287 +@dsl_user_op
288 +def f32_to_bf16_bits(value: Float32, *, loc=None, ip=None) -> Uint32:
289 + return Uint32(
290 + _asm(
291 + T.i32(),
292 + [value.ir_value(loc=loc, ip=ip)],
293 + "{\n\t.reg .b16 bits;\n\tcvt.rn.bf16.f32 bits, $1;\n\t"
294 + "cvt.u32.u16 $0, bits;\n\t}",
295 + "=r,f",
296 + side_effects=False,
297 + loc=loc,
298 + ip=ip,
299 + )
300 + )
301 +
302 +
303 +@dsl_user_op
304 +def shuffle_sync_idx_u32(
305 + value: Uint32, source_lane: Int32, *, loc=None, ip=None
306 +) -> Uint32:
307 + return Uint32(
308 + _asm(
309 + T.i32(),
310 + [value.ir_value(loc=loc, ip=ip), source_lane.ir_value(loc=loc, ip=ip)],
311 + "shfl.sync.idx.b32 $0, $1, $2, 0x1f, 0xffffffff;",
312 + "=r,r,r",
313 + # Preserve full-warp execution across later divergent consumers.
314 + side_effects=True,
315 + loc=loc,
316 + ip=ip,
317 + )
318 + )
319 +
320 +
321 +@dsl_user_op
322 +def stmc_bf16x2(address: Int64, packed: Uint32, *, loc=None, ip=None) -> None:
323 + _asm(
324 + None,
325 + [address.ir_value(loc=loc, ip=ip), packed.ir_value(loc=loc, ip=ip)],
326 + "multimem.st.relaxed.sys.global.bf16x2 [$0], $1;",
327 + "l,r",
328 + side_effects=True,
329 + loc=loc,
330 + ip=ip,
331 + )
332 +
333 +
334 +@dsl_user_op
335 +def stmc_bf16x4(address: Int64, values, *, loc=None, ip=None) -> None:
336 + words = [values[index].ir_value(loc=loc, ip=ip) for index in range(2)]
337 + _asm(
338 + None,
339 + [address.ir_value(loc=loc, ip=ip), *words],
340 + "multimem.st.relaxed.sys.global.v2.bf16x2 [$0], {$1, $2};",
341 + "l,r,r",
342 + side_effects=True,
343 + loc=loc,
344 + ip=ip,
345 + )
346 +
347 +
348 +@dsl_user_op
349 +def stmc_bf16x8(address: Int64, values, *, loc=None, ip=None) -> None:
350 + words = [values[index].ir_value(loc=loc, ip=ip) for index in range(4)]
351 + _asm(
352 + None,
353 + [address.ir_value(loc=loc, ip=ip), *words],
354 + "multimem.st.relaxed.sys.global.v4.bf16x2 [$0], {$1, $2, $3, $4};",
355 + "l,r,r,r,r",
356 + side_effects=True,
357 + loc=loc,
358 + ip=ip,
359 + )
360 +
361 +
362 +class _GraphSafeDLPack:
363 + __slots__ = ("tensor",)
364 +
365 + def __init__(self, tensor: torch.Tensor) -> None:
366 + self.tensor = tensor
367 +
368 + def __dlpack__(self, stream=None):
369 + # stream=-1 skips producer sync; CuTe launches on the current stream,
370 + # which is the capturing one under CUDA graphs.
371 + return self.tensor.__dlpack__(stream=-1)
372 +
373 + def __dlpack_device__(self):
374 + return self.tensor.__dlpack_device__()
375 +
376 +
377 +def to_cute(tensor: torch.Tensor, alignment: int) -> cute.Tensor:
378 + return from_dlpack(_GraphSafeDLPack(tensor.detach()), assumed_align=alignment)
379 +
380 +
381 +def to_cute_dynamic(
382 + tensor: torch.Tensor, alignment: int, *, divisibility: int
383 +) -> cute.Tensor:
384 + return to_cute(tensor, alignment).mark_compact_shape_dynamic(
385 + mode=0, divisibility=divisibility
386 + )
387 +
388 +
389 +def make_fake_dynamic_compact_tensor(dtype, *, alignment: int, divisibility: int):
390 + return make_fake_compact_tensor(
391 + dtype, (cute.sym_int32(divisibility=divisibility),), assumed_align=alignment
392 + )
393 +
394 +
395 +def current_cu_stream() -> cuda.CUstream:
396 + return cuda.CUstream(torch.cuda.current_stream().cuda_stream)
vllm/models/deepseek_v41/nvidia/ops/mhc.py
+73 -18
@@ -9,10 +9,15 @@ import torch
9
10 from vllm.config import VllmConfig
11 from vllm.distributed import get_tp_group
12 +from vllm.logger import init_logger
13 from vllm.model_executor.kernels.mhc.tilelang import (
14 mhc_fused_post_pre_delayed_tilelang,
15 mhc_post_tilelang,
16 )
17 +from vllm.model_executor.layers.fused_moe.moe_output import (
18 + MoEOutput,
19 + UnfinalizedMoEOutput,
20 +)
21 from vllm.platforms import current_platform
22 from vllm.utils.deep_gemm import is_deep_gemm_supported
23
@@ -24,9 +29,17 @@ from .mega_mhc import (
29 if TYPE_CHECKING:
30 from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator
31
32 + from .cute_dsl import AllReduceMHC
33 +
34 +logger = init_logger(__name__)
35 +
36 # GB200 TP4/FlashInfer improves through 16 tokens; larger screens tie or regress.
37 MHC_OVERLAP_MAX_TOKENS = 16
38
39 +# The fused all-reduce + mHC kernel of the overlap path, built once by
40 +# init_mhc_all_reduce.
41 +_all_reduce_mhc: "AllReduceMHC | None" = None
42 +
43
44 def supports_mhc_overlap(vllm_config: VllmConfig) -> bool:
45 """Check kernel requirements and safety of sharing the coefficient stream."""
@@ -53,10 +66,35 @@ def supports_mhc_all_reduce(vllm_config: VllmConfig) -> bool:
66 or config.hc_mult != 4
67 ):
68 return False
69 + # The custom all-reduce's MNNVL buffers succeed exactly when NVLink
70 + # multicast does, which the fused kernel's own mailbox needs too.
71 comm = cast("CudaCommunicator", get_tp_group().device_communicator).ca_comm
72 return comm is not None and bool(comm.mnnvl_lamport_ag_multicast_ptr)
73
74
75 +def init_mhc_all_reduce(vllm_config: VllmConfig) -> None:
76 + """Build the fused all-reduce + mHC kernel the overlap path reduces with.
77 +
78 + Collective over the TP group, so every rank must call it, and only when
79 + ``supports_mhc_all_reduce`` holds.
80 + """
81 + global _all_reduce_mhc
82 + from .cute_dsl import AllReduceMHC
83 +
84 + config = vllm_config.model_config.hf_config
85 + _all_reduce_mhc = AllReduceMHC(
86 + hidden_size=config.hidden_size,
87 + hc_mult=config.hc_mult,
88 + max_num_tokens=MHC_OVERLAP_MAX_TOKENS,
89 + top_k=config.num_experts_per_tok,
90 + device=current_platform.current_device(),
91 + )
92 + logger.info_once(
93 + "DSV4.1 mHC: CuTe DSL Lamport all-reduce fused with mHC for up to %d tokens.",
94 + MHC_OVERLAP_MAX_TOKENS,
95 + )
96 +
97 +
98 def mhc_pre_delayed_overlap(
99 residual: torch.Tensor,
100 fn: torch.Tensor,
@@ -151,7 +189,7 @@ def mhc_pre_delayed_overlap(
189
190
191 def mhc_shifted_post_pre(
154 - x: torch.Tensor,
192 + x: torch.Tensor | MoEOutput,
193 residual: torch.Tensor,
194 post_layer_mix: torch.Tensor,
195 comb_res_mix: torch.Tensor,
@@ -175,39 +213,55 @@ def mhc_shifted_post_pre(
213 ]:
214 """Dispatch shifted post/pre to overlap, Mega-mHC, or fused TileLang.
215
216 + A MoE output left unfinalized is finalized inside the fused all-reduce.
217 When stream is supplied, join it before consuming the returned coefficients.
218 """
219 layer_input = None
181 - if reduce_results:
182 - tp = get_tp_group()
220 + if isinstance(x, MoEOutput):
221 + assert _all_reduce_mhc is not None, "init_mhc_all_reduce was not called"
222 + assert pre_mix is not None and norm_weight is not None
223 + routed = x.routed
224 + assert isinstance(routed, UnfinalizedMoEOutput)
225 + assert x.shared_output is not None
226 + residual, layer_input = _all_reduce_mhc.finalize(
227 + routed.gemm2_permuted,
228 + routed.expert_weights,
229 + routed.expanded_idx_to_permuted_idx,
230 + x.shared_output,
231 + residual,
232 + post_layer_mix,
233 + comb_res_mix,
234 + pre_mix,
235 + norm_weight,
236 + norm_eps,
237 + )
238 + if stream is None:
239 + # Outside FULL graphs, compute the next mixes inline.
240 + stream = torch.cuda.current_stream()
241 + elif reduce_results:
242 if stream is not None and 0 < x.shape[0] <= MHC_OVERLAP_MAX_TOKENS:
184 - comm = cast("CudaCommunicator", tp.device_communicator).ca_comm
185 - assert comm is not None and comm.mnnvl_lamport_epochs is not None
186 - output = torch.empty_like(residual)
187 - layer_input = torch.empty_like(x)
188 - torch.ops._C_custom_ar.all_reduce_mhc(
243 + assert _all_reduce_mhc is not None, "init_mhc_all_reduce was not called"
244 + assert pre_mix is not None and norm_weight is not None
245 + residual, layer_input = _all_reduce_mhc(
246 x,
247 residual,
248 post_layer_mix,
249 comb_res_mix,
250 pre_mix,
251 norm_weight,
195 - output,
196 - layer_input,
197 - comm.mnnvl_lamport_ag_local_ptr,
198 - comm.mnnvl_lamport_ag_multicast_ptr,
199 - comm.mnnvl_lamport_epochs[0],
200 - comm.rank,
201 - comm.mnnvl_buffer_size,
252 norm_eps,
253 )
204 - residual = output
254 else:
206 - x = tp.all_reduce(x)
255 + x = get_tp_group().all_reduce(x)
256 if stream is not None:
257 if layer_input is None:
258 + assert isinstance(x, torch.Tensor)
259 residual = mhc_post_tilelang(x, residual, post_layer_mix, comb_res_mix)
210 - aux = residual.mean(dim=1) if capture_aux else x.new_empty(0, x.shape[1])
260 + aux = (
261 + residual.mean(dim=1)
262 + if capture_aux
263 + else residual.new_empty(0, residual.shape[-1])
264 + )
265 pre_outputs = mhc_pre_delayed_overlap(
266 residual,
267 fn,
@@ -226,6 +280,7 @@ def mhc_shifted_post_pre(
280 )
281 return residual, *pre_outputs, aux
282
283 + assert isinstance(x, torch.Tensor)
284 if can_use_mega_mhc(x, residual, pre_mix, norm_weight, capture_aux):
285 assert pre_mix is not None and norm_weight is not None
286 outputs = mhc_shifted_post_pre_deep_gemm(