From 9a70c233cd5d3ef3b14d79c582a6f9ddfa03cf80 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Sun, 20 Sep 2026 20:28:06 -0700 Subject: [PATCH] [Perf][DSV4.1] Fuse TP all-reduce with mHC input preparation (#57643) Signed-off-by: Woosuk Kwon Co-authored-by: Codex --- CMakeLists.txt | 2 + csrc/libtorch_stable/all_reduce_mhc.cu | 212 ++++++++++++++++++ tests/distributed/test_custom_all_reduce.py | 111 +++++++++ tests/distributed/test_engram_dp_shard.py | 1 + tests/kernels/test_mhc_kernels.py | 46 +++- vllm/models/deepseek_v4/nvidia/model.py | 3 + vllm/models/deepseek_v41/nvidia/model.py | 21 +- .../deepseek_v41/nvidia/ops/mega_mhc.py | 98 -------- vllm/models/deepseek_v41/nvidia/ops/mhc.py | 160 ++++++++++++- 9 files changed, 549 insertions(+), 105 deletions(-) create mode 100644 csrc/libtorch_stable/all_reduce_mhc.cu diff --git a/CMakeLists.txt b/CMakeLists.txt index fbacfd9091..dbf4ca90e2 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -460,6 +460,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") DEFINED CMAKE_CUDA_COMPILER_VERSION AND CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.0) + list(APPEND VLLM_STABLE_EXT_SRC "csrc/libtorch_stable/all_reduce_mhc.cu") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS "9.0a;10.0f;10.1f;10.3f;10.7f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}") diff --git a/csrc/libtorch_stable/all_reduce_mhc.cu b/csrc/libtorch_stable/all_reduce_mhc.cu new file mode 100644 index 0000000000..dc58371156 --- /dev/null +++ b/csrc/libtorch_stable/all_reduce_mhc.cu @@ -0,0 +1,212 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#include "torch_utils.h" +#include "custom_all_gather_reduce_scatter.cuh" +#include +#include + +namespace { +using Tensor = torch::stable::Tensor; +using Pack = vllm::packed_t::P; +constexpr int kHidden = 5120; +constexpr int kRanks = 4; +constexpr int kThreads = 128; +constexpr int kCluster = 5; +constexpr int kPacks = kHidden / Pack::size; + +__global__ __launch_bounds__(kThreads) + __cluster_dims__(1, kCluster, 1) void all_reduce_mhc_kernel( + const Pack* input, const Pack* residual, const float* post, + const float* comb, const float* pre, const Pack* weight, Pack* output, + Pack* normalized, Pack* local, Pack* multicast, uint32_t* epochs, + int rank, int stage_size, float eps) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaGridDependencySynchronize(); + const int token = blockIdx.x; + const int column = blockIdx.y * kThreads + threadIdx.x; + const int idx = token * kPacks + column; + const int size = gridDim.x * kPacks; + const int stage = epochs[0] % 3; + const int dirty = vllm::mnnvl_lamport_dirty_stage(stage); + const int dirty_size = epochs[2 + dirty]; + Pack value = input[idx]; + vllm::store_multimem_lamport_payload( + multicast + stage * stage_size + rank * size + idx, + vllm::sanitize_lamport_payload(value)); + // Multi-token forwards benefit from launching dependents before mHC finishes. + if (gridDim.x > 1) cudaTriggerProgrammaticLaunchCompletion(); + vllm::lamport_cta_arrive(epochs + 1); + for (int i = idx; i < dirty_size; i += size) { + local[dirty * stage_size + i] = vllm::lamport_sentinel(); + } + Pack peers[kRanks]; + bool pending; + do { + vllm::LamportPack packets[kRanks]; + #pragma unroll + for (int peer = 0; peer < kRanks; ++peer) { + auto ptr = local + stage * stage_size + idx + peer * size; + asm volatile("ld.relaxed.sys.global.v4.u32 {%0,%1,%2,%3}, [%4];" + : "=r"(packets[peer].words[0]), "=r"(packets[peer].words[1]), + "=r"(packets[peer].words[2]), "=r"(packets[peer].words[3]) + : "l"(ptr) + : "memory"); + } + pending = false; + #pragma unroll + for (int peer = 0; peer < kRanks; ++peer) { + pending |= vllm::is_lamport_dirty(packets[peer]); + peers[peer] = packets[peer].packed; + } + } while (pending); + auto sum = vllm::upcast(peers[0]); + #pragma unroll + for (int r = 1; r < kRanks; ++r) { + vllm::packed_assign_add(sum, vllm::upcast(peers[r])); + } + // Keep the BF16 boundary between the collective and mHC post. + auto reduced = vllm::upcast(vllm::downcast(sum)); + decltype(sum) streams[4]; + #pragma unroll + for (int r = 0; r < 4; ++r) { + streams[r] = vllm::upcast(residual[(token * 4 + r) * kPacks + column]); + } + decltype(sum) collapsed{}; + #pragma unroll + for (int c = 0; c < 4; ++c) { + decltype(sum) mixed; + #pragma unroll + for (int e = 0; e < Pack::size; ++e) { + // Match post*x followed by the four residual FMAs. + float acc = __fmul_rn(post[token * 4 + c], reduced.data[e]); + #pragma unroll + for (int r = 0; r < 4; ++r) { + acc = fmaf(comb[token * 16 + r * 4 + c], streams[r].data[e], acc); + } + mixed.data[e] = acc; + } + Pack rounded = vllm::downcast(mixed); + output[(token * 4 + c) * kPacks + column] = rounded; + auto as_float = vllm::upcast(rounded); + #pragma unroll + for (int e = 0; e < Pack::size; ++e) { + collapsed.data[e] = + fmaf(pre[token * 4 + c], as_float.data[e], collapsed.data[e]); + } + } + // RMSNorm consumes the rounded BF16 collapse in the unfused path. + collapsed = vllm::upcast(vllm::downcast(collapsed)); + float sqr = 0; + #pragma unroll + for (int e = 0; e < Pack::size; ++e) { + sqr = fmaf(collapsed.data[e], collapsed.data[e], sqr); + } + #pragma unroll + for (int delta = 16; delta; delta /= 2) { + sqr += __shfl_xor_sync(0xffffffff, sqr, delta); + } + __shared__ float warp_sums[kCluster][kThreads / 32]; + auto cluster = cooperative_groups::this_cluster(); + const int lane = threadIdx.x % 32; + if (lane < kCluster) { + *cluster.map_shared_rank(&warp_sums[blockIdx.y][threadIdx.x / 32], lane) = + sqr; + } + cluster.sync(); + float total = 0; + #pragma unroll + for (int c = 0; c < kCluster; ++c) { + #pragma unroll + for (int w = 0; w < kThreads / 32; ++w) total += warp_sums[c][w]; + } + const float scale = rsqrtf(__fadd_rn(total / kHidden, eps)); + auto weights = vllm::upcast(weight[column]); + #pragma unroll + for (int e = 0; e < Pack::size; ++e) { + collapsed.data[e] = collapsed.data[e] * scale * weights.data[e]; + } + normalized[idx] = vllm::downcast(collapsed); + if (idx == 0) { + while (*reinterpret_cast(epochs + 1) < + gridDim.x * gridDim.y) { + } + epochs[2 + stage] = kRanks * size; + epochs[0] = vllm::mnnvl_lamport_next_stage(stage); + epochs[1] = 0; + } + if (gridDim.x == 1) cudaTriggerProgrammaticLaunchCompletion(); +#else + asm volatile("trap;"); +#endif +} + +void all_reduce_mhc(Tensor input, Tensor residual, Tensor post, Tensor comb, + Tensor pre, Tensor weight, Tensor output, Tensor normalized, + int64_t local, int64_t multicast, Tensor epochs, + int64_t rank, int64_t stage_bytes, double eps) { + using torch::headeronly::ScalarType; + const torch::stable::accelerator::DeviceGuard guard(input.get_device_index()); + const int64_t n = input.size(0); + STD_TORCH_CHECK(input.dim() == 2 && input.size(1) == kHidden && n > 0 && + n <= 16); + for (const auto& t : {input, residual, weight, output, normalized}) { + STD_TORCH_CHECK(t.is_contiguous() && + t.scalar_type() == ScalarType::BFloat16); + STD_TORCH_CHECK(t.get_device_index() == input.get_device_index()); + STD_TORCH_CHECK(reinterpret_cast(t.const_data_ptr()) % 16 == 0); + } + for (const auto& t : {post, comb, pre}) { + STD_TORCH_CHECK(t.is_contiguous() && t.scalar_type() == ScalarType::Float); + STD_TORCH_CHECK(t.get_device_index() == input.get_device_index()); + } + STD_TORCH_CHECK(residual.numel() == n * 4 * kHidden && + output.numel() == residual.numel() && + normalized.numel() == input.numel() && + weight.numel() == kHidden); + STD_TORCH_CHECK(post.numel() == n * 4 && pre.numel() == n * 4 && + comb.numel() == n * 16); + STD_TORCH_CHECK(epochs.is_contiguous() && + epochs.scalar_type() == ScalarType::Int && + epochs.numel() >= 5 && + epochs.get_device_index() == input.get_device_index()); + STD_TORCH_CHECK(local && multicast && rank >= 0 && rank < kRanks && + stage_bytes % 16 == 0 && + stage_bytes >= n * kHidden * 2 * kRanks); + STD_TORCH_CHECK(get_device_prop()->major == 10); + cudaLaunchAttribute attr{}; + attr.id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr.val.programmaticStreamSerializationAllowed = 1; + cudaLaunchConfig_t config{ + .gridDim = dim3(n, kCluster), + .blockDim = dim3(kThreads), + .dynamicSmemBytes = 0, + .stream = get_current_cuda_stream(input.get_device_index()), + .attrs = &attr, + .numAttrs = 1}; + STD_CUDA_CHECK(cudaLaunchKernelEx( + &config, all_reduce_mhc_kernel, + static_cast(input.const_data_ptr()), + static_cast(residual.const_data_ptr()), + static_cast(post.const_data_ptr()), + static_cast(comb.const_data_ptr()), + static_cast(pre.const_data_ptr()), + static_cast(weight.const_data_ptr()), + static_cast(output.mutable_data_ptr()), + static_cast(normalized.mutable_data_ptr()), + reinterpret_cast(local), reinterpret_cast(multicast), + static_cast(epochs.mutable_data_ptr()), static_cast(rank), + static_cast(stage_bytes / sizeof(Pack)), static_cast(eps))); +} +} // namespace + +STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, mhc) { + mhc.def( + "all_reduce_mhc(Tensor input, Tensor residual, Tensor post, Tensor comb, " + "Tensor pre, Tensor weight, Tensor! output, Tensor! normalized, " + "int local, int multicast, Tensor! epochs, int rank, int stage_bytes, " + "float eps) -> ()"); +} +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, mhc) { + mhc.impl("all_reduce_mhc", TORCH_BOX(&all_reduce_mhc)); +} diff --git a/tests/distributed/test_custom_all_reduce.py b/tests/distributed/test_custom_all_reduce.py index 7f91f11d7c..7119e596a0 100644 --- a/tests/distributed/test_custom_all_reduce.py +++ b/tests/distributed/test_custom_all_reduce.py @@ -11,6 +11,7 @@ import torch.distributed as dist from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa from vllm.distributed.device_communicators import custom_all_reduce as car from vllm.distributed.parallel_state import get_tp_group, graph_capture +from vllm.platforms import current_platform from ..utils import ( ensure_model_parallel_initialized, @@ -24,6 +25,116 @@ for i, v in enumerate(test_sizes): test_sizes[i] -= v % 8 +@ray.remote(num_gpus=1, max_calls=1) +def _all_reduce_mhc(monkeypatch, tp_size, pp_size, rank, distributed_init_port): + from vllm.model_executor.kernels.mhc.tilelang import ( + mhc_post_tilelang, + mhc_pre_delayed_tilelang, + ) + + with monkeypatch.context() as m: + m.delenv("CUDA_VISIBLE_DEVICES", raising=False) + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) + ensure_model_parallel_initialized(tp_size, pp_size) + comm = get_tp_group().device_communicator.ca_comm + assert comm is not None and comm.mnnvl_lamport_ag_multicast_ptr + fn = torch.zeros(24, 20480, device=device) + scale = torch.ones(3, device=device) + base = torch.zeros(24, device=device) + + # Changing shapes and interleaving all-gather exercise the shared + # Lamport stages, including cleanup of a larger previous payload. + def run(n): + torch.manual_seed(42 + rank) + x = torch.randn(n, 5120, device=device, dtype=torch.bfloat16) + # Packed +0/-0 pairs collide with the Lamport sentinel. + x[:, :16] = 0 + x[:, 9:16:2] = -0.0 + torch.manual_seed(123) + residual = torch.randn(n, 4, 5120, device=device, dtype=torch.bfloat16) + post = torch.rand(n, 4, device=device) + comb = torch.randn(n, 4, 4, device=device) * 0.1 + pre = torch.rand(n, 4, device=device) + weight = torch.randn(5120, device=device, dtype=torch.bfloat16) + output = torch.empty_like(residual) + normalized = torch.empty_like(x) + + def fused(): + torch.ops._C_custom_ar.all_reduce_mhc( + x, + residual, + post, + comb, + pre, + weight, + output, + normalized, + comm.mnnvl_lamport_ag_local_ptr, + comm.mnnvl_lamport_ag_multicast_ptr, + comm.mnnvl_lamport_epochs[0], + rank, + comm.mnnvl_buffer_size, + 1e-6, + ) + + def check(): + gathered = comm.custom_all_gather(x) + assert gathered is not None + peers = gathered.view(tp_size, n, 5120).float() + reduced = peers[0].clone() + for peer in peers[1:]: + reduced.add_(peer) + expected = mhc_post_tilelang( + reduced.bfloat16(), residual, post.unsqueeze(-1), comb + ) + expected_norm = mhc_pre_delayed_tilelang( + expected, + fn, + scale, + base, + 1e-6, + 1e-6, + 1e-6, + 2.0, + 20, + pre_mix=pre, + norm_weight=weight, + )[2] + torch.testing.assert_close(output, expected, rtol=0.008, atol=1e-6) + torch.testing.assert_close( + normalized, expected_norm, rtol=0.008, atol=0.008 + ) + + fused() + check() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(6): + fused() + for _ in range(20): + graph.replay() + check() + x.mul_(0.5) + residual.mul_(2) + graph.replay() + check() + + # Cover fixed Q6 and shrinking/growing adaptive verification batches. + for n in (1, 6, 12, 8, 16, 3, 5, 2, 4, 1): + run(n) + + +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), reason="Requires SM100" +) +def test_all_reduce_mhc_preserves_bf16_boundaries_and_graph_replay(monkeypatch): + if torch.accelerator.device_count() < 4: + pytest.skip("Requires four GPUs with NVLink multicast") + multi_process_parallel(monkeypatch, 4, 1, _all_reduce_mhc) + + @pytest.mark.parametrize( ("dtype", "expected"), [ diff --git a/tests/distributed/test_engram_dp_shard.py b/tests/distributed/test_engram_dp_shard.py index c2fa0f857f..dd7105a23b 100644 --- a/tests/distributed/test_engram_dp_shard.py +++ b/tests/distributed/test_engram_dp_shard.py @@ -512,6 +512,7 @@ def _check_dummy_hash_model_forward( model = SimpleNamespace( use_mega_moe=False, use_sequence_parallel=False, + fuse_mhc_all_reduce=False, engram_hash=state, engram_swa_prefix="swa", engram_dp_shared_memory=dp_shared_memory, diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 8f7df55602..63648ebf2d 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -50,6 +50,40 @@ from vllm.utils.torch_utils import set_random_seed DEVICE = current_platform.device_type +@pytest.mark.parametrize( + "tp,ep,hidden,hc,multicast,expected", + [ + (4, False, 5120, 4, 1, True), + (2, False, 5120, 4, 1, False), + (4, True, 5120, 4, 1, False), + (4, False, 4096, 4, 1, False), + (4, False, 5120, 2, 1, False), + (4, False, 5120, 4, 0, False), + ], +) +def test_deepseek_v41_all_reduce_fusion_requires_kernel_support( + monkeypatch, tp, ep, hidden, hc, multicast, expected +): + """Eligibility depends on kernel inputs, not serving or attention settings.""" + from vllm.models.deepseek_v41.nvidia.ops import mhc + + config = SimpleNamespace( + parallel_config=SimpleNamespace( + tensor_parallel_size=tp, enable_expert_parallel=ep + ), + model_config=SimpleNamespace( + hf_config=SimpleNamespace(hidden_size=hidden, hc_mult=hc) + ), + ) + group = SimpleNamespace( + device_communicator=SimpleNamespace( + ca_comm=SimpleNamespace(mnnvl_lamport_ag_multicast_ptr=multicast) + ) + ) + monkeypatch.setattr(mhc, "get_tp_group", lambda: group) + assert mhc.supports_mhc_all_reduce(config) is expected + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") @pytest.mark.parametrize( "num_tokens,hc_mult,hidden_size", @@ -523,6 +557,7 @@ def test_deepseek_v41_decoder_mixes_match_torch( decoder.rms_norm_eps = 1e-20 decoder.hc_post_alpha = 2.0 decoder.use_sequence_parallel = False + decoder.fuse_mhc_all_reduce = False decoder.mhc_stream = None if mhc_mode != "disabled": from vllm.utils.deep_gemm import is_deep_gemm_supported @@ -597,7 +632,15 @@ def test_deepseek_v41_decoder_mixes_match_torch( return post, res, decoder.attn_norm(collapsed), pre def fused_reference( - x, residual, post_mix, res_mix, *args, capture_aux=False, stream=None, **kw + x, + residual, + post_mix, + res_mix, + *args, + capture_aux=False, + stream=None, + reduce_results=False, + **kw, ): residual = mhc_post_torch(x, residual, post_mix, res_mix) 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 decoder.rms_norm_eps = 1e-6 decoder.hc_post_alpha = 2.0 decoder.use_sequence_parallel = False + decoder.fuse_mhc_all_reduce = False decoder.mhc_stream = None decoder.engram = None from vllm.model_executor.layers.layernorm import RMSNorm diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 39d3c87bc6..7edff8ae02 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -793,6 +793,7 @@ class DeepseekV4MoE(nn.Module): use_sequence_parallel: bool = False, *, num_hash_layers: int, + reduce_results: bool = True, n_routed_experts: int | None = None, n_activated_experts: int | None = None, image_sentinel_lo: int = IMAGE_SENTINEL_BASE_ID, @@ -804,6 +805,7 @@ class DeepseekV4MoE(nn.Module): quant_config = vllm_config.quant_config self.prefix = prefix self.use_sequence_parallel = use_sequence_parallel + self.reduce_results = reduce_results moe_backend = vllm_config.kernel_config.moe_backend validate_fi_moe_ep_config(vllm_config) self.use_mega_moe = moe_backend in MEGA_MOE_BACKENDS @@ -1010,6 +1012,7 @@ class DeepseekV4MoE(nn.Module): self.physical_expert_end = self.experts_end_idx self.experts = FusedMoEFactory( + reduce_results=self.reduce_results, shared_experts=self.shared_experts, gate=self.gate, num_experts=self.n_routed_experts, diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index b5caceab6f..d7c9ddc128 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -17,6 +17,7 @@ from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, ) from vllm.forward_context import ( get_forward_context, @@ -88,10 +89,11 @@ from vllm.v1.worker.ubatching import dbo_current_ubatch_id from ..common.engram import EngramLayout, NgramHashState from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask from .engram import Engram, gather_engram_hashes -from .ops.mega_mhc import mhc_shifted_post_pre from .ops.mhc import ( MHC_OVERLAP_MAX_TOKENS, mhc_pre_delayed_overlap, + mhc_shifted_post_pre, + supports_mhc_all_reduce, supports_mhc_overlap, ) @@ -107,6 +109,7 @@ class DeepseekV4MoE(DeepseekV4MoEBase): vllm_config: VllmConfig, prefix: str = "", use_sequence_parallel: bool = False, + reduce_results: bool = True, ): config = vllm_config.model_config.hf_config n_routed_experts = config.n_routed_experts @@ -125,6 +128,7 @@ class DeepseekV4MoE(DeepseekV4MoEBase): n_routed_experts=n_routed_experts, n_activated_experts=n_activated_experts, num_hash_layers=0, + reduce_results=reduce_results, image_sentinel_lo=IMAGE_SENTINEL_BASE_ID, ) @@ -199,12 +203,14 @@ class DeepseekV4DecoderLayer(nn.Module): candidate_block_buffer: torch.Tensor | None = None, engram_layout: EngramLayout | None = None, mhc_stream: torch.cuda.Stream | None = None, + fuse_mhc_all_reduce: bool = False, ): super().__init__() config = vllm_config.model_config.hf_config self.hidden_size = config.hidden_size self.mhc_stream = mhc_stream + self.fuse_mhc_all_reduce = fuse_mhc_all_reduce self.use_sequence_parallel = _use_sequence_parallel(vllm_config) self.engram: Engram | None = None @@ -228,12 +234,13 @@ class DeepseekV4DecoderLayer(nn.Module): aux_stream_list=aux_stream_list, candidate_block_buffer=candidate_block_buffer, ) - if self.use_sequence_parallel: + if self.use_sequence_parallel or fuse_mhc_all_reduce: self.attn.wo_b.reduce_results = False self.ffn = DeepseekV4MoE( vllm_config, prefix=f"{prefix}.ffn", use_sequence_parallel=self.use_sequence_parallel, + reduce_results=not fuse_mhc_all_reduce, ) self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) @@ -430,6 +437,8 @@ class DeepseekV4DecoderLayer(nn.Module): # and this block's pre, on the full hc stream, so the mix # coefficients see the injected stream. The injection also keeps # the post out of the pre-norm GEMM's fused prologue. + if self.fuse_mhc_all_reduce: + x = tensor_model_parallel_all_reduce(x) previous_post = mhc_post_tilelang(x, residual, post_mix, res_mix) if capture_previous_aux: previous_aux = previous_post.mean(dim=1) @@ -473,6 +482,7 @@ class DeepseekV4DecoderLayer(nn.Module): norm_eps=self.attn_norm.variance_epsilon, capture_aux=capture_previous_aux, stream=mhc_stream, + reduce_results=self.fuse_mhc_all_reduce, ) if capture_previous_aux: previous_aux = aux @@ -503,6 +513,7 @@ class DeepseekV4DecoderLayer(nn.Module): norm_weight=self.ffn_norm.weight, norm_eps=self.ffn_norm.variance_epsilon, stream=mhc_stream, + reduce_results=self.fuse_mhc_all_reduce, ) x = self.ffn(x, input_ids, mega_gate_metadata) if mhc_stream is not None: @@ -540,6 +551,9 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): aux_stream_list = [torch.cuda.Stream() for _ in range(3)] # Keep mHC independent of the streams used inside attention. mhc_stream = torch.cuda.Stream() if supports_mhc_overlap(vllm_config) else None + self.fuse_mhc_all_reduce = mhc_stream is not None and supports_mhc_all_reduce( + vllm_config + ) # Reserved topk indices buffer for all Indexer layers to reuse. self.topk_indices_buffer = torch.empty( @@ -585,6 +599,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): candidate_block_buffer=self.candidate_block_buffer, engram_layout=self.engram_layout, mhc_stream=mhc_stream, + fuse_mhc_all_reduce=self.fuse_mhc_all_reduce, ), prefix=f"{prefix}.layers", ) @@ -794,6 +809,8 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): aux_hidden_by_layer[idx] = previous_aux if layer is not None: # The last layer has no successor to fold its post into. + if self.fuse_mhc_all_reduce: + hidden_states = tensor_model_parallel_all_reduce(hidden_states) hidden_states = mhc_post_tilelang( hidden_states, residual, post_mix, res_mix ) diff --git a/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py b/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py index caae3d8d9b..ef56b5a823 100644 --- a/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py +++ b/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py @@ -5,10 +5,6 @@ import functools import torch -from vllm.model_executor.kernels.mhc.tilelang import ( - mhc_fused_post_pre_delayed_tilelang, - mhc_post_tilelang, -) from vllm.platforms import current_platform from vllm.utils.deep_gemm import ( _import_deep_gemm, @@ -88,97 +84,3 @@ def mhc_shifted_post_pre_deep_gemm( y_bf16, new_prev_mix.squeeze(-1), ) - - -def mhc_shifted_post_pre( - x: torch.Tensor, - residual: torch.Tensor, - post_layer_mix: torch.Tensor, - comb_res_mix: torch.Tensor, - fn: torch.Tensor, - hc_scale: torch.Tensor, - hc_base: torch.Tensor, - rms_eps: float, - hc_pre_eps: float, - hc_sinkhorn_eps: float, - hc_post_mult_value: float, - sinkhorn_repeat: int, - pre_mix: torch.Tensor | None = None, - norm_weight: torch.Tensor | None = None, - norm_eps: float = 1e-6, - capture_aux: bool = False, - *, - stream: torch.cuda.Stream | None = None, -) -> tuple[ - torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor -]: - """Dispatch shifted post/pre to overlap, Mega-mHC, or fused TileLang. - - When stream is supplied, join it before consuming the returned coefficients. - """ - if stream is not None: - from .mhc import mhc_pre_delayed_overlap - - residual = mhc_post_tilelang(x, residual, post_layer_mix, comb_res_mix) - aux = residual.mean(dim=1) if capture_aux else x.new_empty(0, x.shape[1]) - pre_outputs = mhc_pre_delayed_overlap( - residual, - fn, - hc_scale, - hc_base, - rms_eps, - hc_pre_eps, - hc_sinkhorn_eps, - hc_post_mult_value, - sinkhorn_repeat, - pre_mix=pre_mix, - norm_weight=norm_weight, - norm_eps=norm_eps, - stream=stream, - ) - return residual, *pre_outputs, aux - - if ( - pre_mix is not None - and norm_weight is not None - and not capture_aux - and x.shape[0] <= 1 << 20 - and is_mega_mhc_supported(x.shape[1], residual.shape[1]) - ): - outputs = mhc_shifted_post_pre_deep_gemm( - x, - residual, - pre_mix, - post_layer_mix, - comb_res_mix, - fn, - hc_scale, - hc_base, - rms_eps, - hc_pre_eps, - hc_post_mult_value, - hc_sinkhorn_eps, - sinkhorn_repeat, - norm_weight, - norm_eps, - ) - return *outputs, x.new_empty(0, x.shape[1]) - - return mhc_fused_post_pre_delayed_tilelang( - x, - residual, - post_layer_mix, - comb_res_mix, - fn, - hc_scale, - hc_base, - rms_eps, - hc_pre_eps, - hc_sinkhorn_eps, - hc_post_mult_value, - sinkhorn_repeat, - pre_mix=pre_mix, - norm_weight=norm_weight, - norm_eps=norm_eps, - capture_aux=capture_aux, - ) diff --git a/vllm/models/deepseek_v41/nvidia/ops/mhc.py b/vllm/models/deepseek_v41/nvidia/ops/mhc.py index 29c9c3b956..084a102436 100644 --- a/vllm/models/deepseek_v41/nvidia/ops/mhc.py +++ b/vllm/models/deepseek_v41/nvidia/ops/mhc.py @@ -1,15 +1,26 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Split shifted mHC pre so coefficient generation can overlap the sublayer.""" +"""Dispatch DSV4.1 mHC operations and overlap coefficient generation.""" from functools import partial +from typing import TYPE_CHECKING, cast import torch from vllm.config import VllmConfig +from vllm.distributed import get_tp_group +from vllm.model_executor.kernels.mhc.tilelang import ( + mhc_fused_post_pre_delayed_tilelang, + mhc_post_tilelang, +) from vllm.platforms import current_platform from vllm.utils.deep_gemm import is_deep_gemm_supported +from .mega_mhc import is_mega_mhc_supported, mhc_shifted_post_pre_deep_gemm + +if TYPE_CHECKING: + from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator + # GB200 TP4/FlashInfer improves through 16 tokens; larger screens tie or regress. MHC_OVERLAP_MAX_TOKENS = 16 @@ -29,6 +40,20 @@ def supports_mhc_overlap(vllm_config: VllmConfig) -> bool: ) +def supports_mhc_all_reduce(vllm_config: VllmConfig) -> bool: + parallel = vllm_config.parallel_config + config = vllm_config.model_config.hf_config + if ( + parallel.tensor_parallel_size != 4 + or parallel.enable_expert_parallel + or config.hidden_size != 5120 + or config.hc_mult != 4 + ): + return False + comm = cast("CudaCommunicator", get_tp_group().device_communicator).ca_comm + return comm is not None and bool(comm.mnnvl_lamport_ag_multicast_ptr) + + def mhc_pre_delayed_overlap( residual: torch.Tensor, fn: torch.Tensor, @@ -45,6 +70,7 @@ def mhc_pre_delayed_overlap( norm_eps: float = 1e-6, *, stream: torch.cuda.Stream, + layer_input: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Prepare the input on the caller stream and coefficients on another stream. @@ -67,7 +93,11 @@ def mhc_pre_delayed_overlap( post = torch.empty((n, hc), device=residual.device, dtype=torch.float32) comb = torch.empty((n, hc * hc), device=residual.device, dtype=torch.float32) next_pre = torch.empty_like(post) - layer_input = torch.empty((n, hidden), device=residual.device, dtype=torch.bfloat16) + compute_input = layer_input is None + if layer_input is None: + layer_input = torch.empty( + (n, hidden), device=residual.device, dtype=torch.bfloat16 + ) outputs = post.unsqueeze(-1), comb.view(n, hc, hc), layer_input, next_pre if n == 0: return outputs @@ -104,7 +134,7 @@ def mhc_pre_delayed_overlap( ) main = torch.cuda.current_stream() # Prioritize input readiness for small batches before releasing statistics. - if n <= 8: + if compute_input and n <= 8: epilogue(split_mode="input") stream.wait_stream(main) with torch.cuda.stream(stream): @@ -112,6 +142,128 @@ def mhc_pre_delayed_overlap( epilogue(split_mode="stats") for tensor in (x, mix, sqr): tensor.record_stream(stream) - if n > 8: + if compute_input and n > 8: epilogue(split_mode="input") return outputs + + +def mhc_shifted_post_pre( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + pre_mix: torch.Tensor | None = None, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-6, + capture_aux: bool = False, + *, + stream: torch.cuda.Stream | None = None, + reduce_results: bool = False, +) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor +]: + """Dispatch shifted post/pre to overlap, Mega-mHC, or fused TileLang. + + When stream is supplied, join it before consuming the returned coefficients. + """ + layer_input = None + if reduce_results: + tp = get_tp_group() + if stream is not None and 0 < x.shape[0] <= MHC_OVERLAP_MAX_TOKENS: + comm = cast("CudaCommunicator", tp.device_communicator).ca_comm + assert comm is not None and comm.mnnvl_lamport_epochs is not None + output = torch.empty_like(residual) + layer_input = torch.empty_like(x) + torch.ops._C_custom_ar.all_reduce_mhc( + x, + residual, + post_layer_mix, + comb_res_mix, + pre_mix, + norm_weight, + output, + layer_input, + comm.mnnvl_lamport_ag_local_ptr, + comm.mnnvl_lamport_ag_multicast_ptr, + comm.mnnvl_lamport_epochs[0], + comm.rank, + comm.mnnvl_buffer_size, + norm_eps, + ) + residual = output + else: + x = tp.all_reduce(x) + if stream is not None: + if layer_input is None: + residual = mhc_post_tilelang(x, residual, post_layer_mix, comb_res_mix) + aux = residual.mean(dim=1) if capture_aux else x.new_empty(0, x.shape[1]) + pre_outputs = mhc_pre_delayed_overlap( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + pre_mix=pre_mix, + norm_weight=norm_weight, + norm_eps=norm_eps, + stream=stream, + layer_input=layer_input, + ) + return residual, *pre_outputs, aux + + if ( + pre_mix is not None + and norm_weight is not None + and not capture_aux + and x.shape[0] <= 1 << 20 + and is_mega_mhc_supported(x.shape[1], residual.shape[1]) + ): + outputs = mhc_shifted_post_pre_deep_gemm( + x, + residual, + pre_mix, + post_layer_mix, + comb_res_mix, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_post_mult_value, + hc_sinkhorn_eps, + sinkhorn_repeat, + norm_weight, + norm_eps, + ) + return *outputs, x.new_empty(0, x.shape[1]) + + return mhc_fused_post_pre_delayed_tilelang( + x, + residual, + post_layer_mix, + comb_res_mix, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + pre_mix=pre_mix, + norm_weight=norm_weight, + norm_eps=norm_eps, + capture_aux=capture_aux, + )