diff --git a/CMakeLists.txt b/CMakeLists.txt index 58921b9369..02e098f2d9 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -460,8 +460,6 @@ 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 deleted file mode 100644 index dc58371156..0000000000 --- a/csrc/libtorch_stable/all_reduce_mhc.cu +++ /dev/null @@ -1,212 +0,0 @@ -// 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 1d95801f0b..23d66ee973 100644 --- a/tests/distributed/test_custom_all_reduce.py +++ b/tests/distributed/test_custom_all_reduce.py @@ -25,12 +25,43 @@ for i, v in enumerate(test_sizes): test_sizes[i] -= v % 8 +def _bf16_ulps(a: torch.Tensor, b: torch.Tensor) -> int: + """Largest distance between two BF16 tensors in units in the last place.""" + + def ordered(t: torch.Tensor) -> torch.Tensor: + bits = t.view(torch.int16).int() & 0xFFFF + return torch.where(bits >= 0x8000, 0x8000 - bits, bits) + + return int((ordered(a) - ordered(b)).abs().max()) + + +def _unfused_all_reduce_mhc(peers, residual, post, comb, pre, weight, eps): + """All-reduce, mHC post, collapse and RMSNorm in the fused kernel's FP32 + order, rounding to BF16 where the unfused path returns BF16.""" + reduced = peers[0].float() + for peer in peers[1:]: + reduced = reduced + peer.float() + reduced = reduced.bfloat16().float() + output = torch.empty_like(residual) + collapse = torch.zeros_like(reduced) + for target in range(4): + mixed = reduced * post[:, target : target + 1] + for source in range(4): + mixed = torch.addcmul( + mixed, residual[:, source].float(), comb[:, source, target, None] + ) + output[:, target] = mixed.bfloat16() + collapse = torch.addcmul( + collapse, output[:, target].float(), pre[:, target : target + 1] + ) + prenorm = collapse.bfloat16().float() + inv_rms = torch.rsqrt(prenorm.square().mean(-1, keepdim=True) + eps) + return output, (prenorm * inv_rms * weight.float()).bfloat16() + + @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, - ) + from vllm.models.deepseek_v41.nvidia.ops.cute_dsl import AllReduceMHC with monkeypatch.context() as m: m.delenv("CUDA_VISIBLE_DEVICES", raising=False) @@ -38,98 +69,114 @@ def _all_reduce_mhc(monkeypatch, tp_size, pp_size, rank, distributed_init_port): 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) + op = AllReduceMHC( + hidden_size=5120, hc_mult=4, max_num_tokens=16, top_k=6, 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 + def check_eager_and_replayed(fused, check, halved, doubled): + check(fused()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(5): + fused() + captured = fused() + for _ in range(20): + graph.replay() + check(captured) + # Replays must pick up in-place input changes. + halved.mul_(0.5) + doubled.mul_(2) + graph.replay() + check(captured) + + def mhc_inputs(n): 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) + return residual, post, comb, pre, weight + + 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 + residual, post, comb, pre, weight = mhc_inputs(n) def fused(): - torch.ops._C_custom_ar.all_reduce_mhc( - x, + return op(x, residual, post, comb, pre, weight, 1e-6) + + def check(outputs): + peers = get_tp_group().all_gather(x, dim=0).view(tp_size, n, 5120) + output, normalized = _unfused_all_reduce_mhc( + peers, residual, post, comb, pre, weight, 1e-6 + ) + # The mixed hc streams match bit for bit. The RMSNorm sums the + # squares in another order and uses an approximate rsqrt. + assert torch.equal(outputs[0], output) + assert _bf16_ulps(outputs[1], normalized) <= 1 + + check_eager_and_replayed(fused, check, x, residual) + + def run_finalize(n): + torch.manual_seed(7 + rank) + # A padded permuted GEMM2 buffer, like the MoE's. + rows = n * 6 + 5 + gemm2 = torch.randn(rows, 5120, device=device, dtype=torch.bfloat16) + permuted = torch.randperm(rows, device=device)[: n * 6].view(n, 6).int() + # A route to an expert this rank does not hold. + permuted[0, -1] = -1 + weights = torch.rand(n, 6, device=device) + shared = torch.randn(n, 5120, device=device, dtype=torch.bfloat16) + residual, post, comb, pre, weight = mhc_inputs(n) + + def fused(): + return op.finalize( + gemm2, + weights, + permuted, + shared, 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 - ) + def check(outputs): + # One FP32 FMA per route in route order, the shared add, one + # rounding, then the plain path: only the finalize differs. + acc = torch.zeros(n, 5120, device=device) + for k in range(6): + valid = (permuted[:, k] >= 0).unsqueeze(-1) + rows_k = gemm2[permuted[:, k].clamp_min(0).long()].float() + torch.addcmul( + acc, + torch.where(valid, rows_k, 0.0), + torch.where(valid, weights[:, k : k + 1], 0.0), + out=acc, + ) + x = (acc + shared.float()).bfloat16() + expected = op(x, residual, post, comb, pre, weight, 1e-6) + assert torch.equal(outputs[0], expected[0]) + assert torch.equal(outputs[1], expected[1]) - 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() + check_eager_and_replayed(fused, check, gemm2, shared) - # Cover fixed Q6 and shrinking/growing adaptive verification batches. + # Changing shapes cover shrinking and growing batches in one mailbox. for n in (1, 6, 12, 8, 16, 3, 5, 2, 4, 1): run(n) + run_finalize(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): +def test_all_reduce_mhc_matches_unfused_path(monkeypatch): if torch.accelerator.device_count() < 4: pytest.skip("Requires four GPUs with NVLink multicast") multi_process_parallel(monkeypatch, 4, 1, _all_reduce_mhc) diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 8dff34daf9..9e7a35f0fc 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -33,6 +33,14 @@ from vllm.model_executor.kernels.mhc.triton import hc_collapse_triton from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe import ( + TrtLlmMxfp4ExpertsModular, +) +from vllm.model_executor.layers.fused_moe.moe_output import ( + MoEOutput, + UnfinalizedMoEOutput, +) +from vllm.model_executor.layers.fused_moe.runner.moe_runner import _unpack from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig @@ -94,6 +102,7 @@ from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask from .engram import Engram, can_share_engram_tables, gather_engram_hashes from .ops.mhc import ( MHC_OVERLAP_MAX_TOKENS, + init_mhc_all_reduce, mhc_pre_delayed_overlap, mhc_shifted_post_pre, supports_mhc_all_reduce, @@ -135,6 +144,48 @@ class DeepseekV4MoE(DeepseekV4MoEBase): image_sentinel_lo=IMAGE_SENTINEL_BASE_ID, ) + def defer_finalize(self) -> None: + """Leave the routed top-k reduction to the next fused all-reduce + mHC. + + Only the modular TRTLLM MXFP4 experts can stop after GEMM2; the fused + kernel takes as many tokens as the mHC overlap path. + """ + moe_config = self.experts.moe_config + experts_cls = getattr( + self.experts.routed_experts.quant_method, "experts_cls", None + ) + if ( + experts_cls is TrtLlmMxfp4ExpertsModular + and self.experts.routed_scaling_factor == 1.0 + ): + moe_config.defer_moe_finalize(MHC_OVERLAP_MAX_TOKENS) + if moe_config.use_deferred_moe_finalize: + logger.info_once( + "DSV4.1 mHC: MoE top-k finalize fused into the all-reduce for " + "up to %d tokens.", + MHC_OVERLAP_MAX_TOKENS, + ) + + def defers_finalize(self, num_tokens: int) -> bool: + return ( + not self.use_mega_moe + and self.experts.moe_config.should_defer_moe_finalize(num_tokens) + ) + + def forward_unfinalized( + self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None + ) -> MoEOutput: + """``forward`` with the routed top-k reduction and all-reduce left open.""" + # The runner's custom op returns tensors only, so run its body directly. + shared_output, routed = _unpack( + self.experts._forward_impl( + hidden_states, hidden_states, hidden_states, input_ids + ) + ) + assert shared_output is not None + assert isinstance(routed, UnfinalizedMoEOutput) + return MoEOutput(routed=routed, shared_output=shared_output) + def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]: """Pick the CUDA sparse-MLA attention class for the configured backend. @@ -394,7 +445,7 @@ class DeepseekV4DecoderLayer(nn.Module): def forward( self, - x: torch.Tensor, + x: torch.Tensor | MoEOutput, positions: torch.Tensor, input_ids: torch.Tensor | None, pre_mix: torch.Tensor | None = None, @@ -407,7 +458,7 @@ class DeepseekV4DecoderLayer(nn.Module): capture_previous_aux: bool = False, mega_gate_metadata: MegaGateRoutingMetadata | None = None, ) -> tuple[ - torch.Tensor, + torch.Tensor | MoEOutput, torch.Tensor, torch.Tensor, torch.Tensor, @@ -432,6 +483,7 @@ class DeepseekV4DecoderLayer(nn.Module): # sublayer's pre-mix: attention uses the pre-mix carried in (identity # for the first layer), the FFN uses this layer's attention pre-mix. if residual is None: + assert isinstance(x, torch.Tensor) if x.dim() == 2: # First layer: the stream is the embedding broadcast to hc # copies and the identity pre-mix selects copy 0. @@ -471,7 +523,9 @@ class DeepseekV4DecoderLayer(nn.Module): # Engram injection happens between the previous sublayer's post # 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. + # the post out of the pre-norm GEMM's fused prologue. The layer + # before an engram layer finalizes its own MoE. + assert isinstance(x, torch.Tensor) if self.fuse_mhc_all_reduce: x = tensor_model_parallel_all_reduce(x) previous_post = mhc_post_tilelang(x, residual, post_mix, res_mix) @@ -552,7 +606,11 @@ class DeepseekV4DecoderLayer(nn.Module): stream=mhc_stream, reduce_results=self.fuse_mhc_all_reduce, ) - x = self.ffn(x, input_ids, mega_gate_metadata) + if self.fuse_mhc_all_reduce and self.ffn.defers_finalize(x.shape[0]): + # The next layer's first boundary finalizes it with the all-reduce. + x = self.ffn.forward_unfinalized(x, input_ids) + else: + x = self.ffn(x, input_ids, mega_gate_metadata) if mhc_stream is not None: torch.cuda.current_stream().wait_stream(mhc_stream) return x, residual, post_mix, res_mix, ffn_pre, previous_aux @@ -591,6 +649,8 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): self.fuse_mhc_all_reduce = mhc_stream is not None and supports_mhc_all_reduce( vllm_config ) + if self.fuse_mhc_all_reduce: + init_mhc_all_reduce(vllm_config) # Reserved topk indices buffer for all Indexer layers to reuse. self.topk_indices_buffer = torch.empty( @@ -670,6 +730,16 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): ), prefix=f"{prefix}.layers", ) + if self.fuse_mhc_all_reduce: + # A MoE's top-k finalize folds into the next layer's first mHC + # boundary, which the last local layer lacks and an engram layer + # replaces with its own all-reduce. + local_layers = list(islice(self.layers, self.start_layer, self.end_layer)) + for layer, successor in zip(local_layers, local_layers[1:]): + assert isinstance(layer, DeepseekV4DecoderLayer) + assert isinstance(successor, DeepseekV4DecoderLayer) + if successor.engram is None: + layer.ffn.defer_finalize() # The n-gram hash needs a slot-keyed rolling store of compressed ids # (chunked prefill / decode lookback); key it off the first local @@ -874,6 +944,8 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): if self.use_sequence_parallel: previous_aux = sp_all_gather(previous_aux)[:full_num_tokens] aux_hidden_by_layer[idx] = previous_aux + # Without a successor boundary, the last layer finalized its own MoE. + assert isinstance(hidden_states, torch.Tensor) if layer is not None: # The last layer has no successor to fold its post into. if self.fuse_mhc_all_reduce: diff --git a/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/__init__.py b/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/__init__.py new file mode 100644 index 0000000000..5e126a42d6 --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""CuTe DSL kernels for DSV4.1 mHC.""" + +from .all_reduce_mhc import AllReduceMHC + +__all__ = ["AllReduceMHC"] diff --git a/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_reduce_mhc.py b/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_reduce_mhc.py new file mode 100644 index 0000000000..473ad28102 --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_reduce_mhc.py @@ -0,0 +1,1134 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from FlashInfer's low-latency MNNVL CuTe DSL all-reduce +# (flashinfer/comm/mnnvl_cutedsl/kernel_ll, flashinfer-ai/flashinfer@139af6f6). + +"""Lamport all-reduce fused with DSV4.1's mHC post, collapse and RMSNorm. + +FlashInfer's LL kernels, copied: a publish kernel (or its MoE finalize +variant) multicasts each rank's contribution into a three-generation Lamport +mailbox, and a PDL-chained consumer spins until every rank's fragment lands, +reduces it and normalizes. The consumer's residual add and copy are replaced by +the mHC post-mix of the hc residual streams and their pre-mix collapse, which +the RMSNorm then normalizes. The finalize variant also reads FP32 routing +weights. +""" + +from __future__ import annotations + +import math + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem +from cutlass import BFloat16, Float32, Int32, Int64, Uint32 +from cutlass.cute.runtime import make_fake_compact_tensor + +from vllm.distributed import get_tp_group + +from .primitives import ( + NEGATIVE_ZERO_BF16_BITS, + QUAD_BF16, + VEC_BF16, + WARP_SIZE, + bf16x4_to_packed_u32x2, + bf16x8_to_packed_u32x4, + current_cu_stream, + f32_to_bf16_bits, + fragment_has_negative_zero, + load_global_bf16_as_f32, + load_global_u32x2, + load_global_u32x4, + load_volatile_u32, + make_fake_dynamic_compact_tensor, + map_shared_to_peer, + packed_u32x2_to_bf16x4, + packed_u32x4_to_bf16x8, + sanitize_negative_zero_u32x2, + sanitize_negative_zero_u32x4, + shuffle_sync_idx_u32, + stmc_bf16x2, + stmc_bf16x4, + stmc_bf16x8, + store_global_u32, + store_global_u32x4, + store_lamport_sentinel_u32x4, + store_shared_cluster_f32, + to_cute, + to_cute_dynamic, +) + +LAMPORT_GENERATIONS = 3 +NEXT_STAGE = 0 +ACTIVE_STAGE = 1 + + +@cute.jit +def _group_leader_block_sum( + value: Float32, + warp_sums: cute.Tensor, + warps: cutlass.Constexpr[int], + leader_stride: cutlass.Constexpr[int], +) -> Float32: + lane = cute.arch.lane_idx() + warp = cute.arch.warp_idx() + for offset in cutlass.range_constexpr(1, WARP_SIZE): + if cutlass.const_expr(offset >= leader_stride and (offset & (offset - 1)) == 0): + value = value + cute.arch.shuffle_sync_bfly( + value, + offset=offset, + mask=-1, + mask_and_clamp=31, + ) + if lane == 0: + cute.arch.store((warp_sums + warp).llvm_ptr, value) + cute.arch.barrier() + + result = Float32(0.0) + if warp == 0: + if lane < Int32(warps): + result = cute.arch.load( + (warp_sums + lane).llvm_ptr, + Float32, + ) + result = cute.arch.warp_reduction_sum(result) + if lane == 0: + cute.arch.store(warp_sums.llvm_ptr, result) + cute.arch.barrier() + return cute.arch.load(warp_sums.llvm_ptr, Float32) + + +class _SharedOnlyPublishDeviceKernel: + def __init__( + self, + *, + hidden: int, + tp: int, + rank: int, + capacity_m: int, + elements_per_thread: int, + threads: int, + release_before_store: bool, + enable_pdl: bool, + ) -> None: + if elements_per_thread not in (1, QUAD_BF16, VEC_BF16): + raise ValueError("elements_per_thread must be 1, 4, or 8") + if hidden <= 0 or hidden % elements_per_thread: + raise ValueError("hidden must divide evenly across thread fragments") + self.hidden = hidden + self.tp = tp + self.rank = rank + self.capacity_m = capacity_m + self.elements_per_thread = elements_per_thread + self.threads = threads + self.fragments = hidden // elements_per_thread + self.ctas_per_token = math.ceil(self.fragments / threads) + self.release_before_store = release_before_store + self.enable_pdl = enable_pdl + + @cute.jit + def __call__( + self, + shared_output: cute.Tensor, + stage_state: cute.Tensor, + contribution_mailbox_multicast_address: Int64, + m: Int32, + stream: cuda.CUstream, + ) -> None: + self.kernel( + shared_output, + stage_state, + contribution_mailbox_multicast_address, + ).launch( + grid=(m * self.ctas_per_token, 1, 1), + block=(self.threads, 1, 1), + stream=stream, + use_pdl=self.enable_pdl, + ) + + @cute.kernel + def kernel( + self, + shared_output: cute.Tensor, + stage_state: cute.Tensor, + contribution_mailbox_multicast_address: Int64, + ) -> None: + block, _, _ = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + token = block // self.ctas_per_token + cta_in_token = block % self.ctas_per_token + fragment = cta_in_token * self.threads + tidx + + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() + + if cutlass.const_expr(self.elements_per_thread == 1): + bits = Uint32(0) + if fragment < self.fragments: + element = Int64(token) * self.hidden + Int64(fragment) + bits = f32_to_bf16_bits( + load_global_bf16_as_f32( + Int64((shared_output.iterator + element).toint()) + ) + ) + if bits == Uint32(NEGATIVE_ZERO_BF16_BITS): + bits = Uint32(0) + partner_bits = shuffle_sync_idx_u32( + bits, + Int32(cute.arch.lane_idx() | Int32(1)), + ) + packed_word = bits | (partner_bits << Uint32(16)) + elif cutlass.const_expr(self.elements_per_thread == QUAD_BF16): + packed_words = cute.make_rmem_tensor(cute.make_layout((2,)), Uint32) + packed_words.fill(Uint32(0)) + if fragment < self.fragments: + element = Int64(token) * self.hidden + Int64(fragment) * QUAD_BF16 + pointer = cute.make_ptr( + BFloat16, + (shared_output.iterator + element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=8, + ) + packed_words.store( + sanitize_negative_zero_u32x2(load_global_u32x2(pointer)) + ) + else: + packed_words = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32) + packed_words.fill(Uint32(0)) + if fragment < self.fragments: + element = Int64(token) * self.hidden + Int64(fragment) * VEC_BF16 + pointer = cute.make_ptr( + BFloat16, + (shared_output.iterator + element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + packed_words.store( + sanitize_negative_zero_u32x4(load_global_u32x4(pointer)) + ) + + stage = load_volatile_u32(stage_state.iterator + NEXT_STAGE) + if cutlass.const_expr(self.release_before_store): + if block == 0 and tidx == 0: + store_global_u32( + stage_state.iterator + ACTIVE_STAGE, + stage, + ) + cute.arch.barrier() + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() + + if cutlass.const_expr(self.elements_per_thread == 1): + if (cute.arch.lane_idx() & Int32(1)) == Int32( + 0 + ) and fragment < self.fragments: + mailbox_element = ( + (Int64(stage) * self.tp + self.rank) * self.capacity_m + + Int64(token) + ) * self.hidden + Int64(fragment) + stmc_bf16x2( + contribution_mailbox_multicast_address + mailbox_element * 2, + packed_word, + ) + elif cutlass.const_expr(self.elements_per_thread == QUAD_BF16): + if fragment < self.fragments: + mailbox_element = ( + (Int64(stage) * self.tp + self.rank) * self.capacity_m + + Int64(token) + ) * self.hidden + Int64(fragment) * QUAD_BF16 + stmc_bf16x4( + contribution_mailbox_multicast_address + mailbox_element * 2, + packed_words, + ) + else: + if fragment < self.fragments: + mailbox_element = ( + (Int64(stage) * self.tp + self.rank) * self.capacity_m + + Int64(token) + ) * self.hidden + Int64(fragment) * VEC_BF16 + stmc_bf16x8( + contribution_mailbox_multicast_address + mailbox_element * 2, + packed_words, + ) + + if cutlass.const_expr(not self.release_before_store): + if block == 0 and tidx == 0: + store_global_u32( + stage_state.iterator + ACTIVE_STAGE, + stage, + ) + cute.arch.barrier() + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() + + +class _QuadFinalizePublishDeviceKernel: + def __init__( + self, + *, + hidden: int, + top_k: int, + tp: int, + rank: int, + capacity_m: int, + threads: int, + routed_scaling_factor: float, + include_shared_expert: bool, + load_shared_expert_before_pdl: bool, + enable_pdl: bool, + prefetch_group: int, + fp32_weights: bool = False, + ) -> None: + if hidden <= 0 or hidden % QUAD_BF16: + raise ValueError("hidden must be a positive multiple of 4") + self.hidden = hidden + self.top_k = top_k + self.tp = tp + self.rank = rank + self.capacity_m = capacity_m + self.threads = threads + self.routed_scaling_factor = routed_scaling_factor + self.fragments = hidden // QUAD_BF16 + self.ctas_per_token = math.ceil(self.fragments / threads) + self.include_shared_expert = include_shared_expert + self.load_shared_expert_before_pdl = load_shared_expert_before_pdl + self.enable_pdl = enable_pdl + self.prefetch_group = prefetch_group + self.prefetch_groups = (top_k + prefetch_group - 1) // prefetch_group + self.fp32_weights = fp32_weights + + def smem_size_in_bytes(self) -> int: + return self.top_k * 8 + + @cute.jit + def __call__( + self, + routed_output: cute.Tensor, + expert_weights: cute.Tensor, + permuted_indices: cute.Tensor, + shared_output: cute.Tensor, + stage_state: cute.Tensor, + contribution_mailbox_multicast_address: Int64, + m: Int32, + stream: cuda.CUstream, + ) -> None: + self.kernel( + routed_output, + expert_weights, + permuted_indices, + shared_output, + stage_state, + contribution_mailbox_multicast_address, + ).launch( + grid=(m * self.ctas_per_token, 1, 1), + block=(self.threads, 1, 1), + smem=self.smem_size_in_bytes(), + stream=stream, + use_pdl=self.enable_pdl, + ) + + @cute.kernel + def kernel( + self, + routed_output: cute.Tensor, + expert_weights: cute.Tensor, + permuted_indices: cute.Tensor, + shared_output: cute.Tensor, + stage_state: cute.Tensor, + contribution_mailbox_multicast_address: Int64, + ) -> None: + block, _, _ = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + token = block // self.ctas_per_token + cta_in_token = block % self.ctas_per_token + fragment = cta_in_token * self.threads + tidx + + smem = cutlass.utils.SmemAllocator() + staged_indices = smem.allocate_array(Int32, self.top_k) + staged_weights = smem.allocate_array(Float32, self.top_k) + metadata_index = Int32(tidx) + while metadata_index < Int32(self.top_k): + element = Int64(token) * self.top_k + Int64(metadata_index) + row = cute.arch.load( + (permuted_indices.iterator + element).llvm_ptr, + Int32, + ) + if cutlass.const_expr(self.fp32_weights): + weight = cute.arch.load( + (expert_weights.iterator + element).llvm_ptr, + Float32, + ) + else: + weight = load_global_bf16_as_f32( + Int64((expert_weights.iterator + element).toint()) + ) + if cutlass.const_expr(self.routed_scaling_factor != 1.0): + weight = weight * Float32(self.routed_scaling_factor) + if row == Int32(-1): + weight = Float32(0.0) + cute.arch.store( + (staged_indices + metadata_index).llvm_ptr, + row, + ) + cute.arch.store( + (staged_weights + metadata_index).llvm_ptr, + weight, + ) + metadata_index = metadata_index + self.threads + cute.arch.barrier() + + if cutlass.const_expr(self.include_shared_expert): + shared_values = cute.make_rmem_tensor( + cute.make_layout((QUAD_BF16,)), BFloat16 + ) + shared_values.fill(BFloat16(0.0)) + if cutlass.const_expr(self.load_shared_expert_before_pdl): # noqa: SIM102 + if fragment < self.fragments: + shared_element = ( + Int64(token) * self.hidden + Int64(fragment) * QUAD_BF16 + ) + shared_pointer = cute.make_ptr( + BFloat16, + (shared_output.iterator + shared_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=8, + ) + shared_values.store( + packed_u32x2_to_bf16x4(load_global_u32x2(shared_pointer)) + ) + + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() + + if cutlass.const_expr( # noqa: SIM102 + self.include_shared_expert and not self.load_shared_expert_before_pdl + ): + if fragment < self.fragments: + shared_element = ( + Int64(token) * self.hidden + Int64(fragment) * QUAD_BF16 + ) + shared_pointer = cute.make_ptr( + BFloat16, + (shared_output.iterator + shared_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=8, + ) + shared_values.store( + packed_u32x2_to_bf16x4(load_global_u32x2(shared_pointer)) + ) + + stage = load_volatile_u32(stage_state.iterator + NEXT_STAGE) + if fragment < self.fragments: + accumulator = cute.make_rmem_tensor(cute.make_layout((QUAD_BF16,)), Float32) + accumulator.fill(Float32(0.0)) + if cutlass.const_expr(self.prefetch_group == 1): + for k in cutlass.range_constexpr(self.top_k): + row = cute.arch.load((staged_indices + k).llvm_ptr, Int32) + weight = cute.arch.load( + (staged_weights + k).llvm_ptr, + Float32, + ) + if row != Int32(-1): + source_element = ( + Int64(row) * self.hidden + Int64(fragment) * QUAD_BF16 + ) + source_pointer = cute.make_ptr( + BFloat16, + (routed_output.iterator + source_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=8, + ) + accumulator.store( + accumulator.load() + + packed_u32x2_to_bf16x4( + load_global_u32x2(source_pointer) + ).to(Float32) + * weight + ) + else: + inputs = cute.make_rmem_tensor( + cute.make_layout((self.prefetch_group, 2)), + Uint32, + ) + inputs.fill(Uint32(0)) + for group in cutlass.range_constexpr(self.prefetch_groups): + for item in cutlass.range_constexpr(self.prefetch_group): + k = group * self.prefetch_group + item + if cutlass.const_expr(k < self.top_k): + row = cute.arch.load( + (staged_indices + k).llvm_ptr, + Int32, + ) + for word in cutlass.range_constexpr(2): + inputs[item, word] = Uint32(0) + if row != Int32(-1): + source_element = ( + Int64(row) * self.hidden + + Int64(fragment) * QUAD_BF16 + ) + source_pointer = cute.make_ptr( + BFloat16, + (routed_output.iterator + source_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=8, + ) + source = load_global_u32x2(source_pointer) + for word in cutlass.range_constexpr(2): + inputs[item, word] = source[word] + for item in cutlass.range_constexpr(self.prefetch_group): + k = group * self.prefetch_group + item + if cutlass.const_expr(k < self.top_k): + source = cute.make_rmem_tensor( + cute.make_layout((2,)), + Uint32, + ) + for word in cutlass.range_constexpr(2): + source[word] = inputs[item, word] + accumulator.store( + accumulator.load() + + packed_u32x2_to_bf16x4(source.load()).to(Float32) + * cute.arch.load( + (staged_weights + k).llvm_ptr, + Float32, + ) + ) + result = accumulator.load() + if cutlass.const_expr(self.include_shared_expert): + result = result + shared_values.load().to(Float32) + result_packed = bf16x4_to_packed_u32x2(result.to(BFloat16)) + packed = sanitize_negative_zero_u32x2(result_packed) + mailbox_element = ( + (Int64(stage) * self.tp + self.rank) * self.capacity_m + Int64(token) + ) * self.hidden + Int64(fragment) * QUAD_BF16 + stmc_bf16x4( + contribution_mailbox_multicast_address + mailbox_element * 2, + packed, + ) + + if block == 0 and tidx == 0: + store_global_u32( + stage_state.iterator + ACTIVE_STAGE, + stage, + ) + cute.arch.barrier() + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() + + +class _LamportMHCDeviceKernel: + """FlashInfer's ``_LamportResidualRMSNormDeviceKernel`` with mHC. + + The residual add and residual copy become the mHC post-mix of the hc + streams, stored as the residual output, and their pre-mix collapse, which + the unchanged RMSNorm normalizes. + """ + + def __init__( + self, + *, + hidden: int, + hc: int, + tp: int, + capacity_m: int, + cluster_size: int, + rank_lanes: int, + threads: int, + enable_pdl: bool, + ) -> None: + if rank_lanes not in (1, 2, 4, 8): + raise ValueError("rank_lanes must be 1, 2, 4, or 8") + if tp % rank_lanes: + raise ValueError("tp must be divisible by rank_lanes") + if threads <= 0 or threads % WARP_SIZE or threads % rank_lanes: + raise ValueError("threads must be a positive warp and rank-lane multiple") + if hidden <= 0 or hidden % VEC_BF16: + raise ValueError("hidden must be a positive multiple of 8") + self.hidden = hidden + self.hc = hc + self.tp = tp + self.capacity_m = capacity_m + self.cluster_size = cluster_size + self.rank_lanes = rank_lanes + self.threads = threads + self.enable_pdl = enable_pdl + self.fragments = hidden // VEC_BF16 + self.groups_per_cta = threads // rank_lanes + self.fragment_stride = cluster_size * self.groups_per_cta + self.trips = math.ceil(self.fragments / self.fragment_stride) + self.warps = threads // WARP_SIZE + self.rank_waves = tp // rank_lanes + + def smem_size_in_bytes(self) -> int: + return (self.warps + self.cluster_size) * 4 + + @cute.jit + def __call__( + self, + contribution_mailbox: cute.Tensor, + residual_source: cute.Tensor, + post: cute.Tensor, + comb: cute.Tensor, + pre: cute.Tensor, + gamma: cute.Tensor, + residual_output: cute.Tensor, + norm_output: cute.Tensor, + stage_state: cute.Tensor, + eps: Float32, + m: Int32, + stream: cuda.CUstream, + ) -> None: + self.kernel( + contribution_mailbox, + residual_source, + post, + comb, + pre, + gamma, + residual_output, + norm_output, + stage_state, + eps, + ).launch( + grid=(m, self.cluster_size, 1), + block=(self.threads, 1, 1), + cluster=(1, self.cluster_size, 1), + smem=self.smem_size_in_bytes(), + stream=stream, + use_pdl=self.enable_pdl, + ) + + @cute.kernel + def kernel( + self, + contribution_mailbox: cute.Tensor, + residual_source: cute.Tensor, + post: cute.Tensor, + comb: cute.Tensor, + pre: cute.Tensor, + gamma: cute.Tensor, + residual_output: cute.Tensor, + norm_output: cute.Tensor, + stage_state: cute.Tensor, + eps: Float32, + ) -> None: + hc = self.hc + tidx, _, _ = cute.arch.thread_idx() + token, _, _ = cute.arch.block_idx() + cluster_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + rank_lane = tidx % self.rank_lanes + group = tidx // self.rank_lanes + base_fragment = cluster_rank * self.groups_per_cta + group + + prenorm_fragments = cute.make_rmem_tensor( + cute.make_layout( + (self.trips, VEC_BF16), + stride=(VEC_BF16, 1), + ), + BFloat16, + ) + prenorm_fragments.fill(BFloat16(0.0)) + gamma_fragments = cute.make_rmem_tensor( + cute.make_layout( + (self.trips, VEC_BF16), + stride=(VEC_BF16, 1), + ), + BFloat16, + ) + gamma_fragments.fill(BFloat16(0.0)) + residual_fragments = cute.make_rmem_tensor( + cute.make_layout( + (self.trips, hc, VEC_BF16), + stride=(hc * VEC_BF16, VEC_BF16, 1), + ), + BFloat16, + ) + residual_fragments.fill(BFloat16(0.0)) + + for trip in cutlass.range_constexpr(self.trips): + fragment = base_fragment + trip * self.fragment_stride + if fragment < self.fragments and rank_lane == 0: + gamma_element = Int64(fragment) * VEC_BF16 + gamma_pointer = cute.make_ptr( + BFloat16, + (gamma.iterator + gamma_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + gamma_fragments[trip, None].store( + packed_u32x4_to_bf16x8(load_global_u32x4(gamma_pointer)) + ) + for source in cutlass.range_constexpr(hc): + residual_element = ( + Int64(token) * hc + source + ) * self.hidden + gamma_element + residual_pointer = cute.make_ptr( + BFloat16, + (residual_source.iterator + residual_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + residual_fragments[trip, source, None].store( + packed_u32x4_to_bf16x8(load_global_u32x4(residual_pointer)) + ) + post_mix = cute.make_rmem_tensor(cute.make_layout((hc,)), Float32) + pre_mix = cute.make_rmem_tensor(cute.make_layout((hc,)), Float32) + comb_mix = cute.make_rmem_tensor( + cute.make_layout((hc, hc), stride=(hc, 1)), Float32 + ) + for target in cutlass.range_constexpr(hc): + mix_index = Int64(token) * hc + target + post_mix[target] = cute.arch.load( + (post.iterator + mix_index).llvm_ptr, Float32 + ) + pre_mix[target] = cute.arch.load( + (pre.iterator + mix_index).llvm_ptr, Float32 + ) + for source in cutlass.range_constexpr(hc): + comb_index = (Int64(token) * hc + source) * hc + target + comb_mix[source, target] = cute.arch.load( + (comb.iterator + comb_index).llvm_ptr, Float32 + ) + + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() + + active_stage = load_volatile_u32(stage_state.iterator + ACTIVE_STAGE) + for trip in cutlass.range_constexpr(self.trips): + fragment = base_fragment + trip * self.fragment_stride + lane_packed = cute.make_rmem_tensor( + cute.make_layout((self.rank_waves, 4)), + Uint32, + ) + lane_packed.fill(Uint32(0)) + dirty = fragment < self.fragments + while dirty: + dirty = False + for wave in cutlass.range_constexpr(self.rank_waves): + source_rank = wave * self.rank_lanes + rank_lane + if fragment < self.fragments: + source_element = ( + (Int64(active_stage) * self.tp + Int64(source_rank)) + * self.capacity_m + + Int64(token) + ) * self.hidden + Int64(fragment) * VEC_BF16 + source_pointer = cute.make_ptr( + BFloat16, + (contribution_mailbox.iterator + source_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + packed = load_global_u32x4( + source_pointer, + volatile=True, + ) + dirty = dirty | fragment_has_negative_zero(packed) + for word in cutlass.range_constexpr(4): + lane_packed[wave, word] = packed[word] + + lane_sum = cute.make_rmem_tensor(cute.make_layout((VEC_BF16,)), Float32) + lane_sum.fill(Float32(0.0)) + for wave in cutlass.range_constexpr(self.rank_waves): + packed = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32) + for word in cutlass.range_constexpr(4): + packed[word] = lane_packed[wave, word] + lane_sum.store( + lane_sum.load() + packed_u32x4_to_bf16x8(packed.load()).to(Float32) + ) + for offset in cutlass.range_constexpr(1, 5): + if cutlass.const_expr(offset < self.rank_lanes and offset in (1, 2, 4)): + for element in cutlass.range_constexpr(VEC_BF16): + lane_sum[element] = lane_sum[ + element + ] + cute.arch.shuffle_sync_bfly( + lane_sum[element], + offset=offset, + mask=-1, + mask_and_clamp=31, + ) + + if fragment < self.fragments and rank_lane == 0: + # The all-reduced sublayer output is BF16, as the unfused + # collective returns it, before mHC post mixes it in. + reduced = lane_sum.load().to(BFloat16).to(Float32) + collapse = cute.make_rmem_tensor(cute.make_layout((VEC_BF16,)), Float32) + collapse.fill(Float32(0.0)) + for target in cutlass.range_constexpr(hc): + mixed = reduced * post_mix[target] + for source in cutlass.range_constexpr(hc): + mixed = ( + mixed + + residual_fragments[trip, source, None].load().to(Float32) + * comb_mix[source, target] + ) + mixed_bf16 = mixed.to(BFloat16) + output_element = (Int64(token) * hc + target) * self.hidden + Int64( + fragment + ) * VEC_BF16 + store_global_u32x4( + Int64((residual_output.iterator + output_element).toint()), + bf16x8_to_packed_u32x4(mixed_bf16), + ) + collapse.store( + collapse.load() + mixed_bf16.to(Float32) * pre_mix[target] + ) + prenorm_fragments[trip, None].store(collapse.load().to(BFloat16)) + + if token == 0 and cluster_rank == 0 and tidx == 0: + store_global_u32( + stage_state.iterator + NEXT_STAGE, + (active_stage + Uint32(1)) % Uint32(LAMPORT_GENERATIONS), + ) + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() + + for trip in cutlass.range_constexpr(self.trips): + fragment = base_fragment + trip * self.fragment_stride + for wave in cutlass.range_constexpr(self.rank_waves): + source_rank = wave * self.rank_lanes + rank_lane + if fragment < self.fragments: + source_element = ( + (Int64(active_stage) * self.tp + Int64(source_rank)) + * self.capacity_m + + Int64(token) + ) * self.hidden + Int64(fragment) * VEC_BF16 + store_lamport_sentinel_u32x4( + Int64((contribution_mailbox.iterator + source_element).toint()) + ) + + thread_sum = Float32(0.0) + for trip in cutlass.range_constexpr(self.trips): + fragment = base_fragment + trip * self.fragment_stride + if fragment < self.fragments and rank_lane == 0: + values = prenorm_fragments[trip, None].load().to(Float32) + thread_sum = thread_sum + (values * values).reduce( + cute.ReductionOp.ADD, + init_val=Float32(0.0), + reduction_profile=0, + ) + + smem = cutlass.utils.SmemAllocator() + warp_sums = smem.allocate_array(Float32, self.warps) + cluster_sums = smem.allocate_array(Float32, self.cluster_size) + cta_sum = _group_leader_block_sum( + thread_sum, + warp_sums, + self.warps, + self.rank_lanes, + ) + if tidx < self.cluster_size: + local_slot = cluster_sums + cluster_rank + remote_slot = map_shared_to_peer(local_slot, Int32(tidx)) + store_shared_cluster_f32(remote_slot, cta_sum) + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + full_sum = Float32(0.0) + for peer in cutlass.range_constexpr(self.cluster_size): + full_sum = full_sum + cute.arch.load( + (cluster_sums + peer).llvm_ptr, + Float32, + ) + inv_rms = cute.math.rsqrt( + full_sum / Float32(self.hidden) + eps, + fastmath=True, + ) + for trip in cutlass.range_constexpr(self.trips): + fragment = base_fragment + trip * self.fragment_stride + if fragment < self.fragments and rank_lane == 0: + gamma_values = gamma_fragments[trip, None].load().to(Float32) + result = ( + prenorm_fragments[trip, None].load().to(Float32) + * inv_rms + * gamma_values + ).to(BFloat16) + output_element = Int64(token) * self.hidden + Int64(fragment) * VEC_BF16 + store_global_u32x4( + Int64((norm_output.iterator + output_element).toint()), + bf16x8_to_packed_u32x4(result), + ) + + +class AllReduceMHC: + """Lamport TP all-reduce fused with mHC post, collapse and RMSNorm. + + ``__call__`` reduces ``x`` across the TP group, writes the post-mixed hc + streams and the normalized collapse, and leaves the next mixes to the + caller. ``finalize`` does the same for an unfinalized MoE output, folding + the top-k reduction and the shared-expert add into the publish. Each + instance owns a Lamport mailbox; construction is collective over the TP + group. + """ + + def __init__( + self, + *, + hidden_size: int, + hc_mult: int, + max_num_tokens: int, + top_k: int, + device: torch.device, + ) -> None: + group = get_tp_group().device_group + self.hidden_size = hidden = hidden_size + self.hc_mult = hc = hc_mult + self.capacity = capacity = max_num_tokens + self.top_k = top_k + tp = dist.get_world_size(group) + rank = dist.get_rank(group) + device = torch.device(device) + + # FlashInfer's LL_ALL_REDUCE_GB300_TP{4,8}_H5120 and + # LL_FINALIZE_..._H5120_K* presets: a cluster of 5 x 128 threads covers + # a 640-fragment token in one fully used trip, and the finalize stages + # all top_k rows at once. + with torch.accelerator.device_index(device.index): + self._publish = cute.compile( + _SharedOnlyPublishDeviceKernel( + hidden=hidden, + tp=tp, + rank=rank, + capacity_m=capacity, + elements_per_thread=VEC_BF16, + threads=128, + release_before_store=False, + enable_pdl=True, + ), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=hidden + ), + make_fake_compact_tensor(Int32, (2,), assumed_align=4), + Int64(0), + Int32(capacity), + current_cu_stream(), + ) + self._finalize_publish = cute.compile( + _QuadFinalizePublishDeviceKernel( + hidden=hidden, + top_k=top_k, + tp=tp, + rank=rank, + capacity_m=capacity, + threads=128, + # DSV4.1's router folds the scale into the weights. + routed_scaling_factor=1.0, + include_shared_expert=True, + load_shared_expert_before_pdl=False, + enable_pdl=True, + prefetch_group=top_k, + fp32_weights=True, + ), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=hidden + ), + make_fake_dynamic_compact_tensor( + Float32, alignment=4, divisibility=top_k + ), + make_fake_dynamic_compact_tensor( + Int32, alignment=4, divisibility=top_k + ), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=hidden + ), + make_fake_compact_tensor(Int32, (2,), assumed_align=4), + Int64(0), + Int32(capacity), + current_cu_stream(), + ) + self._collective = cute.compile( + _LamportMHCDeviceKernel( + hidden=hidden, + hc=hc, + tp=tp, + capacity_m=capacity, + cluster_size=5, + rank_lanes=1, + threads=128, + enable_pdl=True, + ), + make_fake_compact_tensor( + BFloat16, + (LAMPORT_GENERATIONS * tp * capacity * hidden,), + assumed_align=16, + ), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=hc * hidden + ), + make_fake_dynamic_compact_tensor(Float32, alignment=4, divisibility=hc), + make_fake_dynamic_compact_tensor( + Float32, alignment=4, divisibility=hc * hc + ), + make_fake_dynamic_compact_tensor(Float32, alignment=4, divisibility=hc), + make_fake_compact_tensor(BFloat16, (hidden,), assumed_align=16), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=hc * hidden + ), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=hidden + ), + make_fake_compact_tensor(Int32, (2,), assumed_align=4), + Float32(0.0), + Int32(capacity), + current_cu_stream(), + ) + + self._mailbox = symm_mem.empty( + (LAMPORT_GENERATIONS, tp, capacity, hidden), + dtype=torch.bfloat16, + device=device, + ) + self._mailbox_handle = symm_mem.rendezvous(self._mailbox, group) + multicast = int(self._mailbox_handle.multicast_ptr or 0) + if not multicast or multicast % 16: + raise RuntimeError("NVLink multicast mapping is unavailable") + self._multicast = multicast + # Every slot starts as the Lamport sentinel, BF16 negative zero. + self._mailbox.view(torch.int16).fill_(-32768) + self._stage_state = torch.zeros(2, dtype=torch.int32, device=device) + torch.accelerator.synchronize() + dist.barrier(group=group) + + def __call__( + self, + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + pre: torch.Tensor, + norm_weight: torch.Tensor, + norm_eps: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return (post-mixed residual streams, normalized layer input). + + Args: + x: [M, hidden] BF16 TP-partial sublayer output. + residual: [M, hc, hidden] BF16 residual streams. + post: [M, hc(, 1)] FP32 post-mix. + comb: [M, hc, hc] FP32 residual mix, indexed [source, target]. + pre: [M, hc] FP32 carried pre-mix for the collapse. + norm_weight: [hidden] BF16 RMSNorm weight. + norm_eps: RMSNorm epsilon. + + """ + m = x.shape[0] + self._check(m, residual, post, comb, pre) + if x.shape != (m, self.hidden_size): + raise ValueError("unexpected x shape") + self._publish( + to_cute_dynamic(x.flatten(), 16, divisibility=self.hidden_size), + to_cute(self._stage_state, 4), + Int64(self._multicast), + Int32(m), + current_cu_stream(), + ) + return self._collect(residual, post, comb, pre, norm_weight, norm_eps, m) + + def finalize( + self, + gemm2_permuted: torch.Tensor, + expert_weights: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + shared_output: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + pre: torch.Tensor, + norm_weight: torch.Tensor, + norm_eps: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + """``__call__`` for a MoE output left unfinalized. + + Args: + gemm2_permuted: [rows, hidden] BF16 unweighted GEMM2 output in the + MoE's permuted order. + expert_weights: [M, top_k] FP32 routing weights, already scaled. + expanded_idx_to_permuted_idx: [M, top_k] int32 row of each route, + -1 for an expert this rank does not hold. + shared_output: [M, hidden] BF16 TP-partial shared-expert output. + residual: [M, hc, hidden] BF16 residual streams. + post: [M, hc(, 1)] FP32 post-mix. + comb: [M, hc, hc] FP32 residual mix, indexed [source, target]. + pre: [M, hc] FP32 carried pre-mix for the collapse. + norm_weight: [hidden] BF16 RMSNorm weight. + norm_eps: RMSNorm epsilon. + + """ + m = shared_output.shape[0] + hidden, top_k = self.hidden_size, self.top_k + self._check(m, residual, post, comb, pre) + if gemm2_permuted.dim() != 2 or gemm2_permuted.shape[1] != hidden: + raise ValueError("gemm2_permuted must be [rows, hidden]") + if shared_output.shape != (m, hidden): + raise ValueError("unexpected shared_output shape") + if expert_weights.dtype != torch.float32 or expert_weights.numel() != m * top_k: + raise ValueError("expert_weights must be FP32 [M, top_k]") + if ( + expanded_idx_to_permuted_idx.dtype != torch.int32 + or expanded_idx_to_permuted_idx.numel() != m * top_k + ): + raise ValueError("expanded_idx_to_permuted_idx must be int32 [M, top_k]") + self._finalize_publish( + to_cute_dynamic(gemm2_permuted.flatten(), 16, divisibility=hidden), + to_cute_dynamic(expert_weights.flatten(), 4, divisibility=top_k), + to_cute_dynamic( + expanded_idx_to_permuted_idx.flatten(), 4, divisibility=top_k + ), + to_cute_dynamic(shared_output.flatten(), 16, divisibility=hidden), + to_cute(self._stage_state, 4), + Int64(self._multicast), + Int32(m), + current_cu_stream(), + ) + return self._collect(residual, post, comb, pre, norm_weight, norm_eps, m) + + def _check( + self, + m: int, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + pre: torch.Tensor, + ) -> None: + hc = self.hc_mult + if not 1 <= m <= self.capacity: + raise ValueError(f"M={m} is outside [1, {self.capacity}]") + if residual.shape != (m, hc, self.hidden_size): + raise ValueError("unexpected residual shape") + if ( + post.numel() != m * hc + or pre.numel() != m * hc + or comb.numel() != m * hc * hc + ): + raise ValueError("unexpected mix shapes") + + def _collect( + self, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + pre: torch.Tensor, + norm_weight: torch.Tensor, + norm_eps: float, + m: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + hidden, hc = self.hidden_size, self.hc_mult + residual_output = torch.empty_like(residual) + layer_input = torch.empty( + (m, hidden), dtype=torch.bfloat16, device=residual.device + ) + self._collective( + to_cute(self._mailbox.flatten(), 16), + to_cute_dynamic(residual.flatten(), 16, divisibility=hc * hidden), + to_cute_dynamic(post.flatten(), 4, divisibility=hc), + to_cute_dynamic(comb.flatten(), 4, divisibility=hc * hc), + to_cute_dynamic(pre.flatten(), 4, divisibility=hc), + to_cute(norm_weight, 16), + to_cute_dynamic(residual_output.flatten(), 16, divisibility=hc * hidden), + to_cute_dynamic(layer_input.flatten(), 16, divisibility=hidden), + to_cute(self._stage_state, 4), + Float32(norm_eps), + Int32(m), + current_cu_stream(), + ) + return residual_output, layer_input diff --git a/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/primitives.py b/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/primitives.py new file mode 100644 index 0000000000..16d5f4721a --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/primitives.py @@ -0,0 +1,396 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from FlashInfer's flashinfer/comm/mnnvl_cutedsl/cute_dsl_primitives.py +# and runtime.py (flashinfer-ai/flashinfer@139af6f6). + +"""PTX building blocks and launch helpers for the DSV4.1 LL all-reduce.""" + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass import Float32, Int32, Int64, Uint16, Uint32 +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm, vector +from cutlass.cute.runtime import from_dlpack, make_fake_compact_tensor +from cutlass.cutlass_dsl import T, dsl_user_op + +WARP_SIZE = 32 +VEC_BF16 = 8 +QUAD_BF16 = 4 +NEGATIVE_ZERO_BF16_BITS = 0x8000 +NEGATIVE_ZERO_BF16_PAIR = 0x80008000 + + +def _asm(result, operands, text, constraints, *, side_effects, loc, ip): + return llvm.inline_asm( + result, + operands, + text, + constraints, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def load_global_u32x4( + pointer: cute.Pointer, + *, + volatile: cutlass.Constexpr[bool] = False, + loc=None, + ip=None, +): + address = pointer.toint(loc=loc, ip=ip) + opcode = "ld.volatile.global.v4.u32" if volatile else "ld.global.v4.u32" + loaded = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [address.ir_value(loc=loc, ip=ip)], + f"{opcode} {{$0, $1, $2, $3}}, [$4];", + "=r,=r,=r,=r,l", + side_effects=volatile, + loc=loc, + ip=ip, + ) + packed = vector.from_elements( + ir.VectorType.get([4], T.i32(), loc=loc), + [ + llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip) + for index in range(4) + ], + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 4, Uint32) + + +@dsl_user_op +def load_global_u32x2(pointer: cute.Pointer, *, loc=None, ip=None): + address = pointer.toint(loc=loc, ip=ip) + loaded = _asm( + llvm.StructType.get_literal([T.i32()] * 2), + [address.ir_value(loc=loc, ip=ip)], + "ld.global.v2.u32 {$0, $1}, [$2];", + "=r,=r,l", + side_effects=False, + loc=loc, + ip=ip, + ) + packed = vector.from_elements( + ir.VectorType.get([2], T.i32(), loc=loc), + [ + llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip) + for index in range(2) + ], + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 2, Uint32) + + +@dsl_user_op +def load_global_bf16_as_f32(address: Int64, *, loc=None, ip=None) -> Float32: + return Float32( + _asm( + T.f32(), + [address.ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b16 bits;\n\tld.global.b16 bits, [$1];\n\t" + "cvt.f32.bf16 $0, bits;\n\t}", + "=f,l", + side_effects=False, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def store_global_u32x4(address: Int64, packed, *, loc=None, ip=None) -> None: + words = [packed[index].ir_value(loc=loc, ip=ip) for index in range(4)] + _asm( + None, + [address.ir_value(loc=loc, ip=ip), *words], + "st.global.v4.u32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def store_lamport_sentinel_u32x4(address: Int64, *, loc=None, ip=None) -> None: + sentinel = Uint32(NEGATIVE_ZERO_BF16_PAIR).ir_value(loc=loc, ip=ip) + _asm( + None, + [address.ir_value(loc=loc, ip=ip), sentinel, sentinel, sentinel, sentinel], + "st.global.v4.u32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def load_volatile_u32(pointer: cute.Pointer, *, loc=None, ip=None) -> Uint32: + address = pointer.toint(loc=loc, ip=ip) + return Uint32( + _asm( + T.i32(), + [address.ir_value(loc=loc, ip=ip)], + "ld.volatile.global.u32 $0, [$1];", + "=r,l", + side_effects=True, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def store_global_u32( + pointer: cute.Pointer, value: Uint32, *, loc=None, ip=None +) -> None: + address = pointer.toint(loc=loc, ip=ip) + _asm( + None, + [address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)], + "st.global.u32 [$0], $1;", + "l,r", + side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def packed_u32x4_to_bf16x8(packed, *, loc=None, ip=None): + values = llvm.bitcast( + ir.VectorType.get([VEC_BF16], cutlass.BFloat16.mlir_type, loc=loc), + packed.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(values, VEC_BF16, cutlass.BFloat16) + + +@dsl_user_op +def packed_u32x2_to_bf16x4(packed, *, loc=None, ip=None): + values = llvm.bitcast( + ir.VectorType.get([QUAD_BF16], cutlass.BFloat16.mlir_type, loc=loc), + packed.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(values, QUAD_BF16, cutlass.BFloat16) + + +@dsl_user_op +def bf16x4_to_packed_u32x2(values, *, loc=None, ip=None): + packed = llvm.bitcast( + ir.VectorType.get([2], T.i32(), loc=loc), + values.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 2, Uint32) + + +@dsl_user_op +def bf16x8_to_packed_u32x4(values, *, loc=None, ip=None): + packed = llvm.bitcast( + ir.VectorType.get([4], T.i32(), loc=loc), + values.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 4, Uint32) + + +@cute.jit +def sanitize_negative_zero_u32(word: Uint32) -> Uint32: + low = Uint16(word & Uint32(0xFFFF)) + high = Uint16(word >> Uint32(16)) + if low == Uint16(NEGATIVE_ZERO_BF16_BITS): + word = word & Uint32(0xFFFF0000) + if high == Uint16(NEGATIVE_ZERO_BF16_BITS): + word = word & Uint32(0x0000FFFF) + return word + + +@cute.jit +def sanitize_negative_zero_u32x4(packed): + sanitized = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32) + for index in cutlass.range_constexpr(4): + sanitized[index] = sanitize_negative_zero_u32(packed[index]) + return sanitized.load() + + +@cute.jit +def sanitize_negative_zero_u32x2(packed): + sanitized = cute.make_rmem_tensor(cute.make_layout((2,)), Uint32) + for index in cutlass.range_constexpr(2): + sanitized[index] = sanitize_negative_zero_u32(packed[index]) + return sanitized.load() + + +@cute.jit +def fragment_has_negative_zero(packed): + dirty = False + for index in cutlass.range_constexpr(4): + word = packed[index] + dirty = ( + dirty + | (Uint16(word & Uint32(0xFFFF)) == Uint16(NEGATIVE_ZERO_BF16_BITS)) + | (Uint16(word >> Uint32(16)) == Uint16(NEGATIVE_ZERO_BF16_BITS)) + ) + return dirty + + +@dsl_user_op +def map_shared_to_peer( + smem_pointer: cute.Pointer, peer_rank: Int32, *, loc=None, ip=None +) -> Int32: + address = smem_pointer.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip) + return Int32( + _asm( + T.i32(), + [address, peer_rank.ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + side_effects=False, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def store_shared_cluster_f32( + remote_address: Int32, value: Float32, *, loc=None, ip=None +) -> None: + _asm( + None, + [remote_address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)], + "st.shared::cluster.f32 [$0], $1;", + "r,f", + side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def f32_to_bf16_bits(value: Float32, *, loc=None, ip=None) -> Uint32: + return Uint32( + _asm( + T.i32(), + [value.ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b16 bits;\n\tcvt.rn.bf16.f32 bits, $1;\n\t" + "cvt.u32.u16 $0, bits;\n\t}", + "=r,f", + side_effects=False, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def shuffle_sync_idx_u32( + value: Uint32, source_lane: Int32, *, loc=None, ip=None +) -> Uint32: + return Uint32( + _asm( + T.i32(), + [value.ir_value(loc=loc, ip=ip), source_lane.ir_value(loc=loc, ip=ip)], + "shfl.sync.idx.b32 $0, $1, $2, 0x1f, 0xffffffff;", + "=r,r,r", + # Preserve full-warp execution across later divergent consumers. + side_effects=True, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def stmc_bf16x2(address: Int64, packed: Uint32, *, loc=None, ip=None) -> None: + _asm( + None, + [address.ir_value(loc=loc, ip=ip), packed.ir_value(loc=loc, ip=ip)], + "multimem.st.relaxed.sys.global.bf16x2 [$0], $1;", + "l,r", + side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def stmc_bf16x4(address: Int64, values, *, loc=None, ip=None) -> None: + words = [values[index].ir_value(loc=loc, ip=ip) for index in range(2)] + _asm( + None, + [address.ir_value(loc=loc, ip=ip), *words], + "multimem.st.relaxed.sys.global.v2.bf16x2 [$0], {$1, $2};", + "l,r,r", + side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def stmc_bf16x8(address: Int64, values, *, loc=None, ip=None) -> None: + words = [values[index].ir_value(loc=loc, ip=ip) for index in range(4)] + _asm( + None, + [address.ir_value(loc=loc, ip=ip), *words], + "multimem.st.relaxed.sys.global.v4.bf16x2 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + side_effects=True, + loc=loc, + ip=ip, + ) + + +class _GraphSafeDLPack: + __slots__ = ("tensor",) + + def __init__(self, tensor: torch.Tensor) -> None: + self.tensor = tensor + + def __dlpack__(self, stream=None): + # stream=-1 skips producer sync; CuTe launches on the current stream, + # which is the capturing one under CUDA graphs. + return self.tensor.__dlpack__(stream=-1) + + def __dlpack_device__(self): + return self.tensor.__dlpack_device__() + + +def to_cute(tensor: torch.Tensor, alignment: int) -> cute.Tensor: + return from_dlpack(_GraphSafeDLPack(tensor.detach()), assumed_align=alignment) + + +def to_cute_dynamic( + tensor: torch.Tensor, alignment: int, *, divisibility: int +) -> cute.Tensor: + return to_cute(tensor, alignment).mark_compact_shape_dynamic( + mode=0, divisibility=divisibility + ) + + +def make_fake_dynamic_compact_tensor(dtype, *, alignment: int, divisibility: int): + return make_fake_compact_tensor( + dtype, (cute.sym_int32(divisibility=divisibility),), assumed_align=alignment + ) + + +def current_cu_stream() -> cuda.CUstream: + return cuda.CUstream(torch.cuda.current_stream().cuda_stream) diff --git a/vllm/models/deepseek_v41/nvidia/ops/mhc.py b/vllm/models/deepseek_v41/nvidia/ops/mhc.py index 88288ead4d..edfdbd1ee4 100644 --- a/vllm/models/deepseek_v41/nvidia/ops/mhc.py +++ b/vllm/models/deepseek_v41/nvidia/ops/mhc.py @@ -9,10 +9,15 @@ import torch from vllm.config import VllmConfig from vllm.distributed import get_tp_group +from vllm.logger import init_logger from vllm.model_executor.kernels.mhc.tilelang import ( mhc_fused_post_pre_delayed_tilelang, mhc_post_tilelang, ) +from vllm.model_executor.layers.fused_moe.moe_output import ( + MoEOutput, + UnfinalizedMoEOutput, +) from vllm.platforms import current_platform from vllm.utils.deep_gemm import is_deep_gemm_supported @@ -24,9 +29,17 @@ from .mega_mhc import ( if TYPE_CHECKING: from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator + from .cute_dsl import AllReduceMHC + +logger = init_logger(__name__) + # GB200 TP4/FlashInfer improves through 16 tokens; larger screens tie or regress. MHC_OVERLAP_MAX_TOKENS = 16 +# The fused all-reduce + mHC kernel of the overlap path, built once by +# init_mhc_all_reduce. +_all_reduce_mhc: "AllReduceMHC | None" = None + def supports_mhc_overlap(vllm_config: VllmConfig) -> bool: """Check kernel requirements and safety of sharing the coefficient stream.""" @@ -53,10 +66,35 @@ def supports_mhc_all_reduce(vllm_config: VllmConfig) -> bool: or config.hc_mult != 4 ): return False + # The custom all-reduce's MNNVL buffers succeed exactly when NVLink + # multicast does, which the fused kernel's own mailbox needs too. comm = cast("CudaCommunicator", get_tp_group().device_communicator).ca_comm return comm is not None and bool(comm.mnnvl_lamport_ag_multicast_ptr) +def init_mhc_all_reduce(vllm_config: VllmConfig) -> None: + """Build the fused all-reduce + mHC kernel the overlap path reduces with. + + Collective over the TP group, so every rank must call it, and only when + ``supports_mhc_all_reduce`` holds. + """ + global _all_reduce_mhc + from .cute_dsl import AllReduceMHC + + config = vllm_config.model_config.hf_config + _all_reduce_mhc = AllReduceMHC( + hidden_size=config.hidden_size, + hc_mult=config.hc_mult, + max_num_tokens=MHC_OVERLAP_MAX_TOKENS, + top_k=config.num_experts_per_tok, + device=current_platform.current_device(), + ) + logger.info_once( + "DSV4.1 mHC: CuTe DSL Lamport all-reduce fused with mHC for up to %d tokens.", + MHC_OVERLAP_MAX_TOKENS, + ) + + def mhc_pre_delayed_overlap( residual: torch.Tensor, fn: torch.Tensor, @@ -151,7 +189,7 @@ def mhc_pre_delayed_overlap( def mhc_shifted_post_pre( - x: torch.Tensor, + x: torch.Tensor | MoEOutput, residual: torch.Tensor, post_layer_mix: torch.Tensor, comb_res_mix: torch.Tensor, @@ -175,39 +213,55 @@ def mhc_shifted_post_pre( ]: """Dispatch shifted post/pre to overlap, Mega-mHC, or fused TileLang. + A MoE output left unfinalized is finalized inside the fused all-reduce. When stream is supplied, join it before consuming the returned coefficients. """ layer_input = None - if reduce_results: - tp = get_tp_group() + if isinstance(x, MoEOutput): + assert _all_reduce_mhc is not None, "init_mhc_all_reduce was not called" + assert pre_mix is not None and norm_weight is not None + routed = x.routed + assert isinstance(routed, UnfinalizedMoEOutput) + assert x.shared_output is not None + residual, layer_input = _all_reduce_mhc.finalize( + routed.gemm2_permuted, + routed.expert_weights, + routed.expanded_idx_to_permuted_idx, + x.shared_output, + residual, + post_layer_mix, + comb_res_mix, + pre_mix, + norm_weight, + norm_eps, + ) + if stream is None: + # Outside FULL graphs, compute the next mixes inline. + stream = torch.cuda.current_stream() + elif reduce_results: 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( + assert _all_reduce_mhc is not None, "init_mhc_all_reduce was not called" + assert pre_mix is not None and norm_weight is not None + residual, layer_input = _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) + x = get_tp_group().all_reduce(x) if stream is not None: if layer_input is None: + assert isinstance(x, torch.Tensor) 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]) + aux = ( + residual.mean(dim=1) + if capture_aux + else residual.new_empty(0, residual.shape[-1]) + ) pre_outputs = mhc_pre_delayed_overlap( residual, fn, @@ -226,6 +280,7 @@ def mhc_shifted_post_pre( ) return residual, *pre_outputs, aux + assert isinstance(x, torch.Tensor) if can_use_mega_mhc(x, residual, pre_mix, norm_weight, capture_aux): assert pre_mix is not None and norm_weight is not None outputs = mhc_shifted_post_pre_deep_gemm(