iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Woosuk Kwon
2026-09-20 20:28:06 -07:00
committed by GitHub
co-authored by Codex
parent eb87980585
commit 9a70c233cd
9 changed files with 549 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CMakeLists.txt
+2
View File
@@ -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}")
csrc/libtorch_stable/all_reduce_mhc.cu
+212
View File
@@ -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 <cooperative_groups.h>
#include <torch/csrc/stable/library.h>
namespace {
using Tensor = torch::stable::Tensor;
using Pack = vllm::packed_t<nv_bfloat16>::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>();
}
Pack peers[kRanks];
bool pending;
do {
vllm::LamportPack<Pack> 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<Pack>(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<Pack>(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<Pack>(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<Pack>(collapsed);
if (idx == 0) {
while (*reinterpret_cast<volatile uint32_t*>(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<uintptr_t>(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<const Pack*>(input.const_data_ptr()),
static_cast<const Pack*>(residual.const_data_ptr()),
static_cast<const float*>(post.const_data_ptr()),
static_cast<const float*>(comb.const_data_ptr()),
static_cast<const float*>(pre.const_data_ptr()),
static_cast<const Pack*>(weight.const_data_ptr()),
static_cast<Pack*>(output.mutable_data_ptr()),
static_cast<Pack*>(normalized.mutable_data_ptr()),
reinterpret_cast<Pack*>(local), reinterpret_cast<Pack*>(multicast),
static_cast<uint32_t*>(epochs.mutable_data_ptr()), static_cast<int>(rank),
static_cast<int>(stage_bytes / sizeof(Pack)), static_cast<float>(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));
}
tests/distributed/test_custom_all_reduce.py
+111
View File
@@ -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"),
[
tests/distributed/test_engram_dp_shard.py
+1
View File
@@ -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,
tests/kernels/test_mhc_kernels.py
+45 -1
View File
@@ -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
vllm/models/deepseek_v4/nvidia/model.py
+3
View File
@@ -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,
vllm/models/deepseek_v41/nvidia/model.py
+19 -2
View File
@@ -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
)
vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py
-98
View File
@@ -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,
)
vllm/models/deepseek_v41/nvidia/ops/mhc.py
+156 -4
View File
@@ -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,
)