[JIT] Trait-driven per_token_group_quant: unify the quant kernel family (flat + masked) (#30924)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
DarkSharpness
co-authored by
Claude Fable 5
Claude Fable 5
parent
024639a372
commit
8bb0d8d005
@@ -0,0 +1,76 @@
|
||||
from sglang.jit_kernel.benchmark import marker
|
||||
from sglang.jit_kernel.benchmark.utils import create_empty, create_random
|
||||
|
||||
# per_token_group_quant_8bit_v2 is DEPRECATED (no production call sites); the
|
||||
# kernel is kept only as the perf baseline for this benchmark.
|
||||
from sglang.jit_kernel.per_token_group_quant import per_token_group_quant
|
||||
from sglang.jit_kernel.per_token_group_quant_8bit_v2 import (
|
||||
per_token_group_quant_8bit_v2,
|
||||
)
|
||||
from sglang.kernels.ops.quantization.fp8_kernel import (
|
||||
create_per_token_group_quant_fp8_output_scale,
|
||||
fp8_dtype,
|
||||
fp8_max,
|
||||
fp8_min,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=25, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
|
||||
)
|
||||
|
||||
HIDDEN = 2048
|
||||
LAYOUTS = {
|
||||
"row_major_fp32": (False, False),
|
||||
"col_major_fp32": (True, False),
|
||||
"col_major_ue8m0": (True, True),
|
||||
}
|
||||
|
||||
|
||||
def _jit_v2(G, x, x_q, x_s, scale_ue8m0):
|
||||
per_token_group_quant_8bit_v2(
|
||||
x,
|
||||
x_q,
|
||||
x_s,
|
||||
G,
|
||||
1e-10,
|
||||
float(fp8_min),
|
||||
float(fp8_max),
|
||||
scale_ue8m0=scale_ue8m0,
|
||||
)
|
||||
|
||||
|
||||
def _current(G, x, x_q, x_s, scale_ue8m0):
|
||||
per_token_group_quant(x, x_q, x_s, G, scale_ue8m0=scale_ue8m0)
|
||||
|
||||
|
||||
FN = {"jit_v2": _jit_v2, "current": _current}
|
||||
|
||||
|
||||
@marker.parametrize("group_size", [32, 64, 128], ci_vals=[128])
|
||||
@marker.parametrize("layout", list(LAYOUTS), ci_vals=["col_major_ue8m0"])
|
||||
@marker.parametrize("num_tokens", [2**n for n in range(0, 14)], ci_vals=[1, 32, 2048])
|
||||
@marker.benchmark("impl", ["jit_v2", "current"])
|
||||
def benchmark(group_size: int, layout: str, num_tokens: int, impl: str):
|
||||
column_major, scale_ue8m0 = LAYOUTS[layout]
|
||||
x = create_random(num_tokens, HIDDEN)
|
||||
x_q = create_empty(num_tokens, HIDDEN, dtype=fp8_dtype)
|
||||
x_s = create_per_token_group_quant_fp8_output_scale(
|
||||
x_shape=(num_tokens, HIDDEN),
|
||||
device="cuda",
|
||||
group_size=group_size,
|
||||
column_major_scales=column_major,
|
||||
scale_tma_aligned=column_major,
|
||||
scale_ue8m0=scale_ue8m0,
|
||||
)
|
||||
return marker.do_bench(
|
||||
FN[impl],
|
||||
input_args=(group_size, x, x_q, x_s, scale_ue8m0),
|
||||
graph_clone_args=(1,),
|
||||
memory_args=(x,),
|
||||
memory_output=(x_q, x_s),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run()
|
||||
@@ -1,311 +0,0 @@
|
||||
import itertools
|
||||
from typing import Any, Dict, List
|
||||
|
||||
import torch
|
||||
import triton
|
||||
from sgl_kernel.test_utils import create_per_token_group_quant_test_data
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import get_benchmark_range
|
||||
from sglang.jit_kernel.per_token_group_quant_8bit import (
|
||||
per_token_group_quant_8bit as sglang_per_token_group_quant_8bit,
|
||||
)
|
||||
from sglang.kernels.ops.quantization.fp8_kernel import (
|
||||
create_per_token_group_quant_fp8_output_scale,
|
||||
)
|
||||
from sglang.kernels.ops.quantization.fp8_kernel import (
|
||||
per_token_group_quant_8bit as triton_per_token_group_quant_8bit,
|
||||
)
|
||||
from sglang.srt.utils import is_hip
|
||||
from sglang.srt.utils.bench_utils import bench_kineto
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.utils import is_in_ci
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=13, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
|
||||
)
|
||||
|
||||
IS_CI = is_in_ci()
|
||||
|
||||
_is_hip = is_hip()
|
||||
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
|
||||
|
||||
NUM_TESTS = 30 if IS_CI else 300
|
||||
|
||||
GROUP_SIZE_RANGE = [128]
|
||||
DST_DTYPE_RANGE = [fp8_type_]
|
||||
|
||||
# ---- GEMM-like branch (num_ranks=None) ----
|
||||
NUM_TOKENS_RANGE_GEMM = get_benchmark_range(
|
||||
full_range=[1, 4, 16, 64, 256, 768, 2048, 8192, 16384],
|
||||
ci_range=[768],
|
||||
)
|
||||
HIDDEN_DIM_RANGE_GEMM = [1536, 7168, 16384]
|
||||
NUM_RANKS_RANGE_GEMM = [None]
|
||||
|
||||
|
||||
FLAGS_GEMM_FULL: List[Dict[str, Any]] = [
|
||||
dict(
|
||||
column_major_scales=False,
|
||||
scale_tma_aligned=False,
|
||||
scale_ue8m0=False,
|
||||
fuse_silu_and_mul=False,
|
||||
masked_layout_mode=None,
|
||||
),
|
||||
dict(
|
||||
column_major_scales=True,
|
||||
scale_tma_aligned=False,
|
||||
scale_ue8m0=False,
|
||||
fuse_silu_and_mul=False,
|
||||
masked_layout_mode=None,
|
||||
),
|
||||
dict(
|
||||
column_major_scales=True,
|
||||
scale_tma_aligned=True,
|
||||
scale_ue8m0=False,
|
||||
fuse_silu_and_mul=False,
|
||||
masked_layout_mode=None,
|
||||
),
|
||||
dict(
|
||||
column_major_scales=True,
|
||||
scale_tma_aligned=True,
|
||||
scale_ue8m0=True,
|
||||
fuse_silu_and_mul=False,
|
||||
masked_layout_mode=None,
|
||||
),
|
||||
]
|
||||
FLAGS_GEMM_CI: List[Dict[str, Any]] = [
|
||||
dict(
|
||||
column_major_scales=True,
|
||||
scale_tma_aligned=True,
|
||||
scale_ue8m0=True,
|
||||
fuse_silu_and_mul=False,
|
||||
masked_layout_mode=None,
|
||||
),
|
||||
]
|
||||
FLAGS_RANGE_GEMM = get_benchmark_range(
|
||||
full_range=FLAGS_GEMM_FULL, ci_range=FLAGS_GEMM_CI
|
||||
)
|
||||
|
||||
CONFIGS_GEMM = list(
|
||||
itertools.product(
|
||||
NUM_TOKENS_RANGE_GEMM,
|
||||
HIDDEN_DIM_RANGE_GEMM,
|
||||
GROUP_SIZE_RANGE,
|
||||
NUM_RANKS_RANGE_GEMM,
|
||||
DST_DTYPE_RANGE,
|
||||
FLAGS_RANGE_GEMM,
|
||||
)
|
||||
)
|
||||
|
||||
# ---- MoE-like / multi-rank branch (hidden_dim=2048, num_ranks in {8,16,32,48}) ----
|
||||
NUM_TOKENS_RANGE_MOE = get_benchmark_range(
|
||||
full_range=[1 * 8, 4 * 8, 64 * 8, 256 * 8, 768 * 8],
|
||||
ci_range=[768 * 8],
|
||||
)
|
||||
HIDDEN_DIM_RANGE_MOE = [2048]
|
||||
NUM_RANKS_RANGE_MOE = get_benchmark_range(
|
||||
full_range=[8, 16, 32, 48],
|
||||
ci_range=[48],
|
||||
)
|
||||
|
||||
FLAGS_MOE: List[Dict[str, Any]] = [
|
||||
dict(
|
||||
column_major_scales=True,
|
||||
scale_tma_aligned=True,
|
||||
scale_ue8m0=True,
|
||||
fuse_silu_and_mul=True,
|
||||
masked_layout_mode=None,
|
||||
),
|
||||
dict(
|
||||
column_major_scales=True,
|
||||
scale_tma_aligned=True,
|
||||
scale_ue8m0=True,
|
||||
fuse_silu_and_mul=True,
|
||||
masked_layout_mode="balanced",
|
||||
),
|
||||
dict(
|
||||
column_major_scales=True,
|
||||
scale_tma_aligned=True,
|
||||
scale_ue8m0=True,
|
||||
fuse_silu_and_mul=True,
|
||||
masked_layout_mode="imbalanced",
|
||||
),
|
||||
dict(
|
||||
column_major_scales=True,
|
||||
scale_tma_aligned=True,
|
||||
scale_ue8m0=True,
|
||||
fuse_silu_and_mul=True,
|
||||
masked_layout_mode="extreme",
|
||||
),
|
||||
]
|
||||
FLAGS_RANGE_MOE = get_benchmark_range(full_range=FLAGS_MOE, ci_range=FLAGS_MOE)
|
||||
|
||||
CONFIGS_MOE = list(
|
||||
itertools.product(
|
||||
NUM_TOKENS_RANGE_MOE,
|
||||
HIDDEN_DIM_RANGE_MOE,
|
||||
GROUP_SIZE_RANGE,
|
||||
NUM_RANKS_RANGE_MOE,
|
||||
DST_DTYPE_RANGE,
|
||||
FLAGS_RANGE_MOE,
|
||||
)
|
||||
)
|
||||
|
||||
# ---- Final configs ----
|
||||
CONFIGS = CONFIGS_GEMM + CONFIGS_MOE
|
||||
|
||||
LINE_VALS = ["triton", "aot_v2", "sglang"]
|
||||
LINE_NAMES = ["Triton (Inaccurate)", "AOT v2 (sgl-kernel)", "JIT (this repo)"]
|
||||
STYLES = [("blue", "-"), ("red", "-"), ("green", "-")]
|
||||
|
||||
|
||||
def _flatten_to_2d(t: torch.Tensor) -> torch.Tensor:
|
||||
"""Reshape a tensor with 3+ dims to 2D by merging all leading dims."""
|
||||
if t.ndim <= 2:
|
||||
return t
|
||||
return t.reshape(-1, t.shape[-1])
|
||||
|
||||
|
||||
def _make_sglang_bench_fn(
|
||||
x: torch.Tensor,
|
||||
group_size: int,
|
||||
dst_dtype: torch.dtype,
|
||||
flags: dict,
|
||||
provider: str = "sglang",
|
||||
):
|
||||
"""
|
||||
Adapter that pre-allocates output tensors and returns a zero-arg callable
|
||||
matching the JIT kernel's signature.
|
||||
|
||||
The JIT kernel does not support fuse_silu_and_mul, so when enabled we
|
||||
pre-compute silu+mul on the input. bench_kineto only times the kernel
|
||||
matching the given name, so the pre-processing is not included.
|
||||
|
||||
The JIT kernel expects 2D tensors, so any higher-dimensional inputs
|
||||
(e.g. from masked_layout_mode) are flattened to 2D.
|
||||
"""
|
||||
fuse_silu_and_mul = flags.get("fuse_silu_and_mul", False)
|
||||
column_major_scales = flags.get("column_major_scales", False)
|
||||
scale_tma_aligned = flags.get("scale_tma_aligned", False)
|
||||
scale_ue8m0 = flags.get("scale_ue8m0", False)
|
||||
|
||||
# JIT kernel does not support fuse_silu_and_mul; pre-compute it
|
||||
if fuse_silu_and_mul:
|
||||
half = x.shape[-1] // 2
|
||||
x_input = torch.nn.functional.silu(x[..., :half]) * x[..., half:]
|
||||
else:
|
||||
x_input = x
|
||||
|
||||
# JIT kernel expects 2D (num_tokens, hidden_dim); flatten if needed
|
||||
x_input = _flatten_to_2d(x_input.contiguous())
|
||||
|
||||
out_shape = x_input.shape
|
||||
output_q = torch.empty(out_shape, device=x.device, dtype=dst_dtype)
|
||||
|
||||
fp8_max = torch.finfo(dst_dtype).max
|
||||
fp8_min = -fp8_max
|
||||
|
||||
output_s = create_per_token_group_quant_fp8_output_scale(
|
||||
x_shape=out_shape,
|
||||
device=x.device,
|
||||
group_size=group_size,
|
||||
column_major_scales=column_major_scales,
|
||||
scale_tma_aligned=scale_tma_aligned,
|
||||
scale_ue8m0=scale_ue8m0,
|
||||
)
|
||||
|
||||
if provider == "aot_v2":
|
||||
from sgl_kernel import sgl_per_token_group_quant_8bit as aot_quant
|
||||
|
||||
def _run():
|
||||
aot_quant(
|
||||
x_input,
|
||||
output_q,
|
||||
output_s,
|
||||
group_size,
|
||||
1e-10,
|
||||
fp8_min,
|
||||
fp8_max,
|
||||
scale_ue8m0,
|
||||
False, # fuse_silu_and_mul (already applied to x_input)
|
||||
None, # masked_m (flattened to 2D)
|
||||
enable_v2=True,
|
||||
)
|
||||
|
||||
else:
|
||||
|
||||
def _run():
|
||||
sglang_per_token_group_quant_8bit(
|
||||
input=x_input,
|
||||
output_q=output_q,
|
||||
output_s=output_s,
|
||||
group_size=group_size,
|
||||
eps=1e-10,
|
||||
fp8_min=fp8_min,
|
||||
fp8_max=fp8_max,
|
||||
scale_ue8m0=scale_ue8m0,
|
||||
)
|
||||
|
||||
return _run
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=[
|
||||
"num_tokens",
|
||||
"hidden_dim",
|
||||
"group_size",
|
||||
"num_ranks",
|
||||
"dst_dtype",
|
||||
"flags",
|
||||
],
|
||||
x_vals=CONFIGS,
|
||||
line_arg="provider",
|
||||
line_vals=LINE_VALS,
|
||||
# Triton has multi kernels and we only report the time for the core one
|
||||
line_names=LINE_NAMES,
|
||||
styles=STYLES,
|
||||
ylabel="us",
|
||||
plot_name="per-token-group-quant-8bit-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(
|
||||
num_tokens, hidden_dim, group_size, num_ranks, dst_dtype, flags, provider
|
||||
):
|
||||
print(
|
||||
f"Testing: {num_tokens=} {hidden_dim=} {group_size=} {num_ranks=} {dst_dtype=} {flags=} {provider=}"
|
||||
)
|
||||
|
||||
x, masked_m = create_per_token_group_quant_test_data(
|
||||
num_tokens=num_tokens, hidden_dim=hidden_dim, num_ranks=num_ranks, flags=flags
|
||||
)
|
||||
|
||||
if provider == "triton":
|
||||
fn = triton_per_token_group_quant_8bit
|
||||
kernel_names = "_per_token_group_quant_8bit|_silu_and_mul_post_quant_kernel"
|
||||
bench_fn = lambda: fn(
|
||||
x=x,
|
||||
masked_m=masked_m,
|
||||
group_size=group_size,
|
||||
dst_dtype=dst_dtype,
|
||||
**{k: v for k, v in flags.items() if k not in ["masked_layout_mode"]},
|
||||
)
|
||||
elif provider in ("sglang", "aot_v2"):
|
||||
kernel_names = "per_token_group_quant_8bit_kernel"
|
||||
bench_fn = _make_sglang_bench_fn(
|
||||
x=x,
|
||||
group_size=group_size,
|
||||
dst_dtype=dst_dtype,
|
||||
flags=flags,
|
||||
provider=provider,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown provider: {provider}")
|
||||
|
||||
time_s = bench_kineto(bench_fn, kernel_names=kernel_names, num_tests=NUM_TESTS)
|
||||
return time_s * 1e6
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run(print_data=True)
|
||||
@@ -0,0 +1,121 @@
|
||||
import math
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.benchmark import marker
|
||||
from sglang.jit_kernel.benchmark.utils import create_empty, create_random
|
||||
|
||||
# per_token_group_quant_8bit_v2 is DEPRECATED (no production call sites); the
|
||||
# kernel is kept only as the perf baseline for this benchmark.
|
||||
from sglang.jit_kernel.per_token_group_quant import per_token_group_quant
|
||||
from sglang.jit_kernel.per_token_group_quant_8bit_v2 import (
|
||||
per_token_group_quant_8bit_v2,
|
||||
)
|
||||
from sglang.kernels.ops.quantization.fp8_kernel import (
|
||||
create_per_token_group_quant_fp8_output_scale,
|
||||
fp8_dtype,
|
||||
fp8_max,
|
||||
fp8_min,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=25, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
|
||||
)
|
||||
|
||||
# name -> (moe_intermediate_size, topk, num_experts, group_size)
|
||||
MODELS = {
|
||||
"deepseek_v4": (3072, 6, 384, 32), # DeepSeek-V4 Pro
|
||||
"deepseek_v3": (2048, 8, 256, 128), # DeepSeek-V3/R1
|
||||
"qwen3_235b": (1536, 8, 128, 128), # Qwen3-235B-A22B
|
||||
}
|
||||
|
||||
|
||||
def _jit_v2(G, x, x_q, x_s, masked_m, expected_m, fuse):
|
||||
per_token_group_quant_8bit_v2(
|
||||
x,
|
||||
x_q,
|
||||
x_s,
|
||||
G,
|
||||
1e-10,
|
||||
float(fp8_min),
|
||||
float(fp8_max),
|
||||
scale_ue8m0=True,
|
||||
fuse_silu_and_mul=fuse,
|
||||
masked_m=masked_m,
|
||||
)
|
||||
|
||||
|
||||
def _current(G, x, x_q, x_s, masked_m, expected_m, fuse):
|
||||
per_token_group_quant(
|
||||
x,
|
||||
x_q,
|
||||
x_s,
|
||||
G,
|
||||
scale_ue8m0=True,
|
||||
fuse_silu_and_mul=fuse,
|
||||
masked_m=masked_m,
|
||||
expected_m=expected_m,
|
||||
)
|
||||
|
||||
|
||||
FN = {"jit_v2": _jit_v2, "current": _current}
|
||||
|
||||
|
||||
@marker.parametrize("model", list(MODELS), ci_vals=["deepseek_v3"])
|
||||
@marker.parametrize("num_gpus", [4, 8], ci_vals=[4])
|
||||
@marker.parametrize("fuse_silu", [True, False], ci_vals=[False])
|
||||
@marker.parametrize("balanced", [True, False], ci_vals=[True])
|
||||
@marker.parametrize("num_tokens", [2**n for n in range(8)], ci_vals=[1, 128])
|
||||
@marker.benchmark("impl", ["jit_v2", "current"], unit="us")
|
||||
def benchmark(
|
||||
model: str,
|
||||
fuse_silu: bool,
|
||||
num_gpus: int,
|
||||
num_tokens: int,
|
||||
balanced: bool,
|
||||
impl: str,
|
||||
) -> marker.BenchResult:
|
||||
torch.cuda.random.manual_seed(42)
|
||||
max_tokens = 128 # TODO: test other size
|
||||
hidden_size, topk, num_experts, group_size = MODELS[model]
|
||||
if num_experts % num_gpus != 0 or topk * num_gpus > num_experts:
|
||||
marker.skip("Incompatible model configuration")
|
||||
if impl == "jit_v2" and (hidden_size // group_size) % 16 != 0:
|
||||
marker.skip("v2 masked requires num_groups % 16 == 0")
|
||||
if num_tokens > max_tokens:
|
||||
marker.skip("num_tokens exceeds max_tokens")
|
||||
|
||||
num_local_experts = num_experts // num_gpus
|
||||
padded_tokens = max_tokens * num_gpus
|
||||
expected_m = math.ceil(max_tokens * topk / num_local_experts)
|
||||
in_hidden = hidden_size * (2 if fuse_silu else 1)
|
||||
x = create_random(num_local_experts, padded_tokens, in_hidden)
|
||||
x_q = create_empty(num_local_experts, padded_tokens, hidden_size, dtype=fp8_dtype)
|
||||
x_s = create_per_token_group_quant_fp8_output_scale(
|
||||
x_shape=(num_local_experts, padded_tokens, hidden_size),
|
||||
device="cuda",
|
||||
group_size=group_size,
|
||||
column_major_scales=True,
|
||||
scale_tma_aligned=True,
|
||||
scale_ue8m0=True,
|
||||
)
|
||||
if balanced: # simulation
|
||||
topk_ids = torch.randint(0, num_local_experts, (num_tokens * topk,))
|
||||
masked_m = torch.bincount(topk_ids, minlength=num_local_experts)
|
||||
masked_m = masked_m.cuda().int()
|
||||
else: # only the last few experts receive all tokens
|
||||
masked_m = create_empty(num_local_experts, dtype=torch.int32)
|
||||
masked_m[:-topk].zero_()
|
||||
masked_m[-topk:].fill_(num_tokens)
|
||||
return marker.do_bench(
|
||||
FN[impl],
|
||||
input_args=(group_size, x, x_q, x_s, masked_m, expected_m, fuse_silu),
|
||||
graph_clone_args=(0,),
|
||||
memory_args=(x[:topk, :num_tokens], masked_m),
|
||||
memory_output=(x_q[:topk, :num_tokens], x_s[:topk, :num_tokens]),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run()
|
||||
@@ -0,0 +1,473 @@
|
||||
"""Correctness tests for the trait-driven per_token_group_quant JIT kernel.
|
||||
|
||||
The reference is computed in pure PyTorch (the quantization math itself), NOT by
|
||||
calling the v2 / minimax kernels -- those are being deprecated, so the tests
|
||||
must outlive them.
|
||||
|
||||
Two guard strengths, chosen by what the kernel's numerics can actually pin:
|
||||
- UE8M0 paths: the quant multiplier is an exact power of two (a bit shift, no
|
||||
division), so codes and packed exponent bytes are compared BIT-EXACT
|
||||
against the torch reference. These are the production paths (DeepGEMM dense,
|
||||
EP-MoE), so this is where bit-exactness matters.
|
||||
- fp32 / int8 scale paths: the kernel divides under ``--use_fast_math`` (fast
|
||||
reciprocal), so codes are not bit-reproducible from an exact torch divide.
|
||||
Those tests pin the exactly-reproducible parts -- the stored scale (a single
|
||||
multiply) -- and the dequant round-trip error, which is what downstream
|
||||
actually consumes.
|
||||
"""
|
||||
|
||||
import itertools
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.per_token_group_quant import per_token_group_quant
|
||||
from sglang.jit_kernel.utils import get_ci_test_range
|
||||
from sglang.kernels.ops.quantization.fp8_kernel import (
|
||||
create_per_token_group_quant_fp8_output_scale,
|
||||
fp8_dtype,
|
||||
fp8_max,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
|
||||
|
||||
G = 128
|
||||
FMAX = float(fp8_max) # 448 for e4m3
|
||||
I8_MAX, I8_MIN = 127.0, -128.0
|
||||
EPS = 1e-10
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Pure-torch references (match the kernel's expression order).
|
||||
# --------------------------------------------------------------------------- #
|
||||
def _group_amax(x: torch.Tensor, gs: int) -> torch.Tensor:
|
||||
"""Per-group absmax over the last dim, floored at EPS. Returns [..., ng]."""
|
||||
xf = x.float().unflatten(-1, (-1, gs))
|
||||
return xf.abs().amax(-1).clamp_min(EPS)
|
||||
|
||||
|
||||
def _quantize(x: torch.Tensor, gs: int, quant_scale: torch.Tensor, out_dtype, lo, hi):
|
||||
xf = x.float().unflatten(-1, (-1, gs))
|
||||
q = (xf * quant_scale.unsqueeze(-1)).clamp(lo, hi).to(out_dtype)
|
||||
return q.flatten(-2)
|
||||
|
||||
|
||||
def ref_fp8_fp32_scale(x, gs):
|
||||
"""fp8 codes + fp32 stored scale (scale = amax / FMAX, a single multiply)."""
|
||||
amax = _group_amax(x, gs)
|
||||
scale_inv = amax * (1.0 / FMAX)
|
||||
q = _quantize(x, gs, FMAX / amax, fp8_dtype, -FMAX, FMAX)
|
||||
return q, scale_inv
|
||||
|
||||
|
||||
def ref_int8(x, gs):
|
||||
amax = _group_amax(x, gs)
|
||||
scale_inv = amax * (1.0 / I8_MAX)
|
||||
q = _quantize(x, gs, I8_MAX / amax, torch.int8, I8_MIN, I8_MAX)
|
||||
return q, scale_inv
|
||||
|
||||
|
||||
def ref_fp8_ue8m0(x, gs):
|
||||
"""fp8 codes + UE8M0 exponent bytes [..., ng]. The multiplier 2^-e is exact
|
||||
in fp32, so codes are bit-reproducible (unlike the fp32-scale path)."""
|
||||
amax = _group_amax(x, gs)
|
||||
raw = (amax / FMAX).contiguous()
|
||||
bits = raw.view(torch.int32)
|
||||
exp = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(
|
||||
torch.int32
|
||||
) # ceil to ue8m0
|
||||
quant_scale = ((127 + 127 - exp) << 23).view(torch.float32) # 2^(127 - (exp-127))
|
||||
q = _quantize(x, gs, quant_scale, fp8_dtype, -FMAX, FMAX)
|
||||
return q, exp.to(torch.uint8)
|
||||
|
||||
|
||||
def _decode_packed_exp(s_int32: torch.Tensor, ng: int) -> torch.Tensor:
|
||||
"""Decode an int32 packed-UE8M0 scale (logical [..., ceil(ng/4)]) to the
|
||||
[..., ng] exponent grid, independent of the physical (row/col-major) layout:
|
||||
exponent[..., g] = byte (g % 4) of int32[..., g // 4]."""
|
||||
g = torch.arange(ng, device=s_int32.device)
|
||||
col = s_int32.index_select(-1, g // 4)
|
||||
return ((col >> (8 * (g % 4))) & 0xFF).to(torch.uint8)
|
||||
|
||||
|
||||
def _dequant_rel_err(q, scale_inv, x, gs) -> float:
|
||||
deq = (q.float().unflatten(-1, (-1, gs)) * scale_inv.unsqueeze(-1)).flatten(-2)
|
||||
return ((x.float() - deq).abs() / (x.float().abs() + 1e-6)).mean().item()
|
||||
|
||||
|
||||
def _packed_exp_to_dequant_scale(x_s, ng) -> torch.Tensor:
|
||||
"""Decode a packed-UE8M0 scale buffer to the fp32 dequant scale 2^(e-127)."""
|
||||
exp = _decode_packed_exp(x_s, ng).to(torch.int32)
|
||||
return torch.exp2(exp.float() - 127.0)
|
||||
|
||||
|
||||
def _alloc_scale(x_shape, *, column_major, scale_ue8m0):
|
||||
s = create_per_token_group_quant_fp8_output_scale(
|
||||
x_shape=x_shape,
|
||||
device="cuda",
|
||||
group_size=G,
|
||||
column_major_scales=column_major,
|
||||
scale_tma_aligned=column_major,
|
||||
scale_ue8m0=scale_ue8m0,
|
||||
)
|
||||
s.zero_()
|
||||
return s
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# UE8M0 paths: bit-exact vs the torch reference.
|
||||
# --------------------------------------------------------------------------- #
|
||||
# hidden 768 (Qwen3-30B-A3B moe_intermediate: 6 groups) exercises the non-4
|
||||
# aligned col-packed tail; 128 is a single group.
|
||||
UE8M0_CASES = get_ci_test_range(
|
||||
list(
|
||||
itertools.product(
|
||||
[torch.bfloat16, torch.float16],
|
||||
[1, 7, 38, 333],
|
||||
[128, 768, 2048, 7168],
|
||||
)
|
||||
),
|
||||
[
|
||||
(torch.bfloat16, 1, 128),
|
||||
(torch.bfloat16, 38, 768),
|
||||
(torch.bfloat16, 333, 7168),
|
||||
(torch.float16, 7, 2048),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype,num_tokens,hidden", UE8M0_CASES)
|
||||
def test_ue8m0_bitexact(dtype, num_tokens, hidden):
|
||||
"""Col-major packed UE8M0: fp8 codes and decoded exponent bytes are
|
||||
bit-exact with the torch reference (exact pow-2 multiplier). Covers the
|
||||
aligned and non-4-aligned (hidden=768) pack-tail layouts."""
|
||||
torch.manual_seed(hidden * 10 + num_tokens)
|
||||
x = torch.randn(num_tokens, hidden, device="cuda", dtype=dtype)
|
||||
q_ref, exp_ref = ref_fp8_ue8m0(x, G)
|
||||
|
||||
x_q = torch.zeros_like(x, dtype=fp8_dtype)
|
||||
x_s = _alloc_scale((num_tokens, hidden), column_major=True, scale_ue8m0=True)
|
||||
per_token_group_quant(x, x_q, x_s, G, scale_ue8m0=True)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
assert torch.equal(x_q.view(torch.int8), q_ref.view(torch.int8)), "codes differ"
|
||||
exp = _decode_packed_exp(x_s, hidden // G)
|
||||
assert torch.equal(exp, exp_ref), "exponent bytes differ"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("group_size", get_ci_test_range([16, 32, 64, 128], [16, 64]))
|
||||
def test_ue8m0_group_sizes(group_size):
|
||||
"""Group size is a template axis (v2 dispatched a runtime switch). Each size
|
||||
maps a group onto a different subwarp lane count; codes/exponents must stay
|
||||
bit-exact -- a wrong lane span would fold the wrong elements into absmax."""
|
||||
torch.manual_seed(group_size)
|
||||
num_tokens, hidden = 9, 4096
|
||||
x = torch.randn(num_tokens, hidden, device="cuda", dtype=torch.bfloat16)
|
||||
q_ref, exp_ref = ref_fp8_ue8m0(x, group_size)
|
||||
|
||||
x_q = torch.zeros_like(x, dtype=fp8_dtype)
|
||||
x_s = create_per_token_group_quant_fp8_output_scale(
|
||||
x_shape=(num_tokens, hidden),
|
||||
device="cuda",
|
||||
group_size=group_size,
|
||||
column_major_scales=True,
|
||||
scale_tma_aligned=True,
|
||||
scale_ue8m0=True,
|
||||
)
|
||||
x_s.zero_()
|
||||
per_token_group_quant(x, x_q, x_s, group_size, scale_ue8m0=True)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
assert torch.equal(x_q.view(torch.int8), q_ref.view(torch.int8)), "codes differ"
|
||||
exp = _decode_packed_exp(x_s, hidden // group_size)
|
||||
assert torch.equal(exp, exp_ref), "exponent bytes differ"
|
||||
|
||||
|
||||
# hidden 4096 -> 32 groups (aligned); 768 -> 6 groups (6 % 4 = 2, unaligned:
|
||||
# the last int32 holds 2 real exponent bytes + 2 zero-padded tail bytes).
|
||||
@pytest.mark.parametrize("hidden", [4096, 768])
|
||||
def test_ue8m0_row_packed_bitexact(hidden):
|
||||
"""Row-major packed UE8M0 (int32 [T, ceil(G/4)] contiguous, the minimax
|
||||
layout): bit-exact vs the torch reference. The unaligned hidden exercises
|
||||
the row-major pack-tail zeroing (fill_unaligned)."""
|
||||
torch.manual_seed(hidden)
|
||||
num_tokens = 17
|
||||
x = torch.randn(num_tokens, hidden, device="cuda", dtype=torch.bfloat16)
|
||||
q_ref, exp_ref = ref_fp8_ue8m0(x, G)
|
||||
|
||||
x_q = torch.zeros_like(x, dtype=fp8_dtype)
|
||||
x_s = torch.zeros(
|
||||
num_tokens, (hidden // G + 3) // 4, device="cuda", dtype=torch.int32
|
||||
)
|
||||
per_token_group_quant(x, x_q, x_s, G, scale_ue8m0=True)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
assert torch.equal(x_q.view(torch.int8), q_ref.view(torch.int8)), "codes differ"
|
||||
exp = _decode_packed_exp(x_s, hidden // G)
|
||||
assert torch.equal(exp, exp_ref), "exponent bytes differ"
|
||||
# unaligned tail bytes of the last int32 must be zero-padded, not garbage.
|
||||
ng = hidden // G
|
||||
if ng % 4:
|
||||
last_bytes = x_s[:, -1].contiguous().view(torch.uint8).view(num_tokens, 4)
|
||||
assert torch.all(last_bytes[:, ng % 4 :] == 0), "pack-tail bytes not zeroed"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# fp32 / int8 scale paths: exact stored scale + dequant round-trip (the codes
|
||||
# are not bit-reproducible under fast-math division).
|
||||
# --------------------------------------------------------------------------- #
|
||||
@pytest.mark.parametrize("hidden", [4096, 768])
|
||||
@pytest.mark.parametrize("column_major", [False, True])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
def test_fp32_scale(dtype, column_major, hidden):
|
||||
"""fp32 scale (row-major contiguous / col-major TMA view): the stored scale
|
||||
is amax/FMAX (a single multiply, bit-exact) and dequant round-trips within
|
||||
fp8 error.
|
||||
|
||||
hidden=768 (6 groups, ng % 4 != 0) is a bug regression: the host check
|
||||
used to apply the ue8m0 pack-tail alignment requirement to fp32 scales,
|
||||
which have no packing, and rejected this shape outright."""
|
||||
torch.manual_seed(int(column_major) + 2 * (dtype == torch.float16))
|
||||
num_tokens = 128
|
||||
x = torch.randn(num_tokens, hidden, device="cuda", dtype=dtype)
|
||||
_, scale_ref = ref_fp8_fp32_scale(x, G)
|
||||
|
||||
x_q = torch.zeros_like(x, dtype=fp8_dtype)
|
||||
x_s = _alloc_scale(
|
||||
(num_tokens, hidden), column_major=column_major, scale_ue8m0=False
|
||||
)
|
||||
per_token_group_quant(x, x_q, x_s, G)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(x_s, scale_ref, rtol=0, atol=0)
|
||||
assert _dequant_rel_err(x_q, x_s, x, G) < 0.05
|
||||
|
||||
|
||||
@pytest.mark.parametrize("column_major", [False, True])
|
||||
def test_int8_scale(column_major):
|
||||
"""int8 output (row-major / col-major fp32 scale): exact stored scale
|
||||
(amax/127) + dequant round-trip. Pins the multiply-by-inverse family (v1
|
||||
divided by the scale and differed by a ULP) without depending on v2."""
|
||||
torch.manual_seed(2 + int(column_major))
|
||||
num_tokens, hidden = 33, 4096
|
||||
x = torch.randn(num_tokens, hidden, device="cuda", dtype=torch.bfloat16)
|
||||
_, scale_ref = ref_int8(x, G)
|
||||
|
||||
x_q = torch.zeros(num_tokens, hidden, device="cuda", dtype=torch.int8)
|
||||
x_s = _alloc_scale(
|
||||
(num_tokens, hidden), column_major=column_major, scale_ue8m0=False
|
||||
)
|
||||
per_token_group_quant(x, x_q, x_s, G)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(x_s, scale_ref, rtol=0, atol=0)
|
||||
# int8 group-quant is coarser than fp8 (127 vs 448 levels), so its mean
|
||||
# relative round-trip error on randn data sits a little above the fp8 0.05.
|
||||
assert _dequant_rel_err(x_q, x_s, x, G) < 0.08
|
||||
|
||||
|
||||
def test_group_size_256_roundtrip():
|
||||
"""Group 256 (32 lanes on H100, 16 on Blackwell) is above v2's old cap of
|
||||
128. Pin the derived property: the fp32 scale equals absmax/FMAX and dequant
|
||||
round-trips. A mis-mapped wide subwarp would fold the wrong elements into
|
||||
the group absmax and move the scale."""
|
||||
torch.manual_seed(256)
|
||||
num_tokens, hidden, gs = 9, 4096, 256
|
||||
x = torch.randn(num_tokens, hidden, device="cuda", dtype=torch.bfloat16)
|
||||
_, scale_ref = ref_fp8_fp32_scale(x, gs)
|
||||
|
||||
x_q = torch.zeros_like(x, dtype=fp8_dtype)
|
||||
x_s = torch.zeros(num_tokens, hidden // gs, device="cuda", dtype=torch.float32)
|
||||
per_token_group_quant(x, x_q, x_s, gs)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(x_s, scale_ref, rtol=0, atol=0)
|
||||
assert _dequant_rel_err(x_q, x_s, x, gs) < 0.05
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Fused silu+mul.
|
||||
# --------------------------------------------------------------------------- #
|
||||
def _ref_silu_mul(x, hidden):
|
||||
"""silu in fp32, round to the input dtype, multiply in the input dtype --
|
||||
matching the kernel's fused path exactly."""
|
||||
gate, up = x[..., :hidden], x[..., hidden:]
|
||||
return torch.nn.functional.silu(gate.float()).to(x.dtype) * up
|
||||
|
||||
|
||||
@pytest.mark.parametrize("column_major", [True, False])
|
||||
@pytest.mark.parametrize("scale_ue8m0", [True, False])
|
||||
def test_fused_silu(scale_ue8m0, column_major):
|
||||
"""fuse_silu_and_mul quantizes ``silu(x[..., :h]) * x[..., h:]`` (SGLang's
|
||||
SiluAndMul: first half is the gated half). Covered across all four scale
|
||||
layouts so the fused [gate | up] input layout is pinned everywhere.
|
||||
|
||||
The kernel's silu uses the fast ``__tanhf`` intrinsic on Blackwell, which
|
||||
is not bit-reproducible from torch's sigmoid-based silu, so this is a
|
||||
property test: dequant the kernel output through its own scale and check it
|
||||
round-trips to the torch activation within fp8 error. (The quant math is
|
||||
pinned bit-exact by the non-fused ue8m0 tests; a wrong gate/up split or
|
||||
offset would move the round-trip well past tolerance.)"""
|
||||
torch.manual_seed(int(scale_ue8m0) * 2 + int(column_major))
|
||||
num_tokens, hidden = 37, 4096
|
||||
x = torch.randn(num_tokens, hidden * 2, device="cuda", dtype=torch.bfloat16)
|
||||
act = _ref_silu_mul(x, hidden)
|
||||
|
||||
x_q = torch.zeros(num_tokens, hidden, device="cuda", dtype=fp8_dtype)
|
||||
if scale_ue8m0 and not column_major:
|
||||
x_s = torch.zeros(
|
||||
num_tokens, hidden // G // 4, device="cuda", dtype=torch.int32
|
||||
)
|
||||
else:
|
||||
x_s = _alloc_scale(
|
||||
(num_tokens, hidden), column_major=column_major, scale_ue8m0=scale_ue8m0
|
||||
)
|
||||
per_token_group_quant(
|
||||
x, x_q, x_s, G, scale_ue8m0=scale_ue8m0, fuse_silu_and_mul=True
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
deq_scale = _packed_exp_to_dequant_scale(x_s, hidden // G) if scale_ue8m0 else x_s
|
||||
assert _dequant_rel_err(x_q, deq_scale, act, G) < 0.05
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Masked EP-MoE schedule.
|
||||
# --------------------------------------------------------------------------- #
|
||||
MASKED_CASES = get_ci_test_range(
|
||||
list(itertools.product([2, 5], [2048, 4096], [128, 384])),
|
||||
[(2, 2048, 128), (5, 4096, 384)],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("masked_m_dtype", [torch.int32, torch.int64])
|
||||
@pytest.mark.parametrize("expected_m", [None, 4])
|
||||
@pytest.mark.parametrize("num_experts,hidden,tokens_pad", MASKED_CASES)
|
||||
def test_masked(num_experts, hidden, tokens_pad, expected_m, masked_m_dtype):
|
||||
"""Masked EP-MoE schedule (col-packed ue8m0, plain quant -- no silu, so the
|
||||
quant is bit-reproducible): rows < masked_m[e] are bit-exact vs the torch
|
||||
reference; rows >= masked_m[e] stay zero (untouched). Fusion numerics are
|
||||
covered by test_fused_silu; here the schedule is what's under test.
|
||||
|
||||
masked_m is accepted as int32 or int64 (the latter read as its low word),
|
||||
so both dtypes are exercised.
|
||||
|
||||
expected_m=4 shrinks the grid's token axis far below masked_m, so the
|
||||
grid-stride token loop must still cover every valid token -- guards the
|
||||
host-hint-only contract (a wrong hint can never drop tokens)."""
|
||||
torch.manual_seed(num_experts * 1000 + hidden + tokens_pad)
|
||||
x = torch.randn(
|
||||
num_experts, tokens_pad, hidden, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
masked_m = torch.randint(
|
||||
0, tokens_pad + 1, (num_experts,), device="cuda", dtype=masked_m_dtype
|
||||
)
|
||||
out_shape = (num_experts, tokens_pad, hidden)
|
||||
|
||||
x_q = torch.zeros(out_shape, device="cuda", dtype=fp8_dtype)
|
||||
x_s = _alloc_scale(out_shape, column_major=True, scale_ue8m0=True)
|
||||
per_token_group_quant(
|
||||
x, x_q, x_s, G, scale_ue8m0=True, masked_m=masked_m, expected_m=expected_m
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
q_ref, exp_ref = ref_fp8_ue8m0(x, G)
|
||||
exp = _decode_packed_exp(x_s, hidden // G)
|
||||
for e in range(num_experts):
|
||||
m = int(masked_m[e])
|
||||
assert torch.equal(
|
||||
x_q[e, :m].view(torch.int8), q_ref[e, :m].view(torch.int8)
|
||||
), "written codes differ"
|
||||
assert torch.equal(exp[e, :m], exp_ref[e, :m]), "written exponents differ"
|
||||
assert torch.all(x_q[e, m:].view(torch.int8) == 0), "padding codes touched"
|
||||
|
||||
|
||||
def _as_int32(t: torch.Tensor) -> torch.Tensor:
|
||||
return t.view(torch.int32) if t.dtype == torch.int32 else t
|
||||
|
||||
|
||||
# (out_dtype, column_major_scales, scale_ue8m0); ue8m0 implies fp8 output.
|
||||
AUTO_ALLOC_CASES = [
|
||||
(torch.float8_e4m3fn, True, True),
|
||||
(torch.float8_e4m3fn, False, True),
|
||||
(torch.float8_e4m3fn, True, False),
|
||||
(torch.float8_e4m3fn, False, False),
|
||||
(torch.int8, True, False),
|
||||
(torch.int8, False, False),
|
||||
]
|
||||
|
||||
|
||||
def test_masked_fused():
|
||||
"""The production EP-MoE path: masked schedule + fuse_silu_and_mul +
|
||||
col-packed ue8m0. silu is not bit-reproducible, so check the written rows
|
||||
round-trip to the torch activation and padding rows stay zero."""
|
||||
torch.manual_seed(7)
|
||||
num_experts, tokens_pad, hidden = 3, 256, 2048
|
||||
x = torch.randn(
|
||||
num_experts, tokens_pad, hidden * 2, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
masked_m = torch.randint(
|
||||
0, tokens_pad + 1, (num_experts,), device="cuda", dtype=torch.int32
|
||||
)
|
||||
out_shape = (num_experts, tokens_pad, hidden)
|
||||
|
||||
x_q = torch.zeros(out_shape, device="cuda", dtype=fp8_dtype)
|
||||
x_s = _alloc_scale(out_shape, column_major=True, scale_ue8m0=True)
|
||||
per_token_group_quant(
|
||||
x, x_q, x_s, G, scale_ue8m0=True, fuse_silu_and_mul=True, masked_m=masked_m
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
act = _ref_silu_mul(x, hidden)
|
||||
deq_scale = _packed_exp_to_dequant_scale(x_s, hidden // G)
|
||||
for e in range(num_experts):
|
||||
m = int(masked_m[e])
|
||||
if m > 0:
|
||||
assert _dequant_rel_err(x_q[e, :m], deq_scale[e, :m], act[e, :m], G) < 0.05
|
||||
assert torch.all(x_q[e, m:].view(torch.int8) == 0), "padding touched"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("out_dtype,column_major_scales,scale_ue8m0", AUTO_ALLOC_CASES)
|
||||
def test_auto_allocation(out_dtype, column_major_scales, scale_ue8m0):
|
||||
"""Omitting output_q/output_s allocates them per out_dtype / major mode /
|
||||
scale format and returns (q, s). The auto-allocated run must be bit-
|
||||
identical to quantizing into caller-supplied buffers of the same layout --
|
||||
guards that _allocate_outputs picks the layout the kernel decodes."""
|
||||
torch.manual_seed(int(column_major_scales) * 2 + int(scale_ue8m0))
|
||||
num_tokens, hidden = 38, 2048 # 16 groups, %4 == 0 for row-packed ue8m0
|
||||
x = torch.randn(num_tokens, hidden, device="cuda", dtype=torch.bfloat16)
|
||||
|
||||
q_buf = torch.zeros(num_tokens, hidden, device="cuda", dtype=out_dtype)
|
||||
if scale_ue8m0 and not column_major_scales:
|
||||
s_buf = torch.zeros(
|
||||
num_tokens, hidden // G // 4, device="cuda", dtype=torch.int32
|
||||
)
|
||||
else:
|
||||
s_buf = _alloc_scale(
|
||||
(num_tokens, hidden),
|
||||
column_major=column_major_scales,
|
||||
scale_ue8m0=scale_ue8m0,
|
||||
)
|
||||
per_token_group_quant(x, q_buf, s_buf, G, scale_ue8m0=scale_ue8m0)
|
||||
|
||||
q_auto, s_auto = per_token_group_quant(
|
||||
x,
|
||||
group_size=G,
|
||||
scale_ue8m0=scale_ue8m0,
|
||||
column_major_scales=column_major_scales,
|
||||
out_dtype=out_dtype,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
assert q_auto.dtype == out_dtype and q_auto.shape == x.shape
|
||||
assert s_auto.dtype == s_buf.dtype and s_auto.shape == s_buf.shape
|
||||
assert torch.equal(q_auto.view(torch.int8), q_buf.view(torch.int8)), "codes differ"
|
||||
assert torch.equal(_as_int32(s_auto), _as_int32(s_buf)), "scales differ"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -1,334 +0,0 @@
|
||||
import itertools
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.per_token_group_quant_8bit import (
|
||||
per_token_group_quant_8bit as sglang_per_token_group_quant_8bit,
|
||||
)
|
||||
from sglang.jit_kernel.utils import get_ci_test_range
|
||||
from sglang.srt.utils import is_hip
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=16, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
# Nightly is not redundant here: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1 to expand get_ci_test_range sweeps.
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA required", allow_module_level=True)
|
||||
|
||||
from sgl_kernel import ( # noqa: E402
|
||||
sgl_per_token_group_quant_8bit as aot_per_token_group_quant_8bit,
|
||||
)
|
||||
from sgl_kernel.test_utils import ( # noqa: E402
|
||||
assert_all_close_or_tiny_diff,
|
||||
create_per_token_group_quant_test_data,
|
||||
)
|
||||
|
||||
from sglang.jit_kernel.per_token_group_quant_8bit import ( # noqa: E402
|
||||
per_token_group_quant_8bit as jit_per_token_group_quant_8bit,
|
||||
)
|
||||
from sglang.kernels.ops.quantization.fp8_kernel import ( # noqa: E402
|
||||
create_per_token_group_quant_fp8_output_scale,
|
||||
)
|
||||
from sglang.kernels.ops.quantization.fp8_kernel import ( # noqa: E402
|
||||
per_token_group_quant_8bit as triton_per_token_group_quant_8bit,
|
||||
)
|
||||
|
||||
_is_hip = is_hip()
|
||||
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
|
||||
|
||||
BASE_FLAGS = [
|
||||
dict(
|
||||
column_major_scales=False,
|
||||
scale_tma_aligned=False,
|
||||
scale_ue8m0=False,
|
||||
fuse_silu_and_mul=False,
|
||||
masked_layout_mode=None,
|
||||
),
|
||||
dict(
|
||||
column_major_scales=True,
|
||||
scale_tma_aligned=False,
|
||||
scale_ue8m0=False,
|
||||
fuse_silu_and_mul=False,
|
||||
masked_layout_mode=None,
|
||||
),
|
||||
dict(
|
||||
column_major_scales=True,
|
||||
scale_tma_aligned=True,
|
||||
scale_ue8m0=False,
|
||||
fuse_silu_and_mul=False,
|
||||
masked_layout_mode=None,
|
||||
),
|
||||
dict(
|
||||
column_major_scales=True,
|
||||
scale_tma_aligned=True,
|
||||
scale_ue8m0=True,
|
||||
fuse_silu_and_mul=False,
|
||||
masked_layout_mode=None,
|
||||
),
|
||||
]
|
||||
FUSED_FLAGS = [
|
||||
dict(
|
||||
column_major_scales=True,
|
||||
scale_tma_aligned=True,
|
||||
scale_ue8m0=True,
|
||||
fuse_silu_and_mul=True,
|
||||
masked_layout_mode=None,
|
||||
),
|
||||
dict(
|
||||
column_major_scales=True,
|
||||
scale_tma_aligned=True,
|
||||
scale_ue8m0=True,
|
||||
fuse_silu_and_mul=True,
|
||||
masked_layout_mode="balanced",
|
||||
),
|
||||
dict(
|
||||
column_major_scales=True,
|
||||
scale_tma_aligned=True,
|
||||
scale_ue8m0=True,
|
||||
fuse_silu_and_mul=True,
|
||||
masked_layout_mode="imbalanced",
|
||||
),
|
||||
dict(
|
||||
column_major_scales=True,
|
||||
scale_tma_aligned=True,
|
||||
scale_ue8m0=True,
|
||||
fuse_silu_and_mul=True,
|
||||
masked_layout_mode="extreme",
|
||||
),
|
||||
]
|
||||
|
||||
configs = get_ci_test_range(
|
||||
list(
|
||||
itertools.product(
|
||||
[1, 4, 16, 17, 38, 51, 64, 127, 128, 512, 1024, 4096, 8192],
|
||||
[128, 256, 384, 512, 768, 1024, 1536, 1664, 2048, 4096, 7168, 16384],
|
||||
[16, 32, 64, 128],
|
||||
[None],
|
||||
[fp8_type_],
|
||||
BASE_FLAGS,
|
||||
)
|
||||
)
|
||||
+ list(
|
||||
itertools.product(
|
||||
[1, 4, 1 * 8, 4 * 8, 64 * 8, 256 * 8, 768 * 8],
|
||||
[2048],
|
||||
[128],
|
||||
[8, 16, 32, 48],
|
||||
[fp8_type_],
|
||||
FUSED_FLAGS,
|
||||
)
|
||||
),
|
||||
[
|
||||
(1, 128, 128, None, fp8_type_, BASE_FLAGS[0]),
|
||||
(17, 1536, 128, None, fp8_type_, BASE_FLAGS[2]),
|
||||
(38, 4096, 128, None, fp8_type_, BASE_FLAGS[2]),
|
||||
(51, 4096, 128, None, fp8_type_, BASE_FLAGS[2]),
|
||||
(512, 2048, 128, 8, fp8_type_, FUSED_FLAGS[0]),
|
||||
(2048, 2048, 128, 16, fp8_type_, FUSED_FLAGS[1]),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"num_tokens, hidden_dim, group_size, num_ranks, dst_dtype, flags", configs
|
||||
)
|
||||
def test_per_token_group_quant_with_column_major(
|
||||
num_tokens,
|
||||
hidden_dim,
|
||||
group_size,
|
||||
num_ranks,
|
||||
dst_dtype,
|
||||
flags,
|
||||
):
|
||||
arch_major, _ = torch.cuda.get_device_capability(torch.cuda.current_device())
|
||||
if flags["scale_ue8m0"] and (arch_major <= 9):
|
||||
pytest.skip("Only Blackwell need ue8m0 fusion")
|
||||
return
|
||||
|
||||
if (flags["scale_ue8m0"] and (group_size != 128)) or (
|
||||
(dst_dtype == torch.int8) and flags["column_major_scales"]
|
||||
):
|
||||
pytest.skip()
|
||||
return
|
||||
|
||||
x, masked_m = create_per_token_group_quant_test_data(
|
||||
num_tokens=num_tokens, hidden_dim=hidden_dim, num_ranks=num_ranks, flags=flags
|
||||
)
|
||||
|
||||
execute_kwargs = dict(
|
||||
x=x,
|
||||
masked_m=masked_m,
|
||||
group_size=group_size,
|
||||
eps=1e-10,
|
||||
dst_dtype=dst_dtype,
|
||||
**{k: v for k, v in flags.items() if k not in ["masked_layout_mode"]},
|
||||
)
|
||||
|
||||
def _postprocess(x_q, x_s):
|
||||
if masked_m is not None:
|
||||
print(f"Mask tokens after {masked_m} to be zero")
|
||||
for i in range(len(masked_m)):
|
||||
x_q[i, masked_m[i] :, :] = 0
|
||||
x_s[i, masked_m[i] :, :] = 0
|
||||
return x_q, x_s
|
||||
|
||||
x_q_triton, x_s_triton = _postprocess(
|
||||
*triton_per_token_group_quant_8bit(**execute_kwargs)
|
||||
)
|
||||
|
||||
fuse_silu_and_mul = False
|
||||
out_shape = (*x.shape[:-1], x.shape[-1] // (2 if fuse_silu_and_mul else 1))
|
||||
|
||||
fp8_dtype = torch.float8_e4m3fn
|
||||
fp8_max = torch.finfo(fp8_dtype).max
|
||||
fp8_min = -fp8_max
|
||||
x_q = torch.empty(out_shape, device=x.device, dtype=fp8_dtype)
|
||||
x_s = create_per_token_group_quant_fp8_output_scale(
|
||||
x_shape=out_shape,
|
||||
device=x.device,
|
||||
group_size=group_size,
|
||||
column_major_scales=False,
|
||||
scale_tma_aligned=False,
|
||||
scale_ue8m0=False,
|
||||
)
|
||||
|
||||
execute_kwargs = dict(
|
||||
input=x,
|
||||
output_q=x_q,
|
||||
output_s=x_s,
|
||||
group_size=group_size,
|
||||
eps=1e-10,
|
||||
fp8_max=fp8_max,
|
||||
fp8_min=fp8_min,
|
||||
)
|
||||
x_q_sglang, x_s_sglang = _postprocess(
|
||||
*sglang_per_token_group_quant_8bit(**execute_kwargs)
|
||||
)
|
||||
|
||||
try:
|
||||
assert_all_close_or_tiny_diff(x_q_triton, x_q_sglang)
|
||||
torch.testing.assert_close(
|
||||
x_s_triton.contiguous(),
|
||||
x_s_sglang.contiguous(),
|
||||
rtol=1e-3,
|
||||
atol=1e-5,
|
||||
msg=lambda message: message + f" {x_s_triton=} {x_s_sglang=}",
|
||||
)
|
||||
except AssertionError:
|
||||
print(
|
||||
f"{x.shape=} {x_q_triton.shape=} {x_s_triton.shape=} {x_q_sglang.shape=} {x_s_sglang.shape=}"
|
||||
)
|
||||
print(f"{x=}")
|
||||
print(f"{masked_m=}")
|
||||
print(f"{x_q_triton=}")
|
||||
print(f"{x_s_triton=}")
|
||||
print(f"{x_q_sglang=}")
|
||||
print(f"{x_s_sglang=}")
|
||||
|
||||
raise
|
||||
|
||||
|
||||
LAYOUTS = [
|
||||
(False, False, False),
|
||||
(True, False, False),
|
||||
(True, True, False),
|
||||
(True, True, True),
|
||||
]
|
||||
|
||||
CONFIGS = list(
|
||||
itertools.product(
|
||||
[1, 4, 16, 64, 127, 128, 512, 1024, 4096, 8192],
|
||||
[512, 1536, 2048, 4096, 6144, 7168, 16384],
|
||||
[16, 32, 64, 128],
|
||||
LAYOUTS,
|
||||
[fp8_type_],
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"num_tokens, hidden_dim, group_size, layout, dst_dtype", CONFIGS
|
||||
)
|
||||
def test_jit_matches_aot_v2_byte_identical(
|
||||
num_tokens, hidden_dim, group_size, layout, dst_dtype
|
||||
):
|
||||
column_major_scales, scale_tma_aligned, scale_ue8m0 = layout
|
||||
|
||||
arch_major, _ = torch.cuda.get_device_capability(torch.cuda.current_device())
|
||||
if scale_ue8m0 and arch_major <= 9:
|
||||
pytest.skip("UE8M0 fusion is Blackwell-only")
|
||||
if hidden_dim % group_size != 0:
|
||||
pytest.skip("hidden_dim must be divisible by group_size")
|
||||
|
||||
torch.manual_seed(num_tokens * 131 + hidden_dim + group_size)
|
||||
x = (torch.randn(num_tokens, hidden_dim, device="cuda", dtype=torch.bfloat16)) * 3.0
|
||||
|
||||
fp8_max = torch.finfo(dst_dtype).max
|
||||
fp8_min = -fp8_max
|
||||
|
||||
def _alloc():
|
||||
q = torch.empty_like(x, dtype=dst_dtype)
|
||||
s = create_per_token_group_quant_fp8_output_scale(
|
||||
x_shape=x.shape,
|
||||
device=x.device,
|
||||
group_size=group_size,
|
||||
column_major_scales=column_major_scales,
|
||||
scale_tma_aligned=scale_tma_aligned,
|
||||
scale_ue8m0=scale_ue8m0,
|
||||
)
|
||||
return q, s
|
||||
|
||||
q_aot, s_aot = _alloc()
|
||||
aot_per_token_group_quant_8bit(
|
||||
x,
|
||||
q_aot,
|
||||
s_aot,
|
||||
group_size,
|
||||
1e-10,
|
||||
fp8_min,
|
||||
fp8_max,
|
||||
scale_ue8m0,
|
||||
False,
|
||||
None,
|
||||
enable_v2=True,
|
||||
)
|
||||
|
||||
q_jit, s_jit = _alloc()
|
||||
jit_per_token_group_quant_8bit(
|
||||
x, q_jit, s_jit, group_size, 1e-10, fp8_min, fp8_max, scale_ue8m0=scale_ue8m0
|
||||
)
|
||||
|
||||
# AOT v2 uses -use_fast_math reciprocal; this JIT uses precise division, so an
|
||||
# exact fp8 midpoint can round to an adjacent code (1-ULP, JIT more accurate).
|
||||
qj = q_jit.view(torch.uint8)
|
||||
qa = q_aot.view(torch.uint8)
|
||||
if not torch.equal(qj, qa):
|
||||
mism = qj != qa
|
||||
bj = qj[mism].to(torch.int16)
|
||||
ba = qa[mism].to(torch.int16)
|
||||
same_sign = (bj & 0x80) == (ba & 0x80)
|
||||
one_ulp = (bj - ba).abs() == 1
|
||||
assert bool(
|
||||
(same_sign & one_ulp).all()
|
||||
), f"q mismatch > 1 fp8 ULP {num_tokens=} {hidden_dim=} {group_size=} {layout=}"
|
||||
assert mism.float().mean() < 0.01, (
|
||||
f"too many fp8 ties ({int(mism.sum())}/{mism.numel()}) "
|
||||
f"{num_tokens=} {hidden_dim=} {group_size=} {layout=}"
|
||||
)
|
||||
|
||||
if scale_ue8m0:
|
||||
assert torch.equal(
|
||||
s_jit[:num_tokens].reshape(num_tokens, -1).view(torch.int32),
|
||||
s_aot[:num_tokens].reshape(num_tokens, -1).view(torch.int32),
|
||||
), f"ue8m0 scale mismatch {num_tokens=} {hidden_dim=} {group_size=}"
|
||||
else:
|
||||
assert torch.equal(
|
||||
s_jit[:num_tokens].float(), s_aot[:num_tokens].float()
|
||||
), f"float scale mismatch {num_tokens=} {hidden_dim=} {group_size=}"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -26,7 +26,6 @@ from sglang.kernels.ops.quantization.fp8_kernel import ( # noqa: E402
|
||||
fp8_dtype,
|
||||
fp8_max,
|
||||
fp8_min,
|
||||
sglang_per_token_group_quant_fp8,
|
||||
)
|
||||
|
||||
G = 128
|
||||
@@ -115,44 +114,17 @@ def test_v2_jit_matches_aot(dtype, num_tokens, hidden, fuse_silu_and_mul, scale_
|
||||
assert torch.equal(x_s, s_ref), "scales differ"
|
||||
|
||||
|
||||
ROW_MAJOR_UE8M0_CASES = get_ci_test_range(
|
||||
list(
|
||||
itertools.product(
|
||||
[torch.bfloat16, torch.float16], [1, 33, 128], [128, 512, 4096, 7168]
|
||||
)
|
||||
),
|
||||
[
|
||||
(torch.bfloat16, 1, 128),
|
||||
(torch.bfloat16, 17, 1536),
|
||||
(torch.bfloat16, 33, 7168),
|
||||
(torch.bfloat16, 38, 4096),
|
||||
(torch.float16, 128, 4096),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype,num_tokens,hidden", ROW_MAJOR_UE8M0_CASES)
|
||||
def test_sglang_per_token_group_quant_fp8_row_major_ue8m0(dtype, num_tokens, hidden):
|
||||
"""Row-major scale_ue8m0=True quantizes WITH the rounded (power-of-2) scale.
|
||||
Verify: (1) scales are exact powers of 2, (2) dequant ≈ original within FP8 tolerance.
|
||||
"""
|
||||
torch.manual_seed(num_tokens * 1000 + hidden)
|
||||
x = torch.randn(num_tokens, hidden, device="cuda", dtype=dtype)
|
||||
|
||||
x_q, x_s = sglang_per_token_group_quant_fp8(x, G, scale_ue8m0=True)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Scales must be exact powers of 2
|
||||
log2_s = torch.log2(x_s.abs())
|
||||
assert torch.equal(log2_s, log2_s.round()), "scales are not power-of-2"
|
||||
|
||||
# Dequant should approximate original within FP8 precision
|
||||
x_deq = x_q.float().view(num_tokens, -1, G) * x_s.unsqueeze(-1)
|
||||
x_deq = x_deq.view(num_tokens, hidden)
|
||||
rel_err = (x.float() - x_deq).abs() / (x.float().abs() + 1e-6)
|
||||
assert (
|
||||
rel_err.mean() < 0.05
|
||||
), f"mean relative dequant error too large: {rel_err.mean():.4f}"
|
||||
# NOTE: "row-major + scale_ue8m0=True" names two different formats:
|
||||
# 1. packed int32 [T, ceil(G/4)] (4 exponent bytes per int32) -- supported by
|
||||
# the JIT per_token_group_quant kernel and pinned bit-exact in test_per_token_group_quant
|
||||
# (test_v3_ue8m0_row_packed_bitexact);
|
||||
# 2. fp32 [T, G] storing power-of-two VALUES (the deep_gemm.fp8_einsum
|
||||
# format) -- v2-only. No srt caller requests it (production ties
|
||||
# scale_ue8m0 and column_major_scales to the same DEEPGEMM_SCALE_UE8M0
|
||||
# flag), so the srt entry `sglang_per_token_group_quant_fp8`, which now
|
||||
# routes to the JIT kernel, rejects it loudly instead of allocating an fp32
|
||||
# buffer the kernel cannot fill. The v2 JIT kernel itself still implements it and is
|
||||
# covered by test_v2_jit_matches_aot above.
|
||||
|
||||
|
||||
# Masked (EP-MoE) path: the v2 op only has a masked scheduler for the
|
||||
|
||||
Reference in New Issue
Block a user