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

[Kernel] Add FP32 router GEMM support for Step-5 Preview and Muse Spark Open (#62126)

Browse Source 
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
This commit is contained in:
Jee Jee Li
2026-10-16 22:17:55 -07:00
committed by GitHub
parent 41b8c15809
commit 39983a5e7a
4 changed files with 59 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/fp32_router_gemm.cu
+28 -1
View File
@@ -279,6 +279,27 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a,
launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
output, mat_a, mat_b, stream);
}
} else if constexpr (std::is_same_v<InputT, __nv_bfloat16> &&
kNumExperts == 256 && kHiddenDim == 8192) {
// MuseSpark1x: the 8MB weight makes this shape FFMA-bound past M~10, so
// the wider block wins at every M and token groups clear the M >= 28
// accumulator cliff. Thresholds are from a B300 sweep. BS=192/384 are not
// instantiable: VPT=8 requires kBlockSize to divide kHiddenDim / 8 = 1024.
if (!isBlackwellFamily()) {
launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
output, mat_a, mat_b, stream);
return;
}
if constexpr (kNumTokens >= 10 && kNumTokens % 2 == 0) {
launchFp32RouterGemm<InputT, 256, 2, kNumTokens, kNumExperts, kHiddenDim,
2>(output, mat_a, mat_b, stream);
} else if constexpr (kNumTokens >= 13) {
launchFp32RouterGemm<InputT, 256, 2, kNumTokens, kNumExperts, kHiddenDim>(
output, mat_a, mat_b, stream);
} else {
launchFp32RouterGemm<InputT, 256, 1, kNumTokens, kNumExperts, kHiddenDim>(
output, mat_a, mat_b, stream);
}
} else {
launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
output, mat_a, mat_b, stream);
@@ -287,7 +308,8 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a,
// ---------------------------------------------------------------------------
// Explicit instantiations for both input types.
// Supported (E, H) pairs: (256, 3072), (128, 6144), and (256, 5120).
// Supported (E, H) pairs: (256, 3072), (128, 6144), (256, 5120),
// (352, 4096) [Step-5], and (256, 8192) [Muse Spark].
// ---------------------------------------------------------------------------
#define INSTANTIATE(T, M, E, H) \
@@ -336,6 +358,11 @@ INSTANTIATE_ALL(__nv_bfloat16, 128, 6144)
INSTANTIATE_ALL(float, 256, 5120)
INSTANTIATE_ALL(__nv_bfloat16, 256, 5120)
INSTANTIATE_ALL(float, 352, 4096)
INSTANTIATE_ALL(__nv_bfloat16, 352, 4096)
INSTANTIATE_ALL(float, 256, 8192)
INSTANTIATE_ALL(__nv_bfloat16, 256, 8192)
#define INSTANTIATE_LARGE(T, E, H) \
INSTANTIATE(T, 33, E, H) \
INSTANTIATE(T, 34, E, H) \
csrc/fp32_router_gemm_entry.cu
+10 -2
View File
@@ -29,7 +29,9 @@ static constexpr int FP32_MAX_TOKENS = 32;
static inline bool fp32_router_gemm_supported(int hidden_dim, int num_experts) {
return (hidden_dim == 3072 && num_experts == 256) ||
(hidden_dim == 6144 && num_experts == 128) ||
(hidden_dim == 5120 && num_experts == 256);
(hidden_dim == 5120 && num_experts == 256) ||
(hidden_dim == 4096 && num_experts == 352) ||
(hidden_dim == 8192 && num_experts == 256);
}
// Forward declarations — 4 template params must match fp32_router_gemm.cu
@@ -80,6 +82,12 @@ void dispatchFp32RouterGemm(int num_experts, int hidden_dim, int num_tokens,
} else if (num_experts == 256 && hidden_dim == 5120) {
Fp32LoopUnroller<InputT, 256, 5120, 1, 48>::unroll(num_tokens, output,
mat_a, mat_b, stream);
} else if (num_experts == 352 && hidden_dim == 4096) {
Fp32LoopUnroller<InputT, 352, 4096, 1, FP32_MAX_TOKENS>::unroll(
num_tokens, output, mat_a, mat_b, stream);
} else if (num_experts == 256 && hidden_dim == 8192) {
Fp32LoopUnroller<InputT, 256, 8192, 1, FP32_MAX_TOKENS>::unroll(
num_tokens, output, mat_a, mat_b, stream);
} else {
throw std::invalid_argument(
"fp32_router_gemm: unsupported (hidden_dim, num_experts) pair");
@@ -114,7 +122,7 @@ void fp32_router_gemm(
STD_TORCH_CHECK(
fp32_router_gemm_supported(hidden_dim, num_experts),
"fp32_router_gemm: supported (hidden_dim, num_experts) pairs are "
"(3072, 256), (6144, 128), and (5120, 256)");
"(3072, 256), (6144, 128), (5120, 256), (4096, 352), and (8192, 256)");
const int max_tokens =
hidden_dim == 5120 && num_experts == 256 ? 48 : FP32_MAX_TOKENS;
STD_TORCH_CHECK(num_tokens <= max_tokens,
tests/kernels/test_fp32_router_gemm.py
+14 -9
View File
@@ -3,11 +3,11 @@
"""Tests for fp32_router_gemm kernel: activation×weight→fp32.
Supported (hidden_size, num_experts) pairs:
(3072, 256), (6144, 128), (5120, 256)
(3072, 256), (6144, 128), (5120, 256), (4096, 352), (8192, 256)
Correctness baseline: F.linear in float32. Every supported M is covered so
all tuned geometries (wide-block, experts-per-block, token-group; boundaries
at M=4/5, odd/even, M=15/16) are exercised on Blackwell, and the legacy
at M=4/5, M=9/10, M=12/13, odd/even, M=15/16) are exercised on Blackwell, and the legacy
128/1 geometry everywhere else.
"""
@@ -18,7 +18,7 @@ from vllm._custom_ops import fp32_router_gemm
from vllm.platforms import current_platform
# (hidden_size, num_experts)
SHAPES = [(3072, 256), (6144, 128), (5120, 256)]
SHAPES = [(3072, 256), (6144, 128), (5120, 256), (4096, 352), (8192, 256)]
CASES = [
(m, hidden, experts)
for hidden, experts in SHAPES
@@ -126,11 +126,11 @@ def test_topk_routing_consistency(num_tokens: int, hidden_dim: int, num_experts:
)
def test_zero_tokens_returns_empty():
@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
def test_zero_tokens_returns_empty(hidden_dim, num_experts):
"""M=0 is a graceful no-op returning an empty [0, E] tensor."""
_requires_sm90()
device = torch.device("cuda")
hidden_dim, num_experts = SHAPES[0]
mat_a = torch.empty(0, hidden_dim, dtype=torch.float32, device=device)
mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device)
out = fp32_router_gemm(mat_a, mat_b)
@@ -183,8 +183,13 @@ def test_rejects_invalid_inputs(hidden_dim, num_experts):
@pytest.mark.parametrize(
"tokens", [*range(1, 17), 17, 24, 32, 33, 40, 48, 56, 64, 65, 128]
)
@pytest.mark.parametrize(
"hidden_dim,num_experts", [(5120, 256), (4096, 352), (8192, 256)]
)
@torch.inference_mode()
def test_router_retains_fp32_weights_and_logits(default_vllm_config, dist_init, tokens):
def test_router_retains_fp32_weights_and_logits(
default_vllm_config, dist_init, tokens, hidden_dim, num_experts
):
"""Low-latency routing must retain FP32 weight precision and logits."""
from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
from vllm.platforms import current_platform
@@ -193,16 +198,16 @@ def test_router_retains_fp32_weights_and_logits(default_vllm_config, dist_init,
pytest.skip("SM100 required")
torch.manual_seed(42)
router = GateLinear(
5120, 256, params_dtype=torch.float32, out_dtype=torch.float32
hidden_dim, num_experts, params_dtype=torch.float32, out_dtype=torch.float32
).cuda()
router.weight.normal_(std=0.01)
x = torch.randn(tokens, 5120, device="cuda", dtype=torch.bfloat16)
x = torch.randn(tokens, hidden_dim, device="cuda", dtype=torch.bfloat16)
expected = (x.double() @ router.weight.double().T).float()
actual, bias = router(x)
assert bias is None and actual.dtype == torch.float32
torch.testing.assert_close(actual, expected, rtol=2e-5, atol=2e-5)
torch.testing.assert_close(actual.topk(8).indices, expected.topk(8).indices)
if tokens in (1, 16, 48, 56):
if tokens in (1, 16, 32, 33, 48, 56):
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
captured, _ = router(x)
vllm/model_executor/layers/fused_moe/router/gate_linear.py
+7 -1
View File
@@ -47,7 +47,13 @@ class GateLinear(ReplicatedLinear):
"""
# (hidden_size, num_experts) pairs with an instantiated fp32 kernel.
FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128), (5120, 256)}
FP32_SUPPORTED_SHAPES = {
(3072, 256),
(6144, 128),
(5120, 256),
(4096, 352),
(8192, 256),
}
FP32_MAX_TOKENS = 32
# Largest batch for which tier 1 beats the cuBLAS epilogue. Subclasses
# whose router shape stays ahead further up the range may raise it.