[Kernel] Add router GEMM kernels for (5120, 256) and BF16 output to ll_bf16_gemm (#62042)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Lain <fusiyuan2000@gmail.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: vllm-agent <claw@inferact.ai>
This commit is contained in:
Woosuk Kwon
co-authored by
Lain
Codex
Claude Opus 5.5
vllm-agent
Lain
Codex
Claude Opus 5.5
vllm-agent
parent
153a17166e
commit
596e530675
@@ -1,7 +1,7 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// Router GEMM: activation(T) x weight(fp32) -> fp32, M<=32, for the
|
||||
// Router GEMM: activation(T) x weight(fp32) -> fp32, M<=48, for the
|
||||
// supported (E, H) pairs listed at the bottom of this file.
|
||||
// Supports bf16 or fp32 activation; weight is always fp32.
|
||||
|
||||
@@ -214,8 +214,26 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a,
|
||||
// production-fidelity harness (CUDA-graph replay, per-layer cold weights).
|
||||
// Only enabled on the Blackwell family where it was validated; Hopper and
|
||||
// other shapes / fp32 activation keep the legacy geometry.
|
||||
if constexpr (std::is_same_v<InputT, __nv_bfloat16> && kNumExperts == 128 &&
|
||||
kHiddenDim == 6144) {
|
||||
if constexpr (std::is_same_v<InputT, __nv_bfloat16> && kNumExperts == 256 &&
|
||||
kHiddenDim == 5120) {
|
||||
if (!isBlackwellFamily()) {
|
||||
launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
|
||||
output, mat_a, mat_b, stream);
|
||||
return;
|
||||
}
|
||||
if constexpr (kNumTokens > 16 && kNumTokens % 8 == 0) {
|
||||
constexpr int block_size = kNumTokens == 24 ? 160 : 128;
|
||||
launchFp32RouterGemm<InputT, block_size, 2, kNumTokens, kNumExperts,
|
||||
kHiddenDim, 4>(output, mat_a, mat_b, stream);
|
||||
} else {
|
||||
constexpr int block_size = kNumTokens == 1 ? 64 : 128;
|
||||
constexpr int token_groups = kNumTokens % 2 == 0 ? 2 : 1;
|
||||
launchFp32RouterGemm<InputT, block_size, 1, kNumTokens, kNumExperts,
|
||||
kHiddenDim, token_groups>(output, mat_a, mat_b,
|
||||
stream);
|
||||
}
|
||||
} else if constexpr (std::is_same_v<InputT, __nv_bfloat16> &&
|
||||
kNumExperts == 128 && kHiddenDim == 6144) {
|
||||
// MiniMax-M3. Legacy 128/1 only fills 128 blocks and pays the same
|
||||
// accumulator register cliffs; B300 sweep:
|
||||
// even M in [6, 10] : BS=384, EPB=1, 2 token groups (1.26-1.43x)
|
||||
@@ -268,8 +286,8 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Explicit instantiations: M=1..32, for both input types, for the supported
|
||||
// (E, H) pairs: (256, 3072) [MiniMax-M2/M2.5] and (128, 6144) [MiniMax-M3].
|
||||
// Explicit instantiations for both input types.
|
||||
// Supported (E, H) pairs: (256, 3072), (128, 6144), and (256, 5120).
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#define INSTANTIATE(T, M, E, H) \
|
||||
@@ -315,5 +333,31 @@ INSTANTIATE_ALL(__nv_bfloat16, 256, 3072)
|
||||
INSTANTIATE_ALL(float, 128, 6144)
|
||||
INSTANTIATE_ALL(__nv_bfloat16, 128, 6144)
|
||||
|
||||
INSTANTIATE_ALL(float, 256, 5120)
|
||||
INSTANTIATE_ALL(__nv_bfloat16, 256, 5120)
|
||||
|
||||
#define INSTANTIATE_LARGE(T, E, H) \
|
||||
INSTANTIATE(T, 33, E, H) \
|
||||
INSTANTIATE(T, 34, E, H) \
|
||||
INSTANTIATE(T, 35, E, H) \
|
||||
INSTANTIATE(T, 36, E, H) \
|
||||
INSTANTIATE(T, 37, E, H) \
|
||||
INSTANTIATE(T, 38, E, H) \
|
||||
INSTANTIATE(T, 39, E, H) \
|
||||
INSTANTIATE(T, 40, E, H) \
|
||||
INSTANTIATE(T, 41, E, H) \
|
||||
INSTANTIATE(T, 42, E, H) \
|
||||
INSTANTIATE(T, 43, E, H) \
|
||||
INSTANTIATE(T, 44, E, H) \
|
||||
INSTANTIATE(T, 45, E, H) \
|
||||
INSTANTIATE(T, 46, E, H) \
|
||||
INSTANTIATE(T, 47, E, H) \
|
||||
INSTANTIATE(T, 48, E, H)
|
||||
|
||||
INSTANTIATE_LARGE(float, 256, 5120)
|
||||
INSTANTIATE_LARGE(__nv_bfloat16, 256, 5120)
|
||||
|
||||
#undef INSTANTIATE_LARGE
|
||||
|
||||
#undef INSTANTIATE_ALL
|
||||
#undef INSTANTIATE
|
||||
|
||||
@@ -25,10 +25,11 @@ inline int getSMVersion() {
|
||||
static constexpr int FP32_MAX_TOKENS = 32;
|
||||
|
||||
// Supported (hidden_dim, num_experts) pairs (must match the instantiations in
|
||||
// fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3.
|
||||
// fp32_router_gemm.cu).
|
||||
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 == 6144 && num_experts == 128) ||
|
||||
(hidden_dim == 5120 && num_experts == 256);
|
||||
}
|
||||
|
||||
// Forward declarations — 4 template params must match fp32_router_gemm.cu
|
||||
@@ -60,8 +61,7 @@ struct Fp32LoopUnroller<InputT, kNumExperts, kHiddenDim, kEnd, kEnd> {
|
||||
invokeFp32RouterGemm<InputT, kEnd, kNumExperts, kHiddenDim>(
|
||||
output, mat_a, mat_b, stream);
|
||||
} else {
|
||||
throw std::invalid_argument(
|
||||
"fp32_router_gemm: num_tokens must be in [1, 32]");
|
||||
throw std::invalid_argument("fp32_router_gemm: unsupported num_tokens");
|
||||
}
|
||||
}
|
||||
};
|
||||
@@ -77,6 +77,9 @@ void dispatchFp32RouterGemm(int num_experts, int hidden_dim, int num_tokens,
|
||||
} else if (num_experts == 128 && hidden_dim == 6144) {
|
||||
Fp32LoopUnroller<InputT, 128, 6144, 1, FP32_MAX_TOKENS>::unroll(
|
||||
num_tokens, output, mat_a, mat_b, stream);
|
||||
} else if (num_experts == 256 && hidden_dim == 5120) {
|
||||
Fp32LoopUnroller<InputT, 256, 5120, 1, 48>::unroll(num_tokens, output,
|
||||
mat_a, mat_b, stream);
|
||||
} else {
|
||||
throw std::invalid_argument(
|
||||
"fp32_router_gemm: unsupported (hidden_dim, num_experts) pair");
|
||||
@@ -111,9 +114,11 @@ 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) and (6144, 128)");
|
||||
STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS,
|
||||
"fp32_router_gemm: num_tokens must be in [0, 32]");
|
||||
"(3072, 256), (6144, 128), and (5120, 256)");
|
||||
const int max_tokens =
|
||||
hidden_dim == 5120 && num_experts == 256 ? 48 : FP32_MAX_TOKENS;
|
||||
STD_TORCH_CHECK(num_tokens <= max_tokens,
|
||||
"fp32_router_gemm: num_tokens exceeds supported range");
|
||||
STD_TORCH_CHECK(
|
||||
mat_a.scalar_type() == torch::headeronly::ScalarType::Float ||
|
||||
mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16,
|
||||
|
||||
@@ -582,6 +582,20 @@ def test_fused_topk_bias(
|
||||
assert_routing_results_close(topk_weights, topk_ids, baseline_weights, baseline_ids)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scoring_func", [0, 1])
|
||||
def test_fused_grouped_topk_accepts_idle_rank(scoring_func):
|
||||
"""Zero-token routing must leave CUDA ready for the next batch."""
|
||||
if not current_platform.is_cuda():
|
||||
pytest.skip("CUDA required")
|
||||
from vllm import _custom_ops as ops
|
||||
|
||||
scores = torch.empty((0, 8), device="cuda", dtype=torch.float32)
|
||||
bias = torch.zeros(8, device="cuda")
|
||||
weights, ids = ops.grouped_topk(scores, 2, 1, 2, True, 1.0, bias, scoring_func)
|
||||
torch.testing.assert_close(weights, scores.new_empty((0, 2)))
|
||||
torch.testing.assert_close(ids, scores.new_empty((0, 2), dtype=torch.int32))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m,k", MK_S)
|
||||
@pytest.mark.parametrize("top_k", TOP_KS)
|
||||
@pytest.mark.parametrize(
|
||||
|
||||
@@ -3,9 +3,9 @@
|
||||
"""Tests for fp32_router_gemm kernel: activation×weight→fp32.
|
||||
|
||||
Supported (hidden_size, num_experts) pairs:
|
||||
(3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3
|
||||
(3072, 256), (6144, 128), (5120, 256)
|
||||
|
||||
Correctness baseline: F.linear in float32. Every M in [1, 32] is covered so
|
||||
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
|
||||
128/1 geometry everywhere else.
|
||||
@@ -18,8 +18,12 @@ from vllm._custom_ops import fp32_router_gemm
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
# (hidden_size, num_experts)
|
||||
SHAPES = [(3072, 256), (6144, 128)]
|
||||
ALL_M = list(range(1, 33))
|
||||
SHAPES = [(3072, 256), (6144, 128), (5120, 256)]
|
||||
CASES = [
|
||||
(m, hidden, experts)
|
||||
for hidden, experts in SHAPES
|
||||
for m in range(1, 49 if hidden == 5120 else 33)
|
||||
]
|
||||
# Absolute tolerance for fp32 kernel vs float64 reference
|
||||
ATOL_FP32 = 2e-4
|
||||
ATOL_BF16 = 2e-2 # bf16 activation has lower precision
|
||||
@@ -40,8 +44,7 @@ def _ref(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor:
|
||||
return torch.nn.functional.linear(mat_a.float(), mat_b.float())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
|
||||
@pytest.mark.parametrize("num_tokens", ALL_M)
|
||||
@pytest.mark.parametrize("num_tokens,hidden_dim,num_experts", CASES)
|
||||
def test_fp32_activation(num_tokens: int, hidden_dim: int, num_experts: int):
|
||||
"""fp32 activation → fp32 output should match reference closely."""
|
||||
_requires_sm90()
|
||||
@@ -58,8 +61,7 @@ def test_fp32_activation(num_tokens: int, hidden_dim: int, num_experts: int):
|
||||
torch.testing.assert_close(out, ref, atol=ATOL_FP32, rtol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
|
||||
@pytest.mark.parametrize("num_tokens", ALL_M)
|
||||
@pytest.mark.parametrize("num_tokens,hidden_dim,num_experts", CASES)
|
||||
def test_bf16_activation(num_tokens: int, hidden_dim: int, num_experts: int):
|
||||
"""bf16 activation → fp32 output should match reference within bf16 error."""
|
||||
_requires_sm90()
|
||||
@@ -136,17 +138,19 @@ def test_zero_tokens_returns_empty():
|
||||
assert out.dtype == torch.float32
|
||||
|
||||
|
||||
def test_rejects_invalid_inputs():
|
||||
@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
|
||||
def test_rejects_invalid_inputs(hidden_dim, num_experts):
|
||||
"""The entry must fail loudly, never compute silently wrong results."""
|
||||
_requires_sm90()
|
||||
device = torch.device("cuda")
|
||||
hidden_dim, num_experts = SHAPES[0]
|
||||
mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device)
|
||||
|
||||
# num_tokens > 32 (beyond the instantiated range)
|
||||
# Beyond the instantiated range for this shape.
|
||||
invalid_tokens = 49 if hidden_dim == 5120 else 33
|
||||
with pytest.raises(Exception, match="num_tokens"):
|
||||
fp32_router_gemm(
|
||||
torch.randn(33, hidden_dim, dtype=torch.float32, device=device), mat_b
|
||||
torch.randn(invalid_tokens, hidden_dim, dtype=torch.float32, device=device),
|
||||
mat_b,
|
||||
)
|
||||
|
||||
# unsupported (hidden_dim, num_experts) pair
|
||||
@@ -173,3 +177,36 @@ def test_rejects_invalid_inputs():
|
||||
fp32_router_gemm(
|
||||
torch.randn(4, hidden_dim, dtype=torch.float16, device=device), mat_b
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
||||
@pytest.mark.parametrize(
|
||||
"tokens", [*range(1, 17), 17, 24, 32, 33, 40, 48, 56, 64, 65, 128]
|
||||
)
|
||||
@torch.inference_mode()
|
||||
def test_router_retains_fp32_weights_and_logits(default_vllm_config, dist_init, tokens):
|
||||
"""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
|
||||
|
||||
if not current_platform.is_device_capability_family(100):
|
||||
pytest.skip("SM100 required")
|
||||
torch.manual_seed(42)
|
||||
router = GateLinear(
|
||||
5120, 256, 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)
|
||||
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):
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
captured, _ = router(x)
|
||||
x.normal_()
|
||||
graph.replay()
|
||||
expected = (x.double() @ router.weight.double().T).float()
|
||||
torch.testing.assert_close(captured, expected, rtol=2e-5, atol=2e-5)
|
||||
|
||||
@@ -566,7 +566,7 @@ def test_invalid_output_dtype():
|
||||
with pytest.raises(ValueError, match="output_dtype=torch.float32"):
|
||||
from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ll_bf16_gemm
|
||||
|
||||
ll_bf16_gemm(a, b, output_dtype=torch.bfloat16)
|
||||
ll_bf16_gemm(a, b, output_dtype=torch.float16)
|
||||
|
||||
|
||||
def test_cache_miss_compiles_dotprod():
|
||||
@@ -581,5 +581,29 @@ def test_cache_miss_compiles_dotprod():
|
||||
_assert_close(out, _ref(a, b), context="cache miss dotprod")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("M", [1, 2, 8, 16])
|
||||
@pytest.mark.parametrize("N", [1792, 3584])
|
||||
@torch.inference_mode()
|
||||
def test_bf16_output_and_fp32_cache_coexist(M, N):
|
||||
"""BF16 epilogues retain FP32 accumulation and can replay with new inputs."""
|
||||
from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ll_bf16_gemm
|
||||
|
||||
torch.manual_seed(42)
|
||||
x = torch.randn(M, 5120, device="cuda", dtype=torch.bfloat16)
|
||||
w = torch.randn(N, 5120, device="cuda", dtype=torch.bfloat16) * 0.01
|
||||
for dtype in (torch.float32, torch.bfloat16):
|
||||
actual = ll_bf16_gemm(x, w, output_dtype=dtype)
|
||||
assert actual.dtype == dtype
|
||||
reference = (x.double() @ w.double().T).to(dtype)
|
||||
torch.testing.assert_close(actual, reference, rtol=0.008, atol=1e-5)
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
actual = ll_bf16_gemm(x, w, output_dtype=torch.bfloat16)
|
||||
x.normal_()
|
||||
graph.replay()
|
||||
reference = (x.double() @ w.double().T).bfloat16()
|
||||
torch.testing.assert_close(actual, reference, rtol=0.008, atol=1e-5)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
|
||||
@@ -2424,6 +2424,11 @@ def grouped_topk(
|
||||
raise NotImplementedError(
|
||||
"The fused grouped_topk kernel is only available on CUDA and XPU platforms"
|
||||
)
|
||||
if scores.shape[0] == 0:
|
||||
return (
|
||||
scores.new_empty((0, topk), dtype=torch.float32),
|
||||
scores.new_empty((0, topk), dtype=torch.int32),
|
||||
)
|
||||
return torch.ops._moe_C.grouped_topk(
|
||||
scores,
|
||||
num_expert_group,
|
||||
|
||||
@@ -356,10 +356,12 @@ class LLBf16Dotprod:
|
||||
if tidx == 0:
|
||||
for m in cutlass.range_constexpr(M):
|
||||
partials = sm[m, None].load()
|
||||
gC[m, n_idx] = partials.reduce(
|
||||
cute.ReductionOp.ADD,
|
||||
init_val=cutlass.Float32(0.0),
|
||||
reduction_profile=0,
|
||||
)
|
||||
gC[m, n_idx] = cutlass.Float32(
|
||||
partials.reduce(
|
||||
cute.ReductionOp.ADD,
|
||||
init_val=cutlass.Float32(0.0),
|
||||
reduction_profile=0,
|
||||
)
|
||||
).to(gC.element_type)
|
||||
if const_expr(self.use_pdl):
|
||||
cute.arch.griddepcontrol_launch_dependents()
|
||||
|
||||
@@ -47,7 +47,7 @@ class LLBf16SplitK:
|
||||
|
||||
:param ab_dtype: Element type for A and B operands.
|
||||
:param acc_dtype: Accumulator type used by MMA and reductions.
|
||||
:param out_dtype: Output element type. The public wrapper uses fp32.
|
||||
:param out_dtype: Output element type (fp32 or bf16).
|
||||
:param tile_n: CTA tile size in N. M is fixed to 16 for router batches.
|
||||
:type tile_n: int
|
||||
:param tile_k: K tile size staged through shared memory.
|
||||
@@ -580,6 +580,6 @@ class LLBf16SplitK:
|
||||
reduction_profile=0,
|
||||
)
|
||||
)
|
||||
gC[local_coord] = acc
|
||||
gC[local_coord] = cutlass.Float32(acc).to(gC.element_type)
|
||||
|
||||
cute.arch.sync_threads()
|
||||
|
||||
@@ -76,8 +76,8 @@ def make_fake_gemm_tensors(
|
||||
a_dtype: torch.dtype,
|
||||
b_dtype: torch.dtype,
|
||||
divisibility: int,
|
||||
output_dtype: torch.dtype = torch.float32,
|
||||
):
|
||||
from cutlass import Float32
|
||||
from quack.compile_utils import make_fake_tensor
|
||||
|
||||
hidden_states = make_fake_tensor(
|
||||
@@ -86,7 +86,7 @@ def make_fake_gemm_tensors(
|
||||
router_weight = make_fake_tensor(
|
||||
cutlass_dtype(b_dtype), (N, K), divisibility=divisibility
|
||||
)
|
||||
output = make_fake_tensor(Float32, (M, N), divisibility=1)
|
||||
output = make_fake_tensor(cutlass_dtype(output_dtype), (M, N), divisibility=1)
|
||||
return hidden_states, router_weight, output
|
||||
|
||||
|
||||
@@ -97,6 +97,7 @@ def validate_common_gemm_inputs(
|
||||
*,
|
||||
op_name: str,
|
||||
k_multiple: int | None = None,
|
||||
output_dtypes: tuple[torch.dtype, ...] = (torch.float32,),
|
||||
) -> None:
|
||||
if hidden_states.dim() != 2 or router_weight.dim() != 2:
|
||||
raise ValueError("hidden_states and router_weight must be 2D tensors")
|
||||
@@ -106,8 +107,11 @@ def validate_common_gemm_inputs(
|
||||
raise ValueError(
|
||||
"hidden_states and router_weight must be on the same CUDA device"
|
||||
)
|
||||
if output_dtype != torch.float32:
|
||||
raise ValueError(f"{op_name} only supports output_dtype=torch.float32")
|
||||
if output_dtype not in output_dtypes:
|
||||
raise ValueError(
|
||||
f"{op_name} only supports output_dtype="
|
||||
+ " or ".join(str(dtype) for dtype in output_dtypes)
|
||||
)
|
||||
if hidden_states.shape[1] != router_weight.shape[1]:
|
||||
raise ValueError(
|
||||
"hidden_states and router_weight must have matching K dimensions"
|
||||
|
||||
@@ -47,6 +47,8 @@ _SM100F_TUNED_SPLITK_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int]]]
|
||||
**{M: (8, 5) for M in (5, 8)},
|
||||
9: (8, 2),
|
||||
},
|
||||
(5120, 1792): {M: (2, 4) for M in range(5, 17)},
|
||||
(5120, 3584): {M: (2, 4) for M in range(5, 17)},
|
||||
(7168, 256): {14: (8, 2)},
|
||||
(6144, 256): {M: (8, 2) for M in (9, 12, 16)},
|
||||
(7168, 384): {M: (7, 5) for M in (13, 16)},
|
||||
@@ -101,39 +103,52 @@ class LLBf16Gemm:
|
||||
bs: int = 0
|
||||
split_k: int = 0
|
||||
num_stages: int = 0
|
||||
output_dtype: torch.dtype = torch.float32
|
||||
|
||||
def __init__(self, *, prefetch_pdl_weights: bool = False) -> None:
|
||||
self._prefetch_pdl_weights = prefetch_pdl_weights
|
||||
# Dot-prod: keyed on (M, K, bs), because M and K are Constexpr.
|
||||
self._compiled_cache: dict[tuple[int, int, int], Any] = {}
|
||||
# Split-K: keyed on (split_k, num_stages), fully shape-dynamic.
|
||||
self._splitk_cache: dict[tuple[int, int], Any] = {}
|
||||
# M and K are compile-time constants for dot-product.
|
||||
self._compiled_cache: dict[tuple[int, int, int, torch.dtype], Any] = {}
|
||||
# Split-K is shape-dynamic; both caches distinguish output dtype.
|
||||
self._splitk_cache: dict[tuple[int, int, torch.dtype], Any] = {}
|
||||
|
||||
def dispatch(self, *, M: int, K: int, N: int) -> CompileKey:
|
||||
def dispatch(
|
||||
self, *, M: int, K: int, N: int, output_dtype: torch.dtype = torch.float32
|
||||
) -> CompileKey:
|
||||
tuned_bs, tuned_splitk = _arch_tuned_configs()
|
||||
if M <= _DEFAULT_DOTPROD_MAX_M or K < 2048:
|
||||
bs = tuned_bs.get((K, N), {}).get(M, _DEFAULT_DOTPROD_BS)
|
||||
return self.CompileKey(backend="dotprod", M=M, K=K, bs=bs)
|
||||
return self.CompileKey(
|
||||
backend="dotprod", M=M, K=K, bs=bs, output_dtype=output_dtype
|
||||
)
|
||||
|
||||
split_k, num_stages = tuned_splitk.get((K, N), {}).get(
|
||||
M, _DEFAULT_SPLITK_CONFIG
|
||||
)
|
||||
return self.CompileKey(backend="splitk", split_k=split_k, num_stages=num_stages)
|
||||
return self.CompileKey(
|
||||
backend="splitk",
|
||||
split_k=split_k,
|
||||
num_stages=num_stages,
|
||||
output_dtype=output_dtype,
|
||||
)
|
||||
|
||||
def get_warmup_keys(
|
||||
self,
|
||||
*,
|
||||
shapes: Iterable[tuple[int, int]],
|
||||
m_values: Iterable[int],
|
||||
output_dtype: torch.dtype = torch.float32,
|
||||
) -> list[CompileKey]:
|
||||
return list(
|
||||
dict.fromkeys(
|
||||
self.dispatch(M=M, K=K, N=N) for K, N in shapes for M in m_values
|
||||
self.dispatch(M=M, K=K, N=N, output_dtype=output_dtype)
|
||||
for K, N in shapes
|
||||
for M in m_values
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _fake_gemm_tensors(*, M, K, N, divisibility: int):
|
||||
def _fake_gemm_tensors(*, M, K, N, divisibility: int, output_dtype: torch.dtype):
|
||||
return make_fake_gemm_tensors(
|
||||
M=M,
|
||||
K=K,
|
||||
@@ -141,10 +156,13 @@ class LLBf16Gemm:
|
||||
a_dtype=torch.bfloat16,
|
||||
b_dtype=torch.bfloat16,
|
||||
divisibility=divisibility,
|
||||
output_dtype=output_dtype,
|
||||
)
|
||||
|
||||
def _compile_splitk(self, compile_key: CompileKey) -> None:
|
||||
cute, _ = cute_context()
|
||||
from cutlass import BFloat16, Float32
|
||||
|
||||
from ._ll_bf16_splitk import LLBf16SplitK
|
||||
|
||||
hidden_states, router_weight, output = self._fake_gemm_tensors(
|
||||
@@ -152,8 +170,12 @@ class LLBf16Gemm:
|
||||
K=cute.sym_int(),
|
||||
N=cute.sym_int(),
|
||||
divisibility=8,
|
||||
output_dtype=compile_key.output_dtype,
|
||||
)
|
||||
gemm = LLBf16SplitK(
|
||||
out_dtype=BFloat16
|
||||
if compile_key.output_dtype == torch.bfloat16
|
||||
else Float32,
|
||||
tile_n=16,
|
||||
tile_k=256,
|
||||
num_stages=compile_key.num_stages,
|
||||
@@ -169,7 +191,9 @@ class LLBf16Gemm:
|
||||
current_cuda_stream(),
|
||||
options="--enable-tvm-ffi",
|
||||
)
|
||||
self._splitk_cache[(compile_key.split_k, compile_key.num_stages)] = compiled
|
||||
self._splitk_cache[
|
||||
(compile_key.split_k, compile_key.num_stages, compile_key.output_dtype)
|
||||
] = compiled
|
||||
logger.debug(
|
||||
"Compiled ll_bf16_splitk: sk=%d ns=%d",
|
||||
compile_key.split_k,
|
||||
@@ -187,6 +211,7 @@ class LLBf16Gemm:
|
||||
K=compile_key.K,
|
||||
N=N,
|
||||
divisibility=stride_divisibility,
|
||||
output_dtype=compile_key.output_dtype,
|
||||
)
|
||||
gemm = LLBf16Dotprod(
|
||||
k=compile_key.K,
|
||||
@@ -205,7 +230,9 @@ class LLBf16Gemm:
|
||||
current_cuda_stream(),
|
||||
options="--enable-tvm-ffi --ptxas-options -maxrregcount=64",
|
||||
)
|
||||
self._compiled_cache[(compile_key.M, compile_key.K, compile_key.bs)] = compiled
|
||||
self._compiled_cache[
|
||||
(compile_key.M, compile_key.K, compile_key.bs, compile_key.output_dtype)
|
||||
] = compiled
|
||||
logger.debug(
|
||||
"Compiled ll_bf16_dotprod: M=%d, K=%d, bs=%d",
|
||||
compile_key.M,
|
||||
@@ -215,12 +242,21 @@ class LLBf16Gemm:
|
||||
|
||||
def compile(self, compile_key: CompileKey) -> None:
|
||||
if compile_key.backend == "splitk":
|
||||
splitk_cache_key = (compile_key.split_k, compile_key.num_stages)
|
||||
splitk_cache_key = (
|
||||
compile_key.split_k,
|
||||
compile_key.num_stages,
|
||||
compile_key.output_dtype,
|
||||
)
|
||||
if splitk_cache_key not in self._splitk_cache:
|
||||
self._compile_splitk(compile_key)
|
||||
return
|
||||
|
||||
dotprod_cache_key = (compile_key.M, compile_key.K, compile_key.bs)
|
||||
dotprod_cache_key = (
|
||||
compile_key.M,
|
||||
compile_key.K,
|
||||
compile_key.bs,
|
||||
compile_key.output_dtype,
|
||||
)
|
||||
if dotprod_cache_key not in self._compiled_cache:
|
||||
self._compile_dotprod(compile_key)
|
||||
|
||||
@@ -229,8 +265,11 @@ class LLBf16Gemm:
|
||||
*,
|
||||
shapes: Iterable[tuple[int, int]],
|
||||
m_values: Iterable[int],
|
||||
output_dtype: torch.dtype = torch.float32,
|
||||
) -> None:
|
||||
for compile_key in self.get_warmup_keys(shapes=shapes, m_values=m_values):
|
||||
for compile_key in self.get_warmup_keys(
|
||||
shapes=shapes, m_values=m_values, output_dtype=output_dtype
|
||||
):
|
||||
self.compile(compile_key)
|
||||
|
||||
@staticmethod
|
||||
@@ -251,6 +290,7 @@ class LLBf16Gemm:
|
||||
output_dtype,
|
||||
op_name="ll_bf16_gemm",
|
||||
k_multiple=8,
|
||||
output_dtypes=(torch.float32, torch.bfloat16),
|
||||
)
|
||||
|
||||
def __call__(
|
||||
@@ -258,19 +298,28 @@ class LLBf16Gemm:
|
||||
hidden_states: torch.Tensor, # [M, K] bf16
|
||||
router_weight: torch.Tensor, # [N, K] bf16
|
||||
output_dtype: torch.dtype = torch.float32,
|
||||
) -> torch.Tensor: # [M, N] fp32
|
||||
) -> torch.Tensor: # [M, N]
|
||||
self._validate_inputs(hidden_states, router_weight, output_dtype)
|
||||
|
||||
M, K = hidden_states.shape
|
||||
N = router_weight.shape[0]
|
||||
compile_key = self.dispatch(M=M, K=K, N=N)
|
||||
compile_key = self.dispatch(M=M, K=K, N=N, output_dtype=output_dtype)
|
||||
if compile_key.backend == "splitk":
|
||||
splitk_cache_key = (compile_key.split_k, compile_key.num_stages)
|
||||
splitk_cache_key = (
|
||||
compile_key.split_k,
|
||||
compile_key.num_stages,
|
||||
compile_key.output_dtype,
|
||||
)
|
||||
if splitk_cache_key not in self._splitk_cache:
|
||||
self.compile(compile_key)
|
||||
kernel = self._splitk_cache[splitk_cache_key]
|
||||
else:
|
||||
dotprod_cache_key = (compile_key.M, compile_key.K, compile_key.bs)
|
||||
dotprod_cache_key = (
|
||||
compile_key.M,
|
||||
compile_key.K,
|
||||
compile_key.bs,
|
||||
compile_key.output_dtype,
|
||||
)
|
||||
if dotprod_cache_key not in self._compiled_cache:
|
||||
self.compile(compile_key)
|
||||
kernel = self._compiled_cache[dotprod_cache_key]
|
||||
|
||||
@@ -517,6 +517,7 @@ _TILE_CONFIG_OVERRIDES = {
|
||||
256: (128, 32),
|
||||
512: (128, 19),
|
||||
},
|
||||
(5120, 256): {64: (64, 64)}, # GB200
|
||||
(6144, 128): { # MiniMax-M3
|
||||
64: (32, 64),
|
||||
128: (64, 64),
|
||||
@@ -533,7 +534,7 @@ def _pick_tile_config(N: int, K: int, M: int, num_sms: int) -> tuple[int, int]:
|
||||
if 32 < N <= 512:
|
||||
token_bucket = triton.next_power_of_2(N)
|
||||
overrides = _TILE_CONFIG_OVERRIDES.get((K, M))
|
||||
if overrides is not None:
|
||||
if overrides is not None and token_bucket in overrides:
|
||||
BN, split_k = overrides[token_bucket]
|
||||
return BN, min(split_k, k_tiles)
|
||||
|
||||
@@ -609,6 +610,8 @@ def _bf16x3_router_gemm(X: torch.Tensor, W: torch.Tensor) -> torch.Tensor:
|
||||
"""Return ``X @ W.T`` using the SM100 BF16x3 router GEMM kernel."""
|
||||
N, K = X.shape
|
||||
M, _ = W.shape
|
||||
if N == 0:
|
||||
return X.new_empty((N, M), dtype=torch.float32)
|
||||
num_sms = torch.cuda.get_device_properties(X.device).multi_processor_count
|
||||
|
||||
BN, split_k = _pick_tile_config(N, K, M, num_sms)
|
||||
|
||||
@@ -24,8 +24,8 @@ class GateLinear(ReplicatedLinear):
|
||||
|
||||
1. cuteDSL ll_bf16_gemm (SM90+, M<=16, bf16 in, fp32 out,
|
||||
any shape w/ K divisible by 8)
|
||||
2. fp32 specialized kernel (bf16/fp32 in, fp32 out, M<=32; SM90+
|
||||
model-specific shapes, gfx950 any shape)
|
||||
2. fp32 specialized kernel (bf16/fp32 in, fp32 out; SM90+
|
||||
model-specific shapes and batch limits, gfx950 any shape with M<=32)
|
||||
3. cuteDSL ll_fp32w_gemm (SM90+, M<=32, bf16/fp32 in, fp32 out,
|
||||
any shape)
|
||||
4. bf16x3 CuteDSL kernel (SM100, bf16 in, fp32 weight)
|
||||
@@ -46,9 +46,8 @@ class GateLinear(ReplicatedLinear):
|
||||
split-K from the token count and skip the batch-invariant linear path.
|
||||
"""
|
||||
|
||||
# (hidden_size, num_experts) pairs with an instantiated fp32 kernel:
|
||||
# (3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3
|
||||
FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128)}
|
||||
# (hidden_size, num_experts) pairs with an instantiated fp32 kernel.
|
||||
FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128), (5120, 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.
|
||||
@@ -454,9 +453,9 @@ def fp32_router_gemm_dispatch_impl(
|
||||
weight: torch.Tensor,
|
||||
allow_bf16x3_router_gemm: bool,
|
||||
) -> torch.Tensor:
|
||||
"""Dynamically run fp32 specialized gemm if num_tokens <= FP32_MAX_TOKENS,
|
||||
otherwise optionally run the experimental BF16x3 kernel for medium/large
|
||||
SM100 router batches, then fall back to F.linear.
|
||||
"""Run the FP32 kernel for small batches, with shape-specific limits.
|
||||
|
||||
Larger SM100 BF16 batches use BF16x3; other inputs fall back to F.linear.
|
||||
This must be wrapped in a custom op because our torch.compile integration
|
||||
does not support runtime dispatching on num_tokens.
|
||||
"""
|
||||
@@ -473,7 +472,15 @@ def fp32_router_gemm_dispatch_impl(
|
||||
if weight.dtype == torch.bfloat16:
|
||||
return torch.mm(x, weight.T, out_dtype=torch.float32)
|
||||
return torch.nn.functional.linear(x.float(), weight)
|
||||
if x.shape[0] <= _FP32_ROUTER_GEMM_MAX_TOKENS:
|
||||
max_tokens = _FP32_ROUTER_GEMM_MAX_TOKENS
|
||||
if (
|
||||
allow_bf16x3_router_gemm
|
||||
and x.dtype == torch.bfloat16
|
||||
and weight.shape == (256, 5120)
|
||||
):
|
||||
# GB200: four token groups win through 48 for multiples of eight.
|
||||
max_tokens = 48 if x.shape[0] % 8 == 0 else 16
|
||||
if x.shape[0] <= max_tokens:
|
||||
return ops.fp32_router_gemm(x, weight)
|
||||
|
||||
if allow_bf16x3_router_gemm and x.dtype == torch.bfloat16:
|
||||
|
||||
Reference in New Issue
Block a user