[Kernel] Fuse DeepEP-v2 decode globalize+align into one launch (#53911)
Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>
This commit is contained in:
Elvir Crnčević
+2
-1
@@ -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")
|
||||
|
||||
@@ -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 <cooperative_groups.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <torch/csrc/stable/macros.h>
|
||||
#include <torch/csrc/stable/accelerator.h>
|
||||
#include <torch/csrc/stable/ops.h>
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
#include <torch/headeronly/core/ScalarType.h>
|
||||
|
||||
#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 <bool COOP>
|
||||
__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 <bool COOP>
|
||||
__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<COOP>(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<false><<<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<true>, 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<true>, 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<torch::stable::Tensor> 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<long*>(topk_idx.mutable_data_ptr());
|
||||
const int* p_psum = reinterpret_cast<const int*>(psum.const_data_ptr());
|
||||
int* p_sorted = reinterpret_cast<int*>(sorted_ids.mutable_data_ptr());
|
||||
int* p_expert = reinterpret_cast<int*>(expert_ids.mutable_data_ptr());
|
||||
int* p_num = reinterpret_cast<int*>(num_tokens_post_pad.mutable_data_ptr());
|
||||
int* p_counts = reinterpret_cast<int*>(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<int*>(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);
|
||||
}
|
||||
@@ -44,6 +44,16 @@ void moe_align_block_size(torch::stable::Tensor topk_ids, int64_t num_experts,
|
||||
std::optional<torch::stable::Tensor> maybe_expert_map,
|
||||
std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> 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,
|
||||
|
||||
@@ -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));
|
||||
|
||||
@@ -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])
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user