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,