@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9a70c233cd

[Perf][DSV4.1] Fuse TP all-reduce with mHC input preparation (#57643)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Codex <noreply@openai.com>

Woosuk Kwon committed Sep 20, 2026 at 20:28 UTC 9a70c233cd5d3ef3b14d79c582a6f9ddfa03cf80
9 files changed +549 -105
CMakeLists.txt
+2
@@ -460,6 +460,8 @@ 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 +
465 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
466 cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
467 "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 new
+212
@@ -0,0 +1,212 @@
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
+111
@@ -11,6 +11,7 @@ import torch.distributed as dist
11 from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa
12 from vllm.distributed.device_communicators import custom_all_reduce as car
13 from vllm.distributed.parallel_state import get_tp_group, graph_capture
14 +from vllm.platforms import current_platform
15
16 from ..utils import (
17 ensure_model_parallel_initialized,
@@ -24,6 +25,116 @@ for i, v in enumerate(test_sizes):
25 test_sizes[i] -= v % 8
26
27
28 +@ray.remote(num_gpus=1, max_calls=1)
29 +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 + )
34 +
35 + with monkeypatch.context() as m:
36 + m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
37 + device = torch.device(f"cuda:{rank}")
38 + torch.accelerator.set_device_index(device)
39 + init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
40 + 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
55 + torch.manual_seed(123)
56 + residual = torch.randn(n, 4, 5120, device=device, dtype=torch.bfloat16)
57 + post = torch.rand(n, 4, device=device)
58 + comb = torch.randn(n, 4, 4, device=device) * 0.1
59 + pre = torch.rand(n, 4, device=device)
60 + weight = torch.randn(5120, device=device, dtype=torch.bfloat16)
61 + output = torch.empty_like(residual)
62 + normalized = torch.empty_like(x)
63 +
64 + def fused():
65 + torch.ops._C_custom_ar.all_reduce_mhc(
66 + x,
67 + residual,
68 + post,
69 + comb,
70 + pre,
71 + 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,
99 + 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
108 + )
109 +
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.
125 + for n in (1, 6, 12, 8, 16, 3, 5, 2, 4, 1):
126 + run(n)
127 +
128 +
129 +@pytest.mark.skipif(
130 + not current_platform.is_device_capability_family(100), reason="Requires SM100"
131 +)
132 +def test_all_reduce_mhc_preserves_bf16_boundaries_and_graph_replay(monkeypatch):
133 + if torch.accelerator.device_count() < 4:
134 + pytest.skip("Requires four GPUs with NVLink multicast")
135 + multi_process_parallel(monkeypatch, 4, 1, _all_reduce_mhc)
136 +
137 +
138 @pytest.mark.parametrize(
139 ("dtype", "expected"),
140 [
tests/distributed/test_engram_dp_shard.py
+1
@@ -512,6 +512,7 @@ def _check_dummy_hash_model_forward(
512 model = SimpleNamespace(
513 use_mega_moe=False,
514 use_sequence_parallel=False,
515 + fuse_mhc_all_reduce=False,
516 engram_hash=state,
517 engram_swa_prefix="swa",
518 engram_dp_shared_memory=dp_shared_memory,
tests/kernels/test_mhc_kernels.py
+45 -1
@@ -50,6 +50,40 @@ from vllm.utils.torch_utils import set_random_seed
50 DEVICE = current_platform.device_type
51
52
53 +@pytest.mark.parametrize(
54 + "tp,ep,hidden,hc,multicast,expected",
55 + [
56 + (4, False, 5120, 4, 1, True),
57 + (2, False, 5120, 4, 1, False),
58 + (4, True, 5120, 4, 1, False),
59 + (4, False, 4096, 4, 1, False),
60 + (4, False, 5120, 2, 1, False),
61 + (4, False, 5120, 4, 0, False),
62 + ],
63 +)
64 +def test_deepseek_v41_all_reduce_fusion_requires_kernel_support(
65 + monkeypatch, tp, ep, hidden, hc, multicast, expected
66 +):
67 + """Eligibility depends on kernel inputs, not serving or attention settings."""
68 + from vllm.models.deepseek_v41.nvidia.ops import mhc
69 +
70 + config = SimpleNamespace(
71 + parallel_config=SimpleNamespace(
72 + tensor_parallel_size=tp, enable_expert_parallel=ep
73 + ),
74 + model_config=SimpleNamespace(
75 + hf_config=SimpleNamespace(hidden_size=hidden, hc_mult=hc)
76 + ),
77 + )
78 + group = SimpleNamespace(
79 + device_communicator=SimpleNamespace(
80 + ca_comm=SimpleNamespace(mnnvl_lamport_ag_multicast_ptr=multicast)
81 + )
82 + )
83 + monkeypatch.setattr(mhc, "get_tp_group", lambda: group)
84 + assert mhc.supports_mhc_all_reduce(config) is expected
85 +
86 +
87 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
88 @pytest.mark.parametrize(
89 "num_tokens,hc_mult,hidden_size",
@@ -523,6 +557,7 @@ def test_deepseek_v41_decoder_mixes_match_torch(
557 decoder.rms_norm_eps = 1e-20
558 decoder.hc_post_alpha = 2.0
559 decoder.use_sequence_parallel = False
560 + decoder.fuse_mhc_all_reduce = False
561 decoder.mhc_stream = None
562 if mhc_mode != "disabled":
563 from vllm.utils.deep_gemm import is_deep_gemm_supported
@@ -597,7 +632,15 @@ def test_deepseek_v41_decoder_mixes_match_torch(
632 return post, res, decoder.attn_norm(collapsed), pre
633
634 def fused_reference(
600 - x, residual, post_mix, res_mix, *args, capture_aux=False, stream=None, **kw
635 + x,
636 + residual,
637 + post_mix,
638 + res_mix,
639 + *args,
640 + capture_aux=False,
641 + stream=None,
642 + reduce_results=False,
643 + **kw,
644 ):
645 residual = mhc_post_torch(x, residual, post_mix, res_mix)
646 aux = residual.mean(dim=1) if capture_aux else residual.new_empty(0)
@@ -638,6 +681,7 @@ def test_deepseek_v41_capture_previous_aux(entry, monkeypatch, default_vllm_conf
681 decoder.rms_norm_eps = 1e-6
682 decoder.hc_post_alpha = 2.0
683 decoder.use_sequence_parallel = False
684 + decoder.fuse_mhc_all_reduce = False
685 decoder.mhc_stream = None
686 decoder.engram = None
687 from vllm.model_executor.layers.layernorm import RMSNorm
vllm/models/deepseek_v4/nvidia/model.py
+3
@@ -793,6 +793,7 @@ class DeepseekV4MoE(nn.Module):
793 use_sequence_parallel: bool = False,
794 *,
795 num_hash_layers: int,
796 + reduce_results: bool = True,
797 n_routed_experts: int | None = None,
798 n_activated_experts: int | None = None,
799 image_sentinel_lo: int = IMAGE_SENTINEL_BASE_ID,
@@ -804,6 +805,7 @@ class DeepseekV4MoE(nn.Module):
805 quant_config = vllm_config.quant_config
806 self.prefix = prefix
807 self.use_sequence_parallel = use_sequence_parallel
808 + self.reduce_results = reduce_results
809 moe_backend = vllm_config.kernel_config.moe_backend
810 validate_fi_moe_ep_config(vllm_config)
811 self.use_mega_moe = moe_backend in MEGA_MOE_BACKENDS
@@ -1010,6 +1012,7 @@ class DeepseekV4MoE(nn.Module):
1012 self.physical_expert_end = self.experts_end_idx
1013
1014 self.experts = FusedMoEFactory(
1015 + reduce_results=self.reduce_results,
1016 shared_experts=self.shared_experts,
1017 gate=self.gate,
1018 num_experts=self.n_routed_experts,
vllm/models/deepseek_v41/nvidia/model.py
+19 -2
@@ -17,6 +17,7 @@ from vllm.distributed import (
17 get_pp_group,
18 get_tensor_model_parallel_rank,
19 get_tensor_model_parallel_world_size,
20 + tensor_model_parallel_all_reduce,
21 )
22 from vllm.forward_context import (
23 get_forward_context,
@@ -88,10 +89,11 @@ from vllm.v1.worker.ubatching import dbo_current_ubatch_id
89 from ..common.engram import EngramLayout, NgramHashState
90 from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask
91 from .engram import Engram, gather_engram_hashes
91 -from .ops.mega_mhc import mhc_shifted_post_pre
92 from .ops.mhc import (
93 MHC_OVERLAP_MAX_TOKENS,
94 mhc_pre_delayed_overlap,
95 + mhc_shifted_post_pre,
96 + supports_mhc_all_reduce,
97 supports_mhc_overlap,
98 )
99
@@ -107,6 +109,7 @@ class DeepseekV4MoE(DeepseekV4MoEBase):
109 vllm_config: VllmConfig,
110 prefix: str = "",
111 use_sequence_parallel: bool = False,
112 + reduce_results: bool = True,
113 ):
114 config = vllm_config.model_config.hf_config
115 n_routed_experts = config.n_routed_experts
@@ -125,6 +128,7 @@ class DeepseekV4MoE(DeepseekV4MoEBase):
128 n_routed_experts=n_routed_experts,
129 n_activated_experts=n_activated_experts,
130 num_hash_layers=0,
131 + reduce_results=reduce_results,
132 image_sentinel_lo=IMAGE_SENTINEL_BASE_ID,
133 )
134
@@ -199,12 +203,14 @@ class DeepseekV4DecoderLayer(nn.Module):
203 candidate_block_buffer: torch.Tensor | None = None,
204 engram_layout: EngramLayout | None = None,
205 mhc_stream: torch.cuda.Stream | None = None,
206 + fuse_mhc_all_reduce: bool = False,
207 ):
208 super().__init__()
209
210 config = vllm_config.model_config.hf_config
211 self.hidden_size = config.hidden_size
212 self.mhc_stream = mhc_stream
213 + self.fuse_mhc_all_reduce = fuse_mhc_all_reduce
214 self.use_sequence_parallel = _use_sequence_parallel(vllm_config)
215
216 self.engram: Engram | None = None
@@ -228,12 +234,13 @@ class DeepseekV4DecoderLayer(nn.Module):
234 aux_stream_list=aux_stream_list,
235 candidate_block_buffer=candidate_block_buffer,
236 )
231 - if self.use_sequence_parallel:
237 + if self.use_sequence_parallel or fuse_mhc_all_reduce:
238 self.attn.wo_b.reduce_results = False
239 self.ffn = DeepseekV4MoE(
240 vllm_config,
241 prefix=f"{prefix}.ffn",
242 use_sequence_parallel=self.use_sequence_parallel,
243 + reduce_results=not fuse_mhc_all_reduce,
244 )
245
246 self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
@@ -430,6 +437,8 @@ class DeepseekV4DecoderLayer(nn.Module):
437 # and this block's pre, on the full hc stream, so the mix
438 # coefficients see the injected stream. The injection also keeps
439 # the post out of the pre-norm GEMM's fused prologue.
440 + if self.fuse_mhc_all_reduce:
441 + x = tensor_model_parallel_all_reduce(x)
442 previous_post = mhc_post_tilelang(x, residual, post_mix, res_mix)
443 if capture_previous_aux:
444 previous_aux = previous_post.mean(dim=1)
@@ -473,6 +482,7 @@ class DeepseekV4DecoderLayer(nn.Module):
482 norm_eps=self.attn_norm.variance_epsilon,
483 capture_aux=capture_previous_aux,
484 stream=mhc_stream,
485 + reduce_results=self.fuse_mhc_all_reduce,
486 )
487 if capture_previous_aux:
488 previous_aux = aux
@@ -503,6 +513,7 @@ class DeepseekV4DecoderLayer(nn.Module):
513 norm_weight=self.ffn_norm.weight,
514 norm_eps=self.ffn_norm.variance_epsilon,
515 stream=mhc_stream,
516 + reduce_results=self.fuse_mhc_all_reduce,
517 )
518 x = self.ffn(x, input_ids, mega_gate_metadata)
519 if mhc_stream is not None:
@@ -540,6 +551,9 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
551 aux_stream_list = [torch.cuda.Stream() for _ in range(3)]
552 # Keep mHC independent of the streams used inside attention.
553 mhc_stream = torch.cuda.Stream() if supports_mhc_overlap(vllm_config) else None
554 + self.fuse_mhc_all_reduce = mhc_stream is not None and supports_mhc_all_reduce(
555 + vllm_config
556 + )
557
558 # Reserved topk indices buffer for all Indexer layers to reuse.
559 self.topk_indices_buffer = torch.empty(
@@ -585,6 +599,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
599 candidate_block_buffer=self.candidate_block_buffer,
600 engram_layout=self.engram_layout,
601 mhc_stream=mhc_stream,
602 + fuse_mhc_all_reduce=self.fuse_mhc_all_reduce,
603 ),
604 prefix=f"{prefix}.layers",
605 )
@@ -794,6 +809,8 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
809 aux_hidden_by_layer[idx] = previous_aux
810 if layer is not None:
811 # The last layer has no successor to fold its post into.
812 + if self.fuse_mhc_all_reduce:
813 + hidden_states = tensor_model_parallel_all_reduce(hidden_states)
814 hidden_states = mhc_post_tilelang(
815 hidden_states, residual, post_mix, res_mix
816 )
vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py
-98
@@ -5,10 +5,6 @@ import functools
5
6 import torch
7
8 -from vllm.model_executor.kernels.mhc.tilelang import (
9 - mhc_fused_post_pre_delayed_tilelang,
10 - mhc_post_tilelang,
11 -)
8 from vllm.platforms import current_platform
9 from vllm.utils.deep_gemm import (
10 _import_deep_gemm,
@@ -88,97 +84,3 @@ def mhc_shifted_post_pre_deep_gemm(
84 y_bf16,
85 new_prev_mix.squeeze(-1),
86 )
91 -
92 -
93 -def mhc_shifted_post_pre(
94 - x: torch.Tensor,
95 - residual: torch.Tensor,
96 - post_layer_mix: torch.Tensor,
97 - comb_res_mix: torch.Tensor,
98 - fn: torch.Tensor,
99 - hc_scale: torch.Tensor,
100 - hc_base: torch.Tensor,
101 - rms_eps: float,
102 - hc_pre_eps: float,
103 - hc_sinkhorn_eps: float,
104 - hc_post_mult_value: float,
105 - sinkhorn_repeat: int,
106 - pre_mix: torch.Tensor | None = None,
107 - norm_weight: torch.Tensor | None = None,
108 - norm_eps: float = 1e-6,
109 - capture_aux: bool = False,
110 - *,
111 - stream: torch.cuda.Stream | None = None,
112 -) -> tuple[
113 - torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
114 -]:
115 - """Dispatch shifted post/pre to overlap, Mega-mHC, or fused TileLang.
116 -
117 - When stream is supplied, join it before consuming the returned coefficients.
118 - """
119 - if stream is not None:
120 - from .mhc import mhc_pre_delayed_overlap
121 -
122 - residual = mhc_post_tilelang(x, residual, post_layer_mix, comb_res_mix)
123 - aux = residual.mean(dim=1) if capture_aux else x.new_empty(0, x.shape[1])
124 - pre_outputs = mhc_pre_delayed_overlap(
125 - residual,
126 - fn,
127 - hc_scale,
128 - hc_base,
129 - rms_eps,
130 - hc_pre_eps,
131 - hc_sinkhorn_eps,
132 - hc_post_mult_value,
133 - sinkhorn_repeat,
134 - pre_mix=pre_mix,
135 - norm_weight=norm_weight,
136 - norm_eps=norm_eps,
137 - stream=stream,
138 - )
139 - return residual, *pre_outputs, aux
140 -
141 - if (
142 - pre_mix is not None
143 - and norm_weight is not None
144 - and not capture_aux
145 - and x.shape[0] <= 1 << 20
146 - and is_mega_mhc_supported(x.shape[1], residual.shape[1])
147 - ):
148 - outputs = mhc_shifted_post_pre_deep_gemm(
149 - x,
150 - residual,
151 - pre_mix,
152 - post_layer_mix,
153 - comb_res_mix,
154 - fn,
155 - hc_scale,
156 - hc_base,
157 - rms_eps,
158 - hc_pre_eps,
159 - hc_post_mult_value,
160 - hc_sinkhorn_eps,
161 - sinkhorn_repeat,
162 - norm_weight,
163 - norm_eps,
164 - )
165 - return *outputs, x.new_empty(0, x.shape[1])
166 -
167 - return mhc_fused_post_pre_delayed_tilelang(
168 - x,
169 - residual,
170 - post_layer_mix,
171 - comb_res_mix,
172 - fn,
173 - hc_scale,
174 - hc_base,
175 - rms_eps,
176 - hc_pre_eps,
177 - hc_sinkhorn_eps,
178 - hc_post_mult_value,
179 - sinkhorn_repeat,
180 - pre_mix=pre_mix,
181 - norm_weight=norm_weight,
182 - norm_eps=norm_eps,
183 - capture_aux=capture_aux,
184 - )
vllm/models/deepseek_v41/nvidia/ops/mhc.py
+156 -4
@@ -1,15 +1,26 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Split shifted mHC pre so coefficient generation can overlap the sublayer."""
3 +"""Dispatch DSV4.1 mHC operations and overlap coefficient generation."""
4
5 from functools import partial
6 +from typing import TYPE_CHECKING, cast
7
8 import torch
9
10 from vllm.config import VllmConfig
11 +from vllm.distributed import get_tp_group
12 +from vllm.model_executor.kernels.mhc.tilelang import (
13 + mhc_fused_post_pre_delayed_tilelang,
14 + mhc_post_tilelang,
15 +)
16 from vllm.platforms import current_platform
17 from vllm.utils.deep_gemm import is_deep_gemm_supported
18
19 +from .mega_mhc import is_mega_mhc_supported, mhc_shifted_post_pre_deep_gemm
20 +
21 +if TYPE_CHECKING:
22 + from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator
23 +
24 # GB200 TP4/FlashInfer improves through 16 tokens; larger screens tie or regress.
25 MHC_OVERLAP_MAX_TOKENS = 16
26
@@ -29,6 +40,20 @@ def supports_mhc_overlap(vllm_config: VllmConfig) -> bool:
40 )
41
42
43 +def supports_mhc_all_reduce(vllm_config: VllmConfig) -> bool:
44 + parallel = vllm_config.parallel_config
45 + config = vllm_config.model_config.hf_config
46 + if (
47 + parallel.tensor_parallel_size != 4
48 + or parallel.enable_expert_parallel
49 + or config.hidden_size != 5120
50 + or config.hc_mult != 4
51 + ):
52 + return False
53 + comm = cast("CudaCommunicator", get_tp_group().device_communicator).ca_comm
54 + return comm is not None and bool(comm.mnnvl_lamport_ag_multicast_ptr)
55 +
56 +
57 def mhc_pre_delayed_overlap(
58 residual: torch.Tensor,
59 fn: torch.Tensor,
@@ -45,6 +70,7 @@ def mhc_pre_delayed_overlap(
70 norm_eps: float = 1e-6,
71 *,
72 stream: torch.cuda.Stream,
73 + layer_input: torch.Tensor | None = None,
74 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
75 """Prepare the input on the caller stream and coefficients on another stream.
76
@@ -67,7 +93,11 @@ def mhc_pre_delayed_overlap(
93 post = torch.empty((n, hc), device=residual.device, dtype=torch.float32)
94 comb = torch.empty((n, hc * hc), device=residual.device, dtype=torch.float32)
95 next_pre = torch.empty_like(post)
70 - layer_input = torch.empty((n, hidden), device=residual.device, dtype=torch.bfloat16)
96 + compute_input = layer_input is None
97 + if layer_input is None:
98 + layer_input = torch.empty(
99 + (n, hidden), device=residual.device, dtype=torch.bfloat16
100 + )
101 outputs = post.unsqueeze(-1), comb.view(n, hc, hc), layer_input, next_pre
102 if n == 0:
103 return outputs
@@ -104,7 +134,7 @@ def mhc_pre_delayed_overlap(
134 )
135 main = torch.cuda.current_stream()
136 # Prioritize input readiness for small batches before releasing statistics.
107 - if n <= 8:
137 + if compute_input and n <= 8:
138 epilogue(split_mode="input")
139 stream.wait_stream(main)
140 with torch.cuda.stream(stream):
@@ -112,6 +142,128 @@ def mhc_pre_delayed_overlap(
142 epilogue(split_mode="stats")
143 for tensor in (x, mix, sqr):
144 tensor.record_stream(stream)
115 - if n > 8:
145 + if compute_input and n > 8:
146 epilogue(split_mode="input")
147 return outputs
148 +
149 +
150 +def mhc_shifted_post_pre(
151 + x: torch.Tensor,
152 + residual: torch.Tensor,
153 + post_layer_mix: torch.Tensor,
154 + comb_res_mix: torch.Tensor,
155 + fn: torch.Tensor,
156 + hc_scale: torch.Tensor,
157 + hc_base: torch.Tensor,
158 + rms_eps: float,
159 + hc_pre_eps: float,
160 + hc_sinkhorn_eps: float,
161 + hc_post_mult_value: float,
162 + sinkhorn_repeat: int,
163 + pre_mix: torch.Tensor | None = None,
164 + norm_weight: torch.Tensor | None = None,
165 + norm_eps: float = 1e-6,
166 + capture_aux: bool = False,
167 + *,
168 + stream: torch.cuda.Stream | None = None,
169 + reduce_results: bool = False,
170 +) -> tuple[
171 + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
172 +]:
173 + """Dispatch shifted post/pre to overlap, Mega-mHC, or fused TileLang.
174 +
175 + When stream is supplied, join it before consuming the returned coefficients.
176 + """
177 + layer_input = None
178 + if reduce_results:
179 + tp = get_tp_group()
180 + if stream is not None and 0 < x.shape[0] <= MHC_OVERLAP_MAX_TOKENS:
181 + comm = cast("CudaCommunicator", tp.device_communicator).ca_comm
182 + assert comm is not None and comm.mnnvl_lamport_epochs is not None
183 + output = torch.empty_like(residual)
184 + layer_input = torch.empty_like(x)
185 + torch.ops._C_custom_ar.all_reduce_mhc(
186 + x,
187 + residual,
188 + post_layer_mix,
189 + comb_res_mix,
190 + pre_mix,
191 + norm_weight,
192 + output,
193 + layer_input,
194 + comm.mnnvl_lamport_ag_local_ptr,
195 + comm.mnnvl_lamport_ag_multicast_ptr,
196 + comm.mnnvl_lamport_epochs[0],
197 + comm.rank,
198 + comm.mnnvl_buffer_size,
199 + norm_eps,
200 + )
201 + residual = output
202 + else:
203 + x = tp.all_reduce(x)
204 + if stream is not None:
205 + if layer_input is None:
206 + residual = mhc_post_tilelang(x, residual, post_layer_mix, comb_res_mix)
207 + aux = residual.mean(dim=1) if capture_aux else x.new_empty(0, x.shape[1])
208 + pre_outputs = mhc_pre_delayed_overlap(
209 + residual,
210 + fn,
211 + hc_scale,
212 + hc_base,
213 + rms_eps,
214 + hc_pre_eps,
215 + hc_sinkhorn_eps,
216 + hc_post_mult_value,
217 + sinkhorn_repeat,
218 + pre_mix=pre_mix,
219 + norm_weight=norm_weight,
220 + norm_eps=norm_eps,
221 + stream=stream,
222 + layer_input=layer_input,
223 + )
224 + return residual, *pre_outputs, aux
225 +
226 + if (
227 + pre_mix is not None
228 + and norm_weight is not None
229 + and not capture_aux
230 + and x.shape[0] <= 1 << 20
231 + and is_mega_mhc_supported(x.shape[1], residual.shape[1])
232 + ):
233 + outputs = mhc_shifted_post_pre_deep_gemm(
234 + x,
235 + residual,
236 + pre_mix,
237 + post_layer_mix,
238 + comb_res_mix,
239 + fn,
240 + hc_scale,
241 + hc_base,
242 + rms_eps,
243 + hc_pre_eps,
244 + hc_post_mult_value,
245 + hc_sinkhorn_eps,
246 + sinkhorn_repeat,
247 + norm_weight,
248 + norm_eps,
249 + )
250 + return *outputs, x.new_empty(0, x.shape[1])
251 +
252 + return mhc_fused_post_pre_delayed_tilelang(
253 + x,
254 + residual,
255 + post_layer_mix,
256 + comb_res_mix,
257 + fn,
258 + hc_scale,
259 + hc_base,
260 + rms_eps,
261 + hc_pre_eps,
262 + hc_sinkhorn_eps,
263 + hc_post_mult_value,
264 + sinkhorn_repeat,
265 + pre_mix=pre_mix,
266 + norm_weight=norm_weight,
267 + norm_eps=norm_eps,
268 + capture_aux=capture_aux,
269 + )