[Kernel][DSV4.1] Fuse MoE finalize into the TP all-reduce + mHC boundary (#58586)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Yongye Zhu
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
379e9a1ea8
commit
77871126f9
@@ -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}")
|
||||
|
||||
@@ -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 <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));
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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"]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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)
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user