[Bugfix] Handle padded routes in CUTLASS MoE permutations (#54747)
Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Codex <codex@openai.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
This commit is contained in:
Kevin H. Luu
co-authored by
Codex
Yongye Zhu
Codex
Yongye Zhu
parent
003e34341a
commit
46c8a161f4
@@ -212,23 +212,30 @@ __global__ void shuffleInputRowsKernel(const T* input,
|
||||
int64_t num_src_rows,
|
||||
int64_t num_dst_rows, int64_t num_cols) {
|
||||
int64_t dest_row_idx = blockIdx.x;
|
||||
int64_t const source_row_idx = dst2src_map[dest_row_idx];
|
||||
|
||||
if (blockIdx.x < num_dst_rows) {
|
||||
int64_t const source_row_idx = dst2src_map[dest_row_idx];
|
||||
// Load 128-bits per thread
|
||||
constexpr int64_t ELEM_PER_THREAD = 128 / sizeof(T) / 8;
|
||||
using DataElem = cutlass::Array<T, ELEM_PER_THREAD>;
|
||||
|
||||
// Duplicate and permute rows
|
||||
auto const* source_row_ptr =
|
||||
reinterpret_cast<DataElem const*>(input + source_row_idx * num_cols);
|
||||
auto* dest_row_ptr =
|
||||
reinterpret_cast<DataElem*>(output + dest_row_idx * num_cols);
|
||||
|
||||
int64_t const start_offset = threadIdx.x;
|
||||
int64_t const stride = blockDim.x;
|
||||
int64_t const num_elems_in_col = num_cols / ELEM_PER_THREAD;
|
||||
|
||||
if (source_row_idx < 0 || source_row_idx >= num_src_rows) {
|
||||
DataElem const zero{};
|
||||
for (int elem_index = start_offset; elem_index < num_elems_in_col;
|
||||
elem_index += stride) {
|
||||
dest_row_ptr[elem_index] = zero;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// Duplicate and permute rows
|
||||
auto const* source_row_ptr =
|
||||
reinterpret_cast<DataElem const*>(input + source_row_idx * num_cols);
|
||||
for (int elem_index = start_offset; elem_index < num_elems_in_col;
|
||||
elem_index += stride) {
|
||||
dest_row_ptr[elem_index] = source_row_ptr[elem_index];
|
||||
@@ -322,4 +329,4 @@ STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) {
|
||||
m.impl("moe_permute", TORCH_BOX(&moe_permute));
|
||||
m.impl("moe_permute_with_scratch", TORCH_BOX(&moe_permute_with_scratch));
|
||||
m.impl("moe_unpermute", TORCH_BOX(&moe_unpermute));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -76,9 +76,12 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
|
||||
int32_t const numValidRows =
|
||||
min(numRows,
|
||||
static_cast<int32_t>(__ldca(&input_offset_by_experts[n_experts])));
|
||||
int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
|
||||
|
||||
for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
|
||||
for (int globalIdx = tid; globalIdx < numValidRows * colsPerRow;
|
||||
globalIdx += gridDim.x * blockDim.x) {
|
||||
int rowIdx = globalIdx / colsPerRow;
|
||||
int colIdx = globalIdx % colsPerRow;
|
||||
@@ -193,9 +196,11 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
|
||||
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
|
||||
int32_t const numValidRows =
|
||||
min(numRows, static_cast<int32_t>(shared_input_offsets[n_experts]));
|
||||
int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
|
||||
|
||||
for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
|
||||
for (int globalIdx = tid; globalIdx < numValidRows * colsPerRow;
|
||||
globalIdx += gridDim.x * blockDim.x) {
|
||||
int rowIdx = globalIdx / colsPerRow;
|
||||
int colIdx = globalIdx % colsPerRow;
|
||||
|
||||
@@ -53,11 +53,14 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
|
||||
int32_t const numValidRows =
|
||||
min(numRows,
|
||||
static_cast<int32_t>(__ldca(&input_offset_by_experts[n_experts])));
|
||||
// When fusing SiLU+Mul, input has gate || up layout (doubled width)
|
||||
int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
|
||||
|
||||
// Each global thread processes one element
|
||||
for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
|
||||
for (int globalIdx = tid; globalIdx < numValidRows * colsPerRow;
|
||||
globalIdx += gridDim.x * blockDim.x) {
|
||||
// Calculate which row and column this global thread should process
|
||||
int rowIdx = globalIdx / colsPerRow;
|
||||
@@ -186,11 +189,13 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
|
||||
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
|
||||
int32_t const numValidRows =
|
||||
min(numRows, static_cast<int32_t>(shared_input_offsets[n_experts]));
|
||||
// When fusing SiLU+Mul, input has gate || up layout (doubled width)
|
||||
int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
|
||||
|
||||
// Each global thread processes one element
|
||||
for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
|
||||
for (int globalIdx = tid; globalIdx < numValidRows * colsPerRow;
|
||||
globalIdx += gridDim.x * blockDim.x) {
|
||||
// Calculate which row and column this global thread should process
|
||||
int rowIdx = globalIdx / colsPerRow;
|
||||
|
||||
@@ -62,6 +62,7 @@ __global__ void compute_expert_offsets(
|
||||
tot_offset += swap_ab ? problem_sizes1[i * 3 + 1] : problem_sizes1[i * 3];
|
||||
expert_offsets[i + 1] = tot_offset;
|
||||
}
|
||||
atomic_buffer[num_experts] = tot_offset;
|
||||
}
|
||||
|
||||
__global__ void compute_expert_blockscale_offsets(
|
||||
@@ -81,28 +82,28 @@ __global__ void compute_expert_blockscale_offsets(
|
||||
tot_offset_round += (cur_offset + (128 - 1)) / 128 * 128;
|
||||
blockscale_offsets[i + 1] = tot_offset_round;
|
||||
}
|
||||
atomic_buffer[num_experts] = tot_offset;
|
||||
}
|
||||
|
||||
__global__ void compute_arg_sorts(const int32_t* __restrict__ topk_ids,
|
||||
const int32_t* __restrict__ expert_offsets,
|
||||
int32_t* input_permutation,
|
||||
int32_t* output_permutation,
|
||||
int32_t* atomic_buffer, const int topk_length,
|
||||
const int topk) {
|
||||
int const blk_expert_id = blockIdx.x;
|
||||
int const num_experts = gridDim.x;
|
||||
int32_t const num_tokens = expert_offsets[num_experts];
|
||||
|
||||
for (int i = threadIdx.x; i < topk_length; i += THREADS_PER_EXPERT) {
|
||||
int const expert_id = topk_ids[i];
|
||||
if (expert_id == -1 && blockIdx.x == 0) {
|
||||
if ((expert_id < 0 || expert_id >= num_experts) && blockIdx.x == 0) {
|
||||
int const start = atomicAdd(&atomic_buffer[num_experts], 1);
|
||||
input_permutation[start] = -1;
|
||||
// output_permutation is used to re-order the moe outputs. It is
|
||||
// used as c2 = c2[c_map], where c2 is a torch.tensor that is the
|
||||
// output of the cutlass kernels and c_map is the output_permutation.
|
||||
// c2 is initialized to zeros, therefore by setting the output_permutation
|
||||
// to num_tokens, we are guaranteed to fill the moe outputs to zero
|
||||
// for "invalid" topk_ids.
|
||||
output_permutation[i] = num_tokens;
|
||||
// output of the cutlass kernels and c_map is the output_permutation. Use
|
||||
// out-of-range sentinels for invalid routes; shuffle_rows converts them
|
||||
// to zero rows.
|
||||
output_permutation[i] = topk_length;
|
||||
} else if (expert_id == blk_expert_id) {
|
||||
int start = atomicAdd(&atomic_buffer[expert_id], 1);
|
||||
input_permutation[start] = i / topk;
|
||||
@@ -246,7 +247,7 @@ void get_cutlass_moe_mm_data_caller(
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(device.index());
|
||||
auto stream = get_current_cuda_stream(device.index());
|
||||
torch::stable::Tensor atomic_buffer = torch::stable::new_zeros(
|
||||
topk_ids, {num_experts}, torch::headeronly::ScalarType::Int);
|
||||
topk_ids, {num_experts + 1}, torch::headeronly::ScalarType::Int);
|
||||
|
||||
int num_threads = min(THREADS_PER_EXPERT, topk_ids.numel());
|
||||
|
||||
@@ -275,7 +276,6 @@ void get_cutlass_moe_mm_data_caller(
|
||||
}
|
||||
compute_arg_sorts<<<num_experts, num_threads, 0, stream>>>(
|
||||
static_cast<const int32_t*>(topk_ids.data_ptr()),
|
||||
static_cast<const int32_t*>(expert_offsets.data_ptr()),
|
||||
static_cast<int32_t*>(input_permutation.data_ptr()),
|
||||
static_cast<int32_t*>(output_permutation.data_ptr()),
|
||||
static_cast<int32_t*>(atomic_buffer.data_ptr()), topk_ids.numel(),
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import copy
|
||||
import dataclasses
|
||||
from collections.abc import Callable
|
||||
from math import prod
|
||||
|
||||
import pytest
|
||||
@@ -33,7 +34,9 @@ from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
|
||||
CutlassExpertsFp8,
|
||||
CutlassExpertsMxfp4,
|
||||
CutlassExpertsW4A8Fp8,
|
||||
run_cutlass_moe_fp4,
|
||||
run_cutlass_moe_fp8,
|
||||
run_cutlass_moe_mxfp4,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.oracle import nvfp4 as nvfp4_oracle
|
||||
from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
|
||||
@@ -131,6 +134,166 @@ def test_nvfp4_clamp_allows_shared_activation_backends(
|
||||
assert selected == expected
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
(lambda x: x is None or not ops.cutlass_group_gemm_supported(x.to_int()))(
|
||||
current_platform.get_device_capability()
|
||||
),
|
||||
reason="Grouped gemm is not supported on this GPU type.",
|
||||
)
|
||||
def test_cutlass_moe_permutation_maps_padding_to_zero():
|
||||
"""Invalid top-k routes must not leave unsafe permutation-map entries."""
|
||||
num_experts = 2
|
||||
topk_ids = torch.tensor(
|
||||
[[0, -1], [1, num_experts]], device="cuda", dtype=torch.int32
|
||||
)
|
||||
num_routes = topk_ids.numel()
|
||||
|
||||
expert_offsets = torch.empty(num_experts + 1, device="cuda", dtype=torch.int32)
|
||||
blockscale_offsets = torch.empty_like(expert_offsets)
|
||||
problem_sizes1 = torch.empty(num_experts, 3, device="cuda", dtype=torch.int32)
|
||||
problem_sizes2 = torch.empty_like(problem_sizes1)
|
||||
input_permutation = torch.empty(num_routes, device="cuda", dtype=torch.int32)
|
||||
output_permutation = torch.empty_like(input_permutation)
|
||||
|
||||
ops.get_cutlass_moe_mm_data(
|
||||
topk_ids,
|
||||
expert_offsets,
|
||||
problem_sizes1,
|
||||
problem_sizes2,
|
||||
input_permutation,
|
||||
output_permutation,
|
||||
num_experts,
|
||||
128,
|
||||
128,
|
||||
blockscale_offsets,
|
||||
)
|
||||
|
||||
input_tensor = torch.randn(2, 128, device="cuda", dtype=torch.bfloat16)
|
||||
permuted = ops.shuffle_rows(input_tensor, input_permutation)
|
||||
restored = ops.shuffle_rows(permuted, output_permutation)
|
||||
|
||||
torch.testing.assert_close(permuted[:2], input_tensor)
|
||||
torch.testing.assert_close(permuted[2:], torch.zeros_like(permuted[2:]))
|
||||
torch.testing.assert_close(restored[0], input_tensor[0])
|
||||
torch.testing.assert_close(restored[1], torch.zeros_like(restored[1]))
|
||||
torch.testing.assert_close(restored[2], input_tensor[1])
|
||||
torch.testing.assert_close(restored[3], torch.zeros_like(restored[3]))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("quantization", ["nvfp4", "mxfp4"])
|
||||
@torch.inference_mode()
|
||||
def test_cutlass_fp4_moe_padded_routes_do_not_change_valid_output(quantization: str):
|
||||
"""Padded routes must not participate in either activation quantization."""
|
||||
experts_cls = CutlassExpertsFp4 if quantization == "nvfp4" else CutlassExpertsMxfp4
|
||||
if not experts_cls._supports_current_device():
|
||||
pytest.skip(f"CUTLASS {quantization} MoE is not supported on this GPU")
|
||||
|
||||
set_random_seed(7)
|
||||
num_experts, n, k = 4, 128, 128
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device("cuda")
|
||||
w1 = torch.full(
|
||||
(num_experts, 2 * n, k // 2), 0x11, device=device, dtype=torch.uint8
|
||||
)
|
||||
w2 = torch.full((num_experts, k, n // 2), 0x11, device=device, dtype=torch.uint8)
|
||||
|
||||
runner: Callable[..., None]
|
||||
quant_args: tuple[torch.Tensor, ...]
|
||||
if quantization == "nvfp4":
|
||||
scale_kwargs = {
|
||||
"device": device,
|
||||
"dtype": torch.float8_e4m3fn,
|
||||
}
|
||||
w1_scale = torch.full((num_experts, 2 * n, k // 16), 1 / 16, **scale_kwargs)
|
||||
w2_scale = torch.full((num_experts, k, n // 16), 1 / 16, **scale_kwargs)
|
||||
gscale = torch.ones(num_experts, device=device, dtype=torch.float32)
|
||||
runner = run_cutlass_moe_fp4
|
||||
quant_args = (
|
||||
gscale,
|
||||
w1,
|
||||
w1_scale,
|
||||
gscale,
|
||||
gscale,
|
||||
w2,
|
||||
w2_scale,
|
||||
gscale,
|
||||
)
|
||||
else:
|
||||
scale_exponent = 127 - 4
|
||||
w1_scale = torch.full(
|
||||
(num_experts, 2 * n, k // 32),
|
||||
scale_exponent,
|
||||
device=device,
|
||||
dtype=torch.uint8,
|
||||
)
|
||||
w2_scale = torch.full(
|
||||
(num_experts, k, n // 32),
|
||||
scale_exponent,
|
||||
device=device,
|
||||
dtype=torch.uint8,
|
||||
)
|
||||
runner = run_cutlass_moe_mxfp4
|
||||
quant_args = (w1, w1_scale, w2, w2_scale)
|
||||
|
||||
def run_moe(
|
||||
hidden_states: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
workspace_value: float,
|
||||
) -> torch.Tensor:
|
||||
m, topk = topk_ids.shape
|
||||
output = torch.empty((m, k), device=device, dtype=dtype)
|
||||
workspace13 = torch.full(
|
||||
(m * topk, max(2 * n, k)),
|
||||
workspace_value,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
workspace2 = torch.full(
|
||||
(m * topk, n), workspace_value, device=device, dtype=dtype
|
||||
)
|
||||
|
||||
runner(
|
||||
output,
|
||||
hidden_states,
|
||||
*quant_args,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
MoEActivation.SILU,
|
||||
workspace13,
|
||||
workspace2,
|
||||
m,
|
||||
n,
|
||||
k,
|
||||
num_experts,
|
||||
device,
|
||||
)
|
||||
return output
|
||||
|
||||
valid_input = torch.randn((1, k), device=device, dtype=dtype)
|
||||
expected = run_moe(
|
||||
valid_input,
|
||||
torch.zeros((1, 1), device=device, dtype=torch.int32),
|
||||
torch.ones((1, 1), device=device, dtype=torch.float32),
|
||||
workspace_value=0,
|
||||
)
|
||||
assert expected.abs().max() > 0
|
||||
|
||||
# Force the quantizers' grid-stride loops to revisit padded tail rows.
|
||||
num_tokens = 4096
|
||||
padded_input = torch.zeros((num_tokens, k), device=device, dtype=dtype)
|
||||
padded_input[0] = valid_input[0]
|
||||
padded_ids = torch.full((num_tokens, 2), -1, device=device, dtype=torch.int32)
|
||||
padded_ids[0, 0] = 0
|
||||
padded_weights = torch.zeros((num_tokens, 2), device=device, dtype=torch.float32)
|
||||
padded_weights[0, 0] = 1
|
||||
|
||||
for workspace_value in (0, 64):
|
||||
actual = run_moe(padded_input, padded_ids, padded_weights, workspace_value)
|
||||
torch.testing.assert_close(actual[0], expected[0], atol=1e-2, rtol=1e-2)
|
||||
torch.testing.assert_close(actual[1:], torch.zeros_like(actual[1:]))
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class MOETensors:
|
||||
a: torch.Tensor
|
||||
|
||||
@@ -103,6 +103,21 @@ def test_shuffle_rows_expansion(num_tokens: int, hidden_size: int):
|
||||
)
|
||||
|
||||
|
||||
def test_shuffle_rows_invalid_indices_are_zeroed():
|
||||
"""Invalid sentinel indices should produce zero rows instead of OOB reads."""
|
||||
if not current_platform.is_cuda():
|
||||
pytest.skip("shuffle_rows requires CUDA")
|
||||
|
||||
input_tensor = torch.randn(3, 128, device="cuda", dtype=torch.bfloat16)
|
||||
dst2src_map = torch.tensor([-1, 1, 3], device="cuda", dtype=torch.int32)
|
||||
|
||||
output = shuffle_rows(input_tensor, dst2src_map)
|
||||
|
||||
torch.testing.assert_close(output[0], torch.zeros_like(output[0]))
|
||||
torch.testing.assert_close(output[1], input_tensor[1])
|
||||
torch.testing.assert_close(output[2], torch.zeros_like(output[2]))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [16, 64])
|
||||
@pytest.mark.parametrize("hidden_size", [128, 512])
|
||||
def test_shuffle_rows_random_permutation(num_tokens: int, hidden_size: int):
|
||||
|
||||
Reference in New Issue
Block a user