From b9b401b2ef1db3148d480a1520d95c95eff9c528 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Elvir=20Crn=C4=8Devi=C4=87?= Date: Tue, 13 Oct 2026 17:10:38 +0200 Subject: [PATCH] [Kernel] Fuse DeepEP-v2 decode globalize+align into one launch (#53911) Signed-off-by: Elvir Crncevic --- CMakeLists.txt | 3 +- csrc/moe/fused_globalize_align_block_size.cu | 266 ++++++++++++++++++ csrc/moe/moe_ops.h | 10 + csrc/moe/torch_bindings.cpp | 16 ++ .../test_fused_globalize_align_block_size.py | 229 +++++++++++++++ vllm/_custom_ops.py | 26 ++ .../fused_moe/experts/fused_humming_moe.py | 50 +++- .../fused_globalize_align_block_size.py | 63 +++++ .../layers/fused_moe/modular_kernel.py | 15 + .../fused_moe/prepare_finalize/deepep_v2.py | 24 +- 10 files changed, 686 insertions(+), 16 deletions(-) create mode 100644 csrc/moe/fused_globalize_align_block_size.cu create mode 100644 tests/kernels/moe/test_fused_globalize_align_block_size.py create mode 100644 vllm/model_executor/layers/fused_moe/fused_globalize_align_block_size.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 3a71d7e1e5..e27a031f26 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1287,7 +1287,8 @@ set(VLLM_MOE_EXT_SRC if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_MOE_EXT_SRC "csrc/moe/moe_wna16.cu" - "csrc/moe/grouped_topk_kernels.cu") + "csrc/moe/grouped_topk_kernels.cu" + "csrc/moe/fused_globalize_align_block_size.cu") endif() if(VLLM_GPU_LANG STREQUAL "CUDA") diff --git a/csrc/moe/fused_globalize_align_block_size.cu b/csrc/moe/fused_globalize_align_block_size.cu new file mode 100644 index 0000000000..2e0429b8a6 --- /dev/null +++ b/csrc/moe/fused_globalize_align_block_size.cu @@ -0,0 +1,266 @@ +// Single-launch fusion of the DeepEP-v2 decode metadata chain +// (_globalize_recv_topk_idx + moe_align_block_size + count_and_sort) for the +// humming INDEXED / CUDA-graph decode path. recv values are LOCAL expert ids; +// the kernel bins by local id while storing back the global id, sized by local +// experts. Hybrid launch: one non-cooperative block for small decode, else a +// cooperative multi-SM grid (launched via cudaLaunchKernelEx so grid.sync() is +// CUDA-graph capturable). Barriers: zero+sentinel -> shared histogram +// (distributed reduce) -> block-0 prefix + expert_ids fill -> +// scatter+globalize. + +#include +#include + +#include +#include +#include +#include +#include + +#include "torch_utils.h" + +namespace cg = cooperative_groups; + +namespace vllm { +namespace moe { + +static constexpr int FGAS_THREADS = 256; +static constexpr long FGAS_SB_MAX_EFFECTIVE_WORK = 4096; + +// COOP selects the phase barrier: grid.sync() (multi-block) vs __syncthreads(). +// topk / block_size are runtime args: any (topk, block_size) is supported, and +// the kernel is barrier/memory-bound so the divisions are not on the hot path. +template +__device__ __forceinline__ void fused_gas_body( + long* __restrict__ topk_idx, int* __restrict__ sorted_ids, + int* __restrict__ expert_ids, int* __restrict__ num_tokens_post_pad, + int* __restrict__ counts, int* __restrict__ scatter_idx, int num_recv, + int rank_expert_offset, int global_num_experts, int numel, + int max_num_tokens_padded, int max_num_m_blocks, int local_num_experts, + int topk, int block_size) { + extern __shared__ int sh[]; // [local_num_experts] per-block histogram + const int tid = blockIdx.x * blockDim.x + threadIdx.x; + const int stride = gridDim.x * blockDim.x; + auto BAR = [&]() { + if constexpr (COOP) + cg::this_grid().sync(); + else + __syncthreads(); + }; + + // Phase 0: zero the histogram + sentinel-fill sorted_ids. + for (int e = tid; e < local_num_experts; e += stride) counts[e] = 0; + for (int i = tid; i < max_num_tokens_padded; i += stride) + sorted_ids[i] = numel; + BAR(); + + // Phase 1: per-block shared histogram, one global atomic per (block, expert). + // Reduce stays distributed across blocks (centralizing it in block 0 is + // slower). + for (int e = threadIdx.x; e < local_num_experts; e += blockDim.x) sh[e] = 0; + __syncthreads(); + for (int i = tid; i < numel; i += stride) { + int val = (int)topk_idx[i]; + int row = i / topk; + bool valid = (val >= 0) && (val < local_num_experts) && + (val + rank_expert_offset < global_num_experts) && + (row < num_recv); + if (valid) atomicAdd(&sh[val], 1); + } + __syncthreads(); + for (int e = threadIdx.x; e < local_num_experts; e += blockDim.x) + if (sh[e]) atomicAdd(&counts[e], sh[e]); + BAR(); + + // Phase 2 (block 0). Thread 0 runs the tiny serial prefix scan; then all of + // block 0's threads split the expert_ids fill grid-stride. + if (blockIdx.x == 0) { + int* sh_start = sh; // [local_num_experts] padded start offsets + int& sh_run = sh[local_num_experts]; + if (threadIdx.x == 0) { + int run = 0; + for (int e = 0; e < local_num_experts; e++) { + int c = counts[e]; + int nb = (c + block_size - 1) / block_size; + sh_start[e] = run; + counts[e] = run; // exclusive prefix; empty experts skip + run += nb * block_size; + } + sh_run = run; + *num_tokens_post_pad = run; + } + __syncthreads(); + int run = sh_run; + // owner(j) = highest expert whose padded start <= j*block_size (skips empty + // experts); blocks past the last token get -1. + for (int j = threadIdx.x; j < max_num_m_blocks; j += blockDim.x) { + int off = j * block_size; + if (off >= run) { + expert_ids[j] = -1; + continue; + } + int owner = 0; + for (int e = 1; e < local_num_experts; e++) + if (sh_start[e] <= off) owner = e; + expert_ids[j] = owner; + } + } + BAR(); + + // Phase 3: scatter token indices at cursor[val]++ and globalize in place. + for (int i = tid; i < numel; i += stride) { + int val = (int)topk_idx[i]; + int g = val + rank_expert_offset; + int row = i / topk; + bool valid = (val >= 0) && (val < local_num_experts) && + (g < global_num_experts) && (row < num_recv); + if (scatter_idx != nullptr) scatter_idx[i] = valid ? i : -1; + if (valid) { + int pos = atomicAdd(&counts[val], 1); + sorted_ids[pos] = i; + } + topk_idx[i] = valid ? g : -1; + } +} + +template +__global__ void fused_gas_kernel( + long* __restrict__ topk_idx, const int* __restrict__ psum, + int* __restrict__ sorted_ids, int* __restrict__ expert_ids, + int* __restrict__ num_tokens_post_pad, int* __restrict__ counts, + int* __restrict__ scatter_idx, int P, int rank_expert_offset, + int global_num_experts, int numel, int max_num_tokens_padded, + int max_num_m_blocks, int local_num_experts, int topk, int block_size) { + // num_recv read on-device (psum[P-1]) -> baked into replay, cudagraph-safe. + fused_gas_body(topk_idx, sorted_ids, expert_ids, num_tokens_post_pad, + counts, scatter_idx, psum[P - 1], rank_expert_offset, + global_num_experts, numel, max_num_tokens_padded, + max_num_m_blocks, local_num_experts, topk, block_size); +} + +static int fgas_sm_count() { + // Device-constant: query the driver once, never on the capture/replay path. + static int sm_count = -1; + if (sm_count < 0) { + int dev = 0; + cudaGetDevice(&dev); + cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev); + if (sm_count < 1) sm_count = 1; + } + return sm_count; +} + +static int fgas_cooperative_blocks(int bps, long work) { + int max_blocks = bps * fgas_sm_count(); + if (max_blocks < 1) max_blocks = 1; + long need = (work + FGAS_THREADS - 1) / FGAS_THREADS; // no over-launch + if (need < 1) need = 1; + return need < (long)max_blocks ? (int)need : max_blocks; +} + +static void fgas_launch(long* p_topk, const int* p_psum, int* p_sorted, + int* p_expert, int* p_num, int* p_counts, + int* p_scatter, int P, int reo, int gne, int numel, + int mntp, int mnmb, int local_e, int topk, + int block_size, cudaStream_t stream) { + long work = numel > mntp ? numel : mntp; + long effective_work = (long)mntp + 2L * numel + (long)mnmb * local_e; + // Histogram [local_e] + one scalar (phase-2 padded-run total). + size_t smem = (size_t)(local_e + 1) * sizeof(int); + + // Small decode: one non-cooperative block, __syncthreads barriers. Account + // for both numel passes and the expert_ids owner scan in the crossover. + if (effective_work <= FGAS_SB_MAX_EFFECTIVE_WORK) { + fused_gas_kernel<<<1, FGAS_THREADS, smem, stream>>>( + p_topk, p_psum, p_sorted, p_expert, p_num, p_counts, p_scatter, P, reo, + gne, numel, mntp, mnmb, local_e, topk, block_size); + STD_CUDA_CHECK(cudaGetLastError()); + return; + } + + // Large decode: cooperative multi-SM. Occupancy (fixed per rank) queried + // once. + static int bps = -1; + if (bps < 0) + cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &bps, (const void*)fused_gas_kernel, FGAS_THREADS, smem); + int blocks = fgas_cooperative_blocks(bps, work); + + // Cooperative launch via cudaLaunchKernelEx + the cooperative launch + // attribute: unlike the legacy cudaLaunchCooperativeKernel, this form IS + // stream-capturable, so grid.sync() works inside a captured CUDA graph. + cudaLaunchConfig_t config = {}; + config.gridDim = dim3(blocks); + config.blockDim = dim3(FGAS_THREADS); + config.dynamicSmemBytes = smem; + config.stream = stream; + cudaLaunchAttribute attr = {}; + attr.id = cudaLaunchAttributeCooperative; + attr.val.cooperative = 1; + config.attrs = &attr; + config.numAttrs = 1; + STD_CUDA_CHECK(cudaLaunchKernelEx(&config, fused_gas_kernel, p_topk, + p_psum, p_sorted, p_expert, p_num, p_counts, + p_scatter, P, reo, gne, numel, mntp, mnmb, + local_e, topk, block_size)); +} + +} // namespace moe +} // namespace vllm + +// recv_topk_idx (int64, [N, topk]) is globalized in place and the align/sort +// outputs (int32) are filled. counts is an internal write-cursor allocated +// here. +void fused_globalize_align_block_size( + torch::stable::Tensor topk_idx, torch::stable::Tensor psum, + int64_t rank_expert_offset, int64_t global_num_experts, + int64_t local_num_experts, int64_t block_size, + torch::stable::Tensor sorted_ids, torch::stable::Tensor expert_ids, + torch::stable::Tensor num_tokens_post_pad, + std::optional scatter_idx) { + STD_TORCH_CHECK(topk_idx.scalar_type() == torch::headeronly::ScalarType::Long, + "fused_globalize_align_block_size: topk_idx must be int64"); + STD_TORCH_CHECK( + local_num_experts <= 1024, + "fused_globalize_align_block_size: local_num_experts <= 1024"); + + const torch::stable::accelerator::DeviceGuard device_guard( + topk_idx.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(topk_idx.get_device_index()); + + int numel = (int)topk_idx.numel(); + int topk = (int)topk_idx.size(1); + int mntp = (int)sorted_ids.size(0); + int mnmb = (int)expert_ids.size(0); + int local_e = (int)local_num_experts; + int reo = (int)rank_expert_offset; + int gne = (int)global_num_experts; + int P = (int)psum.size(0); + + // Internal per-local-expert histogram / write-cursor. + torch::stable::Tensor counts = torch::stable::new_empty( + topk_idx, {local_num_experts}, torch::headeronly::ScalarType::Int); + + long* p_topk = reinterpret_cast(topk_idx.mutable_data_ptr()); + const int* p_psum = reinterpret_cast(psum.const_data_ptr()); + int* p_sorted = reinterpret_cast(sorted_ids.mutable_data_ptr()); + int* p_expert = reinterpret_cast(expert_ids.mutable_data_ptr()); + int* p_num = reinterpret_cast(num_tokens_post_pad.mutable_data_ptr()); + int* p_counts = reinterpret_cast(counts.mutable_data_ptr()); + int* p_scatter = nullptr; + if (scatter_idx.has_value()) { + STD_TORCH_CHECK( + scatter_idx->scalar_type() == torch::headeronly::ScalarType::Int && + scatter_idx->device() == topk_idx.device() && + scatter_idx->is_contiguous() && + scatter_idx->numel() == topk_idx.numel(), + "scatter_idx must be contiguous int32 on the same device and match " + "topk_idx.numel()"); + p_scatter = reinterpret_cast(scatter_idx->mutable_data_ptr()); + } + + vllm::moe::fgas_launch(p_topk, p_psum, p_sorted, p_expert, p_num, p_counts, + p_scatter, P, reo, gne, numel, mntp, mnmb, local_e, + topk, (int)block_size, stream); +} diff --git a/csrc/moe/moe_ops.h b/csrc/moe/moe_ops.h index 549a19e981..8ab24fcb33 100644 --- a/csrc/moe/moe_ops.h +++ b/csrc/moe/moe_ops.h @@ -44,6 +44,16 @@ void moe_align_block_size(torch::stable::Tensor topk_ids, int64_t num_experts, std::optional maybe_expert_map, std::optional scatter_idx); +#ifdef USE_CUDA +void fused_globalize_align_block_size( + torch::stable::Tensor topk_idx, torch::stable::Tensor psum, + int64_t rank_expert_offset, int64_t global_num_experts, + int64_t local_num_experts, int64_t block_size, + torch::stable::Tensor sorted_ids, torch::stable::Tensor expert_ids, + torch::stable::Tensor num_tokens_post_pad, + std::optional scatter_idx); +#endif + void batched_moe_align_block_size( int64_t max_tokens_per_batch, int64_t block_size, const torch::stable::Tensor& expert_num_tokens, diff --git a/csrc/moe/torch_bindings.cpp b/csrc/moe/torch_bindings.cpp index 2feec38389..9e8907afd7 100644 --- a/csrc/moe/torch_bindings.cpp +++ b/csrc/moe/torch_bindings.cpp @@ -42,6 +42,18 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) { " Tensor? maybe_expert_map," " Tensor(a!)? scatter_idx=None) -> ()"); +#ifdef USE_CUDA + // Fused globalize + align + sort for the DeepEP-v2 humming INDEXED decode + // path (one launch; num_recv read on device from psum[-1], cudagraph-safe). + m.def( + "fused_globalize_align_block_size(Tensor! topk_idx, Tensor psum," + " int rank_expert_offset, int global_num_experts," + " int local_num_experts, int block_size," + " Tensor! sorted_ids, Tensor! expert_ids," + " Tensor! num_tokens_post_pad," + " Tensor(a!)? scatter_idx=None) -> ()"); +#endif + // Aligning the number of tokens to be processed by each expert such // that it is divisible by the block size, but for the batched case. m.def( @@ -143,6 +155,10 @@ STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) { m.impl("topk_softplus_sqrt", TORCH_BOX(&topk_softplus_sqrt)); m.impl("moe_sum", TORCH_BOX(&moe_sum)); m.impl("moe_align_block_size", TORCH_BOX(&moe_align_block_size)); +#ifdef USE_CUDA + m.impl("fused_globalize_align_block_size", + TORCH_BOX(&fused_globalize_align_block_size)); +#endif m.impl("batched_moe_align_block_size", TORCH_BOX(&batched_moe_align_block_size)); m.impl("moe_lora_align_block_size", TORCH_BOX(&moe_lora_align_block_size)); diff --git a/tests/kernels/moe/test_fused_globalize_align_block_size.py b/tests/kernels/moe/test_fused_globalize_align_block_size.py new file mode 100644 index 0000000000..d8abf27d65 --- /dev/null +++ b/tests/kernels/moe/test_fused_globalize_align_block_size.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness of the fused globalize + align + sort decode metadata kernel. + +The fused op replaces the DeepEP-v2 INDEXED cudagraph-decode chain +(_globalize_recv_topk_idx + moe_align_block_size + count_and_sort). It is +checked against that exact reference chain over a grid of shapes / EP ranks / +fill distributions. +""" + +import pytest +import torch + +from vllm.model_executor.layers.fused_moe.fused_globalize_align_block_size import ( + fused_globalize_align_block_size, +) +from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( + moe_align_block_size, +) +from vllm.platforms import current_platform +from vllm.utils.import_utils import has_deep_ep_v2 + +if current_platform.is_cuda() and has_deep_ep_v2(): + from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 import ( + _globalize_recv_topk_idx, + ) +else: + pytest.skip("Requires CUDA and DeepEP v2", allow_module_level=True) + + +def _expert_map(global_num_experts, rank_expert_offset, local_num_experts, device): + expert_map = torch.full((global_num_experts,), -1, device=device, dtype=torch.int32) + rng = torch.arange(local_num_experts, device=device, dtype=torch.int32) + expert_map[rank_expert_offset : rank_expert_offset + local_num_experts] = rng + return expert_map + + +def _make_case(n, topk, ep_size, local_e, block_size, seed, frac_recv, rank, fill): + torch.manual_seed(seed) + device = "cuda" + global_num_experts = local_e * ep_size + rank_idx = {"first": 0, "mid": ep_size // 2, "last": ep_size - 1}[rank] + rank_expert_offset = rank_idx * local_e + + if fill == "all_nonlocal": + local_idx = torch.full((n, topk), -1, device=device, dtype=torch.int64) + elif fill == "single": + local_idx = torch.zeros((n, topk), device=device, dtype=torch.int64) + elif fill == "all_local": + local_idx = torch.randint( + 0, local_e, (n, topk), device=device, dtype=torch.int64 + ) + else: + local_idx = torch.randint( + -1, local_e, (n, topk), device=device, dtype=torch.int64 + ) + + num_recv = max(0, min(n, int(round(n * frac_recv)))) + if num_recv < n: + local_idx[num_recv:] = torch.randint( + -1, local_e, (n - num_recv, topk), device=device, dtype=torch.int64 + ) + psum = torch.zeros(4, device=device, dtype=torch.int32) + psum[-1] = num_recv + return local_idx, psum, rank_expert_offset, global_num_experts + + +def _group_by_expert(sorted_ids, expert_ids, block_size, valid_len, total): + groups: dict[int, list[int]] = {} + for b in range(valid_len // block_size): + eid = int(expert_ids[b].item()) + blk = sorted_ids[b * block_size : (b + 1) * block_size] + groups.setdefault(eid, []).extend(blk[blk < total].tolist()) + return groups + + +@pytest.mark.parametrize("n", [1, 7, 64, 257]) +# topk=16 with local_num_experts=28 is the Kimi K3 decode config. topk=6 and +# block_size=16, 128 are outside the old template dispatch table: they guard +# against a regression to fixed (topk, block_size) specialization. +@pytest.mark.parametrize("topk", [1, 6, 8, 16]) +@pytest.mark.parametrize("block_size", [16, 32, 64, 128]) +@pytest.mark.parametrize("local_e", [28, 32]) +@pytest.mark.parametrize("rank", ["first", "mid", "last"]) +@pytest.mark.parametrize("fill", ["rand", "all_local", "all_nonlocal", "single"]) +def test_fused_matches_reference(n, topk, block_size, local_e, rank, fill): + ep_size = 4 + local_idx, psum, reo, gne = _make_case( + n, + topk, + ep_size, + local_e, + block_size, + seed=n * 131 + topk, + frac_recv=0.75, + rank=rank, + fill=fill, + ) + + ref_global = _globalize_recv_topk_idx(local_idx.clone(), psum, reo, gne) + ref_sorted, ref_expert_ids, ref_num = moe_align_block_size( + topk_ids=ref_global, + block_size=block_size, + num_experts=gne, + expert_map=_expert_map(gne, reo, local_e, local_idx.device), + ignore_invalid_experts=True, + ) + + fused_global, fused_sorted, fused_expert_ids, fused_num = ( + fused_globalize_align_block_size( + local_idx.clone(), psum, reo, gne, local_e, block_size + ) + ) + + total = n * topk + assert torch.equal(fused_global, ref_global) + assert torch.equal(fused_num, ref_num) + + valid = int(fused_num.item()) + nb = valid // block_size + assert torch.equal(fused_expert_ids[:nb], ref_expert_ids[:nb]) + assert bool((fused_expert_ids[nb:] == -1).all()) + assert bool((fused_sorted[valid:] == total).all()) + + ref_groups = _group_by_expert( + ref_sorted, ref_expert_ids, block_size, int(ref_num.item()), total + ) + fused_groups = _group_by_expert( + fused_sorted, fused_expert_ids, block_size, valid, total + ) + assert set(ref_groups) == set(fused_groups) + for eid in ref_groups: + assert sorted(ref_groups[eid]) == sorted(fused_groups[eid]) + + +@pytest.mark.parametrize("n", [7, 257]) +def test_fused_scatter_idx_matches_reference(n): + """The optional scatter output masks invalid and padded routes.""" + local_e = 28 + local_idx, psum, reo, gne = _make_case( + n, 16, 4, local_e, 64, seed=n, frac_recv=0.75, rank="mid", fill="rand" + ) + ref_global = _globalize_recv_topk_idx(local_idx.clone(), psum, reo, gne) + *_, ref_scatter = moe_align_block_size( + topk_ids=ref_global, + block_size=64, + num_experts=gne, + expert_map=_expert_map(gne, reo, local_e, local_idx.device), + ignore_invalid_experts=True, + return_scatter_idx=True, + ) + + fused_global, *_, fused_scatter = fused_globalize_align_block_size( + local_idx.clone(), psum, reo, gne, local_e, 64, return_scatter_idx=True + ) + assert torch.equal(fused_global, ref_global) + assert torch.equal(fused_scatter, ref_scatter) + + +# These have effective_work 1504 (single block) and 6068 (cooperative). +@pytest.mark.parametrize("n", [1, 64]) +def test_fused_cuda_graph_replay(n): + """Both launch paths remain capturable and consume replay-time inputs.""" + topk = 16 + local_e = 28 + block_size = 64 + ep_size = 4 + + local_idx, psum, reo, gne = _make_case( + n, + topk, + ep_size, + local_e, + block_size, + seed=n, + frac_recv=0.75, + rank="mid", + fill="rand", + ) + graph_idx = torch.empty_like(local_idx) + graph_psum = torch.empty_like(psum) + + # Initialize lazy CUDA state before capture. + fused_globalize_align_block_size( + local_idx.clone(), psum, reo, gne, local_e, block_size + ) + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_idx.copy_(local_idx) + graph_psum.copy_(psum) + graph_outputs = fused_globalize_align_block_size( + graph_idx, graph_psum, reo, gne, local_e, block_size + ) + + for frac_recv in (0.75, 0.5): + psum[-1] = int(round(n * frac_recv)) + ref_idx = _globalize_recv_topk_idx(local_idx.clone(), psum, reo, gne) + ref_sorted, ref_expert_ids, ref_num = moe_align_block_size( + topk_ids=ref_idx, + block_size=block_size, + num_experts=gne, + expert_map=_expert_map(gne, reo, local_e, local_idx.device), + ignore_invalid_experts=True, + ) + + graph.replay() + torch.accelerator.synchronize() + fused_idx, fused_sorted, fused_expert_ids, fused_num = graph_outputs + + assert torch.equal(fused_idx, ref_idx) + assert torch.equal(fused_num, ref_num) + valid = int(fused_num.item()) + num_blocks = valid // block_size + assert torch.equal(fused_expert_ids[:num_blocks], ref_expert_ids[:num_blocks]) + fused_groups = _group_by_expert( + fused_sorted, fused_expert_ids, block_size, valid, n * topk + ) + ref_groups = _group_by_expert( + ref_sorted, + ref_expert_ids, + block_size, + int(ref_num.item()), + n * topk, + ) + assert set(fused_groups) == set(ref_groups) + for eid in ref_groups: + assert sorted(fused_groups[eid]) == sorted(ref_groups[eid]) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index e786207308..1af864bc74 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2195,6 +2195,32 @@ def moe_align_block_size( ) +def fused_globalize_align_block_size( + topk_idx: torch.Tensor, + psum: torch.Tensor, + rank_expert_offset: int, + global_num_experts: int, + local_num_experts: int, + block_size: int, + sorted_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_pad: torch.Tensor, + scatter_idx: torch.Tensor | None = None, +) -> None: + torch.ops._moe_C.fused_globalize_align_block_size( + topk_idx, + psum, + rank_expert_offset, + global_num_experts, + local_num_experts, + block_size, + sorted_ids, + expert_ids, + num_tokens_post_pad, + scatter_idx, + ) + + def batched_moe_align_block_size( max_tokens_per_batch: int, block_size: int, diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py index f27171ef0d..efc2b3f6c8 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py @@ -22,6 +22,9 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEParallelConfig, FusedMoEQuantConfig, ) +from vllm.model_executor.layers.fused_moe.fused_globalize_align_block_size import ( + fused_globalize_align_block_size, +) from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( moe_align_block_size, ) @@ -710,16 +713,45 @@ class HummingIndexedExperts(HummingExpertsBase): num_sms = num_compute_units(topk_ids.get_device()) use_scatter = topk_ids.numel() > 2 * num_sms - alignment = moe_align_block_size( - topk_ids=topk_ids, - block_size=moe_block_size, - num_experts=self.global_num_experts, - expert_map=expert_map, - ignore_invalid_experts=True, - return_scatter_idx=use_scatter, + # When DeepEP-v2 deferred globalization, fuse globalize + align + sort + # into one launch (topk_ids globalized in place); else align only. + rank_expert_offset = ( + expert_tokens_meta.rank_expert_offset + if expert_tokens_meta is not None + else None ) - sorted_ids, expert_ids, num_tokens_padded = alignment[:3] - scatter_idx = alignment[3] if len(alignment) == 4 else None + psum = ( + expert_tokens_meta.psum_recv_per_rank + if expert_tokens_meta is not None + else None + ) + if ( + current_platform.is_cuda() + and rank_expert_offset is not None + and psum is not None + ): + alignment = fused_globalize_align_block_size( + recv_topk_idx=topk_ids, + psum_recv_per_rank=psum, + rank_expert_offset=rank_expert_offset, + global_num_experts=self.global_num_experts, + local_num_experts=self.num_experts, + block_size=moe_block_size, + return_scatter_idx=use_scatter, + ) + _, sorted_ids, expert_ids, num_tokens_padded = alignment[:4] + scatter_idx = alignment[4] if use_scatter else None + else: + alignment = moe_align_block_size( + topk_ids=topk_ids, + block_size=moe_block_size, + num_experts=self.global_num_experts, + expert_map=expert_map, + ignore_invalid_experts=True, + return_scatter_idx=use_scatter, + ) + sorted_ids, expert_ids, num_tokens_padded = alignment[:3] + scatter_idx = alignment[3] if len(alignment) == 4 else None moe_common_kwargs = { "sorted_ids": sorted_ids, diff --git a/vllm/model_executor/layers/fused_moe/fused_globalize_align_block_size.py b/vllm/model_executor/layers/fused_moe/fused_globalize_align_block_size.py new file mode 100644 index 0000000000..b7623141a8 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/fused_globalize_align_block_size.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm import _custom_ops as ops +from vllm.triton_utils import triton + + +def fused_globalize_align_block_size( + recv_topk_idx: torch.Tensor, + psum_recv_per_rank: torch.Tensor, + rank_expert_offset: int, + global_num_experts: int, + local_num_experts: int, + block_size: int, + return_scatter_idx: bool = False, +) -> tuple[torch.Tensor, ...]: + """Fuse the DeepEP-v2 INDEXED decode metadata chain into one launch. + + Replaces ``_globalize_recv_topk_idx`` + ``moe_align_block_size`` + + ``count_and_sort_expert_tokens``. ``recv_topk_idx`` holds LOCAL expert ids + and is globalized in place; valid-token count is read on-device from + ``psum_recv_per_rank[-1]`` (CUDA-graph safe). + + Returns the globalized ``recv_topk_idx`` and the ``moe_align_block_size`` + outputs ``(sorted_ids, expert_ids, num_tokens_post_pad)``. When requested, + the final output is a masked identity scatter index for valid routes. + """ + assert recv_topk_idx.dtype == torch.int64 + _, topk = recv_topk_idx.shape + numel = recv_topk_idx.numel() + device = recv_topk_idx.device + + max_num_tokens_padded = numel + local_num_experts * (block_size - 1) + if numel < local_num_experts: + max_num_tokens_padded = min(numel * block_size, max_num_tokens_padded) + max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size) + + sorted_ids = torch.empty((max_num_tokens_padded,), dtype=torch.int32, device=device) + expert_ids = torch.empty((max_num_m_blocks,), dtype=torch.int32, device=device) + num_tokens_post_pad = torch.empty((1,), dtype=torch.int32, device=device) + scatter_idx = ( + torch.empty((numel, 1), dtype=torch.int32, device=device) + if return_scatter_idx + else None + ) + + ops.fused_globalize_align_block_size( + recv_topk_idx, + psum_recv_per_rank, + rank_expert_offset, + global_num_experts, + local_num_experts, + block_size, + sorted_ids, + expert_ids, + num_tokens_post_pad, + scatter_idx, + ) + + outputs = (recv_topk_idx, sorted_ids, expert_ids, num_tokens_post_pad) + return (*outputs, scatter_idx) if scatter_idx is not None else outputs diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index d9a75c55ec..158e144e72 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -103,6 +103,9 @@ class ExpertTokensMetadata: expert_num_tokens: torch.Tensor | None expert_num_tokens_cpu: torch.Tensor | None psum_recv_per_rank: torch.Tensor | None = None + # local->global expert-id offset, set when globalization is deferred to the + # experts' fused align kernel. + rank_expert_offset: int | None = None @staticmethod def make_from_list( @@ -1048,6 +1051,18 @@ class FusedMoEKernelModularImpl: ): self.prepare_finalize = prepare_finalize self.fused_experts = fused_experts + # Experts that globalize recv_topk_idx themselves let a capable prepare + # stage skip its standalone launch. + from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( + HummingIndexedExperts, + ) + + if ( + current_platform.is_cuda() + and isinstance(fused_experts, HummingIndexedExperts) + and hasattr(prepare_finalize, "defer_globalize") + ): + prepare_finalize.defer_globalize = True self.shared_experts: SharedExperts | None = None moe_parallel_config = fused_experts.moe_config.moe_parallel_config self.moe_parallel_config = moe_parallel_config diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py index a584ab5f52..59614a4c74 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py @@ -127,6 +127,11 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): self.use_cudagraph = use_cudagraph self.sp_size = sp_size + # Set by FusedMoEKernelModularImpl when the experts globalize + # recv_topk_idx themselves, letting the decode path skip the standalone + # _globalize_recv_topk_idx launch. + self.defer_globalize = False + # DBO microbatching: one handle slot per micro-batch. self.handles: list[deep_ep.EPHandle | None] = [None, None] @@ -275,6 +280,7 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): # during profile_run when CUDA graphs are enabled. expert_tokens_meta = None + globalize_deferred = False if recv_topk_idx is None: # do_expand=True (prefill mode): build topk_ids from # per-expert token counts. @@ -307,12 +313,15 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): # be treated as real routed tokens by experts that build routing # over *all* rows (e.g. triton MoE backend's make_routing_data), # polluting the per-expert token lists and corrupting real tokens. - recv_topk_idx = _globalize_recv_topk_idx( - recv_topk_idx, - psum_recv_per_rank, - self.rank_expert_offset, - self.num_experts, - ) + if self.defer_globalize and self.use_cudagraph: + globalize_deferred = True + else: + recv_topk_idx = _globalize_recv_topk_idx( + recv_topk_idx, + psum_recv_per_rank, + self.rank_expert_offset, + self.num_experts, + ) # Reshape recv_topk_weights to match recv_topk_idx shape [N, 1] if recv_topk_weights is not None and recv_topk_weights.ndim == 1: @@ -329,6 +338,9 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): expert_num_tokens_cpu=None, ) expert_tokens_meta.psum_recv_per_rank = psum_recv_per_rank + if globalize_deferred: + # recv_topk_idx still holds LOCAL ids; experts globalize it. + expert_tokens_meta.rank_expert_offset = self.rank_expert_offset if _quantize_before_dispatch(quant_config, defer_input_quant): if quant_config.quant_dtype == "mxfp8" and expert_x_scale is not None: