[minimax-m3] Split 1/4: sparse attention ops + JIT kernels + config foundation (#28712)
This commit is contained in:
Xinyuan Tong
@@ -0,0 +1,265 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Reference-vs-fused unit tests for the MiniMax-M3 ROCm native MXFP8 ops.
|
||||
|
||||
Each fused kernel has a slow PyTorch / dequant-to-bf16 reference; these assert
|
||||
the two agree within tolerance:
|
||||
|
||||
* Fused MXFP8 activation quant (Triton) -> torch reference
|
||||
* Native MXFP8 linear (tl.dot_scaled) -> dequant-to-bf16 @ matmul
|
||||
* Native MXFP8 MoE (dot_scaled grouped GEMM) -> dequant-to-bf16 MoE math
|
||||
|
||||
ROCm-only. The pure quant test runs on any ROCm arch; the native MXFP8
|
||||
``dot_scaled`` linear/MoE tests are gated to CDNA4 gfx95x (the hardware
|
||||
microscaling matrix cores) -- gfx942 has no native ``dot_scaled`` MX path.
|
||||
|
||||
Run: pytest python/sglang/jit_kernel/tests/test_minimax_m3_mxfp8.py -v
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.utils import is_hip
|
||||
|
||||
if not is_hip():
|
||||
pytest.skip(
|
||||
"MiniMax-M3 native MXFP8 ops are the ROCm path.", allow_module_level=True
|
||||
)
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("Requires a GPU.", allow_module_level=True)
|
||||
|
||||
from sglang.srt.layers.quantization.mxfp8_amd_gfx95 import ( # noqa: E402
|
||||
_mxfp8_dot_scaled_linear,
|
||||
_mxfp8_e4m3_quantize_torch,
|
||||
_mxfp8_e4m3_quantize_triton,
|
||||
dequant_mxfp8_to_bf16,
|
||||
)
|
||||
|
||||
DEVICE = "cuda"
|
||||
|
||||
|
||||
def _gcn_arch() -> str:
|
||||
try:
|
||||
return torch.cuda.get_device_properties(0).gcnArchName
|
||||
except Exception: # pragma: no cover - no device / non-AMD
|
||||
return ""
|
||||
|
||||
|
||||
requires_gfx950 = pytest.mark.skipif(
|
||||
"gfx95" not in _gcn_arch(),
|
||||
reason="native MXFP8 dot_scaled is a CDNA4 (gfx95x) feature; "
|
||||
"gfx942 has no native dot_scaled MX path.",
|
||||
)
|
||||
|
||||
|
||||
def _relerr(a: torch.Tensor, b: torch.Tensor) -> float:
|
||||
a = a.float()
|
||||
b = b.float()
|
||||
return ((a - b).norm() / (b.norm() + 1e-8)).item()
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Fused MXFP8 activation quant (Triton vs torch reference)
|
||||
# --------------------------------------------------------------------------- #
|
||||
@pytest.mark.parametrize("shape", [(64, 4096), (1, 6144), (333, 2048)])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@torch.inference_mode()
|
||||
def test_mxfp8_quant_triton_matches_torch(shape, dtype):
|
||||
torch.manual_seed(0)
|
||||
x = torch.randn(*shape, device=DEVICE, dtype=dtype)
|
||||
xq_t, s_t = _mxfp8_e4m3_quantize_torch(x)
|
||||
xq_k, s_k = _mxfp8_e4m3_quantize_triton(x)
|
||||
assert s_k.shape == s_t.shape == (shape[0], shape[1] // 32)
|
||||
# E8M0 block exponents share the round-up ceil(log2(amax/e4m3_max))+127
|
||||
# algorithm; allow at most a 1-step difference at exact powers of two.
|
||||
assert (s_k.int() - s_t.int()).abs().max().item() <= 1
|
||||
# Dequantized values agree to fp8 granularity.
|
||||
deq_t = dequant_mxfp8_to_bf16(xq_t, s_t)
|
||||
deq_k = dequant_mxfp8_to_bf16(xq_k, s_k)
|
||||
assert _relerr(deq_k, deq_t) < 1e-2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m,inter", [(8, 512), (65, 2048)])
|
||||
@torch.inference_mode()
|
||||
def test_minimax_swiglu_mxfp8_quant_matches_unfused_fp32(m, inter):
|
||||
# The fused swiglu+quant kernel keeps the activation in fp32 through the
|
||||
# E8M0 scale selection (no bf16 round-trip; matches the vLLM/ame kernel), so
|
||||
# the reference is the unfused fp32 swiglu followed by MXFP8 quant. Not
|
||||
# bit-identical because the reference quant runs in torch vs the fused triton
|
||||
# path, but numerically equivalent (tight relerr, scales agree within 1 ulp).
|
||||
from sglang.jit_kernel.minimax_m3 import (
|
||||
swiglu_oai_mxfp8_quant,
|
||||
swiglu_oai_split,
|
||||
)
|
||||
from sglang.srt.layers.quantization.mxfp8_amd_gfx95 import mxfp8_e4m3_quantize
|
||||
|
||||
torch.manual_seed(0)
|
||||
alpha, beta, limit = 1.702, 1.0, 7.0
|
||||
gate_up = torch.randn(m, 2 * inter, device=DEVICE, dtype=torch.bfloat16) * 0.5
|
||||
|
||||
act = swiglu_oai_split(
|
||||
gate_up, alpha=alpha, beta=beta, limit=limit, out_dtype=torch.float32
|
||||
)
|
||||
q_ref, s_ref = mxfp8_e4m3_quantize(act)
|
||||
q, s = swiglu_oai_mxfp8_quant(gate_up, alpha=alpha, beta=beta, limit=limit)
|
||||
|
||||
assert q.shape == q_ref.shape
|
||||
assert s.shape == s_ref.shape
|
||||
# E8M0 block scales agree within one exponent step (last-bit amax differences).
|
||||
assert (s.int() - s_ref.int()).abs().max().item() <= 1
|
||||
assert (
|
||||
_relerr(dequant_mxfp8_to_bf16(q, s), dequant_mxfp8_to_bf16(q_ref, s_ref)) < 1e-2
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Native MXFP8 linear (dot_scaled) vs dequant-to-bf16 matmul
|
||||
# --------------------------------------------------------------------------- #
|
||||
@requires_gfx950
|
||||
@pytest.mark.parametrize("m,n,k", [(64, 256, 128), (37, 512, 256), (1, 6144, 4096)])
|
||||
@torch.inference_mode()
|
||||
def test_mxfp8_native_linear(m, n, k):
|
||||
torch.manual_seed(0)
|
||||
w_bf16 = torch.randn(n, k, device=DEVICE, dtype=torch.bfloat16) * 0.1
|
||||
w_fp8, w_scale = _mxfp8_e4m3_quantize_torch(w_bf16)
|
||||
x = torch.randn(m, k, device=DEVICE, dtype=torch.bfloat16) * 0.5
|
||||
|
||||
got = _mxfp8_dot_scaled_linear(x, w_fp8, w_scale)
|
||||
# Reference consumes the SAME quantized weights (isolates activation-quant
|
||||
# noise) -> dequant to bf16, plain matmul.
|
||||
w_deq = dequant_mxfp8_to_bf16(w_fp8, w_scale)
|
||||
ref = torch.nn.functional.linear(x, w_deq).to(x.dtype)
|
||||
assert got.shape == (m, n)
|
||||
assert _relerr(got, ref) < 5e-2
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Native MXFP8 MoE (dot_scaled grouped GEMM) vs dequant-to-bf16 MoE math
|
||||
# --------------------------------------------------------------------------- #
|
||||
def _ref_moe(x, w13, w2, topk_weights, topk_ids, alpha, beta, limit):
|
||||
T, H = x.shape
|
||||
inter = w2.shape[-1]
|
||||
top_k = topk_ids.shape[1]
|
||||
out = torch.zeros(T, H, device=x.device, dtype=torch.float32)
|
||||
for t in range(T):
|
||||
for j in range(top_k):
|
||||
e = int(topk_ids[t, j].item())
|
||||
if e < 0 or e >= w13.shape[0]:
|
||||
continue
|
||||
g1 = x[t].float() @ w13[e].float().T # [2I]
|
||||
gate = g1[:inter]
|
||||
up = g1[inter:]
|
||||
if limit is not None:
|
||||
gate = gate.clamp(max=limit)
|
||||
up = up.clamp(min=-limit, max=limit)
|
||||
act = gate * torch.sigmoid(alpha * gate) * (up + beta)
|
||||
g2 = act @ w2[e].float().T # [H]
|
||||
out[t] += topk_weights[t, j].float() * g2
|
||||
return out.to(x.dtype)
|
||||
|
||||
|
||||
@requires_gfx950
|
||||
@pytest.mark.parametrize(
|
||||
"T,H,inter,E,top_k", [(8, 256, 512, 8, 2), (1, 512, 256, 16, 4)]
|
||||
)
|
||||
@torch.inference_mode()
|
||||
def test_mxfp8_native_moe(T, H, inter, E, top_k):
|
||||
from sglang.srt.layers.moe.moe_runner.triton_utils.mxfp8_moe_amd_gfx95 import (
|
||||
fused_moe_mxfp8_native,
|
||||
)
|
||||
|
||||
torch.manual_seed(0)
|
||||
alpha, beta, limit = 1.702, 1.0, 7.0
|
||||
w13_bf16 = torch.randn(E, 2 * inter, H, device=DEVICE, dtype=torch.bfloat16) * 0.1
|
||||
w2_bf16 = torch.randn(E, H, inter, device=DEVICE, dtype=torch.bfloat16) * 0.1
|
||||
w13_fp8, w13_scale = _mxfp8_e4m3_quantize_torch(w13_bf16)
|
||||
w2_fp8, w2_scale = _mxfp8_e4m3_quantize_torch(w2_bf16)
|
||||
|
||||
x = torch.randn(T, H, device=DEVICE, dtype=torch.bfloat16) * 0.5
|
||||
logits = torch.randn(T, E, device=DEVICE, dtype=torch.float32)
|
||||
topk_weights, topk_ids = logits.softmax(dim=-1).topk(top_k, dim=-1)
|
||||
topk_weights = topk_weights.to(torch.float32)
|
||||
topk_ids = topk_ids.to(torch.int32)
|
||||
|
||||
got = fused_moe_mxfp8_native(
|
||||
x,
|
||||
w13_fp8,
|
||||
w13_scale,
|
||||
w2_fp8,
|
||||
w2_scale,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
alpha=alpha,
|
||||
beta=beta,
|
||||
limit=limit,
|
||||
)
|
||||
# Reference consumes the dequantized weights (same bits the kernel reads).
|
||||
w13_deq = dequant_mxfp8_to_bf16(w13_fp8, w13_scale)
|
||||
w2_deq = dequant_mxfp8_to_bf16(w2_fp8, w2_scale)
|
||||
ref = _ref_moe(x, w13_deq, w2_deq, topk_weights, topk_ids, alpha, beta, limit)
|
||||
assert got.shape == (T, H)
|
||||
assert _relerr(got, ref) < 5e-2
|
||||
|
||||
|
||||
@requires_gfx950
|
||||
@torch.inference_mode()
|
||||
def test_mxfp8_native_moe_ep_expert_map_filters_non_local_routes():
|
||||
from sglang.srt.layers.moe.moe_runner.triton_utils.mxfp8_moe_amd_gfx95 import (
|
||||
fused_moe_mxfp8_native,
|
||||
)
|
||||
|
||||
torch.manual_seed(0)
|
||||
T, H, inter = 4, 256, 512
|
||||
local_E = 3
|
||||
alpha, beta, limit = 1.702, 1.0, 7.0
|
||||
|
||||
w13_bf16 = (
|
||||
torch.randn(local_E, 2 * inter, H, device=DEVICE, dtype=torch.bfloat16) * 0.1
|
||||
)
|
||||
w2_bf16 = torch.randn(local_E, H, inter, device=DEVICE, dtype=torch.bfloat16) * 0.1
|
||||
w13_fp8, w13_scale = _mxfp8_e4m3_quantize_torch(w13_bf16)
|
||||
w2_fp8, w2_scale = _mxfp8_e4m3_quantize_torch(w2_bf16)
|
||||
|
||||
x = torch.randn(T, H, device=DEVICE, dtype=torch.bfloat16) * 0.5
|
||||
topk_ids_global = torch.tensor(
|
||||
[[0, 1, 4], [2, 3, 5], [4, 0, 3], [5, 1, 2]],
|
||||
device=DEVICE,
|
||||
dtype=torch.int32,
|
||||
)
|
||||
topk_weights = torch.tensor(
|
||||
[
|
||||
[0.50, 0.25, 0.25],
|
||||
[0.40, 0.30, 0.30],
|
||||
[0.70, 0.20, 0.10],
|
||||
[0.60, 0.30, 0.10],
|
||||
],
|
||||
device=DEVICE,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
expert_map = torch.tensor([0, -1, 1, -1, 2, -1], device=DEVICE, dtype=torch.int32)
|
||||
|
||||
got = fused_moe_mxfp8_native(
|
||||
x,
|
||||
w13_fp8,
|
||||
w13_scale,
|
||||
w2_fp8,
|
||||
w2_scale,
|
||||
topk_weights,
|
||||
topk_ids_global,
|
||||
alpha=alpha,
|
||||
beta=beta,
|
||||
limit=limit,
|
||||
expert_map=expert_map,
|
||||
)
|
||||
|
||||
topk_ids_local = expert_map[topk_ids_global.long()]
|
||||
w13_deq = dequant_mxfp8_to_bf16(w13_fp8, w13_scale)
|
||||
w2_deq = dequant_mxfp8_to_bf16(w2_fp8, w2_scale)
|
||||
ref = _ref_moe(x, w13_deq, w2_deq, topk_weights, topk_ids_local, alpha, beta, limit)
|
||||
assert got.shape == (T, H)
|
||||
assert _relerr(got, ref) < 5e-2
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -0,0 +1,78 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Reference tests for MiniMax-M3 ROCm Gemma RMSNorm Triton kernels."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.utils import is_hip
|
||||
|
||||
if not is_hip():
|
||||
pytest.skip(
|
||||
"MiniMax-M3 Gemma RMSNorm Triton kernels are ROCm-only.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("Requires a GPU.", allow_module_level=True)
|
||||
|
||||
from sglang.jit_kernel.minimax_m3.rmsnorm import ( # noqa: E402
|
||||
gemma_fused_add_rmsnorm,
|
||||
gemma_rmsnorm,
|
||||
)
|
||||
|
||||
DEVICE = "cuda"
|
||||
EPS = 1e-6
|
||||
|
||||
|
||||
def _gemma_rmsnorm_ref(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
|
||||
orig_dtype = x.dtype
|
||||
x_f = x.float()
|
||||
variance = x_f.pow(2).mean(dim=-1, keepdim=True)
|
||||
out = x_f * torch.rsqrt(variance + EPS)
|
||||
out = out * (1.0 + weight.float())
|
||||
return out.to(orig_dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", [(1, 512), (64, 6144), (257, 6144)])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@torch.inference_mode()
|
||||
def test_gemma_rmsnorm_matches_reference(shape, dtype):
|
||||
torch.manual_seed(0)
|
||||
x = torch.randn(*shape, device=DEVICE, dtype=dtype)
|
||||
weight = torch.randn(shape[-1], device=DEVICE, dtype=torch.float32)
|
||||
|
||||
got = gemma_rmsnorm(x, weight, EPS)
|
||||
ref = _gemma_rmsnorm_ref(x, weight)
|
||||
|
||||
torch.testing.assert_close(got, ref, atol=2e-2, rtol=2e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@torch.inference_mode()
|
||||
def test_gemma_rmsnorm_accepts_strided_2d_input(dtype):
|
||||
torch.manual_seed(0)
|
||||
base = torch.randn(128, 1024, device=DEVICE, dtype=dtype)
|
||||
x = base[:, ::2]
|
||||
weight = torch.randn(x.shape[-1], device=DEVICE, dtype=torch.float32)
|
||||
|
||||
assert not x.is_contiguous()
|
||||
got = gemma_rmsnorm(x, weight, EPS)
|
||||
ref = _gemma_rmsnorm_ref(x, weight)
|
||||
|
||||
torch.testing.assert_close(got, ref, atol=2e-2, rtol=2e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", [(1, 512), (64, 6144), (257, 6144)])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@torch.inference_mode()
|
||||
def test_gemma_fused_add_rmsnorm_matches_reference(shape, dtype):
|
||||
torch.manual_seed(0)
|
||||
x = torch.randn(*shape, device=DEVICE, dtype=dtype)
|
||||
residual = torch.randn(*shape, device=DEVICE, dtype=dtype)
|
||||
weight = torch.randn(shape[-1], device=DEVICE, dtype=torch.float32)
|
||||
|
||||
got, residual_out = gemma_fused_add_rmsnorm(x, residual, weight, EPS)
|
||||
ref_residual = x + residual
|
||||
ref = _gemma_rmsnorm_ref(ref_residual, weight)
|
||||
|
||||
torch.testing.assert_close(residual_out, ref_residual, atol=2e-2, rtol=2e-2)
|
||||
torch.testing.assert_close(got, ref, atol=2e-2, rtol=2e-2)
|
||||
@@ -0,0 +1,493 @@
|
||||
"""
|
||||
Correctness tests for the moe_topk_sigmoid JIT kernel.
|
||||
|
||||
Validates against a pure-PyTorch reference and, when sgl_kernel is available,
|
||||
cross-checks against the AOT implementation.
|
||||
"""
|
||||
|
||||
import itertools
|
||||
import os
|
||||
import sys
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.moe_topk_sigmoid import topk_sigmoid
|
||||
|
||||
try:
|
||||
from sgl_kernel import topk_sigmoid as topk_sigmoid_aot
|
||||
|
||||
AOT_AVAILABLE = True
|
||||
except ImportError:
|
||||
AOT_AVAILABLE = False
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CI / full-range helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_is_ci = (
|
||||
os.getenv("CI", "false").lower() == "true"
|
||||
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
|
||||
)
|
||||
|
||||
# Power-of-2 configs covered by static dispatch (num_experts 1–256)
|
||||
# Plus 48 (non-power-of-2) to exercise the fallback path
|
||||
NUM_TOKENS_FULL = [1, 16, 128, 512, 1024, 2048]
|
||||
NUM_TOKENS_CI = [1, 128, 1024]
|
||||
|
||||
NUM_EXPERTS_FULL = [16, 32, 64, 128, 256, 48] # 48 = fallback path
|
||||
NUM_EXPERTS_CI = [16, 64, 48]
|
||||
|
||||
TOPK_FULL = [1, 2, 4, 8]
|
||||
TOPK_CI = [1, 4]
|
||||
|
||||
DTYPES_FULL = [torch.float32]
|
||||
DTYPES_CI = [torch.float32, torch.bfloat16]
|
||||
|
||||
NUM_TOKENS = NUM_TOKENS_CI if _is_ci else NUM_TOKENS_FULL
|
||||
NUM_EXPERTS = NUM_EXPERTS_CI if _is_ci else NUM_EXPERTS_FULL
|
||||
TOPK_LIST = TOPK_CI if _is_ci else TOPK_FULL
|
||||
DTYPES = DTYPES_CI if _is_ci else DTYPES_FULL
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pure-PyTorch reference
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def grouped_topk_gpu(
|
||||
gating_output: torch.Tensor,
|
||||
topk: int,
|
||||
renormalize: bool,
|
||||
num_expert_group: Optional[int] = None,
|
||||
topk_group: Optional[int] = None,
|
||||
num_fused_shared_experts: int = 0,
|
||||
routed_scaling_factor: Optional[float] = None,
|
||||
apply_routed_scaling_factor_on_output: Optional[bool] = False,
|
||||
scoring_func: str = "softmax",
|
||||
):
|
||||
|
||||
# Scoring function: softmax or sigmoid
|
||||
if scoring_func == "softmax":
|
||||
scores = torch.softmax(gating_output, dim=-1)
|
||||
elif scoring_func == "sigmoid":
|
||||
scores = gating_output.sigmoid()
|
||||
else:
|
||||
raise ValueError(f"Unsupported scoring function: {scoring_func}")
|
||||
|
||||
num_token = scores.shape[0]
|
||||
num_experts = scores.shape[1]
|
||||
group_scores = (
|
||||
scores.view(num_token, num_expert_group, -1).max(dim=-1).values
|
||||
) # [n, n_group]
|
||||
|
||||
group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[
|
||||
1
|
||||
] # [n, top_k_group]
|
||||
|
||||
group_mask = torch.zeros_like(group_scores) # [n, n_group]
|
||||
group_mask.scatter_(1, group_idx, 1) # [n, n_group]
|
||||
score_mask = (
|
||||
group_mask.unsqueeze(-1)
|
||||
.expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group)
|
||||
.reshape(num_token, -1)
|
||||
) # [n, e]
|
||||
tmp_scores = scores.masked_fill(
|
||||
~score_mask.bool(), float("-inf")
|
||||
) # [n, e] - use -inf like VLLM
|
||||
|
||||
topk_weights, topk_ids = torch.topk(
|
||||
tmp_scores,
|
||||
k=topk,
|
||||
dim=-1,
|
||||
sorted=(True if num_fused_shared_experts > 0 else True),
|
||||
)
|
||||
|
||||
if num_fused_shared_experts:
|
||||
topk_ids[:, -1] = torch.randint(
|
||||
low=num_experts,
|
||||
high=num_experts + num_fused_shared_experts,
|
||||
size=(topk_ids.size(0),),
|
||||
dtype=topk_ids.dtype,
|
||||
device=topk_ids.device,
|
||||
)
|
||||
if routed_scaling_factor is not None:
|
||||
topk_weights[:, -1] = (
|
||||
topk_weights[:, :-1].sum(dim=-1) / routed_scaling_factor
|
||||
)
|
||||
|
||||
if renormalize:
|
||||
topk_weights_sum = (
|
||||
topk_weights.sum(dim=-1, keepdim=True)
|
||||
if num_fused_shared_experts == 0
|
||||
else topk_weights[:, :-1].sum(dim=-1, keepdim=True)
|
||||
)
|
||||
topk_weights = topk_weights / topk_weights_sum
|
||||
if apply_routed_scaling_factor_on_output:
|
||||
topk_weights *= routed_scaling_factor
|
||||
|
||||
topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to(torch.int32)
|
||||
return topk_weights, topk_ids
|
||||
|
||||
|
||||
def biased_grouped_topk_impl(
|
||||
gating_output: torch.Tensor,
|
||||
correction_bias: torch.Tensor,
|
||||
topk: int,
|
||||
renormalize: bool,
|
||||
num_expert_group: Optional[int] = None,
|
||||
topk_group: Optional[int] = None,
|
||||
num_fused_shared_experts: int = 0,
|
||||
routed_scaling_factor: Optional[float] = None,
|
||||
apply_routed_scaling_factor_on_output: Optional[bool] = False,
|
||||
):
|
||||
scores = gating_output.sigmoid()
|
||||
|
||||
num_token = scores.shape[0]
|
||||
num_experts = scores.shape[1]
|
||||
scores_for_choice = scores.view(num_token, -1) + correction_bias.unsqueeze(0)
|
||||
group_scores = (
|
||||
scores_for_choice.view(num_token, num_expert_group, -1)
|
||||
.topk(2, dim=-1)[0]
|
||||
.sum(dim=-1)
|
||||
) # [n, n_group]
|
||||
|
||||
group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[
|
||||
1
|
||||
] # [n, top_k_group]
|
||||
|
||||
group_mask = torch.zeros_like(group_scores) # [n, n_group]
|
||||
group_mask.scatter_(1, group_idx, 1) # [n, n_group]
|
||||
score_mask = (
|
||||
group_mask.unsqueeze(-1)
|
||||
.expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group)
|
||||
.reshape(num_token, -1)
|
||||
) # [n, e]
|
||||
tmp_scores = scores_for_choice.masked_fill(
|
||||
~score_mask.bool(), float("-inf")
|
||||
) # [n, e]
|
||||
|
||||
_, topk_ids = torch.topk(
|
||||
tmp_scores,
|
||||
k=topk,
|
||||
dim=-1,
|
||||
sorted=(True if num_fused_shared_experts > 0 else True),
|
||||
)
|
||||
topk_weights = scores.gather(1, topk_ids)
|
||||
|
||||
if num_fused_shared_experts:
|
||||
topk_ids[:, -1] = torch.randint(
|
||||
low=num_experts,
|
||||
high=num_experts + num_fused_shared_experts,
|
||||
size=(topk_ids.size(0),),
|
||||
dtype=topk_ids.dtype,
|
||||
device=topk_ids.device,
|
||||
)
|
||||
if routed_scaling_factor is not None:
|
||||
topk_weights[:, -1] = (
|
||||
topk_weights[:, :-1].sum(dim=-1) / routed_scaling_factor
|
||||
)
|
||||
if renormalize:
|
||||
topk_weights_sum = (
|
||||
topk_weights.sum(dim=-1, keepdim=True)
|
||||
if num_fused_shared_experts == 0
|
||||
else topk_weights[:, :-1].sum(dim=-1, keepdim=True)
|
||||
)
|
||||
topk_weights = topk_weights / topk_weights_sum
|
||||
if apply_routed_scaling_factor_on_output:
|
||||
topk_weights *= routed_scaling_factor
|
||||
|
||||
topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to(torch.int32)
|
||||
return topk_weights, topk_ids
|
||||
|
||||
|
||||
def topk_sigmoid_torch_ref(
|
||||
gating_output: torch.Tensor,
|
||||
topk: int,
|
||||
renormalize: bool,
|
||||
correction_bias: torch.Tensor | None,
|
||||
num_fused_shared_experts: int = 0,
|
||||
routed_scaling_factor: float = 1.0,
|
||||
apply_routed_scaling_factor_on_output: bool = True,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Reference: sigmoid → (add bias) → topk → (renormalize).
|
||||
Indices are selected on biased scores; weights are the unbiased sigmoid values.
|
||||
"""
|
||||
num_experts = gating_output.shape[1]
|
||||
scores = gating_output.float().sigmoid()
|
||||
biased = scores if correction_bias is None else scores + correction_bias.float()
|
||||
_, ref_ids = torch.topk(biased, k=topk, dim=-1)
|
||||
ref_weights = scores.gather(1, ref_ids)
|
||||
if num_fused_shared_experts > 0:
|
||||
ref_ids[:, -1] = torch.randint(
|
||||
low=num_experts,
|
||||
high=num_experts + num_fused_shared_experts,
|
||||
size=(ref_ids.size(0),),
|
||||
dtype=ref_ids.dtype,
|
||||
device=ref_ids.device,
|
||||
)
|
||||
ref_weights[:, -1] = ref_weights[:, :-1].sum(dim=-1) / routed_scaling_factor
|
||||
if renormalize:
|
||||
topk_weights_sum = (
|
||||
ref_weights.sum(dim=-1, keepdim=True)
|
||||
if num_fused_shared_experts == 0
|
||||
else ref_weights[:, :-1].sum(dim=-1, keepdim=True)
|
||||
)
|
||||
ref_weights = ref_weights / topk_weights_sum
|
||||
if apply_routed_scaling_factor_on_output:
|
||||
ref_weights *= routed_scaling_factor
|
||||
return ref_weights.float(), ref_ids.int()
|
||||
|
||||
|
||||
def topk_sigmoid_grouped_ref(
|
||||
gating_output: torch.Tensor,
|
||||
topk: int,
|
||||
renormalize: bool,
|
||||
correction_bias: torch.Tensor | None,
|
||||
num_fused_shared_experts: int = 0,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
if correction_bias is not None:
|
||||
return biased_grouped_topk_impl(
|
||||
gating_output,
|
||||
correction_bias,
|
||||
topk,
|
||||
renormalize,
|
||||
num_expert_group=1,
|
||||
topk_group=1,
|
||||
num_fused_shared_experts=num_fused_shared_experts,
|
||||
routed_scaling_factor=1.0,
|
||||
apply_routed_scaling_factor_on_output=True,
|
||||
)
|
||||
else:
|
||||
return grouped_topk_gpu(
|
||||
gating_output,
|
||||
topk,
|
||||
renormalize,
|
||||
num_expert_group=1,
|
||||
topk_group=1,
|
||||
num_fused_shared_experts=num_fused_shared_experts,
|
||||
routed_scaling_factor=1.0,
|
||||
apply_routed_scaling_factor_on_output=True,
|
||||
scoring_func="sigmoid",
|
||||
)
|
||||
|
||||
|
||||
def topk_sigmoid_ref(
|
||||
gating_output: torch.Tensor,
|
||||
topk: int,
|
||||
renormalize: bool,
|
||||
correction_bias: torch.Tensor | None,
|
||||
num_fused_shared_experts: int = 0,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
return topk_sigmoid_torch_ref(
|
||||
gating_output, topk, renormalize, correction_bias, num_fused_shared_experts
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Correctness: JIT vs PyTorch reference
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"num_tokens, num_experts, topk",
|
||||
list(itertools.product(NUM_TOKENS, NUM_EXPERTS, TOPK_LIST)),
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("renormalize", [False, True])
|
||||
def test_topk_sigmoid_vs_ref(num_tokens, num_experts, topk, dtype, renormalize):
|
||||
if topk > num_experts:
|
||||
pytest.skip("topk > num_experts")
|
||||
|
||||
torch.manual_seed(num_tokens * num_experts)
|
||||
gating = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
|
||||
|
||||
topk_w = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda")
|
||||
topk_i = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
|
||||
topk_sigmoid(topk_w, topk_i, gating, renormalize=renormalize)
|
||||
|
||||
ref_w, ref_i = topk_sigmoid_ref(gating, topk, renormalize, correction_bias=None)
|
||||
|
||||
# Compare sorted weights (indices may differ for ties when dtype != float32)
|
||||
assert torch.allclose(
|
||||
topk_w.sort(dim=-1)[0],
|
||||
ref_w.sort(dim=-1)[0],
|
||||
atol=1e-3,
|
||||
rtol=1e-3,
|
||||
), f"Weight mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk}, renorm={renormalize})"
|
||||
# Exact index match is only reliable for float32 (fp16/bf16 tie-breaking may differ)
|
||||
if dtype == torch.float32:
|
||||
assert torch.equal(
|
||||
topk_i, ref_i
|
||||
), f"Index mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk})"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Correctness: with correction_bias
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"num_tokens, num_experts, topk",
|
||||
list(itertools.product(NUM_TOKENS, NUM_EXPERTS, TOPK_LIST)),
|
||||
)
|
||||
@pytest.mark.parametrize("renormalize", [False, True])
|
||||
def test_topk_sigmoid_with_correction_bias(num_tokens, num_experts, topk, renormalize):
|
||||
if topk > num_experts:
|
||||
pytest.skip("topk > num_experts")
|
||||
|
||||
torch.manual_seed(num_tokens + num_experts + topk)
|
||||
gating = torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda")
|
||||
bias = torch.randn(num_experts, dtype=torch.float32, device="cuda")
|
||||
|
||||
topk_w = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda")
|
||||
topk_i = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
|
||||
topk_sigmoid(topk_w, topk_i, gating, renormalize=renormalize, correction_bias=bias)
|
||||
|
||||
ref_w, ref_i = topk_sigmoid_ref(gating, topk, renormalize, correction_bias=bias)
|
||||
|
||||
assert torch.allclose(
|
||||
topk_w, ref_w, atol=1e-3, rtol=1e-3
|
||||
), f"Weight mismatch with bias (n_exp={num_experts}, topk={topk}, renorm={renormalize})"
|
||||
assert torch.equal(
|
||||
topk_i, ref_i
|
||||
), f"Index mismatch with bias (n_exp={num_experts}, topk={topk})"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Correctness: with fused shared experts
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"num_tokens, num_experts, topk",
|
||||
list(itertools.product(NUM_TOKENS, NUM_EXPERTS, TOPK_LIST)),
|
||||
)
|
||||
@pytest.mark.parametrize("renormalize", [False, True])
|
||||
def test_topk_sigmoid_with_fused_shared_experts(
|
||||
num_tokens, num_experts, topk, renormalize
|
||||
):
|
||||
if topk + 1 > num_experts:
|
||||
pytest.skip("topk > num_experts")
|
||||
|
||||
torch.manual_seed(num_tokens + num_experts)
|
||||
gating = torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda")
|
||||
bias = torch.randn(num_experts, dtype=torch.float32, device="cuda")
|
||||
|
||||
topk_w = torch.empty((num_tokens, topk + 1), dtype=torch.float32, device="cuda")
|
||||
topk_i = torch.empty((num_tokens, topk + 1), dtype=torch.int32, device="cuda")
|
||||
topk_sigmoid(
|
||||
topk_w,
|
||||
topk_i,
|
||||
gating,
|
||||
renormalize=renormalize,
|
||||
correction_bias=bias,
|
||||
num_fused_shared_experts=1,
|
||||
)
|
||||
|
||||
ref_w, ref_i = topk_sigmoid_ref(
|
||||
gating, topk + 1, renormalize, correction_bias=bias, num_fused_shared_experts=1
|
||||
)
|
||||
|
||||
assert torch.allclose(
|
||||
topk_w, ref_w, atol=1e-3, rtol=1e-3
|
||||
), f"Weight mismatch with bias (n_exp={num_experts}, topk={topk}, renorm={renormalize})"
|
||||
assert torch.equal(
|
||||
topk_i, ref_i
|
||||
), f"Index mismatch with bias (n_exp={num_experts}, topk={topk})"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Renormalization: weights should sum to 1 per row
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens, num_experts, topk", [(128, 64, 4), (1, 8, 2)])
|
||||
def test_renormalize_sums_to_one(num_tokens, num_experts, topk):
|
||||
gating = torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda")
|
||||
topk_w = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda")
|
||||
topk_i = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
|
||||
topk_sigmoid(topk_w, topk_i, gating, renormalize=True)
|
||||
row_sums = topk_w.sum(dim=-1)
|
||||
torch.testing.assert_close(
|
||||
row_sums, torch.ones(num_tokens, device="cuda"), rtol=1e-4, atol=1e-4
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Output shape and dtype
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_output_shapes_and_dtypes():
|
||||
num_tokens, num_experts, topk = 64, 128, 4
|
||||
gating = torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda")
|
||||
topk_w = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda")
|
||||
topk_i = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
|
||||
topk_sigmoid(topk_w, topk_i, gating)
|
||||
|
||||
assert topk_w.shape == (num_tokens, topk)
|
||||
assert topk_i.shape == (num_tokens, topk)
|
||||
assert topk_w.dtype == torch.float32
|
||||
assert topk_i.dtype == torch.int32
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fallback path (non-power-of-2 experts)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_experts", [48, 96])
|
||||
def test_fallback_non_power_of_two(num_experts):
|
||||
num_tokens, topk = 64, 2
|
||||
gating = torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda")
|
||||
topk_w = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda")
|
||||
topk_i = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
|
||||
topk_sigmoid(topk_w, topk_i, gating, renormalize=True)
|
||||
|
||||
# Weights should be positive and sum to 1
|
||||
assert torch.all(topk_w > 0)
|
||||
torch.testing.assert_close(
|
||||
topk_w.sum(dim=-1), torch.ones(num_tokens, device="cuda"), rtol=1e-4, atol=1e-4
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cross-validation against AOT sgl_kernel
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.skipif(not AOT_AVAILABLE, reason="sgl_kernel not available")
|
||||
@pytest.mark.parametrize(
|
||||
"num_tokens, num_experts, topk",
|
||||
list(itertools.product([1, 128, 1024], [8, 64, 128], [1, 4])),
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
|
||||
@pytest.mark.parametrize("renormalize", [False, True])
|
||||
def test_topk_sigmoid_vs_aot(num_tokens, num_experts, topk, dtype, renormalize):
|
||||
if topk > num_experts:
|
||||
pytest.skip("topk > num_experts")
|
||||
|
||||
torch.manual_seed(42)
|
||||
gating = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
|
||||
|
||||
topk_w_jit = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda")
|
||||
topk_i_jit = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
|
||||
topk_sigmoid(topk_w_jit, topk_i_jit, gating, renormalize=renormalize)
|
||||
|
||||
topk_w_aot = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda")
|
||||
topk_i_aot = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
|
||||
topk_sigmoid_aot(topk_w_aot, topk_i_aot, gating, renormalize=renormalize)
|
||||
|
||||
assert torch.allclose(
|
||||
topk_w_jit, topk_w_aot, atol=1e-3, rtol=1e-3
|
||||
), f"JIT vs AOT weight mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk})"
|
||||
assert torch.equal(
|
||||
topk_i_jit, topk_i_aot
|
||||
), f"JIT vs AOT index mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk})"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
Reference in New Issue
Block a user