[Kernel] Reclassify kernel tests by ops group + move helpers out of the package (RFC #29630) (#32128)
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
a2935ce329
commit
2d1a7be8c4
@@ -19,21 +19,28 @@ sglang/kernels/
|
||||
selector.py # heuristic select_kernel() and cached get_kernel()
|
||||
fused_op.py # BaseFusedOp: per-operator multi-backend contract
|
||||
ops/
|
||||
<group>/ # one subpackage per operator group
|
||||
<group>/ # one subpackage per operator group (see list below)
|
||||
jit/ # shared JIT CUDA build/runtime infra: utils/, csrc/,
|
||||
# include/, __main__ (KERNEL_PATH resolves here)
|
||||
```
|
||||
|
||||
Groups populated in this phase: `activation`, `gemm`, `kvcache`, `layernorm`,
|
||||
`moe`, `quantization`. The remaining groups (`attention`, `communication`,
|
||||
`diffusion`, `grammar`, `mamba`, `memory`, `sampling`, `spatial`,
|
||||
`speculative`) are reserved package placeholders whose implementations still
|
||||
live in `sglang.kernels.jit` / `sgl_kernel` / `triton_ops` and will migrate in
|
||||
later phases.
|
||||
Operator groups (all populated): `activation`, `attention`, `communication`,
|
||||
`diffusion`, `embeddings`, `gemm`, `grammar`, `kv_canary`, `kvcache`,
|
||||
`layernorm`, `lplb`, `mamba`, `memory`, `model`, `moe`, `quantization`,
|
||||
`sampling`, `spatial`, `speculative`.
|
||||
|
||||
As of the RFC #29630 finale (#32072) the legacy `sglang.jit_kernel` package has
|
||||
been **removed**: its shared build/runtime infra moved to `sglang.kernels.jit`
|
||||
and each JIT-backed operator into its group as
|
||||
`sglang.kernels.ops.<group>._jit_<op>`. Tests and benchmarks live under
|
||||
`test/registered/kernels/` (`ops/<group>/` for tests, `benchmark/<group>/` for
|
||||
benchmarks); shared test helpers are in `sglang.test.kernels`.
|
||||
|
||||
## How it works
|
||||
|
||||
Implementations are not moved yet. Each `ops.<group>` function is a thin
|
||||
wrapper that forwards to a chosen backend, and every backend is described by a
|
||||
`KernelSpec` in the registry so alternatives can be inventoried and compared:
|
||||
Each `ops.<group>` function is a thin wrapper that forwards to a chosen
|
||||
backend, and every backend is described by a `KernelSpec` in the registry so
|
||||
alternatives can be inventoried and compared:
|
||||
|
||||
- `register_kernel(KernelSpec(...))` records metadata only — an operator id
|
||||
(`"<group>.<name>"`), a backend, and an import path (`"module:attr"`). No
|
||||
@@ -85,8 +92,8 @@ What this buys (see the
|
||||
|
||||
- **Unified correctness testing** — a generic harness enumerates
|
||||
`available_backends()` and asserts each one matches `forward_native`
|
||||
(`test/registered/kernels/test_fused_op_gpu_parity.py`); new backends are
|
||||
picked up automatically.
|
||||
(`test/registered/kernels/ops/layernorm/test_fused_op_gpu_parity.py`); new
|
||||
backends are picked up automatically.
|
||||
- **One-switch debugging** — `SGLANG_FORCE_FUSED_OP_BACKEND=torch` (or
|
||||
`set_fused_op_backend(KernelBackend.TORCH)`) flips *every* fused op to its
|
||||
reference implementation for numerical-bug bisection.
|
||||
|
||||
@@ -1,265 +0,0 @@
|
||||
# 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/kernels/jit/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.kernels.ops.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.kernels.ops.moe.minimax_m3_swiglu import (
|
||||
swiglu_oai_mxfp8_quant,
|
||||
swiglu_oai_split,
|
||||
)
|
||||
from sglang.kernels.ops.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.kernels.ops.moe.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.kernels.ops.moe.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"]))
|
||||
@@ -1,78 +0,0 @@
|
||||
# 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.kernels.ops.layernorm.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)
|
||||
@@ -1,93 +0,0 @@
|
||||
"""fused_moe_preprocess must be bit-identical to the torch.sort-based path,
|
||||
and the grouped GEMM must produce identical results under both block_size_m
|
||||
configs (the block schedule and kernel config are chosen together).
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.moe.moe_runner.triton_utils.inkling_moe import (
|
||||
SMALL_M_BLOCK_SIZE_M,
|
||||
compute_grouped_gemm_metadata,
|
||||
fused_moe_preprocess,
|
||||
get_src2dst,
|
||||
grouped_gemm_triton,
|
||||
)
|
||||
|
||||
requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA only")
|
||||
|
||||
E = 256
|
||||
TOPK = 6
|
||||
|
||||
|
||||
def _reference(topk_ids_flat: torch.Tensor):
|
||||
reorder_topk_ids, reorder_ids = torch.sort(
|
||||
topk_ids_flat.to(torch.int16), stable=True
|
||||
)
|
||||
src2dst = get_src2dst(reorder_ids)
|
||||
meta = compute_grouped_gemm_metadata(
|
||||
reorder_topk_ids, E, block_size_m=SMALL_M_BLOCK_SIZE_M
|
||||
)
|
||||
return (src2dst, *meta, reorder_topk_ids)
|
||||
|
||||
|
||||
def _ids(tokens: int, seed: int, skew: bool = False) -> torch.Tensor:
|
||||
torch.manual_seed(seed)
|
||||
if skew: # all tokens on few experts (stresses multi-block experts)
|
||||
return torch.randint(0, 3, (tokens * TOPK,), dtype=torch.int32, device="cuda")
|
||||
return (
|
||||
torch.stack([torch.randperm(E, device="cuda")[:TOPK] for _ in range(tokens)])
|
||||
.view(-1)
|
||||
.to(torch.int32)
|
||||
)
|
||||
|
||||
|
||||
@requires_cuda
|
||||
@pytest.mark.parametrize("tokens", [1, 2, 7, 32, 64, 170, 341]) # n = 6*T <= 2048
|
||||
@pytest.mark.parametrize("skew", [False, True])
|
||||
def test_matches_sort_path(tokens: int, skew: bool):
|
||||
ids = _ids(tokens, seed=tokens, skew=skew)
|
||||
ref = _reference(ids)
|
||||
got = fused_moe_preprocess(ids, E)
|
||||
names = [
|
||||
"src2dst",
|
||||
"num_tokens_per_expert",
|
||||
"expert_token_offs",
|
||||
"expert_block_offs",
|
||||
"expert_block_schedule",
|
||||
"reorder_topk_ids",
|
||||
]
|
||||
for tag, g, r in zip(names, got, ref):
|
||||
assert g.shape == r.shape, (tag, g.shape, r.shape)
|
||||
assert torch.equal(g.long(), r.long()), (
|
||||
tag,
|
||||
g[: min(16, g.numel())],
|
||||
r[: min(16, r.numel())],
|
||||
)
|
||||
|
||||
|
||||
@requires_cuda
|
||||
@pytest.mark.parametrize("tokens", [1, 16, 64])
|
||||
def test_grouped_gemm_small_config_matches(tokens: int):
|
||||
"""GEMM output must be identical whichever (block_size_m, config) runs."""
|
||||
torch.manual_seed(tokens)
|
||||
ids = _ids(tokens, seed=tokens)
|
||||
m, k, n = tokens * TOPK, 768, 1024
|
||||
a = (torch.randn(m, k, device="cuda") * 0.05).to(torch.bfloat16)
|
||||
b = (torch.randn(E, n, k, device="cuda") * 0.02).to(torch.bfloat16)
|
||||
|
||||
sorted_ids, _ = torch.sort(ids.to(torch.int16), stable=True)
|
||||
meta128 = compute_grouped_gemm_metadata(sorted_ids, E)
|
||||
out128 = grouped_gemm_triton(a, b, E, *meta128)
|
||||
|
||||
pre = fused_moe_preprocess(ids, E)
|
||||
out16 = grouped_gemm_triton(a, b, E, *pre[1:5], block_size_m=SMALL_M_BLOCK_SIZE_M)
|
||||
# both are fp32-accumulated bf16 tensor-core dots; BLOCK_K differs so
|
||||
# accumulation grouping may differ by a few ulp
|
||||
torch.testing.assert_close(out16.float(), out128.float(), atol=1e-3, rtol=1e-3)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v", "-x"]))
|
||||
@@ -1,493 +0,0 @@
|
||||
"""
|
||||
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.kernels.ops.moe.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"]))
|
||||
@@ -1,70 +0,0 @@
|
||||
"""fused_decode_sconv_metadata must be bit-identical to the unfused prep.
|
||||
|
||||
The unfused reference is the exact op sequence `_prepare_decode_sconv_metadata`
|
||||
used to launch: two arange calls + ones + precompute_helion_decode_metadata
|
||||
(!= PAD, &, clamp, long, arange x2).
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.models.inkling_common.kernels.sconv import (
|
||||
PAD_SLOT_ID,
|
||||
fused_decode_sconv_metadata,
|
||||
precompute_helion_decode_metadata,
|
||||
)
|
||||
|
||||
requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA only")
|
||||
|
||||
# cross the BLOCK=1024 grid boundary and hit odd sizes
|
||||
BATCH_SIZES = [1, 2, 3, 17, 64, 160, 257, 1023, 1024, 1025]
|
||||
|
||||
|
||||
def _reference(B: int, cache_indices: torch.Tensor):
|
||||
device = cache_indices.device
|
||||
query_start_loc = torch.arange(B + 1, dtype=torch.int32, device=device)
|
||||
has_initial_state = torch.ones(B, dtype=torch.bool, device=device)
|
||||
precomputed = precompute_helion_decode_metadata(
|
||||
B=B, W=4, cache_indices=cache_indices, has_initial_state=has_initial_state
|
||||
)
|
||||
return query_start_loc, has_initial_state, precomputed
|
||||
|
||||
|
||||
@requires_cuda
|
||||
@pytest.mark.parametrize("b", BATCH_SIZES)
|
||||
@pytest.mark.parametrize("idx_dtype", [torch.int32, torch.int64])
|
||||
def test_matches_unfused(b: int, idx_dtype: torch.dtype):
|
||||
torch.manual_seed(b)
|
||||
cache_indices = torch.randint(0, 4096, (b,), dtype=idx_dtype, device="cuda")
|
||||
# sprinkle PAD slots (cudagraph padding lanes)
|
||||
pad = torch.rand(b, device="cuda") < 0.25
|
||||
cache_indices[pad] = PAD_SLOT_ID
|
||||
|
||||
ref_qsl, ref_his, ref_meta = _reference(b, cache_indices)
|
||||
qsl, his, meta = fused_decode_sconv_metadata(B=b, cache_indices=cache_indices)
|
||||
|
||||
for tag, got, ref in (
|
||||
("query_start_loc", qsl, ref_qsl),
|
||||
("has_initial_state", his, ref_his),
|
||||
("cache_mask", meta["cache_mask"], ref_meta["cache_mask"]),
|
||||
("safe_idx", meta["safe_idx"], ref_meta["safe_idx"]),
|
||||
("cu", meta["cu"], ref_meta["cu"]),
|
||||
("si", meta["si"], ref_meta["si"]),
|
||||
):
|
||||
assert got.dtype == ref.dtype, (tag, got.dtype, ref.dtype)
|
||||
assert got.shape == ref.shape, (tag, got.shape, ref.shape)
|
||||
assert torch.equal(got, ref), tag
|
||||
|
||||
|
||||
@requires_cuda
|
||||
def test_all_pad():
|
||||
cache_indices = torch.full((8,), PAD_SLOT_ID, dtype=torch.int32, device="cuda")
|
||||
_, _, meta = fused_decode_sconv_metadata(B=8, cache_indices=cache_indices)
|
||||
assert not meta["cache_mask"].any()
|
||||
assert (meta["safe_idx"] == 0).all()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v", "-x"]))
|
||||
@@ -1,174 +0,0 @@
|
||||
"""fused_extend_sconv_metadata must be bit-identical to the unfused prep.
|
||||
|
||||
The unfused reference is the exact op sequence _prepare_extend_common_metadata
|
||||
+ precompute_helion_extend_metadata used to launch: zeros + cumsum + slice-copy
|
||||
(or arange + ones for verify) + the has_initial_state compare, then != PAD, &,
|
||||
clamp, long, to(int64), arange, searchsorted, clamp, to(int32).
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.models.inkling_common.kernels.sconv import (
|
||||
HIS_ONES,
|
||||
HIS_PREFIX,
|
||||
HIS_SEQ_MINUS_EXT,
|
||||
HIS_ZEROS,
|
||||
PAD_SLOT_ID,
|
||||
fused_extend_sconv_metadata,
|
||||
precompute_helion_extend_metadata,
|
||||
)
|
||||
|
||||
requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA only")
|
||||
|
||||
# cross si tiles (BLOCK_T=256) and the single-tile B bound
|
||||
BATCH_SIZES = [1, 2, 7, 64, 257, 1023]
|
||||
|
||||
|
||||
def _ref_extend(B, extend_seq_lens, his_mode, his_src, cache_indices, T):
|
||||
device = cache_indices.device
|
||||
query_start_loc = torch.zeros(B + 1, dtype=torch.int32, device=device)
|
||||
query_start_loc[1:] = extend_seq_lens.cumsum(dim=0)
|
||||
if his_mode == HIS_ZEROS:
|
||||
has_initial_state = torch.zeros(B, dtype=torch.bool, device=device)
|
||||
elif his_mode == HIS_PREFIX:
|
||||
has_initial_state = his_src > 0
|
||||
else: # HIS_SEQ_MINUS_EXT
|
||||
has_initial_state = (his_src[:B] - extend_seq_lens) > 0
|
||||
meta = precompute_helion_extend_metadata(
|
||||
B=B,
|
||||
T=T,
|
||||
W=4,
|
||||
cache_indices=cache_indices,
|
||||
has_initial_state=has_initial_state,
|
||||
query_start_loc=query_start_loc,
|
||||
)
|
||||
return query_start_loc, has_initial_state, meta
|
||||
|
||||
|
||||
def _ref_verify(B, draft_token_num, cache_indices):
|
||||
device = cache_indices.device
|
||||
query_start_loc = torch.arange(
|
||||
0, (B + 1) * draft_token_num, draft_token_num, dtype=torch.int32, device=device
|
||||
)
|
||||
has_initial_state = torch.ones(B, dtype=torch.bool, device=device)
|
||||
meta = precompute_helion_extend_metadata(
|
||||
B=B,
|
||||
T=B * draft_token_num,
|
||||
W=4,
|
||||
cache_indices=cache_indices,
|
||||
has_initial_state=has_initial_state,
|
||||
query_start_loc=query_start_loc,
|
||||
)
|
||||
return query_start_loc, has_initial_state, meta
|
||||
|
||||
|
||||
def _assert_equal(got, ref):
|
||||
for tag, g, r in (
|
||||
("query_start_loc", got[0], ref[0]),
|
||||
("has_initial_state", got[1], ref[1]),
|
||||
("cache_mask", got[2]["cache_mask"], ref[2]["cache_mask"]),
|
||||
("safe_idx", got[2]["safe_idx"], ref[2]["safe_idx"]),
|
||||
("cu", got[2]["cu"], ref[2]["cu"]),
|
||||
("si", got[2]["si"], ref[2]["si"]),
|
||||
):
|
||||
assert g.dtype == r.dtype, (tag, g.dtype, r.dtype)
|
||||
assert g.shape == r.shape, (tag, g.shape, r.shape)
|
||||
assert torch.equal(g, r), tag
|
||||
|
||||
|
||||
def _cache_indices(b, idx_dtype):
|
||||
ci = torch.randint(0, 4096, (b,), dtype=idx_dtype, device="cuda")
|
||||
pad = torch.rand(b, device="cuda") < 0.25
|
||||
ci[pad] = PAD_SLOT_ID
|
||||
return ci
|
||||
|
||||
|
||||
@requires_cuda
|
||||
@pytest.mark.parametrize("b", BATCH_SIZES)
|
||||
@pytest.mark.parametrize("his_mode", [HIS_ZEROS, HIS_PREFIX, HIS_SEQ_MINUS_EXT])
|
||||
@pytest.mark.parametrize("lens_dtype", [torch.int32, torch.int64])
|
||||
def test_extend_matches_unfused(b, his_mode, lens_dtype):
|
||||
torch.manual_seed(b * 10 + his_mode)
|
||||
lens = torch.randint(0, 33, (b,), dtype=lens_dtype, device="cuda")
|
||||
lens[torch.rand(b, device="cuda") < 0.2] = 0 # zero-length sequences
|
||||
T = int(lens.sum().item())
|
||||
cache_indices = _cache_indices(b, torch.int32)
|
||||
if his_mode == HIS_PREFIX:
|
||||
his_src = torch.randint(0, 3, (b,), dtype=lens_dtype, device="cuda")
|
||||
elif his_mode == HIS_SEQ_MINUS_EXT:
|
||||
his_src = lens + torch.randint(0, 2, (b,), dtype=lens_dtype, device="cuda")
|
||||
else:
|
||||
his_src = None
|
||||
|
||||
ref = _ref_extend(b, lens, his_mode, his_src, cache_indices, T)
|
||||
got = fused_extend_sconv_metadata(
|
||||
B=b,
|
||||
T=T,
|
||||
cache_indices=cache_indices,
|
||||
his_mode=his_mode,
|
||||
extend_seq_lens=lens,
|
||||
his_src=his_src,
|
||||
)
|
||||
assert got is not None
|
||||
_assert_equal(got, ref)
|
||||
|
||||
|
||||
@requires_cuda
|
||||
@pytest.mark.parametrize("b", BATCH_SIZES)
|
||||
@pytest.mark.parametrize("draft_token_num", [1, 9])
|
||||
def test_verify_matches_unfused(b, draft_token_num):
|
||||
torch.manual_seed(b)
|
||||
cache_indices = _cache_indices(b, torch.int64)
|
||||
ref = _ref_verify(b, draft_token_num, cache_indices)
|
||||
got = fused_extend_sconv_metadata(
|
||||
B=b,
|
||||
T=b * draft_token_num,
|
||||
cache_indices=cache_indices,
|
||||
his_mode=HIS_ONES,
|
||||
draft_token_num=draft_token_num,
|
||||
)
|
||||
assert got is not None
|
||||
_assert_equal(got, ref)
|
||||
|
||||
|
||||
@requires_cuda
|
||||
def test_cu_not_spanning_T():
|
||||
"""Dummy capture sequences: cu stops short of T; trailing si rows clamp to
|
||||
B-1 exactly like the reference's searchsorted + clamp."""
|
||||
b = 5
|
||||
lens = torch.tensor([3, 0, 4, 0, 2], dtype=torch.int64, device="cuda")
|
||||
T = int(lens.sum().item()) + 17
|
||||
cache_indices = _cache_indices(b, torch.int32)
|
||||
seq_lens = lens + 1
|
||||
ref = _ref_extend(b, lens, HIS_SEQ_MINUS_EXT, seq_lens, cache_indices, T)
|
||||
got = fused_extend_sconv_metadata(
|
||||
B=b,
|
||||
T=T,
|
||||
cache_indices=cache_indices,
|
||||
his_mode=HIS_SEQ_MINUS_EXT,
|
||||
extend_seq_lens=lens,
|
||||
his_src=seq_lens,
|
||||
)
|
||||
assert got is not None
|
||||
_assert_equal(got, ref)
|
||||
|
||||
|
||||
@requires_cuda
|
||||
def test_fallback_past_batch_bound():
|
||||
b = 1024 # > _FUSED_EXTEND_MAX_B
|
||||
lens = torch.ones(b, dtype=torch.int64, device="cuda")
|
||||
got = fused_extend_sconv_metadata(
|
||||
B=b,
|
||||
T=b,
|
||||
cache_indices=_cache_indices(b, torch.int32),
|
||||
his_mode=HIS_ZEROS,
|
||||
extend_seq_lens=lens,
|
||||
)
|
||||
assert got is None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v", "-x"]))
|
||||
@@ -1,6 +1,6 @@
|
||||
"""Multi-process / multi-GPU launching utilities (torchrun-based).
|
||||
|
||||
Shared `multigpu_launch` helper that both `sglang.kernels.jit.tests.utils` and
|
||||
Shared `multigpu_launch` helper that both `sglang.test.kernels.utils` and
|
||||
`sglang.kernels.jit.benchmark.utils` build their domain-specific entry points on
|
||||
top of (`multigpu_pytest_main`, `multigpu_bench_main`).
|
||||
|
||||
|
||||
+6
-6
@@ -5,21 +5,21 @@ from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.tests.kv_canary._constants import (
|
||||
from sglang.kernels.ops.kv_canary import consts
|
||||
from sglang.kernels.ops.kv_canary.consts import splitmix64, splitmix64_mix3
|
||||
from sglang.kernels.ops.kv_canary.verify import VerifyPlan
|
||||
from sglang.kernels.ops.kv_canary.write import WritePlan
|
||||
from sglang.test.kernels.kv_canary._constants import (
|
||||
_I64_SIGN_BIT,
|
||||
_U64_MASK,
|
||||
DEFAULT_NUM_SLOTS,
|
||||
DEFAULT_RING_CAPACITY,
|
||||
DEFAULT_SLOT_STRIDE_BYTES,
|
||||
)
|
||||
from sglang.kernels.jit.tests.kv_canary._fixtures import (
|
||||
from sglang.test.kernels.kv_canary._fixtures import (
|
||||
make_real_kv_source,
|
||||
make_real_kv_sources,
|
||||
)
|
||||
from sglang.kernels.ops.kv_canary import consts
|
||||
from sglang.kernels.ops.kv_canary.consts import splitmix64, splitmix64_mix3
|
||||
from sglang.kernels.ops.kv_canary.verify import VerifyPlan
|
||||
from sglang.kernels.ops.kv_canary.write import WritePlan
|
||||
|
||||
__all__ = [
|
||||
"FakeViolationLog",
|
||||
+6
-6
@@ -5,12 +5,6 @@ from typing import Any, Callable, Iterator, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.tests.kv_canary._canary_helpers import (
|
||||
FakeViolationLog,
|
||||
assert_canary_buf_equal,
|
||||
assert_canary_state_equal,
|
||||
make_log_pair,
|
||||
)
|
||||
from sglang.kernels.ops.kv_canary import consts
|
||||
from sglang.kernels.ops.kv_canary.plan import launch_canary_plan_kernels
|
||||
from sglang.kernels.ops.kv_canary.plan_ref import (
|
||||
@@ -30,6 +24,12 @@ from sglang.kernels.ops.kv_canary.write import WritePlan, launch_canary_write_ke
|
||||
from sglang.kernels.ops.kv_canary.write_ref import (
|
||||
launch_canary_write_kernel_torch_reference,
|
||||
)
|
||||
from sglang.test.kernels.kv_canary._canary_helpers import (
|
||||
FakeViolationLog,
|
||||
assert_canary_buf_equal,
|
||||
assert_canary_state_equal,
|
||||
make_log_pair,
|
||||
)
|
||||
|
||||
_DEVICE = torch.device("cuda")
|
||||
|
||||
+1
-1
@@ -5,12 +5,12 @@ from typing import Literal, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.tests.kv_canary._constants import DEFAULT_NUM_SLOTS
|
||||
from sglang.kernels.ops.kv_canary.verify import (
|
||||
RealKvSource,
|
||||
VerifyPlan,
|
||||
)
|
||||
from sglang.kernels.ops.kv_canary.write import WritePlan
|
||||
from sglang.test.kernels.kv_canary._constants import DEFAULT_NUM_SLOTS
|
||||
|
||||
_DEVICE = torch.device("cuda")
|
||||
|
||||
+1
-1
@@ -3,7 +3,7 @@ from __future__ import annotations
|
||||
import random
|
||||
from typing import Any, Callable
|
||||
|
||||
from sglang.kernels.jit.tests.kv_canary._differential import (
|
||||
from sglang.test.kernels.kv_canary._differential import (
|
||||
ShrinkResult,
|
||||
shrink_inputs,
|
||||
)
|
||||
+1
-1
@@ -11,10 +11,10 @@ from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.tests.kv_canary._canary_helpers import FakeViolationLog
|
||||
from sglang.kernels.ops.kv_canary import consts
|
||||
from sglang.kernels.ops.kv_canary.verify import CanaryLaunchTag, VerifyPlan
|
||||
from sglang.kernels.ops.kv_canary.write import WritePlan
|
||||
from sglang.test.kernels.kv_canary._canary_helpers import FakeViolationLog
|
||||
|
||||
|
||||
class PlanInvariants:
|
||||
Reference in New Issue
Block a user