[CI] Move JIT kernel tests + benchmarks to test/registered/jit; add in-package guard (#27644)
This commit is contained in:
Liangsheng Yin
@@ -0,0 +1,108 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from sgl_kernel import gelu_and_mul as gelu_and_mul_aot
|
||||
from sgl_kernel import gelu_tanh_and_mul as gelu_tanh_and_mul_aot
|
||||
from sgl_kernel import silu_and_mul as silu_and_mul_aot
|
||||
|
||||
from sglang.jit_kernel.activation import gelu_and_mul as gelu_and_mul_jit
|
||||
from sglang.jit_kernel.activation import gelu_tanh_and_mul as gelu_tanh_and_mul_jit
|
||||
from sglang.jit_kernel.activation import relu2 as relu2_jit
|
||||
from sglang.jit_kernel.activation import silu_and_mul as silu_and_mul_jit
|
||||
from sglang.jit_kernel.benchmark import marker
|
||||
from sglang.jit_kernel.benchmark.utils import create_random
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
|
||||
@torch.compile
|
||||
def silu_and_mul(input: torch.Tensor) -> torch.Tensor:
|
||||
lhs, rhs = input.split(input.shape[-1] // 2, dim=-1)
|
||||
return F.silu(lhs) * rhs
|
||||
|
||||
|
||||
@torch.compile
|
||||
def gelu_and_mul(input: torch.Tensor) -> torch.Tensor:
|
||||
lhs, rhs = input.split(input.shape[-1] // 2, dim=-1)
|
||||
return F.gelu(lhs, approximate="none") * rhs
|
||||
|
||||
|
||||
@torch.compile
|
||||
def gelu_tanh_and_mul(input: torch.Tensor) -> torch.Tensor:
|
||||
lhs, rhs = input.split(input.shape[-1] // 2, dim=-1)
|
||||
return F.gelu(lhs, approximate="tanh") * rhs
|
||||
|
||||
|
||||
OPS = {
|
||||
"silu": (silu_and_mul_aot, silu_and_mul_jit, silu_and_mul),
|
||||
"gelu": (gelu_and_mul_aot, gelu_and_mul_jit, gelu_and_mul),
|
||||
"gelu_tanh": (gelu_tanh_and_mul_aot, gelu_tanh_and_mul_jit, gelu_tanh_and_mul),
|
||||
}
|
||||
|
||||
|
||||
@marker.parametrize("op_name", ["silu", "gelu", "gelu_tanh"])
|
||||
@marker.parametrize("dim", [1024, 4096, 6144, 8192], [4096])
|
||||
@marker.parametrize("batch_size", [2**x for x in range(0, 15)], [8, 512])
|
||||
@marker.benchmark("impl", ["aot", "jit", "torch"])
|
||||
def benchmark(op_name: str, dim: int, batch_size: int, impl: str):
|
||||
x = create_random(batch_size, dim * 2)
|
||||
aot_op, jit_op, torch_op = OPS[op_name]
|
||||
fn = {"aot": aot_op, "jit": jit_op, "torch": torch_op}[impl]
|
||||
return marker.do_bench(fn, input_args=(x,))
|
||||
|
||||
|
||||
def _make_expert_ids(num_tokens: int, skip_ratio: float) -> torch.Tensor:
|
||||
expert_ids = torch.randint(low=0, high=8, size=(num_tokens,), dtype=torch.int32)
|
||||
if skip_ratio > 0:
|
||||
skip = torch.rand(num_tokens) < skip_ratio
|
||||
expert_ids[skip] = -1
|
||||
return expert_ids
|
||||
|
||||
|
||||
@marker.parametrize("op_name", ["silu", "gelu"])
|
||||
@marker.parametrize("dim", [1024, 4096, 8192], [4096])
|
||||
@marker.parametrize("batch_size", [64, 256, 1024, 4096, 16384], [1024])
|
||||
@marker.parametrize("skip_ratio", [0.0, 0.25, 0.5], [0.25])
|
||||
@marker.benchmark("impl", ["unfiltered", "filtered"])
|
||||
def benchmark_filter(
|
||||
op_name: str, dim: int, batch_size: int, skip_ratio: float, impl: str
|
||||
):
|
||||
torch.random.manual_seed(42)
|
||||
x = create_random(batch_size, dim * 2)
|
||||
jit_fn = silu_and_mul_jit if op_name == "silu" else gelu_and_mul_jit
|
||||
extra_kwargs = {}
|
||||
expert_ids = _make_expert_ids(batch_size, skip_ratio)
|
||||
if impl == "filtered":
|
||||
extra_kwargs = {"expert_ids": expert_ids.to(x.device), "expert_step": 1}
|
||||
|
||||
# NOTE: get the unmasked part from `experts_ids`
|
||||
real_skip_ratio = (expert_ids == -1).sum().item() / batch_size
|
||||
effective_bytes = int(x.nbytes * (1 - real_skip_ratio) * 1.5)
|
||||
return marker.do_bench(
|
||||
jit_fn,
|
||||
input_args=(x,),
|
||||
input_kwargs=extra_kwargs,
|
||||
memory_args=None, # x is dynamic (counted in extra_memory_footprint)
|
||||
memory_output=None, # same, output is dynamic
|
||||
extra_memory_footprint=effective_bytes,
|
||||
)
|
||||
|
||||
|
||||
@torch.compile
|
||||
def relu2_torch(input: torch.Tensor) -> torch.Tensor:
|
||||
return F.relu(input).pow(2)
|
||||
|
||||
|
||||
@marker.parametrize("dim", [1024, 4096, 6144, 8192], [4096])
|
||||
@marker.parametrize("batch_size", [2**x for x in range(0, 15)], [8, 512])
|
||||
@marker.benchmark("impl", ["jit", "torch"])
|
||||
def benchmark_unary(dim: int, batch_size: int, impl: str):
|
||||
x = create_random(batch_size, dim)
|
||||
fn = {"jit": relu2_jit, "torch": relu2_torch}[impl]
|
||||
return marker.do_bench(fn, input_args=(x,))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run()
|
||||
benchmark_filter.run()
|
||||
benchmark_unary.run()
|
||||
@@ -0,0 +1,59 @@
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.add_constant import _jit_add_constant_module, add_constant
|
||||
from sglang.jit_kernel.benchmark.utils import (
|
||||
DEFAULT_DEVICE,
|
||||
get_benchmark_range,
|
||||
run_benchmark_no_cudagraph,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=15, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
CONSTANT = 7
|
||||
SIZE_LIST = get_benchmark_range(
|
||||
full_range=[128, 1024, 1025, 4096, 4097, 65536, 2**20, 2**22, 2**24],
|
||||
ci_range=[4096, 2**20],
|
||||
)
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["size"],
|
||||
x_vals=SIZE_LIST,
|
||||
line_arg="provider",
|
||||
line_vals=["jit_module", "jit_wrapper", "torch"],
|
||||
line_names=["JIT module", "JIT wrapper", "PyTorch"],
|
||||
styles=[("blue", "-"), ("orange", "-"), ("green", "--")],
|
||||
ylabel="us",
|
||||
plot_name="add-constant-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(size: int, provider: str):
|
||||
src = torch.arange(size, dtype=torch.int32, device=DEFAULT_DEVICE)
|
||||
|
||||
if provider == "jit_module":
|
||||
dst = torch.empty_like(src)
|
||||
module = _jit_add_constant_module(CONSTANT)
|
||||
|
||||
def fn():
|
||||
module.add_constant(dst, src)
|
||||
|
||||
elif provider == "jit_wrapper":
|
||||
|
||||
def fn():
|
||||
add_constant(src, CONSTANT)
|
||||
|
||||
else:
|
||||
|
||||
def fn():
|
||||
src + CONSTANT
|
||||
|
||||
return run_benchmark_no_cudagraph(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run(print_data=True)
|
||||
@@ -0,0 +1,122 @@
|
||||
import itertools
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.awq_dequantize import awq_dequantize as jit_awq_dequantize
|
||||
from sglang.jit_kernel.benchmark.utils import run_benchmark
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.utils import is_in_ci
|
||||
|
||||
register_cuda_ci(est_time=5, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
try:
|
||||
from sgl_kernel import awq_dequantize as aot_awq_dequantize
|
||||
|
||||
AOT_AVAILABLE = True
|
||||
except ImportError:
|
||||
AOT_AVAILABLE = False
|
||||
|
||||
IS_CI = is_in_ci()
|
||||
|
||||
if IS_CI:
|
||||
qweight_row_range = [128]
|
||||
qweight_cols_range = [16]
|
||||
else:
|
||||
qweight_row_range = [128, 256, 512, 1024, 3584]
|
||||
qweight_cols_range = [16, 32, 64, 128, 448]
|
||||
|
||||
configs = list(itertools.product(qweight_row_range, qweight_cols_range))
|
||||
|
||||
|
||||
def check_correctness():
|
||||
if not AOT_AVAILABLE:
|
||||
print("sgl_kernel AOT not available, skipping correctness check")
|
||||
return
|
||||
|
||||
qweight_row, qweight_col = 128, 16
|
||||
device = torch.device("cuda")
|
||||
qweight = torch.randint(
|
||||
0,
|
||||
torch.iinfo(torch.int32).max,
|
||||
(qweight_row, qweight_col),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
group_size = qweight_row
|
||||
scales_row = qweight_row // group_size
|
||||
scales_col = qweight_col * 8
|
||||
scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device)
|
||||
qzeros = torch.randint(
|
||||
0,
|
||||
torch.iinfo(torch.int32).max,
|
||||
(scales_row, qweight_col),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
jit_out = jit_awq_dequantize(qweight, scales, qzeros)
|
||||
aot_out = aot_awq_dequantize(qweight, scales, qzeros)
|
||||
torch.cuda.synchronize()
|
||||
torch.testing.assert_close(jit_out, aot_out, rtol=0, atol=0)
|
||||
print("Correctness check passed (JIT vs AOT)")
|
||||
|
||||
|
||||
if AOT_AVAILABLE:
|
||||
line_vals = ["jit", "aot"]
|
||||
line_names = ["JIT Kernel", "AOT Kernel"]
|
||||
styles = [("blue", "-"), ("green", "-")]
|
||||
else:
|
||||
line_vals = ["jit"]
|
||||
line_names = ["JIT Kernel"]
|
||||
styles = [("blue", "-")]
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["qweight_row", "qweight_col"],
|
||||
x_vals=configs,
|
||||
line_arg="provider",
|
||||
line_vals=line_vals,
|
||||
line_names=line_names,
|
||||
styles=styles,
|
||||
ylabel="us",
|
||||
plot_name="awq-dequantize-jit-vs-aot",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(qweight_row, qweight_col, provider):
|
||||
device = torch.device("cuda")
|
||||
qweight = torch.randint(
|
||||
0,
|
||||
torch.iinfo(torch.int32).max,
|
||||
(qweight_row, qweight_col),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
group_size = qweight_row
|
||||
scales_row = qweight_row // group_size
|
||||
scales_col = qweight_col * 8
|
||||
scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device)
|
||||
qzeros = torch.randint(
|
||||
0,
|
||||
torch.iinfo(torch.int32).max,
|
||||
(scales_row, qweight_col),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
if provider == "jit":
|
||||
fn = lambda: jit_awq_dequantize(qweight, scales, qzeros)
|
||||
elif provider == "aot":
|
||||
fn = lambda: aot_awq_dequantize(qweight, scales, qzeros)
|
||||
else:
|
||||
raise ValueError(f"Unknown provider: {provider}")
|
||||
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
check_correctness()
|
||||
benchmark.run(print_data=True)
|
||||
@@ -0,0 +1,65 @@
|
||||
import itertools
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import (
|
||||
DEFAULT_DEVICE,
|
||||
get_benchmark_range,
|
||||
run_benchmark,
|
||||
)
|
||||
from sglang.jit_kernel.clamp_position import clamp_position_cuda
|
||||
from sglang.srt.utils import get_compiler_backend
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=13, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
register_amd_ci(est_time=16, suite="jit-kernel-unit-test-amd")
|
||||
|
||||
SIZE_LIST = get_benchmark_range(
|
||||
full_range=[2**n for n in range(4, 16)],
|
||||
ci_range=[256, 4096],
|
||||
)
|
||||
|
||||
configs = list(itertools.product(SIZE_LIST))
|
||||
|
||||
|
||||
def _torch_clamp_position(seq_lens):
|
||||
return torch.clamp(seq_lens - 1, min=0).to(torch.int64)
|
||||
|
||||
|
||||
_compiled_clamp_position = torch.compile(
|
||||
_torch_clamp_position, dynamic=True, backend=get_compiler_backend()
|
||||
)
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["size"],
|
||||
x_vals=configs,
|
||||
line_arg="provider",
|
||||
line_vals=["jit", "torch_compile", "torch"],
|
||||
line_names=["SGL JIT Kernel", "torch.compile", "PyTorch"],
|
||||
styles=[("blue", "-"), ("green", "-."), ("red", "--")],
|
||||
ylabel="us",
|
||||
plot_name="clamp-position-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(size: int, provider: str):
|
||||
seq_lens = torch.randint(
|
||||
0, 10000, (size,), dtype=torch.int64, device=DEFAULT_DEVICE
|
||||
)
|
||||
|
||||
if provider == "jit":
|
||||
fn = lambda: clamp_position_cuda(seq_lens)
|
||||
elif provider == "torch_compile":
|
||||
fn = lambda: _compiled_clamp_position(seq_lens)
|
||||
else:
|
||||
fn = lambda: _torch_clamp_position(seq_lens)
|
||||
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run(print_data=True)
|
||||
@@ -0,0 +1,162 @@
|
||||
import itertools
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
from sgl_kernel import concat_mla_absorb_q as aot_absorb_q
|
||||
from sgl_kernel import concat_mla_k as aot_k
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import run_benchmark
|
||||
from sglang.jit_kernel.concat_mla import concat_mla_absorb_q as jit_absorb_q
|
||||
from sglang.jit_kernel.concat_mla import concat_mla_k as jit_k
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.utils import is_in_ci
|
||||
|
||||
register_cuda_ci(est_time=6, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
IS_CI = is_in_ci()
|
||||
|
||||
NUM_LOCAL_HEADS = 128
|
||||
QK_NOPE_HEAD_DIM = 128
|
||||
QK_ROPE_HEAD_DIM = 64
|
||||
K_HEAD_DIM = QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM
|
||||
|
||||
A_LAST_DIM = 512
|
||||
B_LAST_DIM = 64
|
||||
|
||||
DTYPE = torch.bfloat16
|
||||
DEVICE = "cuda"
|
||||
|
||||
|
||||
def aot_concat_mla_k(k, k_nope, k_rope):
|
||||
aot_k(k, k_nope, k_rope)
|
||||
|
||||
|
||||
def jit_concat_mla_k(k, k_nope, k_rope):
|
||||
jit_k(k, k_nope, k_rope)
|
||||
|
||||
|
||||
def torch_concat_mla_k(k, k_nope, k_rope):
|
||||
nope_head_dim = k_nope.shape[-1]
|
||||
k[:, :, :nope_head_dim] = k_nope
|
||||
k[:, :, nope_head_dim:] = k_rope.expand(-1, k.shape[1], -1)
|
||||
|
||||
|
||||
def aot_concat_mla_absorb_q(a, b):
|
||||
return aot_absorb_q(a, b)
|
||||
|
||||
|
||||
def jit_concat_mla_absorb_q(a, b):
|
||||
return jit_absorb_q(a, b)
|
||||
|
||||
|
||||
def torch_concat_mla_absorb_q(a, b, out):
|
||||
a_last_dim = a.shape[-1]
|
||||
out[:, :, :a_last_dim] = a
|
||||
out[:, :, a_last_dim:] = b
|
||||
|
||||
|
||||
if IS_CI:
|
||||
NUM_TOKENS_VALS = [256, 1024]
|
||||
else:
|
||||
NUM_TOKENS_VALS = [256, 512, 1024, 2048, 4096, 8192, 16384, 32768]
|
||||
|
||||
K_LINE_VALS = ["aot", "jit", "torch"]
|
||||
K_LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "PyTorch"]
|
||||
K_STYLES = [("orange", "-"), ("blue", "--"), ("green", "-.")]
|
||||
|
||||
|
||||
def _create_concat_mla_k_data(num_tokens):
|
||||
"""Allocate oversized containers and slice to produce non-contiguous tensors."""
|
||||
k_nope_container = torch.randn(
|
||||
(num_tokens, NUM_LOCAL_HEADS, QK_NOPE_HEAD_DIM + 128),
|
||||
dtype=DTYPE,
|
||||
device=DEVICE,
|
||||
)
|
||||
k_nope = k_nope_container[:, :, :QK_NOPE_HEAD_DIM]
|
||||
|
||||
k_rope_container = torch.randn(
|
||||
(num_tokens, 1, 128 + QK_ROPE_HEAD_DIM),
|
||||
dtype=DTYPE,
|
||||
device=DEVICE,
|
||||
)
|
||||
k_rope = k_rope_container[:, :, -QK_ROPE_HEAD_DIM:]
|
||||
|
||||
k = torch.empty(
|
||||
(num_tokens, NUM_LOCAL_HEADS, K_HEAD_DIM),
|
||||
dtype=DTYPE,
|
||||
device=DEVICE,
|
||||
)
|
||||
return k, k_nope, k_rope
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["num_tokens"],
|
||||
x_vals=NUM_TOKENS_VALS,
|
||||
line_arg="provider",
|
||||
line_vals=K_LINE_VALS,
|
||||
line_names=K_LINE_NAMES,
|
||||
styles=K_STYLES,
|
||||
ylabel="us",
|
||||
plot_name="concat-mla-k-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def bench_concat_mla_k(num_tokens: int, provider: str):
|
||||
k, k_nope, k_rope = _create_concat_mla_k_data(num_tokens)
|
||||
|
||||
FN_MAP = {
|
||||
"aot": aot_concat_mla_k,
|
||||
"jit": jit_concat_mla_k,
|
||||
"torch": torch_concat_mla_k,
|
||||
}
|
||||
fn = lambda: FN_MAP[provider](k, k_nope, k_rope)
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if IS_CI:
|
||||
ABSORB_Q_VALS = list(itertools.product([4, 16], [16]))
|
||||
else:
|
||||
ABSORB_Q_VALS = list(itertools.product([1, 4, 8, 16, 32], [1, 8, 32, 128]))
|
||||
|
||||
Q_LINE_VALS = ["aot", "jit", "torch"]
|
||||
Q_LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "PyTorch"]
|
||||
Q_STYLES = [("orange", "-"), ("blue", "--"), ("green", "-.")]
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["dim_0", "dim_1"],
|
||||
x_vals=ABSORB_Q_VALS,
|
||||
line_arg="provider",
|
||||
line_vals=Q_LINE_VALS,
|
||||
line_names=Q_LINE_NAMES,
|
||||
styles=Q_STYLES,
|
||||
ylabel="us",
|
||||
plot_name="concat-mla-absorb-q-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def bench_concat_mla_absorb_q(dim_0: int, dim_1: int, provider: str):
|
||||
a = torch.randn(dim_0, dim_1, A_LAST_DIM, dtype=DTYPE, device=DEVICE)
|
||||
b = torch.randn(dim_0, dim_1, B_LAST_DIM, dtype=DTYPE, device=DEVICE)
|
||||
|
||||
if provider == "torch":
|
||||
out = torch.empty(
|
||||
dim_0, dim_1, A_LAST_DIM + B_LAST_DIM, dtype=DTYPE, device=DEVICE
|
||||
)
|
||||
fn = lambda: torch_concat_mla_absorb_q(a, b, out)
|
||||
else:
|
||||
FN_MAP = {
|
||||
"aot": aot_concat_mla_absorb_q,
|
||||
"jit": jit_concat_mla_absorb_q,
|
||||
}
|
||||
fn = lambda: FN_MAP[provider](a, b)
|
||||
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
bench_concat_mla_k.run(print_data=True)
|
||||
bench_concat_mla_absorb_q.run(print_data=True)
|
||||
@@ -0,0 +1,384 @@
|
||||
"""
|
||||
Benchmark JIT custom all-reduce (v2) vs NCCL vs AOT custom all-reduce (v1).
|
||||
|
||||
Usage (torchrun required for multi-GPU):
|
||||
torchrun --nproc_per_node=2 bench_custom_all_reduce.py
|
||||
torchrun --nproc_per_node=4 bench_custom_all_reduce.py --dtype float16
|
||||
torchrun --nproc_per_node=8 bench_custom_all_reduce.py --warmup 10 --iters 100
|
||||
|
||||
The script initializes all three backends, then benchmarks each over a sweep
|
||||
of message sizes. Results are printed as a comparison table on rank 0.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import contextlib
|
||||
import gc
|
||||
import logging
|
||||
import os
|
||||
from math import isnan
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import is_in_ci
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=120,
|
||||
suite="base-b-kernel-benchmark-1-gpu-large",
|
||||
disabled="requires multi-GPU, self-skips in CI",
|
||||
)
|
||||
|
||||
DTYPE_MAP = {
|
||||
"float16": torch.float16,
|
||||
"bfloat16": torch.bfloat16,
|
||||
"float32": torch.float32,
|
||||
}
|
||||
|
||||
MESSAGE_SIZES_BYTES = [
|
||||
4 * 1024, # 4K
|
||||
16 * 1024, # 16K
|
||||
64 * 1024, # 64K
|
||||
128 * 1024, # 128K
|
||||
3 * 64 * 1024, # 192K
|
||||
4 * 64 * 1024, # 256K
|
||||
3 * 128 * 1024, # 384K
|
||||
4 * 128 * 1024, # 512K
|
||||
5 * 128 * 1024, # 640K
|
||||
6 * 128 * 1024, # 768K
|
||||
7 * 128 * 1024, # 896K
|
||||
1 * 1024 * 1024, # 1M
|
||||
2 * 1024 * 1024, # 2M
|
||||
3 * 1024 * 1024, # 2M
|
||||
4 * 1024 * 1024, # 4M
|
||||
8 * 1024 * 1024, # 8M
|
||||
16 * 1024 * 1024, # 16M
|
||||
32 * 1024 * 1024, # 32M
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Backend wrappers - each exposes a uniform interface:
|
||||
# .name - display name
|
||||
# .capture() - context manager for CUDA-graph recording
|
||||
# .all_reduce() - perform an all-reduce and return the result tensor
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class NCCLAllReduceBackend:
|
||||
name = "NCCL"
|
||||
|
||||
def __init__(self, group: dist.ProcessGroup):
|
||||
self.group = group
|
||||
|
||||
def capture(self, register_input: bool):
|
||||
return contextlib.nullcontext()
|
||||
|
||||
def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor:
|
||||
dist.all_reduce(tensor, group=self.group)
|
||||
return tensor
|
||||
|
||||
|
||||
class AOTAllReduceBackend:
|
||||
name = "AOT"
|
||||
|
||||
def __init__(self, group: dist.ProcessGroup, device: torch.device):
|
||||
from sglang.srt.distributed.device_communicators.custom_all_reduce import (
|
||||
CustomAllreduce,
|
||||
)
|
||||
|
||||
max_size = max(MESSAGE_SIZES_BYTES)
|
||||
self.comm = CustomAllreduce(group, device, max_size=max_size)
|
||||
if self.comm.disabled:
|
||||
raise RuntimeError("AOT CustomAllreduce is disabled on this system")
|
||||
|
||||
def capture(self, register_input: bool):
|
||||
return self.comm.capture() # ignore register_input since v1 always requires it
|
||||
|
||||
def all_reduce(self, tensor: torch.Tensor) -> Optional[torch.Tensor]:
|
||||
assert self.comm.should_custom_ar(tensor), str(tensor.shape)
|
||||
return self.comm.custom_all_reduce(tensor)
|
||||
|
||||
|
||||
class JITAllReduceBackend:
|
||||
name = "JIT"
|
||||
|
||||
def __init__(self, group: dist.ProcessGroup, device: torch.device):
|
||||
from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
|
||||
CustomAllReduceV2,
|
||||
)
|
||||
|
||||
max_size = max(MESSAGE_SIZES_BYTES)
|
||||
self.comm = CustomAllReduceV2(group, device, max_pull_size=max_size)
|
||||
if self.comm.disabled:
|
||||
raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system")
|
||||
|
||||
def capture(self, register_input: bool):
|
||||
return self.comm.capture() if register_input else contextlib.nullcontext()
|
||||
|
||||
def all_reduce(self, tensor: torch.Tensor) -> Optional[torch.Tensor]:
|
||||
assert self.comm.should_custom_ar(tensor), str(tensor.shape)
|
||||
return self.comm.custom_all_reduce(tensor)
|
||||
|
||||
|
||||
class FlashInferAllReduceBackend:
|
||||
name = "FI"
|
||||
|
||||
def __init__(self, group: dist.ProcessGroup, dtype: torch.dtype):
|
||||
import flashinfer.comm as comm
|
||||
|
||||
rank = torch.distributed.get_rank(group=group)
|
||||
world_size = torch.distributed.get_world_size(group=group)
|
||||
max_size = max(MESSAGE_SIZES_BYTES)
|
||||
hidden_dim = min(MESSAGE_SIZES_BYTES) // 2
|
||||
num_tokens = max_size // hidden_dim
|
||||
self.comm = comm
|
||||
self.hidden_dim = hidden_dim
|
||||
self.workspace = comm.create_allreduce_fusion_workspace(
|
||||
backend="trtllm",
|
||||
world_size=world_size,
|
||||
rank=rank,
|
||||
max_token_num=num_tokens,
|
||||
hidden_dim=hidden_dim,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
def capture(self, *_):
|
||||
return contextlib.nullcontext()
|
||||
|
||||
def all_reduce(self, tensor: torch.Tensor) -> Optional[torch.Tensor]:
|
||||
return self.comm.allreduce_fusion(
|
||||
input=tensor.view(-1, self.hidden_dim),
|
||||
workspace=self.workspace,
|
||||
pattern=self.comm.AllReduceFusionPattern.kAllReduce,
|
||||
launch_with_pdl=True,
|
||||
fp32_acc=True,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmarking helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser(description=__doc__)
|
||||
p.add_argument("--dtype", choices=DTYPE_MAP.keys(), default="bfloat16")
|
||||
p.add_argument("--warmup", type=int, default=5)
|
||||
p.add_argument("--iters", type=int, default=50)
|
||||
p.add_argument("--no-inplace", dest="register_input", action="store_false")
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def bench_one(
|
||||
backend,
|
||||
inp: torch.Tensor,
|
||||
warmup: int,
|
||||
iters: int,
|
||||
group: dist.ProcessGroup,
|
||||
register_input: bool,
|
||||
) -> float:
|
||||
"""
|
||||
Run *warmup* iterations of all-reduce first.
|
||||
Return the average time for *iters* iterations of all-reduce.
|
||||
"""
|
||||
dist.barrier(group=group)
|
||||
for _ in range(warmup):
|
||||
backend.all_reduce(inp)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Capture a CUDA graph with *iters* all-reduce calls.
|
||||
inp_batch = torch.stack([inp] * 4)
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with backend.capture(register_input):
|
||||
with torch.cuda.graph(graph):
|
||||
for i in range(iters):
|
||||
backend.all_reduce(inp_batch[i % 4])
|
||||
|
||||
torch.cuda.synchronize()
|
||||
# Warm up the graph once.
|
||||
graph.replay()
|
||||
|
||||
# Timed replay.
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
torch.cuda.synchronize()
|
||||
dist.barrier(group=group)
|
||||
graph.replay() # make the stream busy
|
||||
start.record()
|
||||
graph.replay()
|
||||
end.record()
|
||||
torch.cuda.synchronize()
|
||||
return start.elapsed_time(end) / iters
|
||||
|
||||
|
||||
def bench_sweep(
|
||||
backend,
|
||||
sizes_bytes: List[int],
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
warmup: int,
|
||||
iters: int,
|
||||
group: dist.ProcessGroup,
|
||||
register_input: bool,
|
||||
) -> Dict[int, float]:
|
||||
"""Benchmark one backend over all message sizes."""
|
||||
elem_size = torch.tensor([], dtype=dtype).element_size()
|
||||
results: Dict[int, float] = {}
|
||||
for sz in sizes_bytes:
|
||||
numel = sz // elem_size
|
||||
inp = torch.zeros(numel, dtype=dtype, device=device)
|
||||
try:
|
||||
elapsed_ms = bench_one(backend, inp, warmup, iters, group, register_input)
|
||||
results[sz] = elapsed_ms * 1000 # convert to us per iter
|
||||
except AssertionError:
|
||||
results[sz] = float("nan")
|
||||
return results
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Result printing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def print_results(
|
||||
backends: list,
|
||||
all_results: Dict[str, Dict[int, float]],
|
||||
sizes_bytes: List[int],
|
||||
) -> None:
|
||||
"""Print a comparison table on rank 0."""
|
||||
|
||||
def human_bytes(n: int) -> str:
|
||||
for suffix, unit in [("M", 1 << 20), ("K", 1 << 10)]:
|
||||
if n >= unit and n % unit == 0:
|
||||
return f"{n // unit}{suffix}"
|
||||
return f"{n}B"
|
||||
|
||||
def fmt_us(v: float) -> str:
|
||||
return f"{v:13.1f}" if not isnan(v) else " n/a"
|
||||
|
||||
names = [b.name for b in backends]
|
||||
nccl_name = "NCCL"
|
||||
|
||||
# Header
|
||||
header_cols = [f"{n:>13}" for n in names]
|
||||
speedup_cols = [f"{n:>13}/NCCL" for n in names if n != nccl_name]
|
||||
header = f"{'Size':>8} " + " ".join(header_cols)
|
||||
for sc in speedup_cols:
|
||||
header += f" {sc}"
|
||||
header += " "
|
||||
print()
|
||||
print(header)
|
||||
print("-" * len(header))
|
||||
|
||||
# Rows
|
||||
for sz in sizes_bytes:
|
||||
row = f"{human_bytes(sz):>8}"
|
||||
nccl_lat = all_results[nccl_name][sz]
|
||||
for n in names:
|
||||
row += f" {fmt_us(all_results[n][sz])}"
|
||||
for n in names:
|
||||
if n == nccl_name:
|
||||
continue
|
||||
lat = all_results[n][sz]
|
||||
if not isnan(lat):
|
||||
row += f" {nccl_lat / lat:17.2f}x"
|
||||
else:
|
||||
row += f" {'n/a':>17}"
|
||||
print(row)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Distributed setup
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def init_distributed():
|
||||
"""Initialize distributed groups using torchrun env vars.
|
||||
|
||||
Returns (rank, world_size, device, cpu_group, nccl_group).
|
||||
"""
|
||||
import sglang.srt.distributed.parallel_state as ps
|
||||
|
||||
local_rank = int(os.environ.get("LOCAL_RANK", "0"))
|
||||
world_size = int(os.environ.get("WORLD_SIZE", "1"))
|
||||
rank = local_rank
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.cuda.set_device(device)
|
||||
torch.cuda.set_stream(torch.cuda.Stream()) # use a non-default stream
|
||||
|
||||
torch.distributed.init_process_group(backend="gloo")
|
||||
ps._WORLD = coord = ps.init_world_group(
|
||||
ranks=list(range(world_size)),
|
||||
local_rank=local_rank,
|
||||
backend="nccl",
|
||||
)
|
||||
|
||||
cpu_group = coord.cpu_group
|
||||
nccl_group = coord.device_group
|
||||
assert nccl_group is not None
|
||||
return rank, world_size, device, cpu_group, nccl_group
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def main():
|
||||
logging.basicConfig(level=logging.WARNING)
|
||||
args = parse_args()
|
||||
dtype = DTYPE_MAP[args.dtype]
|
||||
|
||||
rank, world_size, device, cpu_group, nccl_group = init_distributed()
|
||||
|
||||
# Instantiate backends.
|
||||
backends = [
|
||||
NCCLAllReduceBackend(nccl_group),
|
||||
JITAllReduceBackend(cpu_group, device),
|
||||
]
|
||||
if world_size in [2, 4, 6, 8]:
|
||||
backends.insert(1, AOTAllReduceBackend(cpu_group, device))
|
||||
if world_size in [2, 4, 8]:
|
||||
backends.append(FlashInferAllReduceBackend(cpu_group, dtype))
|
||||
|
||||
# Run benchmarks.
|
||||
all_results: Dict[str, Dict[int, float]] = {}
|
||||
torch.cuda.synchronize()
|
||||
for backend in backends:
|
||||
if rank == 0:
|
||||
print(f"Benchmarking {backend.name} ...")
|
||||
all_results[backend.name] = bench_sweep(
|
||||
backend,
|
||||
MESSAGE_SIZES_BYTES,
|
||||
dtype,
|
||||
device,
|
||||
args.warmup,
|
||||
args.iters,
|
||||
cpu_group,
|
||||
args.register_input,
|
||||
)
|
||||
|
||||
# Aggregate across ranks (use max to reflect the slowest rank).
|
||||
for name in list(all_results):
|
||||
for sz in MESSAGE_SIZES_BYTES:
|
||||
val = all_results[name].get(sz)
|
||||
if val is None:
|
||||
continue
|
||||
t = torch.tensor([val], dtype=torch.float64, device=device)
|
||||
dist.all_reduce(t, op=dist.ReduceOp.MAX, group=nccl_group)
|
||||
all_results[name][sz] = t.item()
|
||||
|
||||
# Print results on rank 0.
|
||||
if rank == 0:
|
||||
print_results(backends, all_results, MESSAGE_SIZES_BYTES)
|
||||
|
||||
del backends, all_results
|
||||
gc.collect()
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
if __name__ == "__main__" and not is_in_ci():
|
||||
main()
|
||||
@@ -0,0 +1,174 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.benchmark.bench_utils import run_bench
|
||||
from sglang.jit_kernel.benchmark.utils import get_benchmark_range
|
||||
from sglang.srt.utils import is_sm100_supported
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=5, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
try:
|
||||
import deep_gemm
|
||||
from deep_gemm.utils import per_token_cast_to_fp4
|
||||
except Exception:
|
||||
deep_gemm = None
|
||||
per_token_cast_to_fp4 = None
|
||||
|
||||
HEAD_DIM = 128
|
||||
NUM_HEADS = 64
|
||||
BLOCK_KV = 64
|
||||
NEXT_N = 1
|
||||
|
||||
shape_range = get_benchmark_range(
|
||||
full_range=[(256, 8192), (256, 32768)],
|
||||
ci_range=[(256, 8192)],
|
||||
)
|
||||
|
||||
|
||||
def _pack_fp8_cache(k: torch.Tensor, *, num_blocks: int) -> torch.Tensor:
|
||||
k = k.view(num_blocks, BLOCK_KV, 1, HEAD_DIM)
|
||||
scale = k.abs().float().amax(dim=3, keepdim=True).clamp(1.0e-4) / 448.0
|
||||
k_fp8 = (k * (1.0 / scale)).to(torch.float8_e4m3fn)
|
||||
buf = torch.empty(
|
||||
(num_blocks, BLOCK_KV * (HEAD_DIM + 4)), dtype=torch.uint8, device="cuda"
|
||||
)
|
||||
buf[:, : BLOCK_KV * HEAD_DIM].copy_(
|
||||
k_fp8.view(num_blocks, BLOCK_KV * HEAD_DIM).view(torch.uint8)
|
||||
)
|
||||
buf[:, BLOCK_KV * HEAD_DIM :].copy_(
|
||||
scale.view(num_blocks, BLOCK_KV).view(torch.uint8)
|
||||
)
|
||||
return buf.view(num_blocks, BLOCK_KV, 1, HEAD_DIM + 4)
|
||||
|
||||
|
||||
def _pack_fp4_cache(
|
||||
k_fp4: torch.Tensor,
|
||||
k_sf: torch.Tensor,
|
||||
*,
|
||||
num_blocks: int,
|
||||
) -> torch.Tensor:
|
||||
buf = torch.empty((num_blocks, BLOCK_KV * 68), dtype=torch.uint8, device="cuda")
|
||||
buf[:, : BLOCK_KV * 64].view(num_blocks, BLOCK_KV, 64).copy_(
|
||||
k_fp4.view(torch.uint8).view(num_blocks, BLOCK_KV, 64)
|
||||
)
|
||||
buf[:, BLOCK_KV * 64 :].view(num_blocks, BLOCK_KV, 4).copy_(
|
||||
k_sf.contiguous().view(torch.uint8).view(num_blocks, BLOCK_KV, 4)
|
||||
)
|
||||
return buf.view(num_blocks, BLOCK_KV, 1, 68)
|
||||
|
||||
|
||||
def _make_case(batch: int, seq_len_kv: int):
|
||||
if deep_gemm is None or per_token_cast_to_fp4 is None:
|
||||
raise RuntimeError("DeepGEMM is required for this benchmark.")
|
||||
|
||||
blocks_per_seq = triton.cdiv(seq_len_kv, BLOCK_KV)
|
||||
padded_len = blocks_per_seq * BLOCK_KV
|
||||
num_blocks = batch * blocks_per_seq
|
||||
num_cache_tokens = num_blocks * BLOCK_KV
|
||||
page_table = torch.arange(num_blocks, dtype=torch.int32, device="cuda").view(
|
||||
batch, blocks_per_seq
|
||||
)
|
||||
context_lens = torch.full(
|
||||
(batch, NEXT_N), seq_len_kv, dtype=torch.int32, device="cuda"
|
||||
)
|
||||
schedule = deep_gemm.get_paged_mqa_logits_metadata(
|
||||
context_lens, BLOCK_KV, deep_gemm.get_num_sms(), indices=None
|
||||
)
|
||||
|
||||
q = torch.randn(
|
||||
batch, NEXT_N, NUM_HEADS, HEAD_DIM, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
k = torch.randn(num_cache_tokens, HEAD_DIM, device="cuda", dtype=torch.bfloat16)
|
||||
weights = torch.randn(batch * NEXT_N, NUM_HEADS, device="cuda", dtype=torch.float32)
|
||||
|
||||
q_scale = q.abs().float().amax(dim=-1, keepdim=True).clamp(1.0e-4) / 448.0
|
||||
q_fp8 = (q.float() / q_scale).clamp(-448.0, 448.0).to(torch.float8_e4m3fn)
|
||||
weights_fp8 = (
|
||||
weights.view(batch, NEXT_N, NUM_HEADS)[:, :, :, None] * q_scale
|
||||
).view(batch * NEXT_N, NUM_HEADS)
|
||||
k_cache_fp8 = _pack_fp8_cache(k, num_blocks=num_blocks)
|
||||
|
||||
q_fp4_flat, q_sf_flat = per_token_cast_to_fp4(
|
||||
q.view(-1, HEAD_DIM), use_ue8m0=True, gran_k=32, use_packed_ue8m0=True
|
||||
)
|
||||
q_fp4 = q_fp4_flat.view(batch, NEXT_N, NUM_HEADS, HEAD_DIM // 2)
|
||||
q_sf = q_sf_flat.view(batch, NEXT_N, NUM_HEADS)
|
||||
k_fp4, k_sf = per_token_cast_to_fp4(
|
||||
k, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True
|
||||
)
|
||||
k_cache_fp4 = _pack_fp4_cache(k_fp4, k_sf, num_blocks=num_blocks)
|
||||
|
||||
return {
|
||||
"padded_len": padded_len,
|
||||
"page_table": page_table,
|
||||
"context_lens": context_lens,
|
||||
"schedule": schedule,
|
||||
"q_fp8": q_fp8,
|
||||
"weights_fp8": weights_fp8,
|
||||
"k_cache_fp8": k_cache_fp8,
|
||||
"q_fp4": q_fp4,
|
||||
"q_sf": q_sf,
|
||||
"weights": weights,
|
||||
"k_cache_fp4": k_cache_fp4,
|
||||
}
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["batch", "seq_len_kv"],
|
||||
x_vals=shape_range,
|
||||
x_log=False,
|
||||
line_arg="provider",
|
||||
line_vals=["fp8", "fp4"],
|
||||
line_names=["Default FP8 indexer", "FP4 indexer"],
|
||||
styles=[("blue", "-"), ("green", "-")],
|
||||
ylabel="us",
|
||||
plot_name="dsv4-fp4-indexer-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(batch: int, seq_len_kv: int, provider: str):
|
||||
case = _make_case(batch, seq_len_kv)
|
||||
if provider == "fp8":
|
||||
fn = lambda: deep_gemm.fp8_paged_mqa_logits(
|
||||
case["q_fp8"],
|
||||
case["k_cache_fp8"],
|
||||
case["weights_fp8"],
|
||||
case["context_lens"],
|
||||
case["page_table"],
|
||||
case["schedule"],
|
||||
case["padded_len"],
|
||||
clean_logits=False,
|
||||
indices=None,
|
||||
)
|
||||
elif provider == "fp4":
|
||||
fn = lambda: deep_gemm.fp8_fp4_paged_mqa_logits(
|
||||
(case["q_fp4"], case["q_sf"]),
|
||||
case["k_cache_fp4"],
|
||||
case["weights"],
|
||||
case["context_lens"],
|
||||
case["page_table"],
|
||||
case["schedule"],
|
||||
case["padded_len"],
|
||||
clean_logits=False,
|
||||
logits_dtype=torch.float32,
|
||||
indices=None,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown provider: {provider}")
|
||||
return tuple(t * 1000 for t in run_bench(fn, use_cuda_graph=False))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if not is_sm100_supported():
|
||||
print("[skip] DeepSeek V4 FP4 indexer benchmark requires SM100 CUDA.")
|
||||
sys.exit(0)
|
||||
if deep_gemm is None or per_token_cast_to_fp4 is None:
|
||||
print("[skip] DeepGEMM is unavailable.")
|
||||
sys.exit(0)
|
||||
benchmark.run(print_data=True)
|
||||
@@ -0,0 +1,267 @@
|
||||
"""
|
||||
Benchmark: fused_qknorm_rope JIT vs AOT (sgl_kernel)
|
||||
|
||||
Measures throughput (µs) for fused_qk_norm_rope across typical
|
||||
LLM configurations (head_dim × num_heads × num_tokens).
|
||||
|
||||
Run:
|
||||
python test/registered/jit/benchmark/bench_fused_qknorm_rope.py
|
||||
"""
|
||||
|
||||
import itertools
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
|
||||
from sglang.jit_kernel.fused_qknorm_rope import (
|
||||
fused_qk_norm_rope as fused_qk_norm_rope_jit,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=6, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
try:
|
||||
from sgl_kernel import fused_qk_norm_rope as fused_qk_norm_rope_aot
|
||||
|
||||
AOT_AVAILABLE = True
|
||||
except ImportError:
|
||||
fused_qk_norm_rope_aot = None
|
||||
AOT_AVAILABLE = False
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark configuration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
NUM_TOKENS_RANGE = get_benchmark_range(
|
||||
full_range=[1, 64, 256, 1024, 4096],
|
||||
ci_range=[64, 512],
|
||||
)
|
||||
|
||||
# (head_dim, num_heads_q, num_heads_k, num_heads_v) — typical MoE/dense configs
|
||||
MODEL_CONFIGS = get_benchmark_range(
|
||||
full_range=[
|
||||
(64, 32, 8, 8), # small
|
||||
(128, 32, 8, 8), # typical (e.g. Qwen3-8B)
|
||||
(256, 16, 4, 4), # large head_dim
|
||||
],
|
||||
ci_range=[(128, 32, 8, 8)],
|
||||
)
|
||||
|
||||
# Real production shapes (self-attention; num_heads_k == num_heads_v == num_heads_q).
|
||||
# Format: (name, num_tokens, num_heads_q, num_heads_k, num_heads_v, head_dim, rotary_dim)
|
||||
PRODUCTION_SHAPES = [
|
||||
("flux_1024", 4096, 24, 24, 24, 128, 128),
|
||||
("qwen_image_1024", 4096, 32, 32, 32, 128, 128),
|
||||
("qwen_image_partial", 4096, 32, 32, 32, 128, 64),
|
||||
("zimage_1024", 4096, 30, 30, 30, 128, 128),
|
||||
("batch2_medium", 4096, 24, 24, 24, 128, 128), # B=2, T=2048
|
||||
]
|
||||
|
||||
LINE_VALS = ["jit", "aot"] if AOT_AVAILABLE else ["jit"]
|
||||
LINE_NAMES = ["JIT (new)", "AOT sgl_kernel"] if AOT_AVAILABLE else ["JIT (new)"]
|
||||
STYLES = [("blue", "--"), ("orange", "-")] if AOT_AVAILABLE else [("blue", "--")]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark: fused_qk_norm_rope (interleave style, no YaRN)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["num_tokens", "head_dim", "num_heads_q", "num_heads_k", "num_heads_v"],
|
||||
x_vals=[
|
||||
(nt, hd, nq, nk, nv)
|
||||
for nt, (hd, nq, nk, nv) in itertools.product(
|
||||
NUM_TOKENS_RANGE, MODEL_CONFIGS
|
||||
)
|
||||
],
|
||||
line_arg="provider",
|
||||
line_vals=LINE_VALS,
|
||||
line_names=LINE_NAMES,
|
||||
styles=STYLES,
|
||||
ylabel="us",
|
||||
plot_name="fused-qknorm-rope-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def bench_fused_qknorm_rope(
|
||||
num_tokens: int,
|
||||
head_dim: int,
|
||||
num_heads_q: int,
|
||||
num_heads_k: int,
|
||||
num_heads_v: int,
|
||||
provider: str,
|
||||
):
|
||||
device = "cuda"
|
||||
total_heads = num_heads_q + num_heads_k + num_heads_v
|
||||
|
||||
qkv = torch.randn(
|
||||
(num_tokens, total_heads * head_dim), dtype=torch.bfloat16, device=device
|
||||
)
|
||||
q_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device)
|
||||
k_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device)
|
||||
position_ids = torch.arange(num_tokens, dtype=torch.int32, device=device)
|
||||
|
||||
common_kwargs = dict(
|
||||
num_heads_q=num_heads_q,
|
||||
num_heads_k=num_heads_k,
|
||||
num_heads_v=num_heads_v,
|
||||
head_dim=head_dim,
|
||||
eps=1e-5,
|
||||
q_weight=q_weight,
|
||||
k_weight=k_weight,
|
||||
base=10000.0,
|
||||
is_neox=False,
|
||||
position_ids=position_ids,
|
||||
factor=1.0,
|
||||
low=1.0,
|
||||
high=32.0,
|
||||
attention_factor=1.0,
|
||||
rotary_dim=head_dim,
|
||||
)
|
||||
|
||||
if provider == "jit":
|
||||
fn = lambda: fused_qk_norm_rope_jit(qkv.clone(), **common_kwargs)
|
||||
elif provider == "aot":
|
||||
fn = lambda: fused_qk_norm_rope_aot(qkv.clone(), **common_kwargs)
|
||||
else:
|
||||
raise ValueError(f"Unknown provider: {provider}")
|
||||
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark: fused_qk_norm_rope — real production shapes (with speedup column)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def bench_fused_qknorm_rope_production():
|
||||
device = "cuda"
|
||||
header = f"{'name':<22} {'tokens':>6} {'nq':>4} {'nk':>4} {'nv':>4} {'hd':>4} {'rdim':>5} {'JIT(us)':>9} {'AOT(us)':>9} {'speedup':>8}"
|
||||
sep = "-" * len(header)
|
||||
print("\nfused-qknorm-rope-production-shapes:")
|
||||
print(sep)
|
||||
print(header)
|
||||
print(sep)
|
||||
|
||||
for (
|
||||
name,
|
||||
num_tokens,
|
||||
num_heads_q,
|
||||
num_heads_k,
|
||||
num_heads_v,
|
||||
head_dim,
|
||||
rotary_dim,
|
||||
) in PRODUCTION_SHAPES:
|
||||
total_heads = num_heads_q + num_heads_k + num_heads_v
|
||||
qkv = torch.randn(
|
||||
(num_tokens, total_heads * head_dim), dtype=torch.bfloat16, device=device
|
||||
)
|
||||
q_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device)
|
||||
k_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device)
|
||||
position_ids = torch.arange(num_tokens, dtype=torch.int32, device=device)
|
||||
|
||||
common_kwargs = dict(
|
||||
num_heads_q=num_heads_q,
|
||||
num_heads_k=num_heads_k,
|
||||
num_heads_v=num_heads_v,
|
||||
head_dim=head_dim,
|
||||
eps=1e-5,
|
||||
q_weight=q_weight,
|
||||
k_weight=k_weight,
|
||||
base=10000.0,
|
||||
is_neox=False,
|
||||
position_ids=position_ids,
|
||||
factor=1.0,
|
||||
low=1.0,
|
||||
high=32.0,
|
||||
attention_factor=1.0,
|
||||
rotary_dim=rotary_dim,
|
||||
)
|
||||
|
||||
jit_us, _, _ = run_benchmark(
|
||||
lambda: fused_qk_norm_rope_jit(qkv.clone(), **common_kwargs)
|
||||
)
|
||||
if AOT_AVAILABLE:
|
||||
aot_us, _, _ = run_benchmark(
|
||||
lambda: fused_qk_norm_rope_aot(qkv.clone(), **common_kwargs)
|
||||
)
|
||||
speedup = f"{aot_us / jit_us:.2f}x"
|
||||
aot_str = f"{aot_us:9.3f}"
|
||||
else:
|
||||
aot_str = f"{'N/A':>9}"
|
||||
speedup = "N/A"
|
||||
|
||||
print(
|
||||
f"{name:<22} {num_tokens:>6} {num_heads_q:>4} {num_heads_k:>4} {num_heads_v:>4}"
|
||||
f" {head_dim:>4} {rotary_dim:>5} {jit_us:9.3f} {aot_str} {speedup:>8}"
|
||||
)
|
||||
print(sep)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Quick correctness diff
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def calculate_diff():
|
||||
if not AOT_AVAILABLE:
|
||||
print("sgl_kernel not available — skipping AOT diff check")
|
||||
return
|
||||
|
||||
device = "cuda"
|
||||
print("Correctness diff (JIT vs AOT):")
|
||||
|
||||
for head_dim, is_neox in [(64, False), (128, False), (128, True), (256, False)]:
|
||||
num_tokens = 32
|
||||
num_heads_q, num_heads_k, num_heads_v = 4, 2, 2
|
||||
total_heads = num_heads_q + num_heads_k + num_heads_v
|
||||
|
||||
qkv = torch.randn(
|
||||
(num_tokens, total_heads * head_dim), dtype=torch.bfloat16, device=device
|
||||
)
|
||||
q_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device)
|
||||
k_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device)
|
||||
position_ids = torch.arange(num_tokens, dtype=torch.int32, device=device)
|
||||
|
||||
common = dict(
|
||||
num_heads_q=num_heads_q,
|
||||
num_heads_k=num_heads_k,
|
||||
num_heads_v=num_heads_v,
|
||||
head_dim=head_dim,
|
||||
eps=1e-5,
|
||||
q_weight=q_weight,
|
||||
k_weight=k_weight,
|
||||
base=10000.0,
|
||||
is_neox=is_neox,
|
||||
position_ids=position_ids,
|
||||
factor=1.0,
|
||||
low=1.0,
|
||||
high=32.0,
|
||||
attention_factor=1.0,
|
||||
rotary_dim=head_dim,
|
||||
)
|
||||
|
||||
qkv_jit = qkv.clone()
|
||||
fused_qk_norm_rope_jit(qkv_jit, **common)
|
||||
qkv_aot = qkv.clone()
|
||||
fused_qk_norm_rope_aot(qkv_aot, **common)
|
||||
|
||||
match = torch.allclose(qkv_jit.float(), qkv_aot.float(), atol=1e-2, rtol=1e-2)
|
||||
status = "OK" if match else "MISMATCH"
|
||||
max_err = (qkv_jit.float() - qkv_aot.float()).abs().max().item()
|
||||
print(
|
||||
f" head_dim={head_dim:3d} is_neox={str(is_neox):5s} "
|
||||
f"max_err={max_err:.2e} [{status}]"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
calculate_diff()
|
||||
print()
|
||||
bench_fused_qknorm_rope.run(print_data=True)
|
||||
print()
|
||||
bench_fused_qknorm_rope_production()
|
||||
@@ -0,0 +1,119 @@
|
||||
import itertools
|
||||
import math
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import (
|
||||
DEFAULT_DEVICE,
|
||||
DEFAULT_DTYPE,
|
||||
get_benchmark_range,
|
||||
run_benchmark,
|
||||
)
|
||||
from sglang.jit_kernel.hadamard import hadamard_transform
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=5, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
# AOT kernel: might not be available in all environments.
|
||||
# This is used for performance baseline comparison.
|
||||
try:
|
||||
from sgl_kernel import hadamard_transform as hadamard_transform_aot
|
||||
|
||||
AOT_AVAILABLE = True
|
||||
except Exception:
|
||||
AOT_AVAILABLE = False
|
||||
|
||||
# Naive reference implementation using scipy hadamard matrix.
|
||||
try:
|
||||
from scipy.linalg import hadamard
|
||||
|
||||
SCIPY_AVAILABLE = True
|
||||
except ImportError:
|
||||
SCIPY_AVAILABLE = False
|
||||
|
||||
# CI environment uses simplified parameters
|
||||
batch_sizes = get_benchmark_range(
|
||||
full_range=[1, 16, 64, 256],
|
||||
ci_range=[16],
|
||||
)
|
||||
dim_range = get_benchmark_range(
|
||||
full_range=[64, 256, 1024, 4096, 8192, 16384, 32768],
|
||||
ci_range=[1024],
|
||||
)
|
||||
|
||||
|
||||
# Naive reference implementation using precomputed scipy hadamard matrix.
|
||||
def torch_hadamard_transform(x, scale, H, dim, dim_padded):
|
||||
flat = x.reshape(-1, dim)
|
||||
if dim != dim_padded:
|
||||
flat = F.pad(flat, (0, dim_padded - dim))
|
||||
out = F.linear(flat, H) * scale
|
||||
return out[..., :dim].reshape(x.shape)
|
||||
|
||||
|
||||
available_providers = ["jit_kernel"]
|
||||
available_names = ["JIT Kernel"]
|
||||
available_styles = [("red", "-")]
|
||||
|
||||
if AOT_AVAILABLE:
|
||||
available_providers.insert(0, "aot_kernel")
|
||||
available_names.insert(0, "AOT Kernel")
|
||||
available_styles.insert(0, ("green", "-"))
|
||||
|
||||
if SCIPY_AVAILABLE:
|
||||
available_providers.append("naive")
|
||||
available_names.append("Naive (scipy)")
|
||||
available_styles.append(("blue", "-"))
|
||||
|
||||
configs = list(itertools.product(batch_sizes, dim_range))
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["batch_size", "dim"],
|
||||
x_vals=[list(c) for c in configs],
|
||||
line_arg="provider",
|
||||
line_vals=available_providers,
|
||||
line_names=available_names,
|
||||
styles=available_styles,
|
||||
ylabel="us",
|
||||
plot_name="hadamard-transform-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(batch_size: int, dim: int, provider: str) -> Tuple[float, float, float]:
|
||||
scale = 1.0 / math.sqrt(dim)
|
||||
x = torch.randn(batch_size, dim, device=DEFAULT_DEVICE, dtype=DEFAULT_DTYPE)
|
||||
|
||||
FN_MAP = {
|
||||
"jit_kernel": lambda: hadamard_transform(x.clone(), scale=scale),
|
||||
}
|
||||
if AOT_AVAILABLE:
|
||||
FN_MAP["aot_kernel"] = lambda: hadamard_transform_aot(x.clone(), scale=scale)
|
||||
if SCIPY_AVAILABLE:
|
||||
# Precompute Hadamard matrix on GPU to avoid CPU-GPU transfer
|
||||
# during CUDA graph capture.
|
||||
log_dim = math.ceil(math.log2(dim)) if dim > 0 else 0
|
||||
dim_padded = 2**log_dim if dim > 0 else 1
|
||||
H = torch.tensor(
|
||||
hadamard(dim_padded, dtype=float),
|
||||
dtype=DEFAULT_DTYPE,
|
||||
device=DEFAULT_DEVICE,
|
||||
)
|
||||
FN_MAP["naive"] = lambda: torch_hadamard_transform(
|
||||
x.clone(), scale, H, dim, dim_padded
|
||||
)
|
||||
|
||||
fn = FN_MAP[provider]
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("=" * 80)
|
||||
print("Benchmarking Fast Hadamard Transform")
|
||||
print("=" * 80)
|
||||
benchmark.run(print_data=True)
|
||||
@@ -0,0 +1,421 @@
|
||||
"""Benchmark for HiCache JIT kernel performance.
|
||||
|
||||
This benchmark tests the performance of KV cache transfer operations
|
||||
between GPU and CPU (host pinned memory), comparing:
|
||||
- SGL AOT Kernel: Pre-compiled transfer_kv kernels from sgl_kernel
|
||||
- SGL JIT Kernel: JIT-compiled hicache kernels
|
||||
- PyTorch Indexing: Plain PyTorch index copy
|
||||
- PyTorch 2 Stream: PyTorch implementation using 2 CUDA streams
|
||||
|
||||
Tests cover:
|
||||
- One Layer: CPU->GPU
|
||||
- All Layer: GPU->CPU
|
||||
|
||||
Note: Uses do_bench instead of do_bench_cudagraph since CUDA graph
|
||||
capture doesn't support CPU-GPU memory transfers.
|
||||
"""
|
||||
|
||||
import itertools
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
from sgl_kernel import transfer_kv_all_layer, transfer_kv_per_layer
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import DEFAULT_QUANTILES, get_benchmark_range
|
||||
from sglang.jit_kernel.hicache import (
|
||||
can_use_hicache_jit_kernel,
|
||||
transfer_hicache_all_layer,
|
||||
transfer_hicache_one_layer,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=29, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
DISABLE_TORCH = os.environ.get("DISABLE_TORCH", "0") == "1"
|
||||
PAGE_SIZE = 1
|
||||
ENABLE_SORT = True
|
||||
GPU_CACHE_SIZE = 256 * 1024 # 256K tokens on GPU
|
||||
HOST_CACHE_SIZE = 512 * 1024 # 512K tokens on CPU
|
||||
NUM_LAYERS = 8
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class HiCacheCache:
|
||||
k_cache_cuda: torch.Tensor
|
||||
v_cache_cuda: torch.Tensor
|
||||
k_cache_host: torch.Tensor
|
||||
v_cache_host: torch.Tensor
|
||||
|
||||
def get_slice(self, num_layers: int, element_size: int) -> "HiCacheCache":
|
||||
def slice_cuda(t: torch.Tensor) -> torch.Tensor:
|
||||
needed_cuda = num_layers * GPU_CACHE_SIZE
|
||||
return t.view(-1, element_size)[:needed_cuda].unflatten(0, (num_layers, -1))
|
||||
|
||||
def slice_host(t: torch.Tensor) -> torch.Tensor:
|
||||
needed_host = num_layers * HOST_CACHE_SIZE
|
||||
return t.view(-1, element_size)[:needed_host].unflatten(0, (num_layers, -1))
|
||||
|
||||
return HiCacheCache(
|
||||
k_cache_cuda=slice_cuda(self.k_cache_cuda),
|
||||
v_cache_cuda=slice_cuda(self.v_cache_cuda),
|
||||
k_cache_host=slice_host(self.k_cache_host),
|
||||
v_cache_host=slice_host(self.v_cache_host),
|
||||
)
|
||||
|
||||
|
||||
def gen_indices(
|
||||
size: int, max_size: int, *, page_size: int = PAGE_SIZE
|
||||
) -> torch.Tensor:
|
||||
def align(x: int) -> int:
|
||||
return (x + page_size - 1) // page_size
|
||||
|
||||
assert size <= max_size and max_size % page_size == 0
|
||||
indices = torch.randperm(align(max_size))[: align(size)]
|
||||
offsets = torch.arange(page_size)
|
||||
return (indices[:, None] * page_size + offsets).flatten().cuda()[:size]
|
||||
|
||||
|
||||
def sglang_aot_transfer_one(
|
||||
k_cache_dst: torch.Tensor,
|
||||
v_cache_dst: torch.Tensor,
|
||||
indices_dst: torch.Tensor,
|
||||
k_cache_src: torch.Tensor,
|
||||
v_cache_src: torch.Tensor,
|
||||
indices_src: torch.Tensor,
|
||||
item_size: int,
|
||||
) -> None:
|
||||
"""SGL AOT Kernel for single layer transfer."""
|
||||
transfer_kv_per_layer(
|
||||
k_cache_src,
|
||||
k_cache_dst,
|
||||
v_cache_src,
|
||||
v_cache_dst,
|
||||
indices_src,
|
||||
indices_dst,
|
||||
item_size,
|
||||
)
|
||||
|
||||
|
||||
def sglang_jit_transfer_one(
|
||||
k_cache_dst: torch.Tensor,
|
||||
v_cache_dst: torch.Tensor,
|
||||
indices_dst: torch.Tensor,
|
||||
k_cache_src: torch.Tensor,
|
||||
v_cache_src: torch.Tensor,
|
||||
indices_src: torch.Tensor,
|
||||
element_dim: int,
|
||||
) -> None:
|
||||
"""SGL JIT Kernel for single layer transfer."""
|
||||
transfer_hicache_one_layer(
|
||||
k_cache_dst,
|
||||
v_cache_dst,
|
||||
indices_dst,
|
||||
k_cache_src,
|
||||
v_cache_src,
|
||||
indices_src,
|
||||
element_dim=element_dim,
|
||||
)
|
||||
|
||||
|
||||
def sglang_aot_transfer_all(
|
||||
k_ptrs_dst: torch.Tensor,
|
||||
v_ptrs_dst: torch.Tensor,
|
||||
indices_dst: torch.Tensor,
|
||||
k_ptrs_src: torch.Tensor,
|
||||
v_ptrs_src: torch.Tensor,
|
||||
indices_src: torch.Tensor,
|
||||
item_size: int,
|
||||
num_layers: int,
|
||||
) -> None:
|
||||
"""SGL AOT Kernel for all layer transfer."""
|
||||
transfer_kv_all_layer(
|
||||
k_ptrs_src,
|
||||
k_ptrs_dst,
|
||||
v_ptrs_src,
|
||||
v_ptrs_dst,
|
||||
indices_src,
|
||||
indices_dst,
|
||||
item_size,
|
||||
num_layers,
|
||||
)
|
||||
|
||||
|
||||
def sglang_jit_transfer_all(
|
||||
k_ptrs_dst: torch.Tensor,
|
||||
v_ptrs_dst: torch.Tensor,
|
||||
indices_dst: torch.Tensor,
|
||||
k_ptrs_src: torch.Tensor,
|
||||
v_ptrs_src: torch.Tensor,
|
||||
indices_src: torch.Tensor,
|
||||
stride_bytes: int,
|
||||
element_size: int,
|
||||
) -> None:
|
||||
"""SGL JIT Kernel for all layer transfer."""
|
||||
transfer_hicache_all_layer(
|
||||
k_ptrs_dst,
|
||||
v_ptrs_dst,
|
||||
indices_dst,
|
||||
k_ptrs_src,
|
||||
v_ptrs_src,
|
||||
indices_src,
|
||||
kv_cache_src_stride_bytes=stride_bytes,
|
||||
kv_cache_dst_stride_bytes=stride_bytes,
|
||||
element_size=element_size,
|
||||
)
|
||||
|
||||
|
||||
def pytorch_transfer(
|
||||
k_cache_dst: torch.Tensor,
|
||||
v_cache_dst: torch.Tensor,
|
||||
indices_dst_on_dst: torch.Tensor,
|
||||
k_cache_src: torch.Tensor,
|
||||
v_cache_src: torch.Tensor,
|
||||
indices_src_on_src: torch.Tensor,
|
||||
) -> None:
|
||||
"""PyTorch indexing baseline."""
|
||||
dst_device = k_cache_dst.device
|
||||
k_cache_dst[indices_dst_on_dst] = k_cache_src[indices_src_on_src].to(dst_device)
|
||||
v_cache_dst[indices_dst_on_dst] = v_cache_src[indices_src_on_src].to(dst_device)
|
||||
|
||||
|
||||
# Benchmark configuration
|
||||
|
||||
BS_RANGE = get_benchmark_range(
|
||||
full_range=[2**n for n in range(0, 16)],
|
||||
ci_range=[16],
|
||||
)
|
||||
ELEMENT_SIZE_RANGE = get_benchmark_range(
|
||||
full_range=[64, 128, 256, 512, 1024],
|
||||
ci_range=[1024],
|
||||
)
|
||||
|
||||
LINE_VALS = ["aot", "jit", "torch"]
|
||||
LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "PyTorch"]
|
||||
STYLES = [("orange", "-"), ("blue", "--"), ("red", ":")]
|
||||
|
||||
CONFIGS = list(itertools.product(ELEMENT_SIZE_RANGE, BS_RANGE))
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# One Layer Benchmarks
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["element_size", "batch_size"],
|
||||
x_vals=CONFIGS,
|
||||
line_arg="provider",
|
||||
line_vals=LINE_VALS,
|
||||
line_names=LINE_NAMES,
|
||||
styles=STYLES,
|
||||
ylabel="us",
|
||||
plot_name="hicache-one-layer-h2d",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark_one_layer_h2d(
|
||||
element_size: int, batch_size: int, provider: str
|
||||
) -> Tuple[float, float, float]:
|
||||
"""One Layer: Host (CPU) -> Device (GPU)."""
|
||||
global cache
|
||||
cache_local = cache.get_slice(num_layers=NUM_LAYERS, element_size=element_size)
|
||||
k_cache_src = cache_local.k_cache_host
|
||||
v_cache_src = cache_local.v_cache_host
|
||||
k_cache_dst = cache_local.k_cache_cuda
|
||||
v_cache_dst = cache_local.v_cache_cuda
|
||||
torch.manual_seed(batch_size * 65536 + element_size)
|
||||
indices_src_gpu = gen_indices(batch_size, HOST_CACHE_SIZE)
|
||||
indices_dst_gpu = gen_indices(batch_size, GPU_CACHE_SIZE)
|
||||
|
||||
if ENABLE_SORT:
|
||||
indices_src_gpu, mapping = indices_src_gpu.sort()
|
||||
indices_dst_gpu = indices_dst_gpu[mapping]
|
||||
indices_src_cpu = indices_src_gpu.cpu()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
element_bytes = element_size * k_cache_src.element_size()
|
||||
|
||||
FN_MAP = {
|
||||
"aot": lambda: [
|
||||
sglang_aot_transfer_one(
|
||||
k_cache_dst[i],
|
||||
v_cache_dst[i],
|
||||
indices_dst_gpu,
|
||||
k_cache_src[i],
|
||||
v_cache_src[i],
|
||||
indices_src_gpu,
|
||||
element_bytes,
|
||||
)
|
||||
for i in range(NUM_LAYERS)
|
||||
],
|
||||
"jit": lambda: [
|
||||
sglang_jit_transfer_one(
|
||||
k_cache_dst[i],
|
||||
v_cache_dst[i],
|
||||
indices_dst_gpu,
|
||||
k_cache_src[i],
|
||||
v_cache_src[i],
|
||||
indices_src_gpu,
|
||||
element_size,
|
||||
)
|
||||
for i in range(NUM_LAYERS)
|
||||
],
|
||||
"torch": lambda: [
|
||||
pytorch_transfer(
|
||||
k_cache_dst[i],
|
||||
v_cache_dst[i],
|
||||
indices_dst_gpu,
|
||||
k_cache_src[i],
|
||||
v_cache_src[i],
|
||||
indices_src_cpu,
|
||||
)
|
||||
for i in range(NUM_LAYERS)
|
||||
],
|
||||
}
|
||||
|
||||
if provider == "jit" and not can_use_hicache_jit_kernel(element_size=element_bytes):
|
||||
return (float("nan"), float("nan"), float("nan"))
|
||||
|
||||
if DISABLE_TORCH and provider in ["torch"]:
|
||||
return (float("nan"), float("nan"), float("nan"))
|
||||
|
||||
ms, min_ms, max_ms = triton.testing.do_bench( # type: ignore
|
||||
FN_MAP[provider], quantiles=DEFAULT_QUANTILES, warmup=5, rep=25
|
||||
)
|
||||
return (
|
||||
1000 * ms / NUM_LAYERS,
|
||||
1000 * max_ms / NUM_LAYERS,
|
||||
1000 * min_ms / NUM_LAYERS,
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# All Layer Benchmarks
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _create_ptr_tensor(tensors, device="cuda"):
|
||||
"""Create a tensor of data pointers."""
|
||||
return torch.tensor(
|
||||
[t.data_ptr() for t in tensors],
|
||||
dtype=torch.uint64,
|
||||
device=device,
|
||||
)
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["element_size", "batch_size"],
|
||||
x_vals=CONFIGS,
|
||||
line_arg="provider",
|
||||
line_vals=LINE_VALS,
|
||||
line_names=LINE_NAMES,
|
||||
styles=STYLES,
|
||||
ylabel="us",
|
||||
plot_name="hicache-all-layer-d2h",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark_all_layer_d2h(
|
||||
element_size: int, batch_size: int, provider: str
|
||||
) -> Tuple[float, float, float]:
|
||||
"""All Layer: Device (GPU) -> Host (CPU)."""
|
||||
global cache
|
||||
cache_local = cache.get_slice(num_layers=NUM_LAYERS, element_size=element_size)
|
||||
k_caches_src = cache_local.k_cache_cuda
|
||||
v_caches_src = cache_local.v_cache_cuda
|
||||
k_caches_dst = cache_local.k_cache_host
|
||||
v_caches_dst = cache_local.v_cache_host
|
||||
torch.manual_seed(batch_size * 65536 + element_size)
|
||||
|
||||
indices_src_gpu = gen_indices(batch_size, GPU_CACHE_SIZE)
|
||||
indices_dst_gpu = gen_indices(batch_size, HOST_CACHE_SIZE)
|
||||
if ENABLE_SORT:
|
||||
indices_dst_gpu, mapping = indices_dst_gpu.sort()
|
||||
indices_src_gpu = indices_src_gpu[mapping]
|
||||
indices_dst_cpu = indices_dst_gpu.cpu()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
element_bytes = element_size * k_caches_src.element_size()
|
||||
|
||||
k_ptrs_src = _create_ptr_tensor([k_caches_src[i] for i in range(NUM_LAYERS)])
|
||||
v_ptrs_src = _create_ptr_tensor([v_caches_src[i] for i in range(NUM_LAYERS)])
|
||||
k_ptrs_dst = _create_ptr_tensor([k_caches_dst[i] for i in range(NUM_LAYERS)])
|
||||
v_ptrs_dst = _create_ptr_tensor([v_caches_dst[i] for i in range(NUM_LAYERS)])
|
||||
|
||||
FN_MAP = {
|
||||
"aot": lambda: sglang_aot_transfer_all(
|
||||
k_ptrs_dst,
|
||||
v_ptrs_dst,
|
||||
indices_dst_gpu,
|
||||
k_ptrs_src,
|
||||
v_ptrs_src,
|
||||
indices_src_gpu,
|
||||
element_bytes,
|
||||
NUM_LAYERS,
|
||||
),
|
||||
"jit": lambda: sglang_jit_transfer_all(
|
||||
k_ptrs_dst,
|
||||
v_ptrs_dst,
|
||||
indices_dst_gpu,
|
||||
k_ptrs_src,
|
||||
v_ptrs_src,
|
||||
indices_src_gpu,
|
||||
element_bytes,
|
||||
element_bytes,
|
||||
),
|
||||
"torch": lambda: [
|
||||
pytorch_transfer(
|
||||
k_caches_dst[i],
|
||||
v_caches_dst[i],
|
||||
indices_dst_cpu,
|
||||
k_caches_src[i],
|
||||
v_caches_src[i],
|
||||
indices_src_gpu,
|
||||
)
|
||||
for i in range(NUM_LAYERS)
|
||||
],
|
||||
}
|
||||
|
||||
if provider == "jit" and not can_use_hicache_jit_kernel(element_size=element_bytes):
|
||||
return (float("nan"), float("nan"), float("nan"))
|
||||
|
||||
if DISABLE_TORCH and provider in ["torch"]:
|
||||
return (float("nan"), float("nan"), float("nan"))
|
||||
|
||||
ms, min_ms, max_ms = triton.testing.do_bench( # type: ignore
|
||||
FN_MAP[provider], quantiles=DEFAULT_QUANTILES, warmup=5, rep=25
|
||||
)
|
||||
return (
|
||||
1000 * ms / NUM_LAYERS,
|
||||
1000 * max_ms / NUM_LAYERS,
|
||||
1000 * min_ms / NUM_LAYERS,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
MAX_SIZE = max(ELEMENT_SIZE_RANGE)
|
||||
DEVICE_SHAPE = (NUM_LAYERS * GPU_CACHE_SIZE, MAX_SIZE)
|
||||
HOST_SHAPE = (NUM_LAYERS * HOST_CACHE_SIZE, MAX_SIZE)
|
||||
|
||||
cache = HiCacheCache(
|
||||
k_cache_cuda=torch.empty(DEVICE_SHAPE, dtype=torch.bfloat16, device="cuda"),
|
||||
v_cache_cuda=torch.empty(DEVICE_SHAPE, dtype=torch.bfloat16, device="cuda"),
|
||||
k_cache_host=torch.empty(HOST_SHAPE, dtype=torch.bfloat16, pin_memory=True),
|
||||
v_cache_host=torch.empty(HOST_SHAPE, dtype=torch.bfloat16, pin_memory=True),
|
||||
)
|
||||
|
||||
print("=" * 60)
|
||||
print("One Layer: Host -> Device (CPU -> GPU)")
|
||||
print("=" * 60)
|
||||
benchmark_one_layer_h2d.run(print_data=True)
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("All Layer: Device -> Host (GPU -> CPU) [per-layer avg]")
|
||||
print("=" * 60)
|
||||
benchmark_all_layer_d2h.run(print_data=True)
|
||||
@@ -0,0 +1,190 @@
|
||||
import itertools
|
||||
from typing import Dict, Tuple
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE, DEFAULT_DTYPE
|
||||
from sglang.jit_kernel.hisparse import load_cache_to_device_buffer_mla
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=12, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
DEVICE = DEFAULT_DEVICE
|
||||
DTYPE = DEFAULT_DTYPE
|
||||
TOP_K = 2048
|
||||
ITEM_SIZE_BYTES = 512
|
||||
MISS_RATES = [0.2, 0.001]
|
||||
ROUNDS = 5
|
||||
WARMUP_ROUNDS = 5
|
||||
BATCH_SIZES = [1, 10, 100]
|
||||
HOT_BUFFER_SIZES = [4096, 8192]
|
||||
CONFIGS = [
|
||||
(
|
||||
batch_size,
|
||||
hot_buffer_size,
|
||||
miss_rate,
|
||||
batch_size * round(TOP_K * miss_rate),
|
||||
)
|
||||
for batch_size, hot_buffer_size, miss_rate in itertools.product(
|
||||
BATCH_SIZES, HOT_BUFFER_SIZES, MISS_RATES
|
||||
)
|
||||
]
|
||||
|
||||
LINE_VALS = ["jit"]
|
||||
LINE_NAMES = ["SGL JIT Kernel"]
|
||||
STYLES = [("blue", "--")]
|
||||
|
||||
|
||||
def _make_top_k_tokens(
|
||||
num_hits: int, num_misses: int, hot_buffer_size: int
|
||||
) -> torch.Tensor:
|
||||
hit_tokens = torch.arange(num_hits, dtype=torch.int32, device=DEVICE)
|
||||
miss_tokens = hot_buffer_size + torch.arange(
|
||||
num_misses, dtype=torch.int32, device=DEVICE
|
||||
)
|
||||
return torch.cat([hit_tokens, miss_tokens])
|
||||
|
||||
|
||||
def _miss_tokens_per_req(miss_rate: float) -> int:
|
||||
return round(TOP_K * miss_rate)
|
||||
|
||||
|
||||
def _build_inputs(
|
||||
batch_size: int, hot_buffer_size: int, miss_rate: float
|
||||
) -> Dict[str, torch.Tensor | int]:
|
||||
dtype_bytes = torch.empty((), dtype=DTYPE).element_size()
|
||||
kv_dim = ITEM_SIZE_BYTES // dtype_bytes
|
||||
padded_buffer_size = hot_buffer_size + 1
|
||||
seq_len = hot_buffer_size + TOP_K + 1
|
||||
num_misses = _miss_tokens_per_req(miss_rate)
|
||||
num_hits = TOP_K - num_misses
|
||||
|
||||
top_k_row = _make_top_k_tokens(num_hits, num_misses, hot_buffer_size)
|
||||
top_k_tokens = top_k_row.view(1, -1).repeat(batch_size, 1).contiguous()
|
||||
|
||||
host_stride = seq_len
|
||||
total_host_tokens = batch_size * host_stride
|
||||
host_cache = torch.empty(
|
||||
(total_host_tokens, 1, kv_dim), dtype=DTYPE, device="cpu", pin_memory=True
|
||||
)
|
||||
host_cache.copy_(torch.randn_like(host_cache))
|
||||
|
||||
total_device_tokens = batch_size * padded_buffer_size
|
||||
device_buffer = torch.empty(
|
||||
(total_device_tokens, 1, kv_dim), dtype=DTYPE, device=DEVICE
|
||||
)
|
||||
device_buffer.normal_()
|
||||
|
||||
device_buffer_locs = torch.arange(
|
||||
total_device_tokens, dtype=torch.int32, device=DEVICE
|
||||
).view(batch_size, padded_buffer_size)
|
||||
device_buffer_tokens = torch.full(
|
||||
(batch_size, padded_buffer_size), -1, dtype=torch.int32, device=DEVICE
|
||||
)
|
||||
device_buffer_tokens[:, :hot_buffer_size] = torch.arange(
|
||||
hot_buffer_size, dtype=torch.int32, device=DEVICE
|
||||
)
|
||||
|
||||
lru_slots = (
|
||||
torch.arange(hot_buffer_size, dtype=torch.int16, device=DEVICE)
|
||||
.view(1, -1)
|
||||
.repeat(batch_size, 1)
|
||||
)
|
||||
|
||||
return {
|
||||
"top_k_tokens": top_k_tokens,
|
||||
"device_buffer_tokens": device_buffer_tokens,
|
||||
"initial_device_buffer_tokens": device_buffer_tokens.clone(),
|
||||
"host_cache_locs": torch.arange(
|
||||
total_host_tokens, dtype=torch.int64, device=DEVICE
|
||||
).view(batch_size, host_stride),
|
||||
"device_buffer_locs": device_buffer_locs,
|
||||
"host_cache": host_cache,
|
||||
"device_buffer": device_buffer,
|
||||
"top_k_device_locs": torch.empty(
|
||||
(batch_size, TOP_K), dtype=torch.int32, device=DEVICE
|
||||
),
|
||||
"req_pool_indices": torch.arange(batch_size, dtype=torch.int64, device=DEVICE),
|
||||
"seq_lens": torch.full(
|
||||
(batch_size,), seq_len, dtype=torch.int32, device=DEVICE
|
||||
),
|
||||
"lru_slots": lru_slots,
|
||||
"initial_lru_slots": lru_slots.clone(),
|
||||
"num_real_reqs": torch.tensor([batch_size], dtype=torch.int32, device=DEVICE),
|
||||
}
|
||||
|
||||
|
||||
def _time_kernel(batch_size: int, hot_buffer_size: int, miss_rate: float) -> float:
|
||||
state = _build_inputs(batch_size, hot_buffer_size, miss_rate)
|
||||
|
||||
def run_once():
|
||||
state["device_buffer_tokens"].copy_(state["initial_device_buffer_tokens"])
|
||||
state["lru_slots"].copy_(state["initial_lru_slots"])
|
||||
state["top_k_device_locs"].fill_(-1)
|
||||
load_cache_to_device_buffer_mla(
|
||||
top_k_tokens=state["top_k_tokens"],
|
||||
device_buffer_tokens=state["device_buffer_tokens"],
|
||||
host_cache_locs=state["host_cache_locs"],
|
||||
device_buffer_locs=state["device_buffer_locs"],
|
||||
host_cache=state["host_cache"],
|
||||
device_buffer=state["device_buffer"],
|
||||
top_k_device_locs=state["top_k_device_locs"],
|
||||
req_pool_indices=state["req_pool_indices"],
|
||||
seq_lens=state["seq_lens"],
|
||||
lru_slots=state["lru_slots"],
|
||||
item_size_bytes=ITEM_SIZE_BYTES,
|
||||
num_top_k=TOP_K,
|
||||
hot_buffer_size=hot_buffer_size,
|
||||
block_size=1024,
|
||||
num_real_reqs=state["num_real_reqs"],
|
||||
)
|
||||
|
||||
run_once()
|
||||
torch.cuda.synchronize()
|
||||
for _ in range(WARMUP_ROUNDS):
|
||||
run_once()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
for _ in range(ROUNDS):
|
||||
run_once()
|
||||
end.record()
|
||||
torch.cuda.synchronize()
|
||||
return start.elapsed_time(end) * 1000.0 / ROUNDS
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["batch_size", "hot_buffer_size", "miss_rate", "miss_tokens_cnt"],
|
||||
x_vals=CONFIGS,
|
||||
line_arg="provider",
|
||||
line_vals=LINE_VALS,
|
||||
line_names=LINE_NAMES,
|
||||
styles=STYLES,
|
||||
ylabel="us",
|
||||
plot_name="hisparse-latency",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark_latency(
|
||||
batch_size: int,
|
||||
hot_buffer_size: int,
|
||||
miss_rate: float,
|
||||
miss_tokens_cnt: int,
|
||||
provider: str,
|
||||
) -> Tuple[float, float, float]:
|
||||
assert provider == "jit"
|
||||
batch_size = int(batch_size)
|
||||
hot_buffer_size = int(hot_buffer_size)
|
||||
miss_rate = float(miss_rate)
|
||||
assert miss_tokens_cnt == batch_size * _miss_tokens_per_req(miss_rate)
|
||||
avg_us = _time_kernel(batch_size, hot_buffer_size, miss_rate)
|
||||
return avg_us, avg_us, avg_us
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark_latency.run(print_data=True)
|
||||
@@ -0,0 +1,214 @@
|
||||
"""Bench the hybrid ``mla_kv_pack_quantize_fp8`` against an inlined naive Triton baseline."""
|
||||
|
||||
import itertools
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import (
|
||||
DEFAULT_DEVICE,
|
||||
DEFAULT_DTYPE,
|
||||
DEFAULT_QUANTILES,
|
||||
get_benchmark_range,
|
||||
)
|
||||
from sglang.jit_kernel.mla_kv_pack_quantize_fp8 import (
|
||||
mla_kv_pack_quantize_fp8 as hybrid_pack,
|
||||
)
|
||||
from sglang.jit_kernel.utils import is_arch_support_pdl
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=15, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _triton_mla_kv_pack_quantize_fp8_kernel(
|
||||
k_nope_ptr,
|
||||
k_pe_ptr,
|
||||
v_ptr,
|
||||
k_out_ptr,
|
||||
v_out_ptr,
|
||||
k_scale_inv,
|
||||
v_scale_inv,
|
||||
s_total,
|
||||
k_nope_stride_t,
|
||||
k_nope_stride_h,
|
||||
k_pe_stride_t,
|
||||
v_stride_t,
|
||||
v_stride_h,
|
||||
k_out_stride_t,
|
||||
k_out_stride_h,
|
||||
v_out_stride_t,
|
||||
v_out_stride_h,
|
||||
QK_NOPE: tl.constexpr,
|
||||
QK_ROPE: tl.constexpr,
|
||||
V_HEAD: tl.constexpr,
|
||||
FP8_DTYPE: tl.constexpr,
|
||||
BLOCK_S: tl.constexpr,
|
||||
ENABLE_PDL: tl.constexpr,
|
||||
):
|
||||
pid_s = tl.program_id(0)
|
||||
pid_h = tl.program_id(1)
|
||||
t_idx = pid_s * BLOCK_S + tl.arange(0, BLOCK_S)
|
||||
t_mask = t_idx < s_total
|
||||
nope_idx = tl.arange(0, QK_NOPE)
|
||||
rope_idx = tl.arange(0, QK_ROPE)
|
||||
v_idx = tl.arange(0, V_HEAD)
|
||||
if ENABLE_PDL:
|
||||
tl.extra.cuda.gdc_wait()
|
||||
nope_off = (
|
||||
t_idx[:, None] * k_nope_stride_t + pid_h * k_nope_stride_h + nope_idx[None, :]
|
||||
)
|
||||
k_nope = tl.load(k_nope_ptr + nope_off, mask=t_mask[:, None])
|
||||
pe_off = t_idx[:, None] * k_pe_stride_t + rope_idx[None, :]
|
||||
k_pe = tl.load(k_pe_ptr + pe_off, mask=t_mask[:, None])
|
||||
v_off = t_idx[:, None] * v_stride_t + pid_h * v_stride_h + v_idx[None, :]
|
||||
v = tl.load(v_ptr + v_off, mask=t_mask[:, None])
|
||||
k_nope_fp8 = (k_nope.to(tl.float32) * k_scale_inv).to(FP8_DTYPE)
|
||||
k_pe_fp8 = (k_pe.to(tl.float32) * k_scale_inv).to(FP8_DTYPE)
|
||||
v_fp8 = (v.to(tl.float32) * v_scale_inv).to(FP8_DTYPE)
|
||||
k_out_base = t_idx[:, None] * k_out_stride_t + pid_h * k_out_stride_h
|
||||
tl.store(
|
||||
k_out_ptr + k_out_base + nope_idx[None, :], k_nope_fp8, mask=t_mask[:, None]
|
||||
)
|
||||
tl.store(
|
||||
k_out_ptr + k_out_base + QK_NOPE + rope_idx[None, :],
|
||||
k_pe_fp8,
|
||||
mask=t_mask[:, None],
|
||||
)
|
||||
v_out_off = (
|
||||
t_idx[:, None] * v_out_stride_t + pid_h * v_out_stride_h + v_idx[None, :]
|
||||
)
|
||||
tl.store(v_out_ptr + v_out_off, v_fp8, mask=t_mask[:, None])
|
||||
if ENABLE_PDL:
|
||||
tl.extra.cuda.gdc_launch_dependents()
|
||||
|
||||
|
||||
def _triton_pack(k_nope, k_pe, v, k_out, v_out):
|
||||
s, num_heads, qk_nope = k_nope.shape
|
||||
qk_rope = k_pe.shape[-1]
|
||||
v_head = v.shape[-1]
|
||||
k_pe_2d = k_pe.squeeze(1) if k_pe.dim() == 3 else k_pe
|
||||
enable_pdl = is_arch_support_pdl()
|
||||
if s < 512:
|
||||
block_s, num_warps, num_stages = 1, 1, 2
|
||||
elif s < 2048:
|
||||
block_s, num_warps, num_stages = 4, 2, 3
|
||||
else:
|
||||
block_s, num_warps, num_stages = 16, 4, 3
|
||||
extra = {"launch_pdl": True} if enable_pdl else {}
|
||||
grid = (triton.cdiv(s, block_s), num_heads)
|
||||
_triton_mla_kv_pack_quantize_fp8_kernel[grid](
|
||||
k_nope,
|
||||
k_pe_2d,
|
||||
v,
|
||||
k_out,
|
||||
v_out,
|
||||
1.0,
|
||||
1.0,
|
||||
s,
|
||||
k_nope.stride(0),
|
||||
k_nope.stride(1),
|
||||
k_pe_2d.stride(0),
|
||||
v.stride(0),
|
||||
v.stride(1),
|
||||
k_out.stride(0),
|
||||
k_out.stride(1),
|
||||
v_out.stride(0),
|
||||
v_out.stride(1),
|
||||
QK_NOPE=qk_nope,
|
||||
QK_ROPE=qk_rope,
|
||||
V_HEAD=v_head,
|
||||
FP8_DTYPE=tl.float8e4nv,
|
||||
BLOCK_S=block_s,
|
||||
ENABLE_PDL=enable_pdl,
|
||||
num_warps=num_warps,
|
||||
num_stages=num_stages,
|
||||
**extra,
|
||||
)
|
||||
|
||||
|
||||
QK_NOPE = 128
|
||||
QK_ROPE = 64
|
||||
V_HEAD = 128
|
||||
NUM_HEADS = 32
|
||||
NUM_LAYERS = 8
|
||||
|
||||
BS_RANGE = get_benchmark_range(
|
||||
full_range=[1, 4, 16, 64, 256, 1024, 4096, 8192, 16384],
|
||||
ci_range=[1, 64, 1024, 4096, 16384],
|
||||
)
|
||||
|
||||
LINE_VALS = ["hybrid", "triton"]
|
||||
LINE_NAMES = ["hybrid (v0+v1_flat)", "naive Triton"]
|
||||
STYLES = [("green", "-"), ("red", "--")]
|
||||
CONFIGS = list(itertools.product(BS_RANGE))
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["batch_size"],
|
||||
x_vals=CONFIGS,
|
||||
line_arg="provider",
|
||||
line_vals=LINE_VALS,
|
||||
line_names=LINE_NAMES,
|
||||
styles=STYLES,
|
||||
ylabel="us",
|
||||
plot_name="mla-kv-pack-quantize-fp8-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(batch_size: int, provider: str) -> Tuple[float, float, float]:
|
||||
k_nope = torch.randn(
|
||||
(NUM_LAYERS, batch_size, NUM_HEADS, QK_NOPE),
|
||||
dtype=DEFAULT_DTYPE,
|
||||
device=DEFAULT_DEVICE,
|
||||
)
|
||||
k_pe = torch.randn(
|
||||
(NUM_LAYERS, batch_size, 1, QK_ROPE),
|
||||
dtype=DEFAULT_DTYPE,
|
||||
device=DEFAULT_DEVICE,
|
||||
)
|
||||
v = torch.randn(
|
||||
(NUM_LAYERS, batch_size, NUM_HEADS, V_HEAD),
|
||||
dtype=DEFAULT_DTYPE,
|
||||
device=DEFAULT_DEVICE,
|
||||
)
|
||||
k_out = torch.empty(
|
||||
(NUM_LAYERS, batch_size, NUM_HEADS, QK_NOPE + QK_ROPE),
|
||||
dtype=torch.float8_e4m3fn,
|
||||
device=DEFAULT_DEVICE,
|
||||
)
|
||||
v_out = torch.empty(
|
||||
(NUM_LAYERS, batch_size, NUM_HEADS, V_HEAD),
|
||||
dtype=torch.float8_e4m3fn,
|
||||
device=DEFAULT_DEVICE,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
if provider == "hybrid":
|
||||
|
||||
def fn():
|
||||
for i in range(NUM_LAYERS):
|
||||
hybrid_pack(k_nope[i], k_pe[i], v[i], k_out=k_out[i], v_out=v_out[i])
|
||||
|
||||
else:
|
||||
|
||||
def fn():
|
||||
for i in range(NUM_LAYERS):
|
||||
_triton_pack(k_nope[i], k_pe[i], v[i], k_out[i], v_out[i])
|
||||
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
fn, quantiles=DEFAULT_QUANTILES
|
||||
)
|
||||
return (
|
||||
1000 * ms / NUM_LAYERS,
|
||||
1000 * max_ms / NUM_LAYERS,
|
||||
1000 * min_ms / NUM_LAYERS,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run(print_data=True)
|
||||
@@ -0,0 +1,290 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
|
||||
from sglang.jit_kernel.mxfp8 import (
|
||||
es_sm100_mxfp8_blockscaled_grouped_quant,
|
||||
es_sm100_mxfp8_blockscaled_moe_grouped_gemm,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=5, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
|
||||
def is_sm100_supported(device=None) -> bool:
|
||||
if not torch.cuda.is_available():
|
||||
return False
|
||||
return (torch.cuda.get_device_capability(device)[0] == 10) and (
|
||||
torch.version.cuda >= "12.8"
|
||||
)
|
||||
|
||||
|
||||
_SM100_SUPPORTED = is_sm100_supported()
|
||||
|
||||
|
||||
def _probe_sgl_kernel_group_mm() -> tuple[bool, str]:
|
||||
if not _SM100_SUPPORTED:
|
||||
return False, "MXFP8 MoE benchmark requires sm100+ with CUDA 12.8+."
|
||||
try:
|
||||
import sgl_kernel # noqa: F401
|
||||
except Exception as e:
|
||||
return False, f"import sgl_kernel failed: {e}"
|
||||
if not hasattr(sgl_kernel, "es_sm100_mxfp8_blockscaled_grouped_mm"):
|
||||
return False, "sgl_kernel.es_sm100_mxfp8_blockscaled_grouped_mm is missing."
|
||||
try:
|
||||
pass
|
||||
|
||||
# We assume if it's imported, it works
|
||||
except Exception as e:
|
||||
return False, f"calling sgl-kernel grouped_mm op failed: {e}"
|
||||
return True, ""
|
||||
|
||||
|
||||
_SGL_KERNEL_AVAILABLE, _SGL_KERNEL_REASON = _probe_sgl_kernel_group_mm()
|
||||
|
||||
|
||||
def align(val: int, alignment: int = 128) -> int:
|
||||
return int((val + alignment - 1) // alignment * alignment)
|
||||
|
||||
|
||||
def _prepare_case(
|
||||
total_tokens: int, n_g: int, k_g: int, num_experts: int, dtype: torch.dtype
|
||||
) -> dict[str, Any]:
|
||||
device = torch.device("cuda")
|
||||
base = total_tokens // num_experts
|
||||
rem = total_tokens % num_experts
|
||||
m_per_expert = [base + (1 if i < rem else 0) for i in range(num_experts)]
|
||||
|
||||
expert_offset = 0
|
||||
expert_offsets = []
|
||||
aux_expert_offset = 0
|
||||
aux_expert_offsets = []
|
||||
a_blockscale_offset = 0
|
||||
a_blockscale_offsets = []
|
||||
b_blockscale_offset = 0
|
||||
b_blockscale_offsets = []
|
||||
tokens_per_expert_list = []
|
||||
expert_ranges = []
|
||||
problem_sizes = []
|
||||
|
||||
a_list = []
|
||||
b_list = []
|
||||
for g in range(num_experts):
|
||||
m_g = m_per_expert[g]
|
||||
tokens_per_expert_list.append(m_g)
|
||||
expert_ranges.append((expert_offset, expert_offset + m_g))
|
||||
expert_offsets.append(expert_offset)
|
||||
expert_offset += m_g
|
||||
|
||||
aux_expert_offsets.append(aux_expert_offset)
|
||||
aux_expert_offset += n_g
|
||||
|
||||
a_blockscale_offsets.append(a_blockscale_offset)
|
||||
a_blockscale_offset += align(m_g, 128)
|
||||
|
||||
b_blockscale_offsets.append(b_blockscale_offset)
|
||||
b_blockscale_offset += n_g # n_g already align to 128 in practice
|
||||
|
||||
problem_sizes.append([m_g, n_g, k_g])
|
||||
|
||||
a = torch.randn((m_g, k_g), device=device, dtype=dtype) * 0.1
|
||||
b = torch.randn((n_g, k_g), device=device, dtype=dtype) * 0.1
|
||||
a_list.append(a)
|
||||
b_list.append(b)
|
||||
|
||||
a = torch.concat(a_list, dim=0)
|
||||
b = torch.concat(b_list, dim=0)
|
||||
|
||||
_expert_offsets = torch.tensor(expert_offsets).to(device=device, dtype=torch.int32)
|
||||
_aux_expert_offsets = torch.tensor(aux_expert_offsets).to(
|
||||
device=device, dtype=torch.int32
|
||||
)
|
||||
_a_blockscale_offsets = torch.tensor(a_blockscale_offsets).to(
|
||||
device=device, dtype=torch.int32
|
||||
)
|
||||
_b_blockscale_offsets = torch.tensor(b_blockscale_offsets).to(
|
||||
device=device, dtype=torch.int32
|
||||
)
|
||||
_tokens_per_expert = torch.tensor(tokens_per_expert_list).to(
|
||||
device=device, dtype=torch.int32
|
||||
)
|
||||
_problem_sizes = torch.tensor(problem_sizes).to(device=device, dtype=torch.int32)
|
||||
|
||||
a_quant = torch.zeros_like(a, dtype=torch.float8_e4m3fn, device=device)
|
||||
a_scale_factor = torch.zeros(
|
||||
(a_blockscale_offset, k_g // 32), dtype=torch.uint8, device=device
|
||||
)
|
||||
|
||||
b_quant = torch.zeros_like(b, dtype=torch.float8_e4m3fn, device=device)
|
||||
b_scale_factor = torch.zeros(
|
||||
(num_experts * n_g, k_g // 32), dtype=torch.uint8, device=device
|
||||
)
|
||||
|
||||
# Use a global workspace to avoid allocating 1GB every time
|
||||
workspace = torch.empty((1024, 1024, 1024), dtype=torch.uint8, device=device)
|
||||
|
||||
es_sm100_mxfp8_blockscaled_grouped_quant(
|
||||
a,
|
||||
_tokens_per_expert,
|
||||
_expert_offsets,
|
||||
_a_blockscale_offsets,
|
||||
a_quant,
|
||||
a_scale_factor,
|
||||
)
|
||||
|
||||
es_sm100_mxfp8_blockscaled_grouped_quant(
|
||||
b,
|
||||
torch.ones_like(_tokens_per_expert) * n_g,
|
||||
_aux_expert_offsets,
|
||||
_b_blockscale_offsets,
|
||||
b_quant,
|
||||
b_scale_factor,
|
||||
)
|
||||
|
||||
b_quant = b_quant.view(num_experts, n_g, k_g)
|
||||
b_scale_factor = b_scale_factor.view(num_experts, n_g, k_g // 32)
|
||||
|
||||
sgl_b_quant = b_quant.transpose(1, 2)
|
||||
sgl_b_scale_factor = b_scale_factor.transpose(1, 2)
|
||||
|
||||
return {
|
||||
"a": a,
|
||||
"b": b.view(num_experts, n_g, k_g),
|
||||
"b_quant": b_quant,
|
||||
"a_quant": a_quant,
|
||||
"b_scale_factor": b_scale_factor,
|
||||
"a_scale_factor": a_scale_factor,
|
||||
"expert_offsets": _expert_offsets,
|
||||
"a_blockscale_offsets": _a_blockscale_offsets,
|
||||
"tokens_per_expert": _tokens_per_expert,
|
||||
"problem_sizes": _problem_sizes,
|
||||
"sgl_b_quant": sgl_b_quant,
|
||||
"sgl_b_scale_factor": sgl_b_scale_factor,
|
||||
"workspace": workspace,
|
||||
"expert_ranges": expert_ranges,
|
||||
"dtype": dtype,
|
||||
}
|
||||
|
||||
|
||||
def _sgl_kernel_group_mm(case: dict[str, Any]) -> torch.Tensor:
|
||||
from sgl_kernel import es_sm100_mxfp8_blockscaled_grouped_mm
|
||||
|
||||
a_quant = case["a_quant"]
|
||||
sgl_b_quant = case["sgl_b_quant"]
|
||||
a_scale_factor = case["a_scale_factor"]
|
||||
sgl_b_scale_factor = case["sgl_b_scale_factor"]
|
||||
problem_sizes = case["problem_sizes"]
|
||||
expert_offsets = case["expert_offsets"]
|
||||
a_blockscale_offsets = case["a_blockscale_offsets"]
|
||||
dtype = case["dtype"]
|
||||
|
||||
total_tokens = a_quant.shape[0]
|
||||
n_g = sgl_b_quant.shape[2]
|
||||
|
||||
# sgl-kernel takes output pre-allocated
|
||||
d = torch.empty((total_tokens, n_g), device=a_quant.device, dtype=dtype)
|
||||
es_sm100_mxfp8_blockscaled_grouped_mm(
|
||||
d,
|
||||
a_quant,
|
||||
sgl_b_quant,
|
||||
a_scale_factor,
|
||||
sgl_b_scale_factor,
|
||||
problem_sizes,
|
||||
expert_offsets,
|
||||
a_blockscale_offsets,
|
||||
)
|
||||
return d
|
||||
|
||||
|
||||
shape_range = get_benchmark_range(
|
||||
full_range=[
|
||||
# (total_tokens, n_g, k_g, num_experts)
|
||||
(1024, 4096, 4096, 64),
|
||||
(2048, 4096, 4096, 64),
|
||||
(4096, 4096, 4096, 64),
|
||||
]
|
||||
+ [
|
||||
(total_tokens, n_g, k_g, num_experts)
|
||||
for total_tokens in [32 * (2**i) for i in range(9)] # 32 to 8192
|
||||
for n_g, k_g, num_experts in [
|
||||
# DeepSeek-V3/R1, gateup, TP = 1, EP = 8
|
||||
(4096, 7168, 32),
|
||||
# DeepSeek-V3/R1, down, TP = 1, EP = 8
|
||||
(7168, 2048, 32),
|
||||
]
|
||||
],
|
||||
ci_range=[(1024, 2048, 2048, 8)],
|
||||
)
|
||||
|
||||
line_vals = ["jit"]
|
||||
line_names = ["JIT MXFP8 MoE GroupMM"]
|
||||
styles = [("green", "-")]
|
||||
|
||||
if _SGL_KERNEL_AVAILABLE:
|
||||
line_vals.append("sgl_kernel")
|
||||
line_names.append("sgl-kernel MXFP8 MoE GroupMM")
|
||||
styles.append(("orange", "-"))
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["total_tokens", "n_g", "k_g", "num_experts"],
|
||||
x_vals=shape_range,
|
||||
x_log=False,
|
||||
line_arg="provider",
|
||||
line_vals=line_vals,
|
||||
line_names=line_names,
|
||||
styles=styles,
|
||||
ylabel="us",
|
||||
plot_name="mxfp8-moe-groupmm-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(total_tokens, n_g, k_g, num_experts, provider):
|
||||
case = _prepare_case(total_tokens, n_g, k_g, num_experts, torch.bfloat16)
|
||||
|
||||
if provider == "jit":
|
||||
fn = lambda: es_sm100_mxfp8_blockscaled_moe_grouped_gemm(
|
||||
case["b_quant"],
|
||||
case["a_quant"],
|
||||
case["b_scale_factor"],
|
||||
case["a_scale_factor"],
|
||||
case["expert_offsets"],
|
||||
case["a_blockscale_offsets"],
|
||||
case["tokens_per_expert"],
|
||||
case["workspace"],
|
||||
case["dtype"],
|
||||
)
|
||||
elif provider == "sgl_kernel":
|
||||
fn = lambda: _sgl_kernel_group_mm(case)
|
||||
else:
|
||||
raise ValueError(f"Unknown provider: {provider}")
|
||||
|
||||
# Warm up
|
||||
fn()
|
||||
|
||||
# Profile
|
||||
if provider == "jit":
|
||||
torch.cuda.nvtx.range_push("jit")
|
||||
fn()
|
||||
torch.cuda.nvtx.range_pop()
|
||||
elif provider == "sgl_kernel":
|
||||
torch.cuda.nvtx.range_push("sgl_kernel")
|
||||
fn()
|
||||
torch.cuda.nvtx.range_pop()
|
||||
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if not _SM100_SUPPORTED:
|
||||
print("[skip] MXFP8 MoE GroupMM benchmark requires sm100+ with CUDA 12.8+.")
|
||||
sys.exit(0)
|
||||
if not _SGL_KERNEL_AVAILABLE:
|
||||
print(f"[info] sgl-kernel baseline unavailable: {_SGL_KERNEL_REASON}")
|
||||
benchmark.run(print_data=True)
|
||||
@@ -0,0 +1,121 @@
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import (
|
||||
DEFAULT_DEVICE,
|
||||
get_benchmark_range,
|
||||
run_benchmark_no_cudagraph,
|
||||
)
|
||||
from sglang.jit_kernel.ngram_embedding import (
|
||||
compute_n_gram_ids,
|
||||
compute_n_gram_ids_decode,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=15, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
NE_N = 8
|
||||
NE_K = 2
|
||||
VOCAB_SIZE = 32000
|
||||
MAX_CONTEXT_LEN = 1024
|
||||
BATCH_SIZE_LIST = get_benchmark_range(
|
||||
full_range=[1, 2, 8, 32, 128, 512, 1024, 2048, 4096],
|
||||
ci_range=[32, 1024],
|
||||
)
|
||||
|
||||
|
||||
def _make_ngram_params():
|
||||
ne_weights = torch.zeros([NE_N - 1, NE_K, NE_N], dtype=torch.int32)
|
||||
ne_mods = torch.zeros([NE_N - 1, NE_K], dtype=torch.int32)
|
||||
exclusive_sums = torch.zeros([(NE_N - 1) * NE_K + 1], dtype=torch.int32)
|
||||
|
||||
for n in range(2, NE_N + 1):
|
||||
for k in range(NE_K):
|
||||
config_id = (n - 2) * NE_K + k
|
||||
mod = 65537 + 2 * config_id
|
||||
ne_mods[n - 2][k] = mod
|
||||
exclusive_sums[config_id + 1] = exclusive_sums[config_id] + mod
|
||||
for delta in range(NE_N):
|
||||
ne_weights[n - 2][k][delta] = pow(VOCAB_SIZE, delta, mod)
|
||||
|
||||
return (
|
||||
ne_weights.to(DEFAULT_DEVICE),
|
||||
ne_mods.to(DEFAULT_DEVICE),
|
||||
exclusive_sums.to(DEFAULT_DEVICE),
|
||||
)
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["batch_size"],
|
||||
x_vals=BATCH_SIZE_LIST,
|
||||
line_arg="provider",
|
||||
line_vals=["general", "decode"],
|
||||
line_names=["general compute_n_gram_ids", "decode fast path"],
|
||||
styles=[("blue", "-"), ("orange", "-")],
|
||||
ylabel="us",
|
||||
plot_name="ngram-compute-decode",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(batch_size: int, provider: str):
|
||||
num_configs = (NE_N - 1) * NE_K
|
||||
max_running_reqs = batch_size + 8
|
||||
ne_weights, ne_mods, exclusive_sums = _make_ngram_params()
|
||||
ne_token_table = torch.randint(
|
||||
0,
|
||||
VOCAB_SIZE,
|
||||
(max_running_reqs, MAX_CONTEXT_LEN),
|
||||
dtype=torch.int32,
|
||||
device=DEFAULT_DEVICE,
|
||||
)
|
||||
row_indices = torch.arange(batch_size, dtype=torch.int64, device=DEFAULT_DEVICE)
|
||||
column_starts = torch.randint(
|
||||
0, MAX_CONTEXT_LEN, (batch_size,), dtype=torch.int32, device=DEFAULT_DEVICE
|
||||
)
|
||||
n_gram_ids = torch.empty(
|
||||
(batch_size, num_configs), dtype=torch.int32, device=DEFAULT_DEVICE
|
||||
)
|
||||
|
||||
if provider == "general":
|
||||
tokens = torch.empty(batch_size, dtype=torch.int32, device=DEFAULT_DEVICE)
|
||||
exclusive_req_len_sums = torch.arange(
|
||||
batch_size + 1, dtype=torch.int32, device=DEFAULT_DEVICE
|
||||
)
|
||||
|
||||
def fn():
|
||||
compute_n_gram_ids(
|
||||
NE_N,
|
||||
NE_K,
|
||||
ne_weights,
|
||||
ne_mods,
|
||||
exclusive_sums,
|
||||
tokens,
|
||||
exclusive_req_len_sums,
|
||||
ne_token_table,
|
||||
row_indices,
|
||||
column_starts,
|
||||
n_gram_ids,
|
||||
)
|
||||
|
||||
else:
|
||||
|
||||
def fn():
|
||||
compute_n_gram_ids_decode(
|
||||
NE_N,
|
||||
NE_K,
|
||||
ne_weights,
|
||||
ne_mods,
|
||||
exclusive_sums,
|
||||
ne_token_table,
|
||||
row_indices,
|
||||
column_starts,
|
||||
n_gram_ids,
|
||||
)
|
||||
|
||||
return run_benchmark_no_cudagraph(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run(print_data=True)
|
||||
@@ -0,0 +1,76 @@
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import (
|
||||
DEFAULT_DEVICE,
|
||||
get_benchmark_range,
|
||||
run_benchmark_no_cudagraph,
|
||||
)
|
||||
from sglang.jit_kernel.ngram_embedding import (
|
||||
update_token_table,
|
||||
update_token_table_decode,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=15, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
MAX_CONTEXT_LEN = 4096
|
||||
BATCH_SIZE_LIST = get_benchmark_range(
|
||||
full_range=[1, 2, 8, 32, 128, 512, 1024, 2048, 4096],
|
||||
ci_range=[32, 1024],
|
||||
)
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["batch_size"],
|
||||
x_vals=BATCH_SIZE_LIST,
|
||||
line_arg="provider",
|
||||
line_vals=["general", "decode"],
|
||||
line_names=["general update_token_table", "decode fast path"],
|
||||
styles=[("blue", "-"), ("orange", "-")],
|
||||
ylabel="us",
|
||||
plot_name="ngram-update-token-table",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(batch_size: int, provider: str):
|
||||
max_running_reqs = batch_size + 8
|
||||
tokens = torch.arange(batch_size, dtype=torch.int32, device=DEFAULT_DEVICE)
|
||||
token_table = torch.empty(
|
||||
(max_running_reqs, MAX_CONTEXT_LEN), dtype=torch.int32, device=DEFAULT_DEVICE
|
||||
)
|
||||
row_indices = torch.arange(batch_size, dtype=torch.int64, device=DEFAULT_DEVICE)
|
||||
column_starts = torch.randint(
|
||||
0, MAX_CONTEXT_LEN, (batch_size,), dtype=torch.int32, device=DEFAULT_DEVICE
|
||||
)
|
||||
req_lens = torch.ones(batch_size, dtype=torch.int32, device=DEFAULT_DEVICE)
|
||||
|
||||
if provider == "general":
|
||||
|
||||
def fn():
|
||||
update_token_table(
|
||||
tokens,
|
||||
token_table,
|
||||
row_indices,
|
||||
column_starts,
|
||||
req_lens,
|
||||
None,
|
||||
)
|
||||
|
||||
else:
|
||||
|
||||
def fn():
|
||||
update_token_table_decode(
|
||||
tokens,
|
||||
token_table,
|
||||
row_indices,
|
||||
column_starts,
|
||||
)
|
||||
|
||||
return run_benchmark_no_cudagraph(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run(print_data=True)
|
||||
@@ -0,0 +1,100 @@
|
||||
import itertools
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
from flashinfer.norm import fused_add_rmsnorm as fi_fused_add_rmsnorm
|
||||
from flashinfer.norm import rmsnorm as fi_rmsnorm
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
|
||||
from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm
|
||||
from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
|
||||
DTYPE = torch.bfloat16
|
||||
DEVICE = "cuda"
|
||||
|
||||
BS_LIST = get_benchmark_range(
|
||||
full_range=[2**n for n in range(0, 14)],
|
||||
ci_range=[16, 32],
|
||||
)
|
||||
HIDDEN_SIZE_LIST = get_benchmark_range(
|
||||
full_range=sorted([1536, *range(1024, 8192 + 1, 1024)]),
|
||||
ci_range=[512, 2048],
|
||||
)
|
||||
|
||||
LINE_VALS = ["flashinfer", "jit"]
|
||||
LINE_NAMES = ["FlashInfer", "SGL JIT Kernel"]
|
||||
STYLES = [("blue", "--"), ("green", "-.")]
|
||||
NUM_LAYERS = 4 # avoid L2 effect
|
||||
|
||||
configs_0 = list(itertools.product(HIDDEN_SIZE_LIST + [16384], BS_LIST))
|
||||
configs_1 = list(itertools.product(HIDDEN_SIZE_LIST, BS_LIST))
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["hidden_size", "batch_size"],
|
||||
x_vals=configs_0,
|
||||
line_arg="provider",
|
||||
line_vals=LINE_VALS,
|
||||
line_names=LINE_NAMES,
|
||||
styles=STYLES,
|
||||
ylabel="us",
|
||||
plot_name="rmsnorm-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark_rmsnorm(hidden_size: int, batch_size: int, provider: str):
|
||||
input = torch.randn(
|
||||
(NUM_LAYERS, batch_size, hidden_size), dtype=DTYPE, device=DEVICE
|
||||
)
|
||||
weight = torch.randn((NUM_LAYERS, hidden_size), dtype=DTYPE, device=DEVICE)
|
||||
FN_MAP = {"jit": jit_rmsnorm, "flashinfer": fi_rmsnorm}
|
||||
|
||||
def f():
|
||||
fn = FN_MAP[provider]
|
||||
for i in range(NUM_LAYERS):
|
||||
fn(input[i], weight[i], out=input[i])
|
||||
|
||||
return run_benchmark(f, scale=NUM_LAYERS)
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["hidden_size", "batch_size"],
|
||||
x_vals=configs_1,
|
||||
line_arg="provider",
|
||||
line_vals=LINE_VALS,
|
||||
line_names=LINE_NAMES,
|
||||
styles=STYLES,
|
||||
ylabel="us",
|
||||
plot_name="fused-add-rmsnorm-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark_fused_add_rmsnorm(hidden_size: int, batch_size: int, provider: str):
|
||||
input = torch.randn(
|
||||
(NUM_LAYERS, batch_size, hidden_size), dtype=DTYPE, device=DEVICE
|
||||
)
|
||||
residual = torch.randn_like(input)
|
||||
weight = torch.randn((NUM_LAYERS, hidden_size), dtype=DTYPE, device=DEVICE)
|
||||
FN_MAP = {"jit": jit_fused_add_rmsnorm, "flashinfer": fi_fused_add_rmsnorm}
|
||||
|
||||
def f():
|
||||
fn = FN_MAP[provider]
|
||||
for i in range(NUM_LAYERS):
|
||||
fn(input[i], residual[i], weight[i])
|
||||
|
||||
return run_benchmark(f, scale=NUM_LAYERS)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("Benchmarking rmsnorm...")
|
||||
benchmark_rmsnorm.run(print_data=True)
|
||||
|
||||
print("Benchmarking fused_add_rmsnorm...")
|
||||
benchmark_fused_add_rmsnorm.run(print_data=True)
|
||||
@@ -0,0 +1,261 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
|
||||
from sglang.jit_kernel.nvfp4 import (
|
||||
cutlass_fp4_group_mm,
|
||||
scaled_fp4_experts_quant,
|
||||
scaled_fp4_quant,
|
||||
)
|
||||
from sglang.srt.utils import is_sm100_supported
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=5, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
FLOAT4_E2M1_MAX = 6.0
|
||||
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
|
||||
_NVFP4_SUPPORTED = is_sm100_supported()
|
||||
|
||||
|
||||
def _round_up(x: int, y: int) -> int:
|
||||
return ((x + y - 1) // y) * y
|
||||
|
||||
|
||||
def _expert_offsets(m_per_expert: list[int], device: torch.device) -> torch.Tensor:
|
||||
offsets = [0]
|
||||
for m in m_per_expert:
|
||||
offsets.append(offsets[-1] + m)
|
||||
return torch.tensor(offsets, dtype=torch.int32, device=device)
|
||||
|
||||
|
||||
def _blockscale_offsets(m_per_expert: list[int], device: torch.device) -> torch.Tensor:
|
||||
offsets = [0]
|
||||
for m in m_per_expert:
|
||||
offsets.append(offsets[-1] + _round_up(m, 128))
|
||||
return torch.tensor(offsets, dtype=torch.int32, device=device)
|
||||
|
||||
|
||||
def _prepare_case(
|
||||
total_tokens: int, n: int, k: int, num_experts: int, dtype: torch.dtype
|
||||
) -> dict[str, Any]:
|
||||
device = torch.device("cuda")
|
||||
base = total_tokens // num_experts
|
||||
rem = total_tokens % num_experts
|
||||
m_per_expert = [base + (1 if i < rem else 0) for i in range(num_experts)]
|
||||
|
||||
expert_offsets_full = _expert_offsets(m_per_expert, device)
|
||||
blockscale_offsets_full = _blockscale_offsets(m_per_expert, device)
|
||||
|
||||
a = torch.randn((total_tokens, k), device=device, dtype=dtype) * 0.1
|
||||
b = torch.randn((num_experts, n, k), device=device, dtype=dtype) * 0.1
|
||||
|
||||
a_global_scale = torch.empty((num_experts,), device=device, dtype=torch.float32)
|
||||
for i in range(num_experts):
|
||||
start = int(expert_offsets_full[i].item())
|
||||
end = int(expert_offsets_full[i + 1].item())
|
||||
a_global_scale[i] = (
|
||||
FLOAT8_E4M3_MAX
|
||||
* FLOAT4_E2M1_MAX
|
||||
/ a[start:end].abs().max().to(torch.float32)
|
||||
)
|
||||
|
||||
b_global_scale = torch.empty((num_experts,), device=device, dtype=torch.float32)
|
||||
for i in range(num_experts):
|
||||
b_global_scale[i] = (
|
||||
FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / b[i].abs().max().to(torch.float32)
|
||||
)
|
||||
|
||||
a_fp4, a_blockscale = scaled_fp4_experts_quant(
|
||||
a,
|
||||
a_global_scale,
|
||||
expert_offsets_full,
|
||||
blockscale_offsets_full,
|
||||
topk=1,
|
||||
)
|
||||
|
||||
b_fp4 = torch.empty((num_experts, n, k // 2), device=device, dtype=torch.uint8)
|
||||
b_blockscale = torch.empty(
|
||||
(num_experts, _round_up(n, 128), _round_up(k // 16, 4)),
|
||||
device=device,
|
||||
dtype=torch.float8_e4m3fn,
|
||||
)
|
||||
for i in range(num_experts):
|
||||
b_fp4_i, b_scale_i = scaled_fp4_quant(b[i], b_global_scale[i])
|
||||
b_fp4[i].copy_(b_fp4_i)
|
||||
b_blockscale[i].copy_(b_scale_i)
|
||||
|
||||
alphas = (1.0 / (a_global_scale * b_global_scale)).to(torch.float32)
|
||||
params = {
|
||||
"ab_strides": torch.full((num_experts,), k, dtype=torch.int64, device=device),
|
||||
"c_strides": torch.full((num_experts,), n, dtype=torch.int64, device=device),
|
||||
"problem_sizes": torch.tensor(
|
||||
[[m, n, k] for m in m_per_expert], dtype=torch.int32, device=device
|
||||
),
|
||||
"expert_offsets": expert_offsets_full[:-1].contiguous(),
|
||||
"blockscale_offsets": blockscale_offsets_full[:-1].contiguous(),
|
||||
"a_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
|
||||
"b_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
|
||||
"out_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
|
||||
"a_scales_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
|
||||
"b_scales_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
|
||||
"alpha_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
|
||||
"layout_sfa": torch.empty((num_experts, 5), dtype=torch.int64, device=device),
|
||||
"layout_sfb": torch.empty((num_experts, 5), dtype=torch.int64, device=device),
|
||||
}
|
||||
|
||||
expert_ranges: list[tuple[int, int]] = []
|
||||
start = 0
|
||||
for m in m_per_expert:
|
||||
end = start + m
|
||||
expert_ranges.append((start, end))
|
||||
start = end
|
||||
|
||||
return {
|
||||
"a": a,
|
||||
"b": b,
|
||||
"a_fp4": a_fp4,
|
||||
"b_fp4": b_fp4,
|
||||
"a_blockscale": a_blockscale,
|
||||
"b_blockscale": b_blockscale,
|
||||
"alphas": alphas,
|
||||
"params": params,
|
||||
"expert_offsets_full": expert_offsets_full,
|
||||
"expert_ranges": expert_ranges,
|
||||
"dtype": dtype,
|
||||
}
|
||||
|
||||
|
||||
def _torch_ref_group_mm(case: dict[str, Any]) -> torch.Tensor:
|
||||
a = case["a"]
|
||||
b = case["b"]
|
||||
dtype = case["dtype"]
|
||||
expert_ranges = case["expert_ranges"]
|
||||
total_tokens = a.shape[0]
|
||||
n = b.shape[1]
|
||||
out = torch.empty((total_tokens, n), device=a.device, dtype=dtype)
|
||||
for i, (start, end) in enumerate(expert_ranges):
|
||||
out[start:end] = torch.matmul(a[start:end], b[i].t())
|
||||
return out
|
||||
|
||||
|
||||
def _aot_cutlass_fp4_group_mm(case: dict[str, Any]) -> torch.Tensor:
|
||||
a_fp4 = case["a_fp4"]
|
||||
b_fp4 = case["b_fp4"]
|
||||
a_blockscale = case["a_blockscale"]
|
||||
b_blockscale = case["b_blockscale"]
|
||||
alphas = case["alphas"]
|
||||
params = case["params"]
|
||||
out_dtype = case["dtype"]
|
||||
|
||||
out = torch.empty(
|
||||
(a_fp4.shape[0], b_fp4.shape[1]), device=a_fp4.device, dtype=out_dtype
|
||||
)
|
||||
torch.ops.sgl_kernel.cutlass_fp4_group_mm.default(
|
||||
out,
|
||||
a_fp4,
|
||||
b_fp4,
|
||||
a_blockscale,
|
||||
b_blockscale,
|
||||
alphas,
|
||||
params["ab_strides"],
|
||||
params["c_strides"],
|
||||
params["problem_sizes"],
|
||||
params["expert_offsets"],
|
||||
params["blockscale_offsets"],
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def _probe_legacy_aot_group_mm() -> tuple[bool, str]:
|
||||
if not torch.cuda.is_available():
|
||||
return False, "CUDA is not available."
|
||||
if not _NVFP4_SUPPORTED:
|
||||
return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+."
|
||||
try:
|
||||
import sgl_kernel # noqa: F401
|
||||
except Exception as e:
|
||||
return False, f"import sgl_kernel failed: {e}"
|
||||
if not hasattr(torch.ops, "sgl_kernel"):
|
||||
return False, "torch.ops.sgl_kernel is not registered."
|
||||
op = getattr(torch.ops.sgl_kernel, "cutlass_fp4_group_mm", None)
|
||||
if op is None or not hasattr(op, "default"):
|
||||
return False, "torch.ops.sgl_kernel.cutlass_fp4_group_mm.default is missing."
|
||||
try:
|
||||
case = _prepare_case(64, 256, 128, 4, torch.bfloat16)
|
||||
_aot_cutlass_fp4_group_mm(case)
|
||||
torch.cuda.synchronize()
|
||||
except Exception as e:
|
||||
return False, f"calling AOT grouped_mm op failed: {e}"
|
||||
return True, ""
|
||||
|
||||
|
||||
_AOT_GROUP_MM_AVAILABLE, _AOT_GROUP_MM_REASON = _probe_legacy_aot_group_mm()
|
||||
|
||||
shape_range = get_benchmark_range(
|
||||
full_range=[(128, 256, 128, 4), (256, 512, 128, 8), (512, 512, 256, 8)],
|
||||
ci_range=[(128, 256, 128, 4)],
|
||||
)
|
||||
|
||||
line_vals = ["jit"]
|
||||
line_names = ["JIT NVFP4 MoE GroupMM"]
|
||||
styles = [("green", "-")]
|
||||
if _AOT_GROUP_MM_AVAILABLE:
|
||||
line_vals.append("aot_sgl_kernel")
|
||||
line_names.append("AOT NVFP4 MoE GroupMM")
|
||||
styles.append(("orange", "-"))
|
||||
line_vals.append("torch_ref")
|
||||
line_names.append("Torch Ref")
|
||||
styles.append(("blue", "-"))
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["total_tokens", "n", "k", "num_experts"],
|
||||
x_vals=shape_range,
|
||||
x_log=False,
|
||||
line_arg="provider",
|
||||
line_vals=line_vals,
|
||||
line_names=line_names,
|
||||
styles=styles,
|
||||
ylabel="us",
|
||||
plot_name="nvfp4-blockwise-moe-groupmm-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(total_tokens, n, k, num_experts, provider):
|
||||
case = _prepare_case(total_tokens, n, k, num_experts, torch.bfloat16)
|
||||
|
||||
if provider == "jit":
|
||||
fn = lambda: cutlass_fp4_group_mm(
|
||||
case["a_fp4"],
|
||||
case["b_fp4"],
|
||||
case["a_blockscale"],
|
||||
case["b_blockscale"],
|
||||
case["alphas"],
|
||||
case["dtype"],
|
||||
case["params"],
|
||||
)
|
||||
elif provider == "aot_sgl_kernel":
|
||||
fn = lambda: _aot_cutlass_fp4_group_mm(case)
|
||||
elif provider == "torch_ref":
|
||||
fn = lambda: _torch_ref_group_mm(case)
|
||||
else:
|
||||
raise ValueError(f"Unknown provider: {provider}")
|
||||
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if not _NVFP4_SUPPORTED:
|
||||
print("[skip] NVFP4 blockwise MoE benchmark requires sm100+ with CUDA 12.8+.")
|
||||
sys.exit(0)
|
||||
if not _AOT_GROUP_MM_AVAILABLE:
|
||||
print(
|
||||
f"[info] legacy AOT grouped_mm baseline unavailable: {_AOT_GROUP_MM_REASON}"
|
||||
)
|
||||
benchmark.run(print_data=True)
|
||||
@@ -0,0 +1,195 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
|
||||
from sglang.jit_kernel.nvfp4 import scaled_fp4_quant
|
||||
from sglang.srt.utils import is_sm100_supported
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=5, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
FLOAT4_E2M1_MAX = 6.0
|
||||
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
|
||||
BLOCK_SIZE = 16
|
||||
_NVFP4_SUPPORTED = is_sm100_supported()
|
||||
|
||||
try:
|
||||
from flashinfer import fp4_quantize as flashinfer_fp4_quantize
|
||||
except Exception:
|
||||
flashinfer_fp4_quantize = None
|
||||
|
||||
|
||||
def _torch_ref_quant(input: torch.Tensor, input_global_scale: torch.Tensor):
|
||||
m, n = input.shape
|
||||
x = input.view(m, n // BLOCK_SIZE, BLOCK_SIZE)
|
||||
vec_max = torch.max(torch.abs(x), dim=-1, keepdim=True)[0].to(torch.float32)
|
||||
scale = input_global_scale * (vec_max / FLOAT4_E2M1_MAX)
|
||||
scale = scale.to(torch.float8_e4m3fn).to(torch.float32)
|
||||
output_scale = torch.where(scale == 0, torch.zeros_like(scale), 1.0 / scale)
|
||||
|
||||
scaled_x = x.to(torch.float32) * output_scale
|
||||
clipped = torch.clamp(scaled_x, -6.0, 6.0).reshape(m, n)
|
||||
|
||||
rounded = clipped.clone()
|
||||
rounded[(rounded >= 0.0) & (rounded <= 0.25)] = 0.0
|
||||
rounded[(rounded > 0.25) & (rounded < 0.75)] = 0.5
|
||||
rounded[(rounded >= 0.75) & (rounded <= 1.25)] = 1.0
|
||||
rounded[(rounded > 1.25) & (rounded < 1.75)] = 1.5
|
||||
rounded[(rounded >= 1.75) & (rounded <= 2.5)] = 2.0
|
||||
rounded[(rounded > 2.5) & (rounded < 3.5)] = 3.0
|
||||
rounded[(rounded >= 3.5) & (rounded <= 5.0)] = 4.0
|
||||
rounded[rounded > 5.0] = 6.0
|
||||
|
||||
# This baseline intentionally keeps work on GPU but does not pack to uint8.
|
||||
return rounded, scale
|
||||
|
||||
|
||||
def _aot_scaled_fp4_quant(input: torch.Tensor, input_global_scale: torch.Tensor):
|
||||
m, n = input.shape
|
||||
output = torch.empty((m, n // 2), device=input.device, dtype=torch.uint8)
|
||||
rounded_m = ((m + 128 - 1) // 128) * 128
|
||||
scale_n = n // BLOCK_SIZE
|
||||
rounded_n = ((scale_n + 4 - 1) // 4) * 4
|
||||
output_scale = torch.empty(
|
||||
(rounded_m, rounded_n // 4), device=input.device, dtype=torch.int32
|
||||
)
|
||||
torch.ops.sgl_kernel.scaled_fp4_quant.default(
|
||||
output, input, output_scale, input_global_scale
|
||||
)
|
||||
return output, output_scale.view(torch.float8_e4m3fn)
|
||||
|
||||
|
||||
def _probe_legacy_aot_quant() -> tuple[bool, str]:
|
||||
if not torch.cuda.is_available():
|
||||
return False, "CUDA is not available."
|
||||
if not _NVFP4_SUPPORTED:
|
||||
return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+."
|
||||
try:
|
||||
import sgl_kernel # noqa: F401
|
||||
except Exception as e:
|
||||
return False, f"import sgl_kernel failed: {e}"
|
||||
if not hasattr(torch.ops, "sgl_kernel"):
|
||||
return False, "torch.ops.sgl_kernel is not registered."
|
||||
op = getattr(torch.ops.sgl_kernel, "scaled_fp4_quant", None)
|
||||
if op is None or not hasattr(op, "default"):
|
||||
return False, "torch.ops.sgl_kernel.scaled_fp4_quant.default is missing."
|
||||
try:
|
||||
x = torch.randn((16, 64), dtype=torch.bfloat16, device="cuda")
|
||||
global_scale = (
|
||||
FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.abs(x).max().to(torch.float32)
|
||||
)
|
||||
_aot_scaled_fp4_quant(x, global_scale)
|
||||
torch.cuda.synchronize()
|
||||
except Exception as e:
|
||||
return False, f"calling AOT quant op failed: {e}"
|
||||
return True, ""
|
||||
|
||||
|
||||
_AOT_QUANT_AVAILABLE, _AOT_QUANT_REASON = _probe_legacy_aot_quant()
|
||||
|
||||
|
||||
def _probe_flashinfer_quant() -> tuple[bool, str]:
|
||||
if flashinfer_fp4_quantize is None:
|
||||
return False, "import flashinfer.fp4_quantize failed."
|
||||
if not torch.cuda.is_available():
|
||||
return False, "CUDA is not available."
|
||||
if not _NVFP4_SUPPORTED:
|
||||
return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+."
|
||||
try:
|
||||
x = torch.randn((16, 64), dtype=torch.bfloat16, device="cuda")
|
||||
global_scale = (
|
||||
FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.abs(x).max().to(torch.float32)
|
||||
)
|
||||
flashinfer_fp4_quantize(
|
||||
x,
|
||||
global_scale,
|
||||
BLOCK_SIZE, # sf_vec_size
|
||||
False, # use_ue8m0
|
||||
True, # is_sf_swizzled_layout
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
except Exception as e:
|
||||
return False, f"calling flashinfer.fp4_quantize failed: {e}"
|
||||
return True, ""
|
||||
|
||||
|
||||
_FLASHINFER_QUANT_AVAILABLE, _FLASHINFER_QUANT_REASON = _probe_flashinfer_quant()
|
||||
|
||||
shape_range = get_benchmark_range(
|
||||
full_range=[(128, 2048), (512, 4096), (1024, 4096), (2048, 8192)],
|
||||
ci_range=[(128, 2048)],
|
||||
)
|
||||
|
||||
line_vals = []
|
||||
line_names = []
|
||||
styles = []
|
||||
if _FLASHINFER_QUANT_AVAILABLE:
|
||||
line_vals.append("flashinfer")
|
||||
line_names.append("FlashInfer FP4 Quant")
|
||||
styles.append(("purple", "-"))
|
||||
line_vals.append("jit")
|
||||
line_names.append("JIT NVFP4 Quant")
|
||||
styles.append(("green", "-"))
|
||||
if _AOT_QUANT_AVAILABLE:
|
||||
line_vals.append("aot_sgl_kernel")
|
||||
line_names.append("AOT NVFP4 Quant")
|
||||
styles.append(("orange", "-"))
|
||||
line_vals.append("torch_ref")
|
||||
line_names.append("Torch Ref")
|
||||
styles.append(("blue", "-"))
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["m", "n"],
|
||||
x_vals=shape_range,
|
||||
x_log=False,
|
||||
line_arg="provider",
|
||||
line_vals=line_vals,
|
||||
line_names=line_names,
|
||||
styles=styles,
|
||||
ylabel="us",
|
||||
plot_name="nvfp4-quant-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(m, n, provider):
|
||||
x = torch.randn((m, n), dtype=torch.bfloat16, device="cuda")
|
||||
tensor_amax = torch.abs(x).max().to(torch.float32)
|
||||
global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
|
||||
|
||||
if provider == "jit":
|
||||
fn = lambda: scaled_fp4_quant(x, global_scale)
|
||||
elif provider == "flashinfer":
|
||||
fn = lambda: flashinfer_fp4_quantize(
|
||||
x,
|
||||
global_scale,
|
||||
BLOCK_SIZE, # sf_vec_size
|
||||
False, # use_ue8m0
|
||||
True, # is_sf_swizzled_layout
|
||||
)
|
||||
elif provider == "aot_sgl_kernel":
|
||||
fn = lambda: _aot_scaled_fp4_quant(x, global_scale)
|
||||
elif provider == "torch_ref":
|
||||
fn = lambda: _torch_ref_quant(x, global_scale)
|
||||
else:
|
||||
raise ValueError(f"Unknown provider: {provider}")
|
||||
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if not _NVFP4_SUPPORTED:
|
||||
print("[skip] NVFP4 quant benchmark requires sm100+ with CUDA 12.8+.")
|
||||
sys.exit(0)
|
||||
if not _FLASHINFER_QUANT_AVAILABLE:
|
||||
print(
|
||||
f"[info] flashinfer quant baseline unavailable: {_FLASHINFER_QUANT_REASON}"
|
||||
)
|
||||
if not _AOT_QUANT_AVAILABLE:
|
||||
print(f"[info] legacy AOT quant baseline unavailable: {_AOT_QUANT_REASON}")
|
||||
benchmark.run(print_data=True)
|
||||
@@ -0,0 +1,187 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
|
||||
from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant
|
||||
from sglang.srt.utils import is_sm100_supported, is_sm120_supported
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=5, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
FLOAT4_E2M1_MAX = 6.0
|
||||
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
|
||||
BLOCK_SIZE = 16
|
||||
_NVFP4_SUPPORTED = is_sm100_supported() or is_sm120_supported()
|
||||
|
||||
K_E2M1_TO_FLOAT = [
|
||||
0.0,
|
||||
0.5,
|
||||
1.0,
|
||||
1.5,
|
||||
2.0,
|
||||
3.0,
|
||||
4.0,
|
||||
6.0,
|
||||
0.0,
|
||||
-0.5,
|
||||
-1.0,
|
||||
-1.5,
|
||||
-2.0,
|
||||
-3.0,
|
||||
-4.0,
|
||||
-6.0,
|
||||
]
|
||||
|
||||
|
||||
def _dequantize_to_fp16(
|
||||
tensor_fp4: torch.Tensor, tensor_sf: torch.Tensor, global_scale: torch.Tensor
|
||||
):
|
||||
m, packed_k = tensor_fp4.shape
|
||||
k = packed_k * 2
|
||||
flat = tensor_fp4.flatten()
|
||||
high = (flat & 0xF0) >> 4
|
||||
low = flat & 0x0F
|
||||
f_h = torch.tensor([K_E2M1_TO_FLOAT[x] for x in high], device=tensor_fp4.device)
|
||||
f_l = torch.tensor([K_E2M1_TO_FLOAT[x] for x in low], device=tensor_fp4.device)
|
||||
val = torch.stack((f_l, f_h), dim=-1).reshape(m, k)
|
||||
|
||||
rounded_m = ((m + 128 - 1) // 128) * 128
|
||||
scale_n = k // BLOCK_SIZE
|
||||
rounded_n = ((scale_n + 4 - 1) // 4) * 4
|
||||
sf = tensor_sf.view(torch.float8_e4m3fn)
|
||||
tmp = torch.reshape(sf, (1, rounded_m // 128, rounded_n // 4, 32, 4, 4))
|
||||
tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5))
|
||||
scale = torch.reshape(tmp, (rounded_m, rounded_n))[:m, :scale_n].to(torch.float32)
|
||||
scale = scale / global_scale
|
||||
|
||||
return (val.view(m, scale_n, BLOCK_SIZE) * scale.unsqueeze(-1)).reshape(m, k)
|
||||
|
||||
|
||||
def _aot_cutlass_scaled_fp4_mm(
|
||||
a: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
block_scale_a: torch.Tensor,
|
||||
block_scale_b: torch.Tensor,
|
||||
alpha: torch.Tensor,
|
||||
out_dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
out = torch.empty((a.shape[0], b.shape[0]), dtype=out_dtype, device=a.device)
|
||||
torch.ops.sgl_kernel.cutlass_scaled_fp4_mm.default(
|
||||
out, a, b, block_scale_a, block_scale_b, alpha
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def _probe_legacy_aot_scaled_mm() -> tuple[bool, str]:
|
||||
if not torch.cuda.is_available():
|
||||
return False, "CUDA is not available."
|
||||
if not _NVFP4_SUPPORTED:
|
||||
return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+."
|
||||
try:
|
||||
import sgl_kernel # noqa: F401
|
||||
except Exception as e:
|
||||
return False, f"import sgl_kernel failed: {e}"
|
||||
if not hasattr(torch.ops, "sgl_kernel"):
|
||||
return False, "torch.ops.sgl_kernel is not registered."
|
||||
op = getattr(torch.ops.sgl_kernel, "cutlass_scaled_fp4_mm", None)
|
||||
if op is None or not hasattr(op, "default"):
|
||||
return False, "torch.ops.sgl_kernel.cutlass_scaled_fp4_mm.default is missing."
|
||||
try:
|
||||
m, n, k = 16, 32, 64
|
||||
a = torch.randn((m, k), dtype=torch.bfloat16, device="cuda")
|
||||
b = torch.randn((n, k), dtype=torch.bfloat16, device="cuda")
|
||||
a_global_scale = (
|
||||
FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.amax(a.flatten(), dim=-1)
|
||||
).to(torch.float32)
|
||||
b_global_scale = (
|
||||
FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.amax(b.flatten(), dim=-1)
|
||||
).to(torch.float32)
|
||||
alpha = 1.0 / (a_global_scale * b_global_scale)
|
||||
a_fp4, a_sf = scaled_fp4_quant(a, a_global_scale)
|
||||
b_fp4, b_sf = scaled_fp4_quant(b, b_global_scale)
|
||||
_aot_cutlass_scaled_fp4_mm(a_fp4, b_fp4, a_sf, b_sf, alpha, torch.bfloat16)
|
||||
torch.cuda.synchronize()
|
||||
except Exception as e:
|
||||
return False, f"calling AOT scaled_mm op failed: {e}"
|
||||
return True, ""
|
||||
|
||||
|
||||
_AOT_SCALED_MM_AVAILABLE, _AOT_SCALED_MM_REASON = _probe_legacy_aot_scaled_mm()
|
||||
|
||||
shape_range = get_benchmark_range(
|
||||
full_range=[(128, 4096, 4096), (512, 4096, 4096), (1024, 8192, 4096)],
|
||||
ci_range=[(128, 4096, 4096)],
|
||||
)
|
||||
|
||||
line_vals = ["jit"]
|
||||
line_names = ["JIT NVFP4 GEMM"]
|
||||
styles = [("green", "-")]
|
||||
if _AOT_SCALED_MM_AVAILABLE:
|
||||
line_vals.append("aot_sgl_kernel")
|
||||
line_names.append("AOT NVFP4 GEMM")
|
||||
styles.append(("orange", "-"))
|
||||
line_vals.append("torch_ref")
|
||||
line_names.append("Torch Ref")
|
||||
styles.append(("blue", "-"))
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["m", "n", "k"],
|
||||
x_vals=shape_range,
|
||||
x_log=False,
|
||||
line_arg="provider",
|
||||
line_vals=line_vals,
|
||||
line_names=line_names,
|
||||
styles=styles,
|
||||
ylabel="us",
|
||||
plot_name="nvfp4-scaled-mm-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(m, n, k, provider):
|
||||
a = torch.randn((m, k), dtype=torch.bfloat16, device="cuda")
|
||||
b = torch.randn((n, k), dtype=torch.bfloat16, device="cuda")
|
||||
|
||||
a_global_scale = (
|
||||
FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.amax(a.flatten(), dim=-1)
|
||||
).to(torch.float32)
|
||||
b_global_scale = (
|
||||
FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.amax(b.flatten(), dim=-1)
|
||||
).to(torch.float32)
|
||||
alpha = 1.0 / (a_global_scale * b_global_scale)
|
||||
|
||||
a_fp4, a_sf = scaled_fp4_quant(a, a_global_scale)
|
||||
b_fp4, b_sf = scaled_fp4_quant(b, b_global_scale)
|
||||
|
||||
if provider == "jit":
|
||||
fn = lambda: cutlass_scaled_fp4_mm(
|
||||
a_fp4, b_fp4, a_sf, b_sf, alpha, torch.bfloat16
|
||||
)
|
||||
elif provider == "aot_sgl_kernel":
|
||||
fn = lambda: _aot_cutlass_scaled_fp4_mm(
|
||||
a_fp4, b_fp4, a_sf, b_sf, alpha, torch.bfloat16
|
||||
)
|
||||
elif provider == "torch_ref":
|
||||
a_ref = _dequantize_to_fp16(a_fp4, a_sf, a_global_scale)
|
||||
b_ref = _dequantize_to_fp16(b_fp4, b_sf, b_global_scale)
|
||||
fn = lambda: torch.matmul(a_ref, b_ref.t())
|
||||
else:
|
||||
raise ValueError(f"Unknown provider: {provider}")
|
||||
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if not _NVFP4_SUPPORTED:
|
||||
print("[skip] NVFP4 scaled_mm benchmark requires sm100/sm120 with CUDA 12.8+.")
|
||||
sys.exit(0)
|
||||
if not _AOT_SCALED_MM_AVAILABLE:
|
||||
print(
|
||||
f"[info] legacy AOT scaled_mm baseline unavailable: {_AOT_SCALED_MM_REASON}"
|
||||
)
|
||||
benchmark.run(print_data=True)
|
||||
@@ -0,0 +1,120 @@
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
|
||||
from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=5, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
try:
|
||||
from vllm import _custom_ops as ops
|
||||
|
||||
VLLM_AVAILABLE = True
|
||||
except ImportError:
|
||||
ops = None
|
||||
VLLM_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from sglang.srt.utils import is_hip
|
||||
|
||||
_is_hip = is_hip()
|
||||
except ImportError:
|
||||
_is_hip = False
|
||||
|
||||
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
|
||||
|
||||
|
||||
def vllm_scaled_fp8_quant(
|
||||
input: torch.Tensor,
|
||||
scale: Optional[torch.Tensor] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
if not VLLM_AVAILABLE:
|
||||
return sglang_scaled_fp8_quant(input, scale)
|
||||
return ops.scaled_fp8_quant(input, scale)
|
||||
|
||||
|
||||
def sglang_scaled_fp8_quant(
|
||||
input: torch.Tensor,
|
||||
scale: Optional[torch.Tensor] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
fp8_type_: torch.dtype = torch.float8_e4m3fn
|
||||
output = torch.empty_like(input, device=input.device, dtype=fp8_type_)
|
||||
is_static = True
|
||||
if scale is None:
|
||||
scale = torch.zeros(1, device=input.device, dtype=torch.float32)
|
||||
is_static = False
|
||||
per_tensor_quant_fp8(input, output, scale, is_static)
|
||||
|
||||
return output, scale
|
||||
|
||||
|
||||
def calculate_diff(batch_size: int, seq_len: int):
|
||||
device = torch.device("cuda")
|
||||
x = torch.rand((batch_size, seq_len), dtype=torch.bfloat16, device=device)
|
||||
|
||||
if not VLLM_AVAILABLE:
|
||||
print("vLLM not available, skipping comparison")
|
||||
return
|
||||
|
||||
vllm_out, vllm_scale = vllm_scaled_fp8_quant(x)
|
||||
sglang_out, sglang_scale = sglang_scaled_fp8_quant(x)
|
||||
|
||||
vllm_out = vllm_out.to(torch.float32)
|
||||
sglang_out = sglang_out.to(torch.float32)
|
||||
|
||||
triton.testing.assert_close(vllm_out, sglang_out, rtol=1e-3, atol=1e-3)
|
||||
triton.testing.assert_close(vllm_scale, sglang_scale, rtol=1e-3, atol=1e-3)
|
||||
|
||||
|
||||
# Benchmark configuration
|
||||
element_range = get_benchmark_range(
|
||||
full_range=[2**n for n in range(10, 20)],
|
||||
ci_range=[16384],
|
||||
)
|
||||
|
||||
if VLLM_AVAILABLE:
|
||||
line_vals = ["vllm", "sglang"]
|
||||
line_names = ["VLLM", "SGL Kernel"]
|
||||
styles = [("blue", "-"), ("green", "-")]
|
||||
else:
|
||||
line_vals = ["sglang"]
|
||||
line_names = ["SGL Kernel"]
|
||||
styles = [("green", "-")]
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["element_count"],
|
||||
x_vals=element_range,
|
||||
line_arg="provider",
|
||||
line_vals=line_vals,
|
||||
line_names=line_names,
|
||||
styles=styles,
|
||||
ylabel="us",
|
||||
plot_name="per-tensor-quant-fp8-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(element_count, provider):
|
||||
dtype = torch.float16
|
||||
device = torch.device("cuda")
|
||||
|
||||
x = torch.randn(element_count, 4096, device=device, dtype=dtype)
|
||||
|
||||
if provider == "vllm":
|
||||
fn = lambda: vllm_scaled_fp8_quant(x.clone())
|
||||
elif provider == "sglang":
|
||||
fn = lambda: sglang_scaled_fp8_quant(x.clone())
|
||||
else:
|
||||
raise ValueError(f"Unknown provider: {provider}")
|
||||
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
calculate_diff(batch_size=4, seq_len=4096)
|
||||
benchmark.run(print_data=True)
|
||||
@@ -0,0 +1,287 @@
|
||||
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.srt.layers.quantization.fp8_kernel import (
|
||||
create_per_token_group_quant_fp8_output_scale,
|
||||
)
|
||||
from sglang.srt.layers.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, suite="base-b-kernel-benchmark-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", "sglang"]
|
||||
LINE_NAMES = ["Triton (Inaccurate)", "SGL Kernel"]
|
||||
STYLES = [("blue", "-"), ("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,
|
||||
):
|
||||
"""
|
||||
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,
|
||||
)
|
||||
|
||||
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 == "sglang":
|
||||
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,
|
||||
)
|
||||
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,75 @@
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.benchmark import marker
|
||||
from sglang.jit_kernel.benchmark.utils import create_random
|
||||
from sglang.jit_kernel.norm import fused_inplace_qknorm
|
||||
from sglang.srt.utils import get_current_device_stream_fast
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=10, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
alt_stream = torch.cuda.Stream()
|
||||
|
||||
torch._dynamo.config.recompile_limit = 100
|
||||
|
||||
|
||||
# NOTE: now aot fallback to flashinfer
|
||||
def sglang_aot_qknorm(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
) -> None:
|
||||
from flashinfer import rmsnorm # lazy import to avoid crash
|
||||
|
||||
current_stream = get_current_device_stream_fast()
|
||||
alt_stream.wait_stream(current_stream)
|
||||
rmsnorm(q, q_weight, out=q)
|
||||
with torch.cuda.stream(alt_stream):
|
||||
rmsnorm(k, k_weight, out=k)
|
||||
current_stream.wait_stream(alt_stream)
|
||||
|
||||
|
||||
@torch.compile()
|
||||
def torch_impl_qknorm(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
eps: float = 1e-6,
|
||||
) -> None:
|
||||
q_mean = q.float().pow(2).mean(dim=-1, keepdim=True)
|
||||
k_mean = k.float().pow(2).mean(dim=-1, keepdim=True)
|
||||
q_norm = (q_mean + eps).rsqrt()
|
||||
k_norm = (k_mean + eps).rsqrt()
|
||||
q.copy_(q.float() * q_norm * q_weight.float())
|
||||
k.copy_(k.float() * k_norm * k_weight.float())
|
||||
|
||||
|
||||
FN_MAP = {
|
||||
"aot": sglang_aot_qknorm,
|
||||
"jit": fused_inplace_qknorm,
|
||||
"torch": torch_impl_qknorm,
|
||||
}
|
||||
|
||||
|
||||
@marker.parametrize("head_dim", [128, 256, 512, 1024], [128])
|
||||
@marker.parametrize("GQA", [4, 8], [4])
|
||||
@marker.parametrize("num_kv_heads", [1, 2, 4, 8], [1])
|
||||
@marker.parametrize("batch_size", [2**n for n in range(0, 14)], [16])
|
||||
@marker.benchmark("impl", ["aot", "jit", "torch"])
|
||||
def benchmark(head_dim: int, GQA: int, num_kv_heads: int, batch_size: int, impl: str):
|
||||
num_qo_heads = GQA * num_kv_heads
|
||||
q = create_random(batch_size, num_qo_heads, head_dim)
|
||||
k = create_random(batch_size, num_kv_heads, head_dim)
|
||||
q_weight = create_random(head_dim)
|
||||
k_weight = create_random(head_dim)
|
||||
return marker.do_bench(
|
||||
FN_MAP[impl],
|
||||
input_args=(q, k, q_weight, k_weight),
|
||||
memory_output=(q, k), # inplace write to q, k
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run()
|
||||
@@ -0,0 +1,123 @@
|
||||
import itertools
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
from sgl_kernel import rmsnorm
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import run_benchmark
|
||||
from sglang.jit_kernel.norm import fused_inplace_qknorm_across_heads
|
||||
from sglang.srt.utils import get_current_device_stream_fast
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.utils import is_in_ci
|
||||
|
||||
register_cuda_ci(est_time=12, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
IS_CI = is_in_ci()
|
||||
|
||||
alt_stream = torch.cuda.Stream()
|
||||
|
||||
|
||||
def sglang_jit_qknorm_across_heads(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
) -> None:
|
||||
|
||||
fused_inplace_qknorm_across_heads(q, k, q_weight, k_weight)
|
||||
|
||||
|
||||
def sglang_aot_qknorm_across_heads(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
) -> None:
|
||||
|
||||
current_stream = get_current_device_stream_fast()
|
||||
alt_stream.wait_stream(current_stream)
|
||||
rmsnorm(q, q_weight, out=q)
|
||||
with torch.cuda.stream(alt_stream):
|
||||
rmsnorm(k, k_weight, out=k)
|
||||
current_stream.wait_stream(alt_stream)
|
||||
|
||||
|
||||
def flashinfer_qknorm_across_heads(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
) -> None:
|
||||
from flashinfer import rmsnorm
|
||||
|
||||
rmsnorm(q, q_weight, out=q)
|
||||
rmsnorm(k, k_weight, out=k)
|
||||
|
||||
|
||||
@torch.compile()
|
||||
def torch_impl_qknorm_across_heads(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
eps: float = 1e-6,
|
||||
) -> None:
|
||||
q_mean = q.float().pow(2).mean(dim=-1, keepdim=True)
|
||||
k_mean = k.float().pow(2).mean(dim=-1, keepdim=True)
|
||||
q_norm = (q_mean + eps).rsqrt()
|
||||
k_norm = (k_mean + eps).rsqrt()
|
||||
q.copy_(q.float() * q_norm * q_weight.float())
|
||||
k.copy_(k.float() * k_norm * k_weight.float())
|
||||
|
||||
|
||||
DTYPE = torch.bfloat16
|
||||
DEVICE = "cuda"
|
||||
|
||||
if IS_CI:
|
||||
BS_RANGE = [16]
|
||||
HIDDEN_DIM_RANGE = [1024]
|
||||
else:
|
||||
BS_RANGE = [2**n for n in range(0, 14)]
|
||||
HIDDEN_DIM_RANGE = [512, 1024, 2048, 4096, 8192]
|
||||
|
||||
LINE_VALS = ["jit", "aot", "flashinfer", "torch"]
|
||||
LINE_NAMES = ["SGL JIT Kernel", "SGL AOT Kernel", "FlashInfer", "PyTorch"]
|
||||
STYLES = [("blue", "-"), ("orange", "--"), ("green", "-."), ("red", ":")]
|
||||
|
||||
configs = list(itertools.product(BS_RANGE, HIDDEN_DIM_RANGE))
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["batch_size", "hidden_dim"],
|
||||
x_vals=configs,
|
||||
line_arg="provider",
|
||||
line_vals=LINE_VALS,
|
||||
line_names=LINE_NAMES,
|
||||
styles=STYLES,
|
||||
ylabel="us",
|
||||
plot_name="qknorm-across-heads-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(
|
||||
batch_size: int, hidden_dim: int, provider: str
|
||||
) -> Tuple[float, float, float]:
|
||||
q = torch.randn((batch_size, hidden_dim), dtype=DTYPE, device=DEVICE)
|
||||
k = torch.randn((batch_size, hidden_dim), dtype=DTYPE, device=DEVICE)
|
||||
q_weight = torch.randn(hidden_dim, dtype=DTYPE, device=DEVICE)
|
||||
k_weight = torch.randn(hidden_dim, dtype=DTYPE, device=DEVICE)
|
||||
FN_MAP = {
|
||||
"jit": sglang_jit_qknorm_across_heads,
|
||||
"aot": sglang_aot_qknorm_across_heads,
|
||||
"flashinfer": flashinfer_qknorm_across_heads,
|
||||
"torch": torch_impl_qknorm_across_heads,
|
||||
}
|
||||
fn = lambda: FN_MAP[provider](q, k, q_weight, k_weight)
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run(print_data=True)
|
||||
@@ -0,0 +1,239 @@
|
||||
import itertools
|
||||
|
||||
import sgl_kernel
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.utils import is_in_ci
|
||||
|
||||
register_cuda_ci(est_time=5, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
|
||||
def torch_top_k_renorm_probs(probs, top_k):
|
||||
"""Vectorized PyTorch implementation of top-k renormalization."""
|
||||
batch_size, vocab_size = probs.shape
|
||||
|
||||
# Handle scalar or tensor k
|
||||
if isinstance(top_k, int):
|
||||
k_val = min(max(top_k, 1), vocab_size)
|
||||
# Get top-k indices for all batches at once
|
||||
_, topk_indices = torch.topk(probs, k_val, dim=1, largest=True)
|
||||
|
||||
# Create mask: batch_size x vocab_size
|
||||
mask = torch.zeros_like(probs)
|
||||
mask.scatter_(1, topk_indices, 1.0)
|
||||
|
||||
# Vectorized renormalization
|
||||
masked_probs = probs * mask
|
||||
renorm_probs = masked_probs / (masked_probs.sum(dim=1, keepdim=True) + 1e-10)
|
||||
return renorm_probs
|
||||
else:
|
||||
# Variable k per batch - need to handle separately
|
||||
renorm_probs = torch.zeros_like(probs)
|
||||
for i in range(batch_size):
|
||||
k_val = min(max(top_k[i].item(), 1), vocab_size)
|
||||
_, topk_indices = torch.topk(probs[i], k_val, largest=True)
|
||||
mask = torch.zeros_like(probs[i])
|
||||
mask[topk_indices] = 1.0
|
||||
masked_probs = probs[i] * mask
|
||||
renorm_probs[i] = masked_probs / (masked_probs.sum() + 1e-10)
|
||||
return renorm_probs
|
||||
|
||||
|
||||
def torch_top_p_renorm_probs(probs, top_p, eps=1e-5):
|
||||
"""Vectorized PyTorch implementation of top-p renormalization."""
|
||||
batch_size, vocab_size = probs.shape
|
||||
|
||||
# Handle scalar or tensor p
|
||||
if isinstance(top_p, float):
|
||||
p_val = top_p
|
||||
# Vectorized implementation for uniform top_p
|
||||
# Sort probs in descending order
|
||||
sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=1)
|
||||
cumsum_probs = torch.cumsum(sorted_probs, dim=1)
|
||||
|
||||
# Find cutoff: where cumsum exceeds top_p
|
||||
cutoff_mask = cumsum_probs <= p_val
|
||||
# Keep at least one token (the highest prob)
|
||||
cutoff_mask[:, 0] = True
|
||||
|
||||
# Create mask in original order
|
||||
mask = torch.zeros_like(probs)
|
||||
mask.scatter_(1, sorted_indices, cutoff_mask.float())
|
||||
|
||||
# Vectorized renormalization
|
||||
masked_probs = probs * mask
|
||||
renorm_probs = masked_probs / (masked_probs.sum(dim=1, keepdim=True) + eps)
|
||||
return renorm_probs
|
||||
else:
|
||||
# Variable p per batch - need to handle separately
|
||||
renorm_probs = torch.zeros_like(probs)
|
||||
for i in range(batch_size):
|
||||
p_val = top_p[i].item()
|
||||
sorted_prob, indices = torch.sort(probs[i], descending=False)
|
||||
cdf = torch.cumsum(sorted_prob, dim=-1)
|
||||
mask = torch.zeros(vocab_size, dtype=torch.float32, device=probs.device)
|
||||
mask.scatter_(0, indices, (cdf >= (1 - p_val) - eps).float())
|
||||
masked_probs = probs[i] * mask
|
||||
renorm_probs[i] = masked_probs / (masked_probs.sum() + eps)
|
||||
return renorm_probs
|
||||
|
||||
|
||||
def calculate_diff_top_k_renorm(batch_size, vocab_size, k):
|
||||
"""Compare Torch reference and SGLang kernel for top-k renorm correctness."""
|
||||
torch.manual_seed(42)
|
||||
device = torch.device("cuda")
|
||||
|
||||
pre_norm_prob = torch.rand(batch_size, vocab_size, device=device)
|
||||
probs = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
|
||||
|
||||
top_k_tensor = torch.full((batch_size,), k, device=device, dtype=torch.int32)
|
||||
|
||||
torch_output = torch_top_k_renorm_probs(probs, top_k_tensor)
|
||||
sglang_output = sgl_kernel.top_k_renorm_prob(probs, top_k_tensor)
|
||||
|
||||
torch.testing.assert_close(torch_output, sglang_output, rtol=1e-3, atol=1e-3)
|
||||
|
||||
|
||||
def calculate_diff_top_p_renorm(batch_size, vocab_size, p):
|
||||
"""Compare Torch reference and SGLang kernel for top-p renorm correctness."""
|
||||
torch.manual_seed(42)
|
||||
device = torch.device("cuda")
|
||||
|
||||
pre_norm_prob = torch.rand(batch_size, vocab_size, device=device)
|
||||
probs = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
|
||||
|
||||
top_p_tensor = torch.full((batch_size,), p, device=device, dtype=torch.float32)
|
||||
|
||||
torch_output = torch_top_p_renorm_probs(probs, top_p_tensor)
|
||||
sglang_output = sgl_kernel.top_p_renorm_prob(probs, top_p_tensor)
|
||||
|
||||
torch.testing.assert_close(torch_output, sglang_output, rtol=1e-3, atol=1e-3)
|
||||
|
||||
|
||||
# Parameter space - simplified for CI
|
||||
if is_in_ci():
|
||||
batch_size_range = [16]
|
||||
vocab_size_range = [111]
|
||||
k_range = [10]
|
||||
p_range = [0.5]
|
||||
else:
|
||||
batch_size_range = [16, 64, 128]
|
||||
vocab_size_range = [111, 32000, 128256]
|
||||
k_range = [10, 100, 500]
|
||||
p_range = [0.1, 0.5, 0.9]
|
||||
|
||||
configs_k = list(itertools.product(batch_size_range, vocab_size_range, k_range))
|
||||
configs_p = list(itertools.product(batch_size_range, vocab_size_range, p_range))
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["batch_size", "vocab_size", "k"],
|
||||
x_vals=configs_k,
|
||||
line_arg="provider",
|
||||
line_vals=["torch", "sglang"],
|
||||
line_names=["Torch Reference", "SGL Kernel"],
|
||||
styles=[("red", "-"), ("green", "-")],
|
||||
ylabel="us",
|
||||
plot_name="top-k-renorm-probs-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark_top_k_renorm(batch_size, vocab_size, k, provider):
|
||||
# Skip invalid configurations
|
||||
if k >= vocab_size:
|
||||
return float("nan"), float("nan"), float("nan")
|
||||
|
||||
torch.manual_seed(42)
|
||||
device = torch.device("cuda")
|
||||
|
||||
pre_norm_prob = torch.rand(batch_size, vocab_size, device=device)
|
||||
probs = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
|
||||
top_k_tensor = torch.full((batch_size,), k, device=device, dtype=torch.int32)
|
||||
|
||||
if provider == "torch":
|
||||
fn = lambda: torch_top_k_renorm_probs(probs.clone(), top_k_tensor)
|
||||
elif provider == "sglang":
|
||||
fn = lambda: sgl_kernel.top_k_renorm_prob(probs.clone(), top_k_tensor)
|
||||
|
||||
return run_benchmark_no_cudagraph(fn)
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["batch_size", "vocab_size", "p"],
|
||||
x_vals=configs_p,
|
||||
line_arg="provider",
|
||||
line_vals=["torch", "sglang"],
|
||||
line_names=["Torch Reference", "SGL Kernel"],
|
||||
styles=[("red", "-"), ("blue", "-")],
|
||||
ylabel="us",
|
||||
plot_name="top-p-renorm-probs-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark_top_p_renorm(batch_size, vocab_size, p, provider):
|
||||
torch.manual_seed(42)
|
||||
device = torch.device("cuda")
|
||||
|
||||
pre_norm_prob = torch.rand(batch_size, vocab_size, device=device)
|
||||
probs = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
|
||||
top_p_tensor = torch.full((batch_size,), p, device=device, dtype=torch.float32)
|
||||
|
||||
if provider == "torch":
|
||||
fn = lambda: torch_top_p_renorm_probs(probs.clone(), top_p_tensor)
|
||||
elif provider == "sglang":
|
||||
fn = lambda: sgl_kernel.top_p_renorm_prob(probs.clone(), top_p_tensor)
|
||||
|
||||
return run_benchmark_no_cudagraph(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("=" * 60)
|
||||
print("Running correctness checks...")
|
||||
print("=" * 60)
|
||||
|
||||
# Correctness checks - simplified for CI
|
||||
if is_in_ci():
|
||||
test_configs_k = [configs_k[0]] if configs_k else [(16, 111, 10)]
|
||||
test_configs_p = [configs_p[0]] if configs_p else [(16, 111, 0.5)]
|
||||
else:
|
||||
test_configs_k = configs_k[:3] # Test first 3 configs
|
||||
test_configs_p = configs_p[:3]
|
||||
|
||||
print("\n1. Testing top_k_renorm_probs...")
|
||||
for cfg in test_configs_k:
|
||||
batch_size, vocab_size, k = cfg
|
||||
if k < vocab_size: # Skip invalid configs
|
||||
calculate_diff_top_k_renorm(batch_size, vocab_size, k)
|
||||
print(
|
||||
f" ✓ Passed: batch_size={batch_size}, vocab_size={vocab_size}, k={k}"
|
||||
)
|
||||
|
||||
print("\n2. Testing top_p_renorm_probs...")
|
||||
for cfg in test_configs_p:
|
||||
calculate_diff_top_p_renorm(*cfg)
|
||||
batch_size, vocab_size, p = cfg
|
||||
print(f" ✓ Passed: batch_size={batch_size}, vocab_size={vocab_size}, p={p}")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("All correctness checks passed!")
|
||||
print("=" * 60)
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Starting performance benchmarks...")
|
||||
print("=" * 60)
|
||||
|
||||
print("\n1. Benchmarking top_k_renorm_probs...")
|
||||
benchmark_top_k_renorm.run(print_data=True)
|
||||
|
||||
print("\n2. Benchmarking top_p_renorm_probs...")
|
||||
benchmark_top_p_renorm.run(print_data=True)
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Benchmarking complete!")
|
||||
print("=" * 60)
|
||||
@@ -0,0 +1,73 @@
|
||||
import itertools
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import (
|
||||
DEFAULT_DEVICE,
|
||||
get_benchmark_range,
|
||||
run_benchmark,
|
||||
)
|
||||
from sglang.jit_kernel.resolve_future_token_ids import resolve_future_token_ids_cuda
|
||||
from sglang.srt.utils import get_compiler_backend
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=10, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
register_amd_ci(est_time=10, suite="jit-kernel-unit-test-amd")
|
||||
|
||||
SIZE_LIST = get_benchmark_range(
|
||||
full_range=[2**n for n in range(4, 16)], # 16 … 32K elements
|
||||
ci_range=[256, 4096],
|
||||
)
|
||||
|
||||
configs = list(itertools.product(SIZE_LIST))
|
||||
|
||||
|
||||
def _torch_resolve(input_ids, future_map):
|
||||
input_ids[:] = torch.where(
|
||||
input_ids < 0,
|
||||
future_map[torch.clamp(-input_ids, min=0)],
|
||||
input_ids,
|
||||
)
|
||||
|
||||
|
||||
_compiled_resolve = torch.compile(
|
||||
_torch_resolve, dynamic=True, backend=get_compiler_backend()
|
||||
)
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["size"],
|
||||
x_vals=configs,
|
||||
line_arg="provider",
|
||||
line_vals=["jit", "torch_compile", "torch"],
|
||||
line_names=["SGL JIT Kernel", "torch.compile", "PyTorch"],
|
||||
styles=[("blue", "-"), ("green", "-."), ("red", "--")],
|
||||
ylabel="us",
|
||||
plot_name="resolve-future-token-ids-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(size: int, provider: str):
|
||||
map_size = 8192
|
||||
future_map = torch.randint(
|
||||
0, 50000, (map_size,), dtype=torch.int64, device=DEFAULT_DEVICE
|
||||
)
|
||||
input_ids = torch.randint(
|
||||
-map_size + 1, 50000, (size,), dtype=torch.int64, device=DEFAULT_DEVICE
|
||||
)
|
||||
|
||||
if provider == "jit":
|
||||
fn = lambda: resolve_future_token_ids_cuda(input_ids.clone(), future_map)
|
||||
elif provider == "torch_compile":
|
||||
fn = lambda: _compiled_resolve(input_ids.clone(), future_map)
|
||||
else:
|
||||
fn = lambda: _torch_resolve(input_ids.clone(), future_map)
|
||||
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run(print_data=True)
|
||||
@@ -0,0 +1,308 @@
|
||||
import itertools
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import (
|
||||
DEFAULT_DEVICE,
|
||||
DEFAULT_DTYPE,
|
||||
get_benchmark_range,
|
||||
run_benchmark,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=6, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
MAX_SEQ_LEN = 131072
|
||||
ROPE_BASE = 10000.0
|
||||
ROPE_DIM = 128
|
||||
CACHE_SIZE = 1024 * 1024
|
||||
|
||||
|
||||
def create_cos_sin_cache(
|
||||
rotary_dim: int = ROPE_DIM,
|
||||
max_position: int = MAX_SEQ_LEN,
|
||||
base: float = ROPE_BASE,
|
||||
) -> torch.Tensor:
|
||||
inv_freq = 1.0 / (
|
||||
base
|
||||
** (
|
||||
torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=DEFAULT_DEVICE)
|
||||
/ rotary_dim
|
||||
)
|
||||
)
|
||||
t = torch.arange(max_position, dtype=torch.float32, device=DEFAULT_DEVICE)
|
||||
freqs = torch.einsum("i,j->ij", t, inv_freq)
|
||||
cos = freqs.cos()
|
||||
sin = freqs.sin()
|
||||
return torch.cat((cos, sin), dim=-1)
|
||||
|
||||
|
||||
# Pre-build the cache once
|
||||
COS_SIN_CACHE = create_cos_sin_cache()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# RoPE-only provider implementations
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def flashinfer_rope(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
is_neox: bool,
|
||||
) -> None:
|
||||
from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace
|
||||
|
||||
head_size = q.shape[-1]
|
||||
apply_rope_with_cos_sin_cache_inplace(
|
||||
positions=positions,
|
||||
query=q.view(q.shape[0], -1),
|
||||
key=k.view(k.shape[0], -1),
|
||||
head_size=head_size,
|
||||
cos_sin_cache=COS_SIN_CACHE,
|
||||
is_neox=is_neox,
|
||||
)
|
||||
|
||||
|
||||
def sglang_pos_enc_rope(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
is_neox: bool,
|
||||
) -> None:
|
||||
from sglang.jit_kernel.rope import rotary_embedding_with_key
|
||||
|
||||
head_size = q.shape[-1]
|
||||
rotary_embedding_with_key(
|
||||
positions=positions,
|
||||
query=q.view(q.shape[0], -1),
|
||||
key=k.view(k.shape[0], -1),
|
||||
head_size=head_size,
|
||||
cos_sin_cache=COS_SIN_CACHE,
|
||||
is_neox=is_neox,
|
||||
)
|
||||
|
||||
|
||||
def sglang_fused_rope(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
is_neox: bool,
|
||||
) -> None:
|
||||
from sglang.jit_kernel.rope import apply_rope_inplace
|
||||
|
||||
apply_rope_inplace(q, k, COS_SIN_CACHE, positions, is_neox=is_neox)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# RoPE + KV cache store provider implementations
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def jit_rope_then_store(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
k_cache: torch.Tensor,
|
||||
v_cache: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
out_loc: torch.Tensor,
|
||||
is_neox: bool,
|
||||
) -> None:
|
||||
from sglang.jit_kernel.kvcache import store_cache
|
||||
from sglang.jit_kernel.rope import apply_rope_inplace
|
||||
|
||||
head_size = q.shape[-1]
|
||||
row_dim = k.shape[-2] * head_size
|
||||
apply_rope_inplace(
|
||||
positions=positions,
|
||||
q=q,
|
||||
k=k,
|
||||
rope_dim=head_size,
|
||||
cos_sin_cache=COS_SIN_CACHE,
|
||||
is_neox=is_neox,
|
||||
)
|
||||
store_cache(
|
||||
k.view(-1, row_dim),
|
||||
v.view(-1, row_dim),
|
||||
k_cache,
|
||||
v_cache,
|
||||
out_loc,
|
||||
)
|
||||
|
||||
|
||||
def jit_fused_rope_store(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
k_cache: torch.Tensor,
|
||||
v_cache: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
out_loc: torch.Tensor,
|
||||
is_neox: bool,
|
||||
) -> None:
|
||||
from sglang.jit_kernel.rope import apply_rope_inplace_with_kvcache
|
||||
|
||||
apply_rope_inplace_with_kvcache(
|
||||
q, k, v, k_cache, v_cache, COS_SIN_CACHE, positions, out_loc, is_neox=is_neox
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark configuration (shared)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
BS_RANGE = get_benchmark_range(
|
||||
full_range=[2**n for n in range(0, 16)],
|
||||
ci_range=[16],
|
||||
)
|
||||
QK_HEAD_RANGE = get_benchmark_range(
|
||||
full_range=[(8, 1), (16, 2), (32, 8)],
|
||||
ci_range=[(16, 2)],
|
||||
)
|
||||
QK_HEAD_RANGE = [f"{q},{k}" for q, k in QK_HEAD_RANGE]
|
||||
IS_NEOX_RANGE = get_benchmark_range(
|
||||
full_range=[True, False],
|
||||
ci_range=[True],
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark 1: RoPE only
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
ROPE_LINE_VALS = ["flashinfer", "jit_pos_enc", "jit_fused_rope"]
|
||||
ROPE_LINE_NAMES = [
|
||||
"FlashInfer",
|
||||
"SGL JIT PosEnc",
|
||||
"SGL JIT Fused RoPE",
|
||||
]
|
||||
ROPE_STYLES = [("green", "-."), ("red", "-"), ("blue", "--")]
|
||||
|
||||
rope_configs = list(itertools.product(QK_HEAD_RANGE, IS_NEOX_RANGE, BS_RANGE))
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["num_q_k_heads", "is_neox", "batch_size"],
|
||||
x_vals=rope_configs,
|
||||
line_arg="provider",
|
||||
line_vals=ROPE_LINE_VALS,
|
||||
line_names=ROPE_LINE_NAMES,
|
||||
styles=ROPE_STYLES,
|
||||
ylabel="us",
|
||||
plot_name="rope-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(batch_size: int, num_q_k_heads: str, is_neox: bool, provider: str):
|
||||
qo, kv = num_q_k_heads.split(",")
|
||||
num_qo_heads = int(qo)
|
||||
num_kv_heads = int(kv)
|
||||
q = torch.randn(
|
||||
(batch_size, num_qo_heads, ROPE_DIM),
|
||||
dtype=DEFAULT_DTYPE,
|
||||
device=DEFAULT_DEVICE,
|
||||
)
|
||||
k = torch.randn(
|
||||
(batch_size, num_kv_heads, ROPE_DIM),
|
||||
dtype=DEFAULT_DTYPE,
|
||||
device=DEFAULT_DEVICE,
|
||||
)
|
||||
seed = batch_size << 16 | num_qo_heads << 8 | num_kv_heads << 4 | is_neox
|
||||
torch.random.manual_seed(seed)
|
||||
positions = torch.randint(
|
||||
MAX_SEQ_LEN, (batch_size,), device=DEFAULT_DEVICE, dtype=torch.int64
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
FN_MAP = {
|
||||
"flashinfer": flashinfer_rope,
|
||||
"jit_pos_enc": sglang_pos_enc_rope,
|
||||
"jit_fused_rope": sglang_fused_rope,
|
||||
}
|
||||
fn = lambda: FN_MAP[provider](q, k, positions, is_neox)
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark 2: RoPE + KV cache store
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
STORE_LINE_VALS = ["jit_rope_then_store", "jit_fused_store"]
|
||||
STORE_LINE_NAMES = [
|
||||
"SGL JIT RoPE + Store",
|
||||
"SGL JIT Fused RoPE + Store",
|
||||
]
|
||||
STORE_STYLES = [("red", "-"), ("blue", "--")]
|
||||
|
||||
store_configs = list(itertools.product(QK_HEAD_RANGE, IS_NEOX_RANGE, BS_RANGE))
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["num_q_k_heads", "is_neox", "batch_size"],
|
||||
x_vals=store_configs,
|
||||
line_arg="provider",
|
||||
line_vals=STORE_LINE_VALS,
|
||||
line_names=STORE_LINE_NAMES,
|
||||
styles=STORE_STYLES,
|
||||
ylabel="us",
|
||||
plot_name="rope-store-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark_store(batch_size: int, num_q_k_heads: str, is_neox: bool, provider: str):
|
||||
qo, kv = num_q_k_heads.split(",")
|
||||
num_qo_heads = int(qo)
|
||||
num_kv_heads = int(kv)
|
||||
q = torch.randn(
|
||||
(batch_size, num_qo_heads, ROPE_DIM),
|
||||
dtype=DEFAULT_DTYPE,
|
||||
device=DEFAULT_DEVICE,
|
||||
)
|
||||
k = torch.randn(
|
||||
(batch_size, num_kv_heads, ROPE_DIM),
|
||||
dtype=DEFAULT_DTYPE,
|
||||
device=DEFAULT_DEVICE,
|
||||
)
|
||||
v = torch.randn(
|
||||
(batch_size, num_kv_heads, ROPE_DIM),
|
||||
dtype=DEFAULT_DTYPE,
|
||||
device=DEFAULT_DEVICE,
|
||||
)
|
||||
row_size = num_kv_heads * ROPE_DIM
|
||||
k_cache = torch.zeros(
|
||||
CACHE_SIZE, row_size, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE
|
||||
)
|
||||
v_cache = torch.zeros(
|
||||
CACHE_SIZE, row_size, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE
|
||||
)
|
||||
out_loc = torch.randperm(CACHE_SIZE, device=DEFAULT_DEVICE, dtype=torch.int64)[
|
||||
:batch_size
|
||||
]
|
||||
seed = batch_size << 16 | num_qo_heads << 8 | num_kv_heads << 4 | is_neox
|
||||
torch.random.manual_seed(seed)
|
||||
positions = torch.randint(
|
||||
MAX_SEQ_LEN, (batch_size,), device=DEFAULT_DEVICE, dtype=torch.int64
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
FN_MAP = {
|
||||
"jit_rope_then_store": jit_rope_then_store,
|
||||
"jit_fused_store": jit_fused_rope_store,
|
||||
}
|
||||
fn = lambda: FN_MAP[provider](
|
||||
q, k, v, k_cache, v_cache, positions, out_loc, is_neox
|
||||
)
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("Running RoPE performance benchmark...")
|
||||
benchmark.run(print_data=True)
|
||||
print("\nRunning RoPE + KV cache store performance benchmark...")
|
||||
benchmark_store.run(print_data=True)
|
||||
@@ -0,0 +1,127 @@
|
||||
"""Benchmark the set_mla_kv_buffer dispatcher.
|
||||
|
||||
Compares three providers across a batch-size sweep:
|
||||
- ``wrapper``: the high-level wrapper exposed by ``set_mla_kv_buffer_triton``
|
||||
(dispatches to TMA on SM90+, Triton fallback otherwise).
|
||||
- ``jit_tma``: the JIT CUDA TMA bulk-store kernel directly.
|
||||
- ``triton``: the BLOCK-tiled Triton kernel (SM<90 fallback path).
|
||||
"""
|
||||
|
||||
import itertools
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import (
|
||||
DEFAULT_DEVICE,
|
||||
DEFAULT_DTYPE,
|
||||
DEFAULT_QUANTILES,
|
||||
get_benchmark_range,
|
||||
)
|
||||
from sglang.jit_kernel.set_mla_kv_buffer import set_mla_kv_buffer as jit_set
|
||||
from sglang.jit_kernel.utils import is_arch_support_pdl
|
||||
from sglang.srt.mem_cache.utils import set_mla_kv_buffer_kernel as sglang_triton_kernel
|
||||
from sglang.srt.mem_cache.utils import set_mla_kv_buffer_triton as sglang_wrapper
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=9, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
|
||||
def _triton_baseline(kv_buffer, loc, cache_k_nope, cache_k_rope):
|
||||
nope_dim = cache_k_nope.shape[-1]
|
||||
rope_dim = cache_k_rope.shape[-1]
|
||||
total_dim = nope_dim + rope_dim
|
||||
BLOCK = 128
|
||||
n_loc = loc.numel()
|
||||
grid = (n_loc, triton.cdiv(total_dim, BLOCK))
|
||||
pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
|
||||
sglang_triton_kernel[grid](
|
||||
kv_buffer,
|
||||
cache_k_nope,
|
||||
cache_k_rope,
|
||||
loc,
|
||||
kv_buffer.stride(0),
|
||||
cache_k_nope.stride(0),
|
||||
cache_k_rope.stride(0),
|
||||
nope_dim,
|
||||
rope_dim,
|
||||
BLOCK=BLOCK,
|
||||
**pdl_kwargs,
|
||||
)
|
||||
|
||||
|
||||
NUM_LAYERS = 8
|
||||
CACHE_SIZE = (2 * 1024 * 1024) // NUM_LAYERS
|
||||
|
||||
NOPE_DIM = 512
|
||||
ROPE_DIM = 64
|
||||
|
||||
BS_RANGE = get_benchmark_range(
|
||||
full_range=[1, 8, 32, 128, 512, 1024, 2048, 4096, 8192, 16384],
|
||||
ci_range=[1, 128, 2048, 4096, 8192],
|
||||
)
|
||||
|
||||
LINE_VALS = ["wrapper", "jit_tma", "triton"]
|
||||
LINE_NAMES = ["Wrapper (auto)", "JIT TMA bulk-store", "Triton (BLOCK=128 baseline)"]
|
||||
STYLES = [("blue", "-"), ("green", "--"), ("red", "-.")]
|
||||
X_NAMES = ["batch_size"]
|
||||
CONFIGS = list(itertools.product(BS_RANGE))
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=X_NAMES,
|
||||
x_vals=CONFIGS,
|
||||
line_arg="provider",
|
||||
line_vals=LINE_VALS,
|
||||
line_names=LINE_NAMES,
|
||||
styles=STYLES,
|
||||
ylabel="us",
|
||||
plot_name="set-mla-kv-buffer-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(batch_size: int, provider: str) -> Tuple[float, float, float]:
|
||||
cache_k_nope = torch.randn(
|
||||
(NUM_LAYERS, batch_size, 1, NOPE_DIM),
|
||||
dtype=DEFAULT_DTYPE,
|
||||
device=DEFAULT_DEVICE,
|
||||
)
|
||||
cache_k_rope = torch.randn(
|
||||
(NUM_LAYERS, batch_size, 1, ROPE_DIM),
|
||||
dtype=DEFAULT_DTYPE,
|
||||
device=DEFAULT_DEVICE,
|
||||
)
|
||||
kv_buffer = torch.randn(
|
||||
(NUM_LAYERS, CACHE_SIZE, 1, NOPE_DIM + ROPE_DIM),
|
||||
dtype=DEFAULT_DTYPE,
|
||||
device=DEFAULT_DEVICE,
|
||||
)
|
||||
loc = torch.randperm(CACHE_SIZE, device=DEFAULT_DEVICE)[:batch_size]
|
||||
torch.cuda.synchronize()
|
||||
|
||||
FN_MAP = {
|
||||
"wrapper": sglang_wrapper,
|
||||
"jit_tma": lambda buf, loc, n, r: jit_set(buf, loc, n, r),
|
||||
"triton": _triton_baseline,
|
||||
}
|
||||
|
||||
def fn():
|
||||
impl = FN_MAP[provider]
|
||||
for i in range(NUM_LAYERS):
|
||||
impl(kv_buffer[i], loc, cache_k_nope[i], cache_k_rope[i])
|
||||
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
fn, quantiles=DEFAULT_QUANTILES
|
||||
)
|
||||
return (
|
||||
1000 * ms / NUM_LAYERS,
|
||||
1000 * max_ms / NUM_LAYERS,
|
||||
1000 * min_ms / NUM_LAYERS,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run(print_data=True)
|
||||
@@ -0,0 +1,73 @@
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.benchmark import marker
|
||||
from sglang.jit_kernel.benchmark.utils import (
|
||||
DEFAULT_DEVICE,
|
||||
create_empty,
|
||||
create_random,
|
||||
)
|
||||
from sglang.jit_kernel.kvcache import store_cache
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=9, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
|
||||
@torch.compile()
|
||||
def torch_compile_store_cache(
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
k_cache: torch.Tensor,
|
||||
v_cache: torch.Tensor,
|
||||
indices: torch.Tensor,
|
||||
) -> None:
|
||||
k_cache[indices] = k
|
||||
v_cache[indices] = v
|
||||
|
||||
|
||||
alt_stream = torch.cuda.Stream()
|
||||
|
||||
|
||||
def torch_streams_store_cache(
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
k_cache: torch.Tensor,
|
||||
v_cache: torch.Tensor,
|
||||
indices: torch.Tensor,
|
||||
) -> None:
|
||||
current_stream = torch.cuda.current_stream()
|
||||
alt_stream.wait_stream(current_stream)
|
||||
k_cache[indices] = k
|
||||
with torch.cuda.stream(alt_stream):
|
||||
v_cache[indices] = v
|
||||
current_stream.wait_stream(alt_stream)
|
||||
|
||||
|
||||
CACHE_SIZE = 2 * 1024 * 1024
|
||||
FN_MAP = {
|
||||
"jit": store_cache,
|
||||
"torch_compile": torch_compile_store_cache,
|
||||
"torch_streams": torch_streams_store_cache,
|
||||
}
|
||||
|
||||
|
||||
@marker.parametrize("item_size", [64, 128, 256, 512, 1024], [1024])
|
||||
@marker.parametrize("batch_size", [2**n for n in range(0, 15)], [16])
|
||||
@marker.benchmark("impl", ["jit", "torch_compile", "torch_streams"])
|
||||
def benchmark(batch_size: int, item_size: int, impl: str):
|
||||
torch.manual_seed(42)
|
||||
k = create_random(batch_size, item_size)
|
||||
k_cache = create_empty(CACHE_SIZE, item_size)
|
||||
v = create_random(batch_size, item_size)
|
||||
v_cache = create_empty(CACHE_SIZE, item_size)
|
||||
indices = torch.randperm(CACHE_SIZE, device=DEFAULT_DEVICE)[:batch_size]
|
||||
return marker.do_bench(
|
||||
FN_MAP[impl],
|
||||
input_args=(k, v, k_cache, v_cache, indices),
|
||||
graph_clone_args=(0, 1, 4), # not need to clone cache, which is large
|
||||
memory_args=(k, v, indices), # k_cache / v_cache excluded
|
||||
memory_output=(k, v), # inplace write, size = k + v
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run()
|
||||
@@ -0,0 +1,170 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
import sglang.srt.distributed.parallel_state as ps
|
||||
from sglang.jit_kernel.all_reduce import (
|
||||
fused_parallel_qknorm,
|
||||
get_fused_parallel_qknorm_max_occupancy,
|
||||
)
|
||||
from sglang.jit_kernel.utils import get_ci_test_range
|
||||
from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
|
||||
CustomAllReduceV2,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=120,
|
||||
suite="base-b-kernel-benchmark-1-gpu-large",
|
||||
disabled="requires multi-GPU, self-skips in CI",
|
||||
)
|
||||
|
||||
Q_K_DIMS = [(6144, 1024)]
|
||||
DTYPE = torch.bfloat16
|
||||
EPS = 1e-6
|
||||
BATCH_SIZES = get_ci_test_range([2**i for i in range(15)], [1, 64, 1024])
|
||||
NUM_LAYERS = 8
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--warmup", type=int, default=10)
|
||||
parser.add_argument("--iters", type=int, default=100)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def init_distributed():
|
||||
local_rank = int(os.environ["LOCAL_RANK"])
|
||||
world_size = int(os.environ["WORLD_SIZE"])
|
||||
rank = local_rank
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.cuda.set_device(device)
|
||||
|
||||
dist.init_process_group(backend="gloo")
|
||||
ps._WORLD = coord = ps.init_world_group(
|
||||
ranks=list(range(world_size)),
|
||||
local_rank=local_rank,
|
||||
backend="nccl",
|
||||
)
|
||||
|
||||
cpu_group = coord.cpu_group
|
||||
max_occupancy = get_fused_parallel_qknorm_max_occupancy(
|
||||
DTYPE, world_size, Q_K_DIMS[0][0], Q_K_DIMS[0][1]
|
||||
)
|
||||
if rank == 0:
|
||||
print(f"Max occupancy for fused_parallel_qknorm: {max_occupancy} blocks/SM")
|
||||
|
||||
props = torch.cuda.get_device_properties(device)
|
||||
comm = CustomAllReduceV2(
|
||||
cpu_group,
|
||||
device,
|
||||
max_pull_size=0,
|
||||
max_push_size=8 * max(BATCH_SIZES),
|
||||
max_push_blocks=props.multi_processor_count * max_occupancy,
|
||||
)
|
||||
comm_ = CustomAllReduceV2(cpu_group, device)
|
||||
if comm.disabled or comm_.disabled:
|
||||
raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system")
|
||||
return rank, world_size, device, cpu_group, comm, comm_
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def bench_one(fn, warmup: int, iters: int) -> float:
|
||||
for _ in range(warmup):
|
||||
fn(0)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
for i in range(NUM_LAYERS):
|
||||
fn(i)
|
||||
|
||||
graph.replay()
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
graph.replay()
|
||||
start.record()
|
||||
for i in range(iters):
|
||||
graph.replay()
|
||||
end.record()
|
||||
torch.cuda.synchronize()
|
||||
return start.elapsed_time(end) * 1000.0 / (iters * NUM_LAYERS)
|
||||
|
||||
|
||||
def rmsnorm_baseline(
|
||||
comm_,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
world_size: int,
|
||||
) -> None:
|
||||
from sglang.srt.models.minimax_m2 import rms_apply_serial, rms_sumsq_serial
|
||||
|
||||
sum_sq = rms_sumsq_serial(q, k)
|
||||
sum_sq = comm_.custom_all_reduce(sum_sq)
|
||||
rms_apply_serial(q, k, q_weight, k_weight, sum_sq, world_size, EPS)
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
rank, world_size, device, _, comm, comm_ = init_distributed()
|
||||
torch.cuda.set_stream(torch.cuda.Stream())
|
||||
|
||||
if rank == 0:
|
||||
print(
|
||||
f"{'q_dim':>8} {'k_dim':>8} {'batch':>8} {'fused_us':>12} {'baseline_us':>12}"
|
||||
)
|
||||
|
||||
for q_dim, k_dim in Q_K_DIMS:
|
||||
local_q_dim = q_dim // world_size
|
||||
local_k_dim = k_dim // world_size
|
||||
for batch_size in BATCH_SIZES:
|
||||
q = torch.randn(
|
||||
NUM_LAYERS, batch_size, local_q_dim, device=device, dtype=DTYPE
|
||||
)
|
||||
k = torch.randn(
|
||||
NUM_LAYERS, batch_size, local_k_dim, device=device, dtype=DTYPE
|
||||
)
|
||||
q_weight = torch.randn(NUM_LAYERS, local_q_dim, device=device, dtype=DTYPE)
|
||||
k_weight = torch.randn(NUM_LAYERS, local_k_dim, device=device, dtype=DTYPE)
|
||||
|
||||
def run_fused(i: int):
|
||||
fused_parallel_qknorm(
|
||||
comm.obj,
|
||||
q[i],
|
||||
k[i],
|
||||
q_weight[i],
|
||||
k_weight[i],
|
||||
EPS,
|
||||
)
|
||||
|
||||
def run_baseline(i: int):
|
||||
rmsnorm_baseline(
|
||||
comm_,
|
||||
q[i],
|
||||
k[i],
|
||||
q_weight[i],
|
||||
k_weight[i],
|
||||
world_size,
|
||||
)
|
||||
|
||||
fused_us = bench_one(run_fused, args.warmup, args.iters)
|
||||
baseline_us = bench_one(run_baseline, args.warmup, args.iters)
|
||||
|
||||
if rank == 0:
|
||||
print(
|
||||
f"{q_dim:8d} {k_dim:8d} {batch_size:8d} "
|
||||
f"{fused_us:12.1f} {baseline_us:12.1f}"
|
||||
)
|
||||
|
||||
comm.close()
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,366 @@
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import statistics
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable
|
||||
|
||||
import flashinfer
|
||||
import sgl_kernel
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import DEFAULT_DTYPE
|
||||
from sglang.jit_kernel.utils import KERNEL_PATH
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.utils import is_in_ci
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=120,
|
||||
suite="base-b-kernel-benchmark-1-gpu-large",
|
||||
disabled="standalone diffusion NVFP4 benchmark",
|
||||
)
|
||||
|
||||
SCRIPT_DIR = Path(__file__).resolve().parent
|
||||
REPO_ROOT = (
|
||||
Path(os.environ["SGLANG_NVFP4_REPO_ROOT"])
|
||||
if os.environ.get("SGLANG_NVFP4_REPO_ROOT")
|
||||
# Anchor on the installed jit_kernel package (python/sglang/jit_kernel) so
|
||||
# this stays correct regardless of where the benchmark file lives.
|
||||
else KERNEL_PATH.parents[2]
|
||||
)
|
||||
DEFAULT_OUTPUT_DIR = REPO_ROOT / "outputs" / "nvfp4_benchmarks"
|
||||
DEFAULT_SHAPE_LIBRARY = SCRIPT_DIR / "diffusion_nvfp4_shapes.json"
|
||||
DTYPE = DEFAULT_DTYPE
|
||||
WARMUP = 8
|
||||
ITERS = 20
|
||||
FLOAT4_E2M1_MAX = 6.0
|
||||
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
|
||||
METHODS = ("cutlass", "flashinfer_auto", "flashinfer_cudnn")
|
||||
|
||||
|
||||
def benchmark_provider(
|
||||
fn: Callable[[], torch.Tensor],
|
||||
warmup: int = WARMUP,
|
||||
iters: int = ITERS,
|
||||
) -> tuple[float, float, float]:
|
||||
for _ in range(warmup):
|
||||
y = fn()
|
||||
del y
|
||||
torch.cuda.synchronize()
|
||||
|
||||
times_ms: list[float] = []
|
||||
for _ in range(iters):
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
y = fn()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
times_ms.append(start.elapsed_time(end))
|
||||
del y
|
||||
return statistics.median(times_ms), max(times_ms), min(times_ms)
|
||||
|
||||
|
||||
def make_global_scale(x: torch.Tensor) -> torch.Tensor:
|
||||
max_abs = torch.amax(x.abs()).clamp_min_(1e-6)
|
||||
return (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / max_abs).to(torch.float32)
|
||||
|
||||
|
||||
def build_quantized_inputs(
|
||||
m: int,
|
||||
n: int,
|
||||
k: int,
|
||||
device: torch.device,
|
||||
seed: int,
|
||||
) -> dict[str, Any]:
|
||||
assert k % 16 == 0, f"NVFP4 requires k % 16 == 0, got k={k}"
|
||||
|
||||
gen = torch.Generator(device=device)
|
||||
gen.manual_seed(seed)
|
||||
x = torch.randn((m, k), device=device, dtype=DTYPE, generator=gen)
|
||||
w = torch.randn((n, k), device=device, dtype=DTYPE, generator=gen)
|
||||
|
||||
x_global_scale = make_global_scale(x)
|
||||
w_global_scale = make_global_scale(w)
|
||||
alpha = (1.0 / (x_global_scale * w_global_scale)).to(torch.float32)
|
||||
|
||||
x_fp4, x_sf = flashinfer.fp4_quantize(x, x_global_scale)
|
||||
w_fp4, w_sf = flashinfer.fp4_quantize(w, w_global_scale)
|
||||
if x_sf.dtype == torch.uint8:
|
||||
x_sf = x_sf.view(torch.float8_e4m3fn)
|
||||
if w_sf.dtype == torch.uint8:
|
||||
w_sf = w_sf.view(torch.float8_e4m3fn)
|
||||
|
||||
return {
|
||||
"x_fp4": x_fp4,
|
||||
"w_fp4": w_fp4,
|
||||
"x_sf": x_sf,
|
||||
"w_sf": w_sf,
|
||||
"alpha": alpha,
|
||||
}
|
||||
|
||||
|
||||
def make_shape_id(
|
||||
model: str, shape_kind: str, prefix: str, m: int, n: int, k: int
|
||||
) -> str:
|
||||
prefix_slug = re.sub(r"[^a-zA-Z0-9]+", "_", prefix).strip("_")
|
||||
return f"{model}_{shape_kind}_{prefix_slug}_{m}x{n}x{k}"
|
||||
|
||||
|
||||
def load_shape_cases(shape_library: Path) -> list[dict[str, Any]]:
|
||||
payload = json.loads(shape_library.read_text(encoding="utf-8"))
|
||||
if not isinstance(payload, dict) or not payload:
|
||||
raise RuntimeError(
|
||||
f"Expected a non-empty model->shape list mapping in {shape_library}."
|
||||
)
|
||||
|
||||
rows: list[dict[str, Any]] = []
|
||||
for model, shapes in payload.items():
|
||||
if not isinstance(shapes, list):
|
||||
raise RuntimeError(
|
||||
f"Expected {model} to map to a list of shapes in {shape_library}."
|
||||
)
|
||||
for shape in shapes:
|
||||
m, n, k = (int(x) for x in shape["shape"])
|
||||
count = int(shape["count"])
|
||||
shape_kind = str(shape.get("kind", "actual_runtime_linear"))
|
||||
prefix = str(shape.get("prefix", ""))
|
||||
rows.append(
|
||||
{
|
||||
"shape_id": make_shape_id(model, shape_kind, prefix, m, n, k),
|
||||
"source_model": model,
|
||||
"shape_kind": shape_kind,
|
||||
"runtime_prefix": prefix,
|
||||
"m": m,
|
||||
"n": n,
|
||||
"k": k,
|
||||
"count": count,
|
||||
"approx_flops": 2 * m * n * k * count,
|
||||
}
|
||||
)
|
||||
|
||||
if not rows:
|
||||
raise RuntimeError(f"No shapes found in {shape_library}.")
|
||||
return rows
|
||||
|
||||
|
||||
def split_csv_arg(text: str | None) -> set[str]:
|
||||
if text is None or not text.strip():
|
||||
return set()
|
||||
return {item.strip() for item in text.split(",") if item.strip()}
|
||||
|
||||
|
||||
def select_shape_cases(
|
||||
rows: list[dict[str, Any]],
|
||||
*,
|
||||
models: set[str],
|
||||
shape_kinds: set[str],
|
||||
top_k: int,
|
||||
rank_by: str,
|
||||
) -> list[dict[str, Any]]:
|
||||
filtered = [
|
||||
row
|
||||
for row in rows
|
||||
if (not models or row["source_model"] in models)
|
||||
and (not shape_kinds or row["shape_kind"] in shape_kinds)
|
||||
]
|
||||
key = "approx_flops" if rank_by == "flops" else "count"
|
||||
return sorted(filtered, key=lambda row: int(row[key]), reverse=True)[:top_k]
|
||||
|
||||
|
||||
def write_csv(rows: list[dict[str, Any]], output_path: Path) -> None:
|
||||
with output_path.open("w", newline="", encoding="utf-8") as f:
|
||||
writer = csv.DictWriter(
|
||||
f,
|
||||
fieldnames=[
|
||||
"shape_id",
|
||||
"source_model",
|
||||
"shape_kind",
|
||||
"runtime_prefix",
|
||||
"m",
|
||||
"n",
|
||||
"k",
|
||||
"count",
|
||||
"approx_flops",
|
||||
"method",
|
||||
"median_ms",
|
||||
"min_ms",
|
||||
"max_ms",
|
||||
"tflops",
|
||||
],
|
||||
)
|
||||
writer.writeheader()
|
||||
writer.writerows(rows)
|
||||
|
||||
|
||||
def write_markdown(rows: list[dict[str, Any]], output_path: Path) -> None:
|
||||
shape_rows = []
|
||||
seen_shape_ids = set()
|
||||
for row in rows:
|
||||
if row["shape_id"] in seen_shape_ids:
|
||||
continue
|
||||
seen_shape_ids.add(row["shape_id"])
|
||||
shape_rows.append(row)
|
||||
|
||||
lines: list[str] = []
|
||||
lines.append("# Diffusion NVFP4 Scaled MM Benchmark")
|
||||
lines.append("")
|
||||
lines.append("## Shape Cases")
|
||||
lines.append("")
|
||||
lines.append("| Shape ID | Model | Shape Kind | Calls | Shape `(M,N,K)` | Prefix |")
|
||||
lines.append("|---|---|---|---:|---|---|")
|
||||
for row in shape_rows:
|
||||
lines.append(
|
||||
f"| {row['shape_id']} | {row['source_model']} | {row['shape_kind']} | {row['count']} | `({row['m']}, {row['n']}, {row['k']})` | `{row['runtime_prefix']}` |"
|
||||
)
|
||||
lines.append("")
|
||||
|
||||
for shape_row in shape_rows:
|
||||
shape_id = shape_row["shape_id"]
|
||||
scoped = [row for row in rows if row["shape_id"] == shape_id]
|
||||
lines.append(f"## {shape_id}")
|
||||
lines.append("")
|
||||
lines.append("| Method | Median ms | TFLOPS |")
|
||||
lines.append("|---|---:|---:|")
|
||||
for row in sorted(scoped, key=lambda item: float(item["median_ms"])):
|
||||
lines.append(
|
||||
f"| {row['method']} | {float(row['median_ms']):.4f} | {float(row['tflops']):.1f} |"
|
||||
)
|
||||
lines.append("")
|
||||
|
||||
output_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
|
||||
|
||||
|
||||
def run_shape_suite(shape_cases: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
device = torch.device("cuda")
|
||||
rows: list[dict[str, Any]] = []
|
||||
for idx, shape in enumerate(shape_cases):
|
||||
m = int(shape["m"])
|
||||
n = int(shape["n"])
|
||||
k = int(shape["k"])
|
||||
quantized = build_quantized_inputs(m, n, k, device, seed=idx)
|
||||
|
||||
metadata = {
|
||||
"shape_id": str(shape["shape_id"]),
|
||||
"source_model": str(shape["source_model"]),
|
||||
"shape_kind": str(shape["shape_kind"]),
|
||||
"runtime_prefix": str(shape["runtime_prefix"]),
|
||||
"m": m,
|
||||
"n": n,
|
||||
"k": k,
|
||||
"count": int(shape["count"]),
|
||||
"approx_flops": int(shape["approx_flops"]),
|
||||
}
|
||||
|
||||
providers: dict[str, Callable[[], torch.Tensor]] = {
|
||||
"cutlass": lambda: sgl_kernel.cutlass_scaled_fp4_mm(
|
||||
quantized["x_fp4"],
|
||||
quantized["w_fp4"],
|
||||
quantized["x_sf"],
|
||||
quantized["w_sf"],
|
||||
quantized["alpha"],
|
||||
DTYPE,
|
||||
),
|
||||
"flashinfer_auto": lambda: flashinfer.mm_fp4(
|
||||
quantized["x_fp4"],
|
||||
quantized["w_fp4"].T,
|
||||
quantized["x_sf"],
|
||||
quantized["w_sf"].T,
|
||||
quantized["alpha"],
|
||||
DTYPE,
|
||||
backend="auto",
|
||||
),
|
||||
"flashinfer_cudnn": lambda: flashinfer.mm_fp4(
|
||||
quantized["x_fp4"],
|
||||
quantized["w_fp4"].T,
|
||||
quantized["x_sf"],
|
||||
quantized["w_sf"].T,
|
||||
quantized["alpha"],
|
||||
DTYPE,
|
||||
backend="cudnn",
|
||||
),
|
||||
}
|
||||
|
||||
for method in METHODS:
|
||||
median_ms, max_ms, min_ms = benchmark_provider(providers[method])
|
||||
rows.append(
|
||||
{
|
||||
**metadata,
|
||||
"method": method,
|
||||
"median_ms": median_ms,
|
||||
"min_ms": min_ms,
|
||||
"max_ms": max_ms,
|
||||
"tflops": (2 * m * n * k) / (median_ms / 1e3) / 1e12,
|
||||
}
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Benchmark diffusion NVFP4 GEMM backends on the captured diffusion shape library."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--models",
|
||||
help="Comma-separated source_model filter. Default: all models in the JSON shape library.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--shape-kinds",
|
||||
help="Comma-separated shape_kind filter. Default: benchmark every shape kind in the JSON shape library.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--top-k",
|
||||
type=int,
|
||||
default=64,
|
||||
help="Benchmark the top-k shapes after filtering and ranking.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--rank-by",
|
||||
choices=["flops", "count"],
|
||||
default="flops",
|
||||
help="How to rank shapes before selecting top-k.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
default=str(DEFAULT_OUTPUT_DIR),
|
||||
help="Directory for CSV/Markdown outputs.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if is_in_ci():
|
||||
print("Skipping bench_diffusion_nvfp4_scaled_mm.py in CI")
|
||||
return
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError("CUDA is required for NVFP4 scaled mm benchmarks.")
|
||||
if not DEFAULT_SHAPE_LIBRARY.exists():
|
||||
raise RuntimeError(
|
||||
f"Shape library not found at {DEFAULT_SHAPE_LIBRARY}. "
|
||||
"Commit or copy the generated diffusion_nvfp4_shapes.json first."
|
||||
)
|
||||
|
||||
shape_cases = load_shape_cases(DEFAULT_SHAPE_LIBRARY)
|
||||
selected_shapes = select_shape_cases(
|
||||
shape_cases,
|
||||
models=split_csv_arg(args.models),
|
||||
shape_kinds=split_csv_arg(args.shape_kinds),
|
||||
top_k=args.top_k,
|
||||
rank_by=args.rank_by,
|
||||
)
|
||||
if not selected_shapes:
|
||||
raise RuntimeError("No shapes matched the requested filters.")
|
||||
rows = run_shape_suite(selected_shapes)
|
||||
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
csv_path = output_dir / "diffusion_nvfp4_scaled_mm.csv"
|
||||
md_path = output_dir / "diffusion_nvfp4_scaled_mm_summary.md"
|
||||
write_csv(rows, csv_path)
|
||||
write_markdown(rows, md_path)
|
||||
print(f"Wrote {csv_path}")
|
||||
print(f"Wrote {md_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,138 @@
|
||||
# Benchmarks SGLang fused layernorm/rmsnorm scale shift kernels
|
||||
# 1. fused_norm_scale_shift
|
||||
# 2. fused_scale_residual_norm_scale_shift
|
||||
import itertools
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph
|
||||
from sglang.multimodal_gen.runtime.layers.layernorm import (
|
||||
LayerNormScaleShift,
|
||||
RMSNormScaleShift,
|
||||
ScaleResidualLayerNormScaleShift,
|
||||
ScaleResidualRMSNormScaleShift,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.utils import is_in_ci
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=17,
|
||||
suite="base-b-kernel-benchmark-1-gpu-large",
|
||||
disabled="Temporarily skipped to unblock flashinfer upgrade. Ref: https://github.com/sgl-project/sglang/actions/runs/23735552939/job/69139238979?pr=21422",
|
||||
)
|
||||
|
||||
if is_in_ci():
|
||||
B_RANGE, S_RANGE, D_RANGE = [1], [128], [1024]
|
||||
else:
|
||||
B_RANGE, S_RANGE, D_RANGE = [1], [128, 1024, 4096], [1024, 3072, 4096]
|
||||
|
||||
NORM_TYPE_RANGE = ["layer", "rms"]
|
||||
AFFINE_RANGE = [True, False]
|
||||
DTYPE = torch.bfloat16
|
||||
DEVICE = "cuda"
|
||||
EPS = 1e-5
|
||||
LINE_VALS = ["native", "cuda"]
|
||||
LINE_NAMES = ["SGLang Native", "SGLang Fused"]
|
||||
STYLES = [("red", "-"), ("blue", "--")]
|
||||
config = list(
|
||||
itertools.product(B_RANGE, S_RANGE, D_RANGE, NORM_TYPE_RANGE, AFFINE_RANGE)
|
||||
)
|
||||
|
||||
|
||||
def preprocess_layer(layer, affine: bool, D: int, DTYPE: torch.dtype):
|
||||
if affine:
|
||||
weight = torch.randn(D, dtype=DTYPE, device=DEVICE)
|
||||
bias = torch.randn(D, dtype=DTYPE, device=DEVICE)
|
||||
with torch.no_grad():
|
||||
layer.norm.weight.copy_(weight)
|
||||
if hasattr(layer.norm, "bias"):
|
||||
layer.norm.bias.copy_(bias)
|
||||
layer.requires_grad_(False)
|
||||
return layer.to(DEVICE)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Benchmark 1: fused_norm_scale_shift
|
||||
# ============================================================================
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["B", "S", "D", "norm_type", "affine"],
|
||||
x_vals=config,
|
||||
line_arg="provider",
|
||||
line_vals=LINE_VALS,
|
||||
line_names=LINE_NAMES,
|
||||
styles=STYLES,
|
||||
ylabel="us",
|
||||
plot_name="fused_norm_scale_shift",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def bench_fused_norm_scale_shift(
|
||||
B: int, S: int, D: int, norm_type, affine: bool, provider: str
|
||||
) -> Tuple[float, float, float]:
|
||||
x = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
|
||||
scale = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
|
||||
shift = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
|
||||
if norm_type == "layer":
|
||||
layer = LayerNormScaleShift(D, EPS, affine, dtype=DTYPE)
|
||||
else:
|
||||
layer = RMSNormScaleShift(D, EPS, affine, dtype=DTYPE)
|
||||
layer = preprocess_layer(layer, affine, D, DTYPE)
|
||||
if provider == "native":
|
||||
fn = lambda: layer.forward_native(x, shift, scale)
|
||||
else:
|
||||
fn = lambda: layer.forward_cuda(x, shift, scale)
|
||||
|
||||
return run_benchmark_no_cudagraph(fn)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Benchmark 2: fused_scale_residual_norm_scale_shift
|
||||
# ============================================================================
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["B", "S", "D", "norm_type", "affine"],
|
||||
x_vals=config,
|
||||
line_arg="provider",
|
||||
line_vals=LINE_VALS,
|
||||
line_names=LINE_NAMES,
|
||||
styles=STYLES,
|
||||
ylabel="us",
|
||||
plot_name="fused_scale_residual_norm_scale_shift",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def bench_fused_scale_residual_norm_scale_shift(
|
||||
B: int, S: int, D: int, norm_type, affine: bool, provider: str
|
||||
) -> Tuple[float, float, float]:
|
||||
residual = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
|
||||
x = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
|
||||
scale = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
|
||||
shift = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
|
||||
gate = torch.randn(B, 1, D, dtype=DTYPE, device=DEVICE)
|
||||
if norm_type == "layer":
|
||||
layer = ScaleResidualLayerNormScaleShift(D, EPS, affine, dtype=DTYPE).to(DEVICE)
|
||||
else:
|
||||
layer = ScaleResidualRMSNormScaleShift(D, EPS, affine, dtype=DTYPE).to(DEVICE)
|
||||
layer = preprocess_layer(layer, affine, D, DTYPE)
|
||||
if provider == "native":
|
||||
fn = lambda: layer.forward_native(residual, x, gate, shift, scale)
|
||||
else:
|
||||
fn = lambda: layer.forward_cuda(residual, x, gate, shift, scale)
|
||||
|
||||
return run_benchmark_no_cudagraph(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(f"\n{'='*80}")
|
||||
print("Benchmark: fused_norm_scale_shift")
|
||||
print(f"{'='*80}\n")
|
||||
bench_fused_norm_scale_shift.run(print_data=True)
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("Benchmark: fused_scale_residual_norm_scale_shift")
|
||||
print(f"{'='*80}\n")
|
||||
bench_fused_scale_residual_norm_scale_shift.run(print_data=True)
|
||||
@@ -0,0 +1,313 @@
|
||||
import argparse
|
||||
import csv
|
||||
import statistics
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.diffusion.triton.group_norm_silu import triton_group_norm_silu
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.utils import is_in_ci
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=45,
|
||||
suite="base-b-kernel-benchmark-1-gpu-large",
|
||||
disabled="standalone benchmark",
|
||||
)
|
||||
|
||||
DEVICE = "cuda"
|
||||
EPS = 1e-5
|
||||
QUANTILES = [0.5, 0.2, 0.8]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Case:
|
||||
name: str
|
||||
shape: tuple[int, ...]
|
||||
num_groups: int
|
||||
|
||||
|
||||
CASES = [
|
||||
Case("token_2d", (4, 128), 32),
|
||||
Case("image_2d", (2, 64, 32, 32), 32),
|
||||
Case("video_3d_small", (1, 64, 4, 16, 16), 32),
|
||||
Case("threshold_3d", (1, 128, 1, 256, 256), 32),
|
||||
Case("hunyuan_video_large", (1, 128, 20, 256, 256), 32),
|
||||
# LTX-2 latent upsampler (`LatentUpsampler` + `ResBlock`) operates on
|
||||
# `[B, mid_channels=512, F, H, W]` tensors with num_groups=32. The
|
||||
# `small` and `pre_720p` cases stay in the default set; the larger
|
||||
# `post_720p` case is opt-in via LARGE_CASES below.
|
||||
Case("ltx2_upsampler_small", (1, 512, 8, 45, 80), 32),
|
||||
Case("ltx2_upsampler_pre_720p", (1, 512, 16, 90, 160), 32),
|
||||
]
|
||||
|
||||
# Cases too large to fit comfortably alongside the native-path intermediates
|
||||
# on consumer GPUs (e.g. 24 GB L4). Opt in with `--cases large` (large only),
|
||||
# `--cases all-large` (default + large), or by name.
|
||||
#
|
||||
# `ltx2_upsampler_post_720p` is ~471M bf16 elements (~940 MB tensor) and the
|
||||
# eager `silu(group_norm(x))` reference materializes mean / variance /
|
||||
# normalized / silu intermediates -- working set lands around 5 GB. On
|
||||
# H100 / H200 this is fine and surfaces the asymptotic ~14x kernel speedup;
|
||||
# on a 24 GB GPU it can OOM, so it's gated out of `--cases all`.
|
||||
LARGE_CASES = [
|
||||
Case("ltx2_upsampler_post_720p", (1, 512, 16, 180, 320), 32),
|
||||
]
|
||||
|
||||
CASE_BY_NAME = {case.name: case for case in CASES + LARGE_CASES}
|
||||
|
||||
|
||||
def dtype_from_name(name: str) -> torch.dtype:
|
||||
mapping = {
|
||||
"bf16": torch.bfloat16,
|
||||
"bfloat16": torch.bfloat16,
|
||||
"fp16": torch.float16,
|
||||
"float16": torch.float16,
|
||||
"fp32": torch.float32,
|
||||
"float32": torch.float32,
|
||||
}
|
||||
return mapping[name]
|
||||
|
||||
|
||||
def dtype_name(dtype: torch.dtype) -> str:
|
||||
mapping = {
|
||||
torch.bfloat16: "bf16",
|
||||
torch.float16: "fp16",
|
||||
torch.float32: "fp32",
|
||||
}
|
||||
return mapping[dtype]
|
||||
|
||||
|
||||
def parse_dtypes(text: str) -> list[torch.dtype]:
|
||||
return [dtype_from_name(item.strip()) for item in text.split(",") if item.strip()]
|
||||
|
||||
|
||||
def parse_cases(text: str) -> list[Case]:
|
||||
if text == "all":
|
||||
return CASES
|
||||
if text == "large":
|
||||
return LARGE_CASES
|
||||
if text == "all-large":
|
||||
return CASES + LARGE_CASES
|
||||
names = [item.strip() for item in text.split(",") if item.strip()]
|
||||
missing = sorted(set(names) - CASE_BY_NAME.keys())
|
||||
if missing:
|
||||
raise ValueError(f"Unknown cases: {missing}")
|
||||
return [CASE_BY_NAME[name] for name in names]
|
||||
|
||||
|
||||
def tolerance(dtype: torch.dtype) -> tuple[float, float]:
|
||||
if dtype == torch.float32:
|
||||
return 1e-5, 1e-5
|
||||
if dtype == torch.bfloat16:
|
||||
return 7e-2, 2e-2
|
||||
return 3e-3, 3e-3
|
||||
|
||||
|
||||
def native_group_norm_silu(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor,
|
||||
num_groups: int,
|
||||
) -> torch.Tensor:
|
||||
return F.silu(F.group_norm(x, num_groups, weight=weight, bias=bias, eps=EPS))
|
||||
|
||||
|
||||
def make_inputs(case: Case, dtype: torch.dtype) -> tuple[torch.Tensor, ...]:
|
||||
generator = torch.Generator(device=DEVICE)
|
||||
generator.manual_seed(len(case.shape) * 1009 + case.shape[1] * 17 + case.num_groups)
|
||||
x = torch.randn(case.shape, device=DEVICE, dtype=dtype, generator=generator)
|
||||
weight = torch.randn(case.shape[1], device=DEVICE, dtype=dtype, generator=generator)
|
||||
bias = torch.randn(case.shape[1], device=DEVICE, dtype=dtype, generator=generator)
|
||||
return x, weight, bias
|
||||
|
||||
|
||||
def do_bench_us(fn: Callable[[], object], warmup: int, rep: int) -> tuple[float, ...]:
|
||||
median_ms, p20_ms, p80_ms = triton.testing.do_bench(
|
||||
fn,
|
||||
quantiles=QUANTILES,
|
||||
warmup=warmup,
|
||||
rep=rep,
|
||||
)
|
||||
return median_ms * 1000.0, p20_ms * 1000.0, p80_ms * 1000.0
|
||||
|
||||
|
||||
def summarize(values: list[float]) -> float:
|
||||
return statistics.median(values)
|
||||
|
||||
|
||||
def run_case(
|
||||
case: Case,
|
||||
dtype: torch.dtype,
|
||||
rounds: int,
|
||||
warmup: int,
|
||||
rep: int,
|
||||
) -> dict[str, object]:
|
||||
x, weight, bias = make_inputs(case, dtype)
|
||||
|
||||
with torch.inference_mode():
|
||||
actual = triton_group_norm_silu(
|
||||
x, weight, bias, num_groups=case.num_groups, eps=EPS
|
||||
)
|
||||
expected = native_group_norm_silu(x, weight, bias, case.num_groups)
|
||||
atol, rtol = tolerance(dtype)
|
||||
torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
|
||||
|
||||
native_stats = []
|
||||
fused_stats = []
|
||||
for _ in range(rounds):
|
||||
native_stats.append(
|
||||
do_bench_us(
|
||||
lambda: native_group_norm_silu(x, weight, bias, case.num_groups),
|
||||
warmup=warmup,
|
||||
rep=rep,
|
||||
)
|
||||
)
|
||||
fused_stats.append(
|
||||
do_bench_us(
|
||||
lambda: triton_group_norm_silu(
|
||||
x, weight, bias, num_groups=case.num_groups, eps=EPS
|
||||
),
|
||||
warmup=warmup,
|
||||
rep=rep,
|
||||
)
|
||||
)
|
||||
|
||||
native_median_us = summarize([stats[0] for stats in native_stats])
|
||||
fused_median_us = summarize([stats[0] for stats in fused_stats])
|
||||
torch.cuda.empty_cache()
|
||||
return {
|
||||
"case": case.name,
|
||||
"shape": "x".join(str(dim) for dim in case.shape),
|
||||
"groups": case.num_groups,
|
||||
"dtype": dtype_name(dtype),
|
||||
"native_median_us": native_median_us,
|
||||
"native_p20_us": summarize([stats[1] for stats in native_stats]),
|
||||
"native_p80_us": summarize([stats[2] for stats in native_stats]),
|
||||
"fused_median_us": fused_median_us,
|
||||
"fused_p20_us": summarize([stats[1] for stats in fused_stats]),
|
||||
"fused_p80_us": summarize([stats[2] for stats in fused_stats]),
|
||||
"speedup": native_median_us / fused_median_us,
|
||||
"rounds": rounds,
|
||||
"warmup": warmup,
|
||||
"rep": rep,
|
||||
}
|
||||
|
||||
|
||||
def run_profile(case: Case, dtype: torch.dtype, provider: str, iters: int) -> None:
|
||||
x, weight, bias = make_inputs(case, dtype)
|
||||
|
||||
if provider == "native":
|
||||
|
||||
def fn() -> torch.Tensor:
|
||||
return native_group_norm_silu(x, weight, bias, case.num_groups)
|
||||
|
||||
elif provider == "fused":
|
||||
|
||||
def fn() -> torch.Tensor:
|
||||
return triton_group_norm_silu(
|
||||
x, weight, bias, num_groups=case.num_groups, eps=EPS
|
||||
)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unknown provider: {provider}")
|
||||
|
||||
with torch.inference_mode():
|
||||
for _ in range(5):
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
for _ in range(iters):
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
|
||||
def write_csv(rows: list[dict[str, object]], output_path: Path) -> None:
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
fieldnames = list(rows[0].keys()) if rows else []
|
||||
with output_path.open("w", newline="", encoding="utf-8") as f:
|
||||
writer = csv.DictWriter(f, fieldnames=fieldnames)
|
||||
writer.writeheader()
|
||||
writer.writerows(rows)
|
||||
|
||||
|
||||
def print_rows(rows: list[dict[str, object]]) -> None:
|
||||
header = (
|
||||
"case",
|
||||
"dtype",
|
||||
"shape",
|
||||
"native_us",
|
||||
"fused_us",
|
||||
"speedup",
|
||||
)
|
||||
print("| " + " | ".join(header) + " |")
|
||||
print("|---|---|---|---:|---:|---:|")
|
||||
for row in rows:
|
||||
print(
|
||||
"| {case} | {dtype} | {shape} | {native:.2f} | {fused:.2f} | {speedup:.3f}x |".format(
|
||||
case=row["case"],
|
||||
dtype=row["dtype"],
|
||||
shape=row["shape"],
|
||||
native=row["native_median_us"],
|
||||
fused=row["fused_median_us"],
|
||||
speedup=row["speedup"],
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Benchmark fused GroupNorm+SiLU against PyTorch GroupNorm+SiLU."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cases",
|
||||
default="all",
|
||||
help=(
|
||||
"Comma-separated case names, or one of: 'all' (default-sized "
|
||||
"cases only), 'large' (high-memory cases only -- requires "
|
||||
"H100/H200-class GPU), 'all-large' (both). See CASES + LARGE_CASES."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--dtypes", default="bf16,fp16")
|
||||
parser.add_argument("--rounds", type=int, default=3)
|
||||
parser.add_argument("--warmup", type=int, default=25)
|
||||
parser.add_argument("--rep", type=int, default=100)
|
||||
parser.add_argument("--output-csv", default="")
|
||||
parser.add_argument("--profile-provider", choices=["native", "fused"], default="")
|
||||
parser.add_argument("--profile-iters", type=int, default=20)
|
||||
args = parser.parse_args()
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError("CUDA is required for this benchmark.")
|
||||
|
||||
cases = parse_cases(args.cases)
|
||||
dtypes = parse_dtypes(args.dtypes)
|
||||
|
||||
if args.profile_provider:
|
||||
if len(cases) != 1 or len(dtypes) != 1:
|
||||
raise ValueError(
|
||||
"--profile-provider requires exactly one case and one dtype"
|
||||
)
|
||||
run_profile(cases[0], dtypes[0], args.profile_provider, args.profile_iters)
|
||||
return
|
||||
|
||||
rows = []
|
||||
for case in cases:
|
||||
for dtype in dtypes:
|
||||
rows.append(run_case(case, dtype, args.rounds, args.warmup, args.rep))
|
||||
|
||||
print_rows(rows)
|
||||
if args.output_csv:
|
||||
write_csv(rows, Path(args.output_csv))
|
||||
print(f"Wrote {args.output_csv}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if is_in_ci():
|
||||
print("Skipping bench_group_norm_silu.py in CI")
|
||||
sys.exit(0)
|
||||
main()
|
||||
@@ -0,0 +1,753 @@
|
||||
import argparse
|
||||
import csv
|
||||
import functools
|
||||
import importlib
|
||||
import math
|
||||
import os
|
||||
import statistics
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE
|
||||
from sglang.jit_kernel.diffusion.triton.norm import norm_infer, rms_norm_fn
|
||||
from sglang.jit_kernel.diffusion.triton.rmsnorm_onepass import triton_one_pass_rms_norm
|
||||
from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm
|
||||
from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm
|
||||
from sglang.jit_kernel.utils import KERNEL_PATH
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.utils import is_in_ci
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=120,
|
||||
suite="base-b-kernel-benchmark-1-gpu-large",
|
||||
disabled="self-skips in CI, standalone tool",
|
||||
)
|
||||
|
||||
os.environ.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1")
|
||||
|
||||
REPO_ROOT = KERNEL_PATH.parents[2]
|
||||
THIRD_PARTY_ROOT = REPO_ROOT / "third_party"
|
||||
|
||||
FLAGGEMS_REPO = "https://github.com/flagos-ai/FlagGems.git"
|
||||
QUACK_REPO = "https://github.com/Dao-AILab/quack.git"
|
||||
|
||||
TORCH_LN = "torch.nn.LayerNorm"
|
||||
SGL_RMS = "sglang.RMSNorm.forward_cuda"
|
||||
SGL_FUSED = "sgl_kernel.fused_add_rmsnorm"
|
||||
SGL_LN = "sglang.LayerNormScaleShift"
|
||||
SGL_RES_LN = "sglang.ScaleResidualLayerNormScaleShift"
|
||||
SGL_LN_PAIR = f"{SGL_LN} / {SGL_RES_LN}"
|
||||
MOVA_LN_MIX = f"{TORCH_LN} / {SGL_LN_PAIR}"
|
||||
|
||||
ACTUAL_DIFFUSION_GROUPS: list[
|
||||
tuple[str, str, list[tuple[str, str, tuple[int, ...], str]]]
|
||||
] = [
|
||||
(
|
||||
"qwen",
|
||||
"1 GPU",
|
||||
[
|
||||
("qwen_ln_4096x3072", "layernorm", (1, 4096, 3072), SGL_LN_PAIR),
|
||||
("qwen_ln_26x3072", "layernorm", (1, 26, 3072), SGL_LN_PAIR),
|
||||
("qwen_ln_6x3072", "layernorm", (1, 6, 3072), SGL_LN_PAIR),
|
||||
("qwen_rms_26x3584", "rmsnorm", (1, 26, 3584), SGL_RMS),
|
||||
("qwen_rms_6x3584", "rmsnorm", (1, 6, 3584), SGL_RMS),
|
||||
],
|
||||
),
|
||||
(
|
||||
"qwen-edit",
|
||||
"1 GPU",
|
||||
[
|
||||
("qwen_edit_ln_200x3072", "layernorm", (1, 200, 3072), SGL_LN_PAIR),
|
||||
("qwen_edit_ln_203x3072", "layernorm", (1, 203, 3072), SGL_LN_PAIR),
|
||||
("qwen_edit_ln_8308x3072", "layernorm", (1, 8308, 3072), TORCH_LN),
|
||||
("qwen_edit_rms_200x3584", "rmsnorm", (1, 200, 3584), SGL_RMS),
|
||||
("qwen_edit_rms_203x3584", "rmsnorm", (1, 203, 3584), SGL_RMS),
|
||||
],
|
||||
),
|
||||
(
|
||||
"flux",
|
||||
"1 GPU",
|
||||
[
|
||||
("flux_ln_77x768", "layernorm", (1, 77, 768), TORCH_LN),
|
||||
("flux_ln_512x3072", "layernorm", (1, 512, 3072), TORCH_LN),
|
||||
("flux_ln_4096x3072", "layernorm", (1, 4096, 3072), TORCH_LN),
|
||||
("flux_ln_4608x3072", "layernorm", (1, 4608, 3072), TORCH_LN),
|
||||
("flux_rms_512x4096", "rmsnorm", (1, 512, 4096), SGL_RMS),
|
||||
],
|
||||
),
|
||||
(
|
||||
"flux2",
|
||||
"1 GPU",
|
||||
[
|
||||
("flux2_ln_512x6144", "layernorm", (1, 512, 6144), TORCH_LN),
|
||||
("flux2_ln_4096x6144", "layernorm", (1, 4096, 6144), TORCH_LN),
|
||||
("flux2_ln_4608x6144", "layernorm", (1, 4608, 6144), TORCH_LN),
|
||||
("flux2_rms_4608x48x128", "rmsnorm", (1, 4608, 48, 128), SGL_RMS),
|
||||
],
|
||||
),
|
||||
(
|
||||
"zimage",
|
||||
"1 GPU",
|
||||
[
|
||||
("zimage_ln_4128x3840", "layernorm", (1, 4128, 3840), TORCH_LN),
|
||||
("zimage_rms_32x3840", "rmsnorm", (1, 32, 3840), SGL_RMS),
|
||||
("zimage_rms_4096x3840", "rmsnorm", (1, 4096, 3840), SGL_RMS),
|
||||
("zimage_rms_4128x3840", "rmsnorm", (1, 4128, 3840), SGL_RMS),
|
||||
("zimage_rms_32x2560", "rmsnorm", (32, 2560), SGL_RMS),
|
||||
],
|
||||
),
|
||||
(
|
||||
"wan-ti2v",
|
||||
"1 GPU",
|
||||
[
|
||||
("wan_ti2v_ln_17850x3072", "layernorm", (1, 17850, 3072), SGL_LN_PAIR),
|
||||
("wan_ti2v_rms_17850x3072", "rmsnorm", (1, 17850, 3072), SGL_RMS),
|
||||
("wan_ti2v_rms_512x3072", "rmsnorm", (1, 512, 3072), SGL_RMS),
|
||||
("wan_ti2v_rms_512x4096", "rmsnorm", (1, 512, 4096), SGL_RMS),
|
||||
],
|
||||
),
|
||||
(
|
||||
"hunyuanvideo",
|
||||
"1 GPU",
|
||||
[
|
||||
("hunyuan_ln_46x768", "layernorm", (1, 46, 768), TORCH_LN),
|
||||
("hunyuan_ln_45x3072", "layernorm", (1, 45, 3072), SGL_LN_PAIR),
|
||||
("hunyuan_ln_27030x3072", "layernorm", (1, 27030, 3072), SGL_LN_PAIR),
|
||||
("hunyuan_ln_27075x3072", "layernorm", (1, 27075, 3072), SGL_LN),
|
||||
("hunyuan_rms_140x4096", "rmsnorm", (1, 140, 4096), SGL_RMS),
|
||||
("hunyuan_rms_45x24x128", "rmsnorm", (1, 45, 24, 128), SGL_RMS),
|
||||
("hunyuan_rms_27030x24x128", "rmsnorm", (1, 27030, 24, 128), SGL_RMS),
|
||||
("hunyuan_rms_27075x24x128", "rmsnorm", (1, 27075, 24, 128), SGL_RMS),
|
||||
("hunyuan_fused_add_140x4096", "fused_add_rmsnorm", (140, 4096), SGL_FUSED),
|
||||
],
|
||||
),
|
||||
(
|
||||
"mova-720p",
|
||||
"4 GPU, ulysses=4, ring=1",
|
||||
[
|
||||
("mova_ln_101x1536", "layernorm", (1, 101, 1536), MOVA_LN_MIX),
|
||||
("mova_ln_403x1536", "layernorm", (1, 403, 1536), TORCH_LN),
|
||||
("mova_ln_44100x5120", "layernorm", (1, 44100, 5120), MOVA_LN_MIX),
|
||||
("mova_ln_176400x5120", "layernorm", (1, 176400, 5120), SGL_LN),
|
||||
("mova_rms_101x1536", "rmsnorm", (1, 101, 1536), SGL_RMS),
|
||||
("mova_rms_101x5120", "rmsnorm", (1, 101, 5120), SGL_RMS),
|
||||
("mova_rms_44100x1536", "rmsnorm", (1, 44100, 1536), SGL_RMS),
|
||||
("mova_rms_44100x5120", "rmsnorm", (1, 44100, 5120), SGL_RMS),
|
||||
("mova_rms_512x1536", "rmsnorm", (1, 512, 1536), SGL_RMS),
|
||||
("mova_rms_512x4096", "rmsnorm", (1, 512, 4096), SGL_RMS),
|
||||
("mova_rms_512x5120", "rmsnorm", (1, 512, 5120), SGL_RMS),
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
ACTUAL_DIFFUSION_SHAPES: list[dict[str, object]] = [
|
||||
{
|
||||
"shape_id": shape_id,
|
||||
"model": model,
|
||||
"gpu_config": gpu_config,
|
||||
"op": op,
|
||||
"input_shape": list(input_shape),
|
||||
"source_impl": source_impl,
|
||||
}
|
||||
for model, gpu_config, cases in ACTUAL_DIFFUSION_GROUPS
|
||||
for shape_id, op, input_shape, source_impl in cases
|
||||
]
|
||||
|
||||
|
||||
def effective_rows_from_shape(input_shape: list[int]) -> int:
|
||||
rows = 1
|
||||
for dim in input_shape[:-1]:
|
||||
rows *= dim
|
||||
return rows
|
||||
|
||||
|
||||
def ensure_repo(repo_name: str, repo_url: str) -> Path:
|
||||
repo_path = THIRD_PARTY_ROOT / repo_name
|
||||
if repo_path.exists():
|
||||
return repo_path
|
||||
repo_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
subprocess.run(
|
||||
["git", "clone", "--depth", "1", repo_url, str(repo_path)],
|
||||
check=True,
|
||||
cwd=REPO_ROOT,
|
||||
)
|
||||
return repo_path
|
||||
|
||||
|
||||
def ensure_python_dep(module_name: str, package_name: str | None = None) -> None:
|
||||
package_name = package_name or module_name
|
||||
try:
|
||||
importlib.import_module(module_name)
|
||||
except ModuleNotFoundError:
|
||||
subprocess.run(
|
||||
[sys.executable, "-m", "pip", "install", package_name],
|
||||
check=True,
|
||||
)
|
||||
|
||||
|
||||
def dtype_from_name(name: str) -> torch.dtype:
|
||||
mapping = {
|
||||
"bf16": torch.bfloat16,
|
||||
"bfloat16": torch.bfloat16,
|
||||
"fp16": torch.float16,
|
||||
"float16": torch.float16,
|
||||
"fp32": torch.float32,
|
||||
"float32": torch.float32,
|
||||
}
|
||||
return mapping[name]
|
||||
|
||||
|
||||
def dtype_name(dtype: torch.dtype) -> str:
|
||||
mapping = {
|
||||
torch.bfloat16: "bf16",
|
||||
torch.float16: "fp16",
|
||||
torch.float32: "fp32",
|
||||
}
|
||||
return mapping[dtype]
|
||||
|
||||
|
||||
def normalize_hidden_sizes(text: str) -> list[int]:
|
||||
return [int(x) for x in text.split(",") if x]
|
||||
|
||||
|
||||
def normalize_dtypes(text: str) -> list[torch.dtype]:
|
||||
return [dtype_from_name(x.strip()) for x in text.split(",") if x.strip()]
|
||||
|
||||
|
||||
def prewarm(fn: Callable[[], object], iters: int = 3) -> None:
|
||||
for _ in range(iters):
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
|
||||
def benchmark_provider(
|
||||
fn: Callable[[], object],
|
||||
setup_fn: Callable[[], None] | None = None,
|
||||
warmup: int = 10,
|
||||
rep: int = 30,
|
||||
) -> tuple[float, float, float]:
|
||||
for _ in range(warmup):
|
||||
if setup_fn is not None:
|
||||
setup_fn()
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
start_event = torch.cuda.Event(enable_timing=True)
|
||||
end_event = torch.cuda.Event(enable_timing=True)
|
||||
times_us: list[float] = []
|
||||
for _ in range(rep):
|
||||
if setup_fn is not None:
|
||||
setup_fn()
|
||||
start_event.record()
|
||||
fn()
|
||||
end_event.record()
|
||||
end_event.synchronize()
|
||||
times_us.append(start_event.elapsed_time(end_event) * 1000.0)
|
||||
|
||||
return statistics.median(times_us), max(times_us), min(times_us)
|
||||
|
||||
|
||||
def geometric_mean(values: list[float]) -> float:
|
||||
if not values:
|
||||
return float("nan")
|
||||
return math.exp(sum(math.log(v) for v in values) / len(values))
|
||||
|
||||
|
||||
@functools.cache
|
||||
def load_flaggems():
|
||||
ensure_python_dep("sqlalchemy")
|
||||
ensure_repo("FlagGems", FLAGGEMS_REPO)
|
||||
src_root = THIRD_PARTY_ROOT / "FlagGems" / "src"
|
||||
if str(src_root) not in sys.path:
|
||||
sys.path.insert(0, str(src_root))
|
||||
from flag_gems.fused.fused_add_rms_norm import fused_add_rms_norm
|
||||
from flag_gems.ops.layernorm import layer_norm
|
||||
from flag_gems.ops.rms_norm import rms_norm
|
||||
|
||||
return rms_norm, layer_norm, fused_add_rms_norm
|
||||
|
||||
|
||||
@functools.cache
|
||||
def load_quack():
|
||||
repo_path = ensure_repo("quack", QUACK_REPO)
|
||||
try:
|
||||
quack_rmsnorm = importlib.import_module("quack.rmsnorm")
|
||||
except ModuleNotFoundError:
|
||||
subprocess.run(
|
||||
[sys.executable, "-m", "pip", "install", "-e", str(repo_path)],
|
||||
check=True,
|
||||
)
|
||||
quack_rmsnorm = importlib.import_module("quack.rmsnorm")
|
||||
|
||||
return quack_rmsnorm.rmsnorm_fwd, quack_rmsnorm.layernorm_fwd
|
||||
|
||||
|
||||
def build_rmsnorm_providers(dtype: torch.dtype, batch_size: int, hidden_size: int):
|
||||
import flashinfer.norm as flashinfer_norm
|
||||
import sgl_kernel
|
||||
|
||||
x = torch.randn((batch_size, hidden_size), device=DEFAULT_DEVICE, dtype=dtype)
|
||||
weight = torch.randn(hidden_size, device=DEFAULT_DEVICE, dtype=dtype)
|
||||
|
||||
jit_out = torch.empty_like(x)
|
||||
sgl_out = torch.empty_like(x)
|
||||
flashinfer_out = torch.empty_like(x)
|
||||
|
||||
flaggems_rms_norm, _, _ = load_flaggems()
|
||||
quack_rmsnorm_fwd, _ = load_quack()
|
||||
|
||||
providers = {
|
||||
"pytorch": lambda: F.rms_norm(x, (hidden_size,), weight, 1e-6),
|
||||
"sgl_kernel": lambda: sgl_kernel.rmsnorm(x, weight, eps=1e-6, out=sgl_out),
|
||||
"flashinfer": lambda: flashinfer_norm.rmsnorm(
|
||||
x, weight, eps=1e-6, out=flashinfer_out
|
||||
),
|
||||
"jit_rmsnorm": lambda: jit_rmsnorm(x, weight, jit_out, 1e-6),
|
||||
"quack": lambda: quack_rmsnorm_fwd(x, weight, eps=1e-6),
|
||||
"triton_rms_norm_fn": lambda: rms_norm_fn(
|
||||
x, weight, bias=None, residual=None, eps=1e-6
|
||||
),
|
||||
"flaggems": lambda: flaggems_rms_norm(x, (hidden_size,), weight, 1e-6),
|
||||
}
|
||||
if hidden_size <= 128:
|
||||
providers["triton_one_pass"] = lambda: triton_one_pass_rms_norm(x, weight, 1e-6)
|
||||
return providers
|
||||
|
||||
|
||||
def build_fused_add_rmsnorm_providers(
|
||||
dtype: torch.dtype, batch_size: int, hidden_size: int
|
||||
):
|
||||
import flashinfer.norm as flashinfer_norm
|
||||
import sgl_kernel
|
||||
|
||||
base_x = torch.randn((batch_size, hidden_size), device=DEFAULT_DEVICE, dtype=dtype)
|
||||
base_residual = torch.randn_like(base_x)
|
||||
weight = torch.randn(hidden_size, device=DEFAULT_DEVICE, dtype=dtype)
|
||||
|
||||
x = base_x.clone()
|
||||
residual = base_residual.clone()
|
||||
|
||||
def reset():
|
||||
x.copy_(base_x)
|
||||
residual.copy_(base_residual)
|
||||
|
||||
_, _, flaggems_fused_add_rms_norm = load_flaggems()
|
||||
quack_rmsnorm_fwd, _ = load_quack()
|
||||
|
||||
def pytorch_impl():
|
||||
out = x + residual
|
||||
return F.rms_norm(out, (hidden_size,), weight, 1e-6)
|
||||
|
||||
providers = {
|
||||
"pytorch": (pytorch_impl, reset),
|
||||
"sgl_kernel": (
|
||||
lambda: sgl_kernel.fused_add_rmsnorm(x, residual, weight, eps=1e-6),
|
||||
reset,
|
||||
),
|
||||
"flashinfer": (
|
||||
lambda: flashinfer_norm.fused_add_rmsnorm(x, residual, weight, eps=1e-6),
|
||||
reset,
|
||||
),
|
||||
"jit_fused_add_rmsnorm": (
|
||||
lambda: jit_fused_add_rmsnorm(x, residual, weight, 1e-6),
|
||||
reset,
|
||||
),
|
||||
"quack": (
|
||||
lambda: quack_rmsnorm_fwd(x, weight, residual=residual, eps=1e-6),
|
||||
reset,
|
||||
),
|
||||
"flaggems": (
|
||||
lambda: flaggems_fused_add_rms_norm(
|
||||
x, residual, (hidden_size,), weight, 1e-6
|
||||
),
|
||||
reset,
|
||||
),
|
||||
}
|
||||
return providers
|
||||
|
||||
|
||||
def build_layernorm_providers(dtype: torch.dtype, batch_size: int, hidden_size: int):
|
||||
import flashinfer.norm as flashinfer_norm
|
||||
|
||||
x = torch.randn((batch_size, hidden_size), device=DEFAULT_DEVICE, dtype=dtype)
|
||||
weight = torch.randn(hidden_size, device=DEFAULT_DEVICE, dtype=dtype)
|
||||
bias = torch.randn(hidden_size, device=DEFAULT_DEVICE, dtype=dtype)
|
||||
flashinfer_weight = torch.randn(
|
||||
hidden_size, device=DEFAULT_DEVICE, dtype=torch.float32
|
||||
)
|
||||
flashinfer_bias = torch.randn(
|
||||
hidden_size, device=DEFAULT_DEVICE, dtype=torch.float32
|
||||
)
|
||||
|
||||
triton_out = torch.empty_like(x)
|
||||
|
||||
_, flaggems_layer_norm, _ = load_flaggems()
|
||||
_, quack_layernorm_fwd = load_quack()
|
||||
|
||||
providers = {
|
||||
"pytorch": lambda: F.layer_norm(x, (hidden_size,), weight, bias, 1e-6),
|
||||
"triton_norm_infer": lambda: norm_infer(
|
||||
x, weight, bias, eps=1e-6, is_rms_norm=False, out=triton_out
|
||||
),
|
||||
"flashinfer": lambda: flashinfer_norm.layernorm(
|
||||
x, flashinfer_weight, flashinfer_bias, 1e-6
|
||||
),
|
||||
"quack": lambda: quack_layernorm_fwd(
|
||||
x, flashinfer_weight, flashinfer_bias, 1e-6
|
||||
),
|
||||
"flaggems": lambda: flaggems_layer_norm(x, (hidden_size,), weight, bias)[0],
|
||||
}
|
||||
return providers
|
||||
|
||||
|
||||
def maybe_benchmark(
|
||||
op_name: str,
|
||||
provider_name: str,
|
||||
fn: Callable[[], object],
|
||||
rows: list[dict[str, object]],
|
||||
dtype: torch.dtype,
|
||||
batch_size: int,
|
||||
hidden_size: int,
|
||||
reset: Callable[[], None] | None = None,
|
||||
metadata: dict[str, object] | None = None,
|
||||
) -> None:
|
||||
metadata = metadata or {}
|
||||
try:
|
||||
median_us, max_us, min_us = benchmark_provider(fn, reset)
|
||||
rows.append(
|
||||
{
|
||||
"op": op_name,
|
||||
"provider": provider_name,
|
||||
"dtype": dtype_name(dtype),
|
||||
"batch_size": batch_size,
|
||||
"hidden_size": hidden_size,
|
||||
"median_us": median_us,
|
||||
"min_us": min_us,
|
||||
"max_us": max_us,
|
||||
"status": "ok",
|
||||
"error": "",
|
||||
**metadata,
|
||||
}
|
||||
)
|
||||
except Exception as exc: # pragma: no cover - benchmark failures are data
|
||||
rows.append(
|
||||
{
|
||||
"op": op_name,
|
||||
"provider": provider_name,
|
||||
"dtype": dtype_name(dtype),
|
||||
"batch_size": batch_size,
|
||||
"hidden_size": hidden_size,
|
||||
"median_us": "",
|
||||
"min_us": "",
|
||||
"max_us": "",
|
||||
"status": "unsupported",
|
||||
"error": str(exc),
|
||||
**metadata,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def write_csv(rows: list[dict[str, object]], output_path: Path) -> None:
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with output_path.open("w", newline="", encoding="utf-8") as f:
|
||||
writer = csv.DictWriter(
|
||||
f,
|
||||
fieldnames=[
|
||||
"op",
|
||||
"provider",
|
||||
"dtype",
|
||||
"batch_size",
|
||||
"hidden_size",
|
||||
"median_us",
|
||||
"min_us",
|
||||
"max_us",
|
||||
"shape_id",
|
||||
"source_model",
|
||||
"source_gpu_config",
|
||||
"source_input_shape",
|
||||
"source_impl",
|
||||
"status",
|
||||
"error",
|
||||
],
|
||||
)
|
||||
writer.writeheader()
|
||||
writer.writerows(rows)
|
||||
|
||||
|
||||
def write_markdown(rows: list[dict[str, object]], output_path: Path) -> None:
|
||||
lines: list[str] = []
|
||||
lines.append("# Norm Benchmark Summary")
|
||||
lines.append("")
|
||||
actual_shape_rows = [row for row in rows if row.get("shape_id")]
|
||||
if actual_shape_rows:
|
||||
seen: set[tuple[str, str, str, str, str, str]] = set()
|
||||
lines.append("## Diffusion Shape Cases")
|
||||
lines.append("")
|
||||
lines.append(
|
||||
"| Shape ID | Op | Model | GPU Config | Input Shape | Source Impl |"
|
||||
)
|
||||
lines.append("|---|---|---|---|---|---|")
|
||||
for row in actual_shape_rows:
|
||||
key = (
|
||||
str(row.get("shape_id", "")),
|
||||
str(row.get("op", "")),
|
||||
str(row.get("source_model", "")),
|
||||
str(row.get("source_gpu_config", "")),
|
||||
str(row.get("source_input_shape", "")),
|
||||
str(row.get("source_impl", "")),
|
||||
)
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
lines.append(
|
||||
f"| {key[0]} | {key[1]} | {key[2]} | {key[3]} | `{key[4]}` | {key[5]} |"
|
||||
)
|
||||
lines.append("")
|
||||
for op_name in ("rmsnorm", "fused_add_rmsnorm", "layernorm"):
|
||||
for dtype in sorted({row["dtype"] for row in rows}):
|
||||
scoped = [
|
||||
row
|
||||
for row in rows
|
||||
if row["op"] == op_name
|
||||
and row["dtype"] == dtype
|
||||
and row["status"] == "ok"
|
||||
]
|
||||
if not scoped:
|
||||
continue
|
||||
provider_to_values: dict[str, list[float]] = {}
|
||||
provider_to_speedups: dict[str, list[float]] = {}
|
||||
by_shape: dict[tuple[str, int, int], dict[str, float]] = {}
|
||||
for row in scoped:
|
||||
provider = str(row["provider"])
|
||||
value = float(row["median_us"])
|
||||
provider_to_values.setdefault(provider, []).append(value)
|
||||
shape = (
|
||||
str(row.get("shape_id", "")),
|
||||
int(row["batch_size"]),
|
||||
int(row["hidden_size"]),
|
||||
)
|
||||
by_shape.setdefault(shape, {})[provider] = value
|
||||
for shape, perf in by_shape.items():
|
||||
if "pytorch" not in perf:
|
||||
continue
|
||||
baseline = perf["pytorch"]
|
||||
for provider, value in perf.items():
|
||||
provider_to_speedups.setdefault(provider, []).append(
|
||||
baseline / value
|
||||
)
|
||||
|
||||
lines.append(f"## {op_name} ({dtype})")
|
||||
lines.append("")
|
||||
lines.append(
|
||||
"| Provider | Geomean Speedup vs PyTorch | Median Latency (us) | Win Count |"
|
||||
)
|
||||
lines.append("|---|---:|---:|---:|")
|
||||
wins: dict[str, int] = {}
|
||||
for perf in by_shape.values():
|
||||
best_provider = min(perf, key=perf.get)
|
||||
wins[best_provider] = wins.get(best_provider, 0) + 1
|
||||
for provider in sorted(provider_to_values):
|
||||
geomean_speedup = geometric_mean(provider_to_speedups.get(provider, []))
|
||||
median_latency = statistics.median(provider_to_values[provider])
|
||||
win_count = wins.get(provider, 0)
|
||||
lines.append(
|
||||
f"| {provider} | {geomean_speedup:.3f}x | {median_latency:.2f} | {win_count} |"
|
||||
)
|
||||
lines.append("")
|
||||
output_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
|
||||
|
||||
|
||||
def run_suite(
|
||||
hidden_sizes: list[int],
|
||||
batch_sizes: list[int],
|
||||
dtypes: list[torch.dtype],
|
||||
ops: list[str],
|
||||
) -> list[dict[str, object]]:
|
||||
rows: list[dict[str, object]] = []
|
||||
for dtype in dtypes:
|
||||
for batch_size in batch_sizes:
|
||||
for hidden_size in hidden_sizes:
|
||||
if "rmsnorm" in ops:
|
||||
rms_providers = build_rmsnorm_providers(
|
||||
dtype, batch_size, hidden_size
|
||||
)
|
||||
for provider_name, fn in rms_providers.items():
|
||||
maybe_benchmark(
|
||||
"rmsnorm",
|
||||
provider_name,
|
||||
fn,
|
||||
rows,
|
||||
dtype,
|
||||
batch_size,
|
||||
hidden_size,
|
||||
)
|
||||
|
||||
if "fused_add_rmsnorm" in ops:
|
||||
fused_providers = build_fused_add_rmsnorm_providers(
|
||||
dtype, batch_size, hidden_size
|
||||
)
|
||||
for provider_name, provider in fused_providers.items():
|
||||
fn, reset = provider
|
||||
maybe_benchmark(
|
||||
"fused_add_rmsnorm",
|
||||
provider_name,
|
||||
fn,
|
||||
rows,
|
||||
dtype,
|
||||
batch_size,
|
||||
hidden_size,
|
||||
reset,
|
||||
)
|
||||
|
||||
if "layernorm" in ops:
|
||||
layernorm_providers = build_layernorm_providers(
|
||||
dtype, batch_size, hidden_size
|
||||
)
|
||||
for provider_name, fn in layernorm_providers.items():
|
||||
maybe_benchmark(
|
||||
"layernorm",
|
||||
provider_name,
|
||||
fn,
|
||||
rows,
|
||||
dtype,
|
||||
batch_size,
|
||||
hidden_size,
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
def run_shape_suite(
|
||||
shape_cases: list[dict[str, object]],
|
||||
dtypes: list[torch.dtype],
|
||||
) -> list[dict[str, object]]:
|
||||
rows: list[dict[str, object]] = []
|
||||
for case in shape_cases:
|
||||
op_name = str(case["op"])
|
||||
input_shape = [int(x) for x in case["input_shape"]]
|
||||
batch_size = effective_rows_from_shape(input_shape)
|
||||
hidden_size = input_shape[-1]
|
||||
metadata = {
|
||||
"shape_id": str(case["shape_id"]),
|
||||
"source_model": str(case["model"]),
|
||||
"source_gpu_config": str(case["gpu_config"]),
|
||||
"source_input_shape": str(input_shape),
|
||||
"source_impl": str(case["source_impl"]),
|
||||
}
|
||||
for dtype in dtypes:
|
||||
if op_name == "rmsnorm":
|
||||
providers = build_rmsnorm_providers(dtype, batch_size, hidden_size)
|
||||
for provider_name, fn in providers.items():
|
||||
maybe_benchmark(
|
||||
op_name,
|
||||
provider_name,
|
||||
fn,
|
||||
rows,
|
||||
dtype,
|
||||
batch_size,
|
||||
hidden_size,
|
||||
metadata=metadata,
|
||||
)
|
||||
elif op_name == "fused_add_rmsnorm":
|
||||
providers = build_fused_add_rmsnorm_providers(
|
||||
dtype, batch_size, hidden_size
|
||||
)
|
||||
for provider_name, provider in providers.items():
|
||||
fn, reset = provider
|
||||
maybe_benchmark(
|
||||
op_name,
|
||||
provider_name,
|
||||
fn,
|
||||
rows,
|
||||
dtype,
|
||||
batch_size,
|
||||
hidden_size,
|
||||
reset,
|
||||
metadata=metadata,
|
||||
)
|
||||
elif op_name == "layernorm":
|
||||
providers = build_layernorm_providers(dtype, batch_size, hidden_size)
|
||||
for provider_name, fn in providers.items():
|
||||
maybe_benchmark(
|
||||
op_name,
|
||||
provider_name,
|
||||
fn,
|
||||
rows,
|
||||
dtype,
|
||||
batch_size,
|
||||
hidden_size,
|
||||
metadata=metadata,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported op in shape preset: {op_name}")
|
||||
return rows
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Benchmark RMSNorm/LayerNorm implementations across providers."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hidden-sizes",
|
||||
default="64,128,256,512,1024,2048,4096,8192,16384",
|
||||
help="Comma-separated hidden sizes.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch-sizes",
|
||||
default="1,16,128,1024",
|
||||
help="Comma-separated batch sizes.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dtypes",
|
||||
default="bf16,fp16",
|
||||
help="Comma-separated dtypes: bf16, fp16, fp32.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
default=str(REPO_ROOT / "outputs" / "norm_benchmarks"),
|
||||
help="Directory for CSV/Markdown outputs.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ops",
|
||||
default="rmsnorm,fused_add_rmsnorm,layernorm",
|
||||
help="Comma-separated ops to benchmark.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--shape-preset",
|
||||
choices=["grid", "diffusion-actual"],
|
||||
default="grid",
|
||||
help="Use the default grid sweep or the captured diffusion workload shapes.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError("CUDA is required for norm benchmarks.")
|
||||
|
||||
hidden_sizes = normalize_hidden_sizes(args.hidden_sizes)
|
||||
batch_sizes = normalize_hidden_sizes(args.batch_sizes)
|
||||
dtypes = normalize_dtypes(args.dtypes)
|
||||
ops = [op.strip() for op in args.ops.split(",") if op.strip()]
|
||||
|
||||
if args.shape_preset == "diffusion-actual":
|
||||
shape_cases = [case for case in ACTUAL_DIFFUSION_SHAPES if case["op"] in ops]
|
||||
rows = run_shape_suite(shape_cases, dtypes)
|
||||
else:
|
||||
rows = run_suite(hidden_sizes, batch_sizes, dtypes, ops)
|
||||
output_dir = Path(args.output_dir)
|
||||
csv_path = output_dir / "norm_impls.csv"
|
||||
md_path = output_dir / "norm_impls_summary.md"
|
||||
write_csv(rows, csv_path)
|
||||
write_markdown(rows, md_path)
|
||||
print(f"Wrote {csv_path}")
|
||||
print(f"Wrote {md_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if is_in_ci():
|
||||
print("Skipping bench_norm_impls.py in CI")
|
||||
sys.exit(0)
|
||||
main()
|
||||
@@ -0,0 +1,190 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import (
|
||||
DEFAULT_DEVICE,
|
||||
DEFAULT_DTYPE,
|
||||
get_benchmark_range,
|
||||
run_benchmark_no_cudagraph,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=13, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
MAX_SEQ_LEN = 131072
|
||||
ROPE_BASE = 10000.0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CaseSpec:
|
||||
name: str
|
||||
batch_size: int
|
||||
num_tokens: int
|
||||
num_heads: int
|
||||
head_dim: int
|
||||
rope_dim: int
|
||||
is_neox: bool
|
||||
|
||||
|
||||
BENCH_CASES = (
|
||||
CaseSpec("flux_1024", 1, 4096, 24, 128, 128, False),
|
||||
CaseSpec("qwen_image_1024", 1, 4096, 32, 128, 128, False),
|
||||
CaseSpec("qwen_image_partial", 1, 4096, 32, 128, 64, False),
|
||||
# Z-Image-Turbo default 1024x1024 config: dim=3840, num_heads=30 -> head_dim=128.
|
||||
CaseSpec("zimage_1024", 1, 4096, 30, 128, 128, False),
|
||||
CaseSpec("batch2_medium", 2, 2048, 24, 128, 128, False),
|
||||
)
|
||||
CASE_BY_NAME = {case.name: case for case in BENCH_CASES}
|
||||
CASE_NAMES = get_benchmark_range(
|
||||
full_range=[case.name for case in BENCH_CASES],
|
||||
ci_range=[case.name for case in BENCH_CASES],
|
||||
)
|
||||
LINE_VALS = ["split", "fused"]
|
||||
LINE_NAMES = ["JIT QKNorm + FlashInfer RoPE", "SGL JIT Fused QKNorm+RoPE"]
|
||||
STYLES = [("red", "-"), ("blue", "--")]
|
||||
|
||||
|
||||
def create_cos_sin_cache(
|
||||
rotary_dim: int,
|
||||
max_position: int = MAX_SEQ_LEN,
|
||||
base: float = ROPE_BASE,
|
||||
) -> torch.Tensor:
|
||||
inv_freq = 1.0 / (
|
||||
base
|
||||
** (
|
||||
torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=DEFAULT_DEVICE)
|
||||
/ rotary_dim
|
||||
)
|
||||
)
|
||||
t = torch.arange(max_position, dtype=torch.float32, device=DEFAULT_DEVICE)
|
||||
freqs = torch.einsum("i,j->ij", t, inv_freq)
|
||||
return torch.cat((freqs.cos(), freqs.sin()), dim=-1)
|
||||
|
||||
|
||||
def make_inputs(case: CaseSpec) -> dict[str, torch.Tensor | bool]:
|
||||
seed = (
|
||||
case.batch_size * 1_000_003
|
||||
+ case.num_tokens * 8191
|
||||
+ case.num_heads * 127
|
||||
+ case.head_dim * 17
|
||||
+ case.rope_dim
|
||||
)
|
||||
generator = torch.Generator(device=DEFAULT_DEVICE)
|
||||
generator.manual_seed(seed)
|
||||
return {
|
||||
"q": torch.randn(
|
||||
case.batch_size * case.num_tokens,
|
||||
case.num_heads,
|
||||
case.head_dim,
|
||||
device=DEFAULT_DEVICE,
|
||||
dtype=DEFAULT_DTYPE,
|
||||
generator=generator,
|
||||
),
|
||||
"k": torch.randn(
|
||||
case.batch_size * case.num_tokens,
|
||||
case.num_heads,
|
||||
case.head_dim,
|
||||
device=DEFAULT_DEVICE,
|
||||
dtype=DEFAULT_DTYPE,
|
||||
generator=generator,
|
||||
),
|
||||
"q_weight": torch.randn(
|
||||
case.head_dim,
|
||||
device=DEFAULT_DEVICE,
|
||||
dtype=DEFAULT_DTYPE,
|
||||
generator=generator,
|
||||
),
|
||||
"k_weight": torch.randn(
|
||||
case.head_dim,
|
||||
device=DEFAULT_DEVICE,
|
||||
dtype=DEFAULT_DTYPE,
|
||||
generator=generator,
|
||||
),
|
||||
"positions": torch.randint(
|
||||
0,
|
||||
MAX_SEQ_LEN,
|
||||
(case.batch_size * case.num_tokens,),
|
||||
device=DEFAULT_DEVICE,
|
||||
dtype=torch.int64,
|
||||
generator=generator,
|
||||
),
|
||||
"cos_sin_cache": create_cos_sin_cache(case.rope_dim),
|
||||
"is_neox": case.is_neox,
|
||||
}
|
||||
|
||||
|
||||
def clone_inputs(
|
||||
inputs: dict[str, torch.Tensor | bool],
|
||||
) -> dict[str, torch.Tensor | bool]:
|
||||
out: dict[str, torch.Tensor | bool] = {}
|
||||
for key, value in inputs.items():
|
||||
out[key] = value.clone() if isinstance(value, torch.Tensor) else value
|
||||
return out
|
||||
|
||||
|
||||
def split_qknorm_rope(inputs: dict[str, torch.Tensor | bool]) -> None:
|
||||
from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace
|
||||
|
||||
from sglang.jit_kernel.norm import fused_inplace_qknorm
|
||||
|
||||
q = inputs["q"]
|
||||
k = inputs["k"]
|
||||
q_weight = inputs["q_weight"]
|
||||
k_weight = inputs["k_weight"]
|
||||
positions = inputs["positions"]
|
||||
cos_sin_cache = inputs["cos_sin_cache"]
|
||||
is_neox = bool(inputs["is_neox"])
|
||||
|
||||
fused_inplace_qknorm(q, k, q_weight, k_weight)
|
||||
apply_rope_with_cos_sin_cache_inplace(
|
||||
positions=positions,
|
||||
query=q.view(q.shape[0], -1),
|
||||
key=k.view(k.shape[0], -1),
|
||||
head_size=q.shape[-1],
|
||||
cos_sin_cache=cos_sin_cache,
|
||||
is_neox=is_neox,
|
||||
)
|
||||
|
||||
|
||||
def fused_qknorm_rope(inputs: dict[str, torch.Tensor | bool]) -> None:
|
||||
from sglang.jit_kernel.diffusion.qknorm_rope import fused_inplace_qknorm_rope
|
||||
|
||||
fused_inplace_qknorm_rope(
|
||||
inputs["q"],
|
||||
inputs["k"],
|
||||
inputs["q_weight"],
|
||||
inputs["k_weight"],
|
||||
inputs["cos_sin_cache"],
|
||||
inputs["positions"],
|
||||
is_neox=bool(inputs["is_neox"]),
|
||||
rope_dim=inputs["cos_sin_cache"].shape[-1],
|
||||
)
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["case_name"],
|
||||
x_vals=CASE_NAMES,
|
||||
line_arg="provider",
|
||||
line_vals=LINE_VALS,
|
||||
line_names=LINE_NAMES,
|
||||
styles=STYLES,
|
||||
ylabel="us",
|
||||
plot_name="diffusion-qknorm-rope-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(case_name: str, provider: str) -> Tuple[float, float, float]:
|
||||
case = CASE_BY_NAME[case_name]
|
||||
inputs = make_inputs(case)
|
||||
fn = split_qknorm_rope if provider == "split" else fused_qknorm_rope
|
||||
return run_benchmark_no_cudagraph(lambda: fn(inputs))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("Running diffusion qknorm + rope performance benchmark...")
|
||||
benchmark.run(print_data=True)
|
||||
@@ -0,0 +1,184 @@
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph
|
||||
from sglang.jit_kernel.diffusion.triton.norm import norm_infer
|
||||
from sglang.jit_kernel.diffusion.triton.scale_shift import (
|
||||
fuse_layernorm_scale_shift_gate_select01_kernel,
|
||||
fuse_residual_layernorm_scale_shift_gate_select01_kernel,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.utils import is_in_ci
|
||||
|
||||
register_cuda_ci(est_time=13, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
if is_in_ci():
|
||||
B_RANGE, S_RANGE, D_RANGE = [1], [128], [3072]
|
||||
else:
|
||||
B_RANGE, S_RANGE, D_RANGE = [1, 2], [128, 512, 2048], [1024, 1536, 3072]
|
||||
|
||||
DTYPE = torch.bfloat16
|
||||
DEVICE = "cuda"
|
||||
EPS = 1e-6
|
||||
LINE_VALS = ["split", "fused"]
|
||||
LINE_NAMES = ["Triton Norm + Torch Select", "Fused Triton"]
|
||||
STYLES = [("red", "-"), ("blue", "--")]
|
||||
CONFIG = [(b, s, d) for b in B_RANGE for s in S_RANGE for d in D_RANGE]
|
||||
|
||||
|
||||
def _make_common_inputs(batch_size: int, seq_len: int, hidden_size: int):
|
||||
x = torch.randn(batch_size, seq_len, hidden_size, dtype=DTYPE, device=DEVICE)
|
||||
weight = torch.randn(hidden_size, dtype=DTYPE, device=DEVICE)
|
||||
bias = torch.randn(hidden_size, dtype=DTYPE, device=DEVICE)
|
||||
index = torch.randint(0, 2, (batch_size, seq_len), dtype=torch.int32, device=DEVICE)
|
||||
scale0 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE)
|
||||
shift0 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE)
|
||||
gate0 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE)
|
||||
scale1 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE)
|
||||
shift1 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE)
|
||||
gate1 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE)
|
||||
return x, weight, bias, index, scale0, shift0, gate0, scale1, shift1, gate1
|
||||
|
||||
|
||||
def _apply_select01_modulation(
|
||||
x: torch.Tensor,
|
||||
scale0: torch.Tensor,
|
||||
shift0: torch.Tensor,
|
||||
gate0: torch.Tensor,
|
||||
scale1: torch.Tensor,
|
||||
shift1: torch.Tensor,
|
||||
gate1: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
):
|
||||
idx = index.bool().unsqueeze(-1)
|
||||
scale = torch.where(idx, scale1.unsqueeze(1), scale0.unsqueeze(1))
|
||||
shift = torch.where(idx, shift1.unsqueeze(1), shift0.unsqueeze(1))
|
||||
gate = torch.where(idx, gate1.unsqueeze(1), gate0.unsqueeze(1))
|
||||
return x * (1 + scale) + shift, gate
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["B", "S", "D"],
|
||||
x_vals=CONFIG,
|
||||
line_arg="provider",
|
||||
line_vals=LINE_VALS,
|
||||
line_names=LINE_NAMES,
|
||||
styles=STYLES,
|
||||
ylabel="us",
|
||||
plot_name="qwen_image_layernorm_scale_shift_gate_select01",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def bench_layernorm_scale_shift_gate_select01(
|
||||
B: int, S: int, D: int, provider: str
|
||||
) -> Tuple[float, float, float]:
|
||||
x, weight, bias, index, scale0, shift0, gate0, scale1, shift1, gate1 = (
|
||||
_make_common_inputs(B, S, D)
|
||||
)
|
||||
|
||||
if provider == "split":
|
||||
|
||||
def fn():
|
||||
normalized = norm_infer(
|
||||
x.view(-1, x.shape[-1]),
|
||||
weight,
|
||||
bias,
|
||||
eps=EPS,
|
||||
is_rms_norm=False,
|
||||
).view_as(x)
|
||||
return _apply_select01_modulation(
|
||||
normalized, scale0, shift0, gate0, scale1, shift1, gate1, index
|
||||
)
|
||||
|
||||
else:
|
||||
|
||||
def fn():
|
||||
return fuse_layernorm_scale_shift_gate_select01_kernel(
|
||||
x,
|
||||
weight=weight,
|
||||
bias=bias,
|
||||
scale0=scale0,
|
||||
shift0=shift0,
|
||||
gate0=gate0,
|
||||
scale1=scale1,
|
||||
shift1=shift1,
|
||||
gate1=gate1,
|
||||
index=index,
|
||||
eps=EPS,
|
||||
)
|
||||
|
||||
return run_benchmark_no_cudagraph(fn)
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["B", "S", "D"],
|
||||
x_vals=CONFIG,
|
||||
line_arg="provider",
|
||||
line_vals=LINE_VALS,
|
||||
line_names=LINE_NAMES,
|
||||
styles=STYLES,
|
||||
ylabel="us",
|
||||
plot_name="qwen_image_residual_layernorm_scale_shift_gate_select01",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def bench_residual_layernorm_scale_shift_gate_select01(
|
||||
B: int, S: int, D: int, provider: str
|
||||
) -> Tuple[float, float, float]:
|
||||
x, weight, bias, index, scale0, shift0, gate0, scale1, shift1, gate1 = (
|
||||
_make_common_inputs(B, S, D)
|
||||
)
|
||||
residual = torch.randn_like(x)
|
||||
residual_gate = torch.randn_like(x)
|
||||
|
||||
if provider == "split":
|
||||
|
||||
def fn():
|
||||
residual_out = residual + residual_gate * x
|
||||
normalized = norm_infer(
|
||||
residual_out.view(-1, residual_out.shape[-1]),
|
||||
weight,
|
||||
bias,
|
||||
eps=EPS,
|
||||
is_rms_norm=False,
|
||||
).view_as(residual_out)
|
||||
return _apply_select01_modulation(
|
||||
normalized, scale0, shift0, gate0, scale1, shift1, gate1, index
|
||||
)
|
||||
|
||||
else:
|
||||
|
||||
def fn():
|
||||
return fuse_residual_layernorm_scale_shift_gate_select01_kernel(
|
||||
x,
|
||||
residual=residual,
|
||||
residual_gate=residual_gate,
|
||||
weight=weight,
|
||||
bias=bias,
|
||||
scale0=scale0,
|
||||
shift0=shift0,
|
||||
gate0=gate0,
|
||||
scale1=scale1,
|
||||
shift1=shift1,
|
||||
gate1=gate1,
|
||||
index=index,
|
||||
eps=EPS,
|
||||
)
|
||||
|
||||
return run_benchmark_no_cudagraph(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(f"\n{'=' * 80}")
|
||||
print("Benchmark: qwen_image layernorm + scale_shift_gate_select01")
|
||||
print(f"{'=' * 80}\n")
|
||||
bench_layernorm_scale_shift_gate_select01.run(print_data=True)
|
||||
|
||||
print(f"\n{'=' * 80}")
|
||||
print("Benchmark: qwen_image residual + layernorm + scale_shift_gate_select01")
|
||||
print(f"{'=' * 80}\n")
|
||||
bench_residual_layernorm_scale_shift_gate_select01.run(print_data=True)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,345 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.kv_canary.utils import (
|
||||
POOL_AXIS,
|
||||
SWA_WINDOW,
|
||||
BenchCase,
|
||||
build_fast_matrix_cases,
|
||||
build_full_matrix_cases,
|
||||
naive_cumsum_fn,
|
||||
)
|
||||
from sglang.jit_kernel.benchmark.utils import (
|
||||
DEFAULT_DEVICE,
|
||||
get_benchmark_range,
|
||||
run_benchmark,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary.plan import launch_canary_plan_kernels
|
||||
from sglang.jit_kernel.kv_canary.verify import VerifyPlan
|
||||
from sglang.jit_kernel.kv_canary.write import WritePlan
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=900, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
_TOTAL_TOKENS_AXIS: list[int] = [256, 4096, 65536, 262144]
|
||||
_TOTAL_TOKENS_BS_AXIS: list[int] = [1, 32, 256]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True, kw_only=True)
|
||||
class _TotalTokensBenchCase:
|
||||
bs: int
|
||||
total_tokens: int
|
||||
pool_kind: str
|
||||
|
||||
|
||||
def _build_total_tokens_cases() -> list[_TotalTokensBenchCase]:
|
||||
cases: list[_TotalTokensBenchCase] = []
|
||||
for bs in _TOTAL_TOKENS_BS_AXIS:
|
||||
for total_tokens in _TOTAL_TOKENS_AXIS:
|
||||
if total_tokens < bs:
|
||||
continue
|
||||
for pool_kind in POOL_AXIS:
|
||||
cases.append(
|
||||
_TotalTokensBenchCase(
|
||||
bs=bs, total_tokens=total_tokens, pool_kind=pool_kind
|
||||
)
|
||||
)
|
||||
return cases
|
||||
|
||||
|
||||
_POOL_CAPACITY_VERIFY_CAP_AXIS: list[int] = [16384, 262144, 1398028]
|
||||
_POOL_CAPACITY_BS_AXIS: list[int] = [1, 4, 32]
|
||||
_POOL_CAPACITY_PREFIX_LEN: int = 512
|
||||
_POOL_CAPACITY_BS_PADDED_AXIS: list[Optional[int]] = [None, 4096]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True, kw_only=True)
|
||||
class _PoolCapacityBenchCase:
|
||||
"""One pool-capacity bench point.
|
||||
|
||||
Attributes:
|
||||
bs: Number of active (non-padding) requests in the launch.
|
||||
bs_padded: Total request-axis size of the input tensors. ``None`` means
|
||||
no padding (``bs_padded == bs``); a concrete value pads ``req_pool_indices``
|
||||
with ``REQ_POOL_IDX_PADDING`` sentinels in rows ``[bs, bs_padded)``.
|
||||
prefix_len: Per-active-request prefix length.
|
||||
verify_capacity: Plan tensor row capacity.
|
||||
pool_kind: "full".
|
||||
"""
|
||||
|
||||
bs: int
|
||||
bs_padded: Optional[int]
|
||||
prefix_len: int
|
||||
verify_capacity: int
|
||||
pool_kind: str
|
||||
|
||||
|
||||
def _build_pool_capacity_cases() -> list[_PoolCapacityBenchCase]:
|
||||
cases: list[_PoolCapacityBenchCase] = []
|
||||
for bs in _POOL_CAPACITY_BS_AXIS:
|
||||
for verify_capacity in _POOL_CAPACITY_VERIFY_CAP_AXIS:
|
||||
for bs_padded in _POOL_CAPACITY_BS_PADDED_AXIS:
|
||||
if bs_padded is not None and bs_padded < bs:
|
||||
continue
|
||||
cases.append(
|
||||
_PoolCapacityBenchCase(
|
||||
bs=bs,
|
||||
bs_padded=bs_padded,
|
||||
prefix_len=_POOL_CAPACITY_PREFIX_LEN,
|
||||
verify_capacity=verify_capacity,
|
||||
pool_kind="full",
|
||||
)
|
||||
)
|
||||
return cases
|
||||
|
||||
|
||||
_X_NAMES_MATRIX = ["scenario", "bs", "prefix_len", "mode", "extend_len", "pool_kind"]
|
||||
|
||||
|
||||
def _cases_to_matrix_x_vals(
|
||||
cases: list[BenchCase],
|
||||
) -> list[tuple[str, int, int, str, int, str]]:
|
||||
return [
|
||||
(c.scenario, c.bs, c.prefix_len, c.mode, c.extend_len, c.pool_kind)
|
||||
for c in cases
|
||||
]
|
||||
|
||||
|
||||
_X_VALS_MATRIX = _cases_to_matrix_x_vals(
|
||||
get_benchmark_range(
|
||||
full_range=build_full_matrix_cases(),
|
||||
ci_range=build_fast_matrix_cases(),
|
||||
)
|
||||
)
|
||||
|
||||
_X_NAMES_TT = ["bs", "total_tokens", "pool_kind"]
|
||||
_X_VALS_TT = [(c.bs, c.total_tokens, c.pool_kind) for c in _build_total_tokens_cases()]
|
||||
|
||||
_X_NAMES_PC = ["bs", "bs_padded", "prefix_len", "verify_capacity", "pool_kind"]
|
||||
_X_VALS_PC = [
|
||||
(
|
||||
c.bs,
|
||||
c.bs_padded if c.bs_padded is not None else c.bs,
|
||||
c.prefix_len,
|
||||
c.verify_capacity,
|
||||
c.pool_kind,
|
||||
)
|
||||
for c in _build_pool_capacity_cases()
|
||||
]
|
||||
|
||||
|
||||
def _build_plan_inputs(
|
||||
*,
|
||||
bs: int,
|
||||
prefix_len: int,
|
||||
extend_len: int,
|
||||
pool_kind: str,
|
||||
device: torch.device,
|
||||
verify_capacity_override: Optional[int] = None,
|
||||
bs_padded: Optional[int] = None,
|
||||
) -> dict:
|
||||
swa_window_size = SWA_WINDOW if pool_kind == "swa_window_128" else 0
|
||||
verify_per_req = min(prefix_len, SWA_WINDOW) if swa_window_size > 0 else prefix_len
|
||||
if verify_capacity_override is not None:
|
||||
verify_capacity = max(1, verify_capacity_override)
|
||||
else:
|
||||
verify_capacity = max(1, bs * verify_per_req)
|
||||
|
||||
effective_bs = bs_padded if bs_padded is not None else bs
|
||||
if effective_bs < bs:
|
||||
raise ValueError(f"kv-canary bench: bs_padded={bs_padded} must be >= bs={bs}")
|
||||
write_req_capacity = max(1, effective_bs)
|
||||
|
||||
verify_plan = VerifyPlan.allocate(verify_capacity=verify_capacity, device=device)
|
||||
write_plan = WritePlan.allocate(
|
||||
write_req_capacity=write_req_capacity, device=device
|
||||
)
|
||||
|
||||
req_pool_indices = torch.zeros(effective_bs, dtype=torch.int64, device=device)
|
||||
req_pool_indices[:bs] = torch.arange(1, bs + 1, dtype=torch.int64, device=device)
|
||||
prefix_lens = torch.zeros(effective_bs, dtype=torch.int64, device=device)
|
||||
prefix_lens[:bs] = prefix_len
|
||||
extend_seq_lens = torch.zeros(effective_bs, dtype=torch.int64, device=device)
|
||||
extend_seq_lens[:bs] = extend_len
|
||||
|
||||
max_seq_len = max(prefix_len + extend_len, 1)
|
||||
req_to_token_rows = effective_bs + 1
|
||||
req_to_token = torch.zeros(
|
||||
req_to_token_rows,
|
||||
max_seq_len,
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
if bs > 0:
|
||||
row_idx = torch.arange(1, bs + 1, dtype=torch.int32, device=device).unsqueeze(1)
|
||||
col_idx = torch.arange(max_seq_len, dtype=torch.int32, device=device).unsqueeze(
|
||||
0
|
||||
)
|
||||
req_to_token[1 : bs + 1] = (row_idx - 1) * max_seq_len + col_idx
|
||||
|
||||
if swa_window_size > 0:
|
||||
full_pool_size = effective_bs * max_seq_len + 1
|
||||
full_to_swa: Optional[torch.Tensor] = torch.arange(
|
||||
full_pool_size + 1, dtype=torch.int64, device=device
|
||||
)
|
||||
full_to_swa[-1] = -1
|
||||
else:
|
||||
full_to_swa = None
|
||||
|
||||
return dict(
|
||||
verify_plan_out=verify_plan,
|
||||
write_plan_out=write_plan,
|
||||
req_pool_indices=req_pool_indices,
|
||||
prefix_lens=prefix_lens,
|
||||
extend_seq_lens=extend_seq_lens,
|
||||
req_to_token=req_to_token,
|
||||
swa_window_size=swa_window_size,
|
||||
full_to_swa_index_mapping=full_to_swa,
|
||||
verify_capacity=int(verify_plan.verify_slot_indices.shape[0]),
|
||||
)
|
||||
|
||||
|
||||
def _make_plan_callable(inputs: dict):
|
||||
def fn() -> None:
|
||||
launch_canary_plan_kernels(
|
||||
verify_plan_out=inputs["verify_plan_out"],
|
||||
write_plan_out=inputs["write_plan_out"],
|
||||
req_pool_indices=inputs["req_pool_indices"],
|
||||
prefix_lens=inputs["prefix_lens"],
|
||||
extend_seq_lens=inputs["extend_seq_lens"],
|
||||
req_to_token=inputs["req_to_token"],
|
||||
swa_window_size=inputs["swa_window_size"],
|
||||
full_to_swa_index_mapping=inputs["full_to_swa_index_mapping"],
|
||||
verify_capacity=inputs["verify_capacity"],
|
||||
req_to_verify_expected_tokens=None,
|
||||
req_to_verify_expected_tokens_valid_lens=None,
|
||||
kv_token_id_vs_position_offset=0,
|
||||
)
|
||||
|
||||
return fn
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=_X_NAMES_MATRIX,
|
||||
x_vals=_X_VALS_MATRIX,
|
||||
line_arg="provider",
|
||||
line_vals=["canary", "naive"],
|
||||
line_names=["canary_plan_step", "naive torch.cumsum"],
|
||||
styles=[("blue", "-"), ("red", "--")],
|
||||
ylabel="us",
|
||||
plot_name="kv-canary-plan-matrix-perf",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark_matrix(
|
||||
scenario: str,
|
||||
bs: int,
|
||||
prefix_len: int,
|
||||
mode: str,
|
||||
extend_len: int,
|
||||
pool_kind: str,
|
||||
provider: str,
|
||||
) -> Tuple[float, float, float]:
|
||||
del scenario
|
||||
del mode
|
||||
|
||||
device = torch.device(DEFAULT_DEVICE)
|
||||
if provider == "canary":
|
||||
inputs = _build_plan_inputs(
|
||||
bs=bs,
|
||||
prefix_len=prefix_len,
|
||||
extend_len=extend_len,
|
||||
pool_kind=pool_kind,
|
||||
device=device,
|
||||
)
|
||||
fn = _make_plan_callable(inputs)
|
||||
else:
|
||||
fn = naive_cumsum_fn(bs=bs, device=device)
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=_X_NAMES_TT,
|
||||
x_vals=_X_VALS_TT,
|
||||
line_arg="provider",
|
||||
line_vals=["canary", "naive"],
|
||||
line_names=["canary_plan_step", "naive torch.cumsum"],
|
||||
styles=[("blue", "-"), ("red", "--")],
|
||||
ylabel="us",
|
||||
plot_name="kv-canary-plan-total-tokens-perf",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark_total_tokens(
|
||||
bs: int,
|
||||
total_tokens: int,
|
||||
pool_kind: str,
|
||||
provider: str,
|
||||
) -> Tuple[float, float, float]:
|
||||
device = torch.device(DEFAULT_DEVICE)
|
||||
per_req_prefix = max(1, total_tokens // max(bs, 1))
|
||||
if provider == "canary":
|
||||
inputs = _build_plan_inputs(
|
||||
bs=bs,
|
||||
prefix_len=per_req_prefix,
|
||||
extend_len=1,
|
||||
pool_kind=pool_kind,
|
||||
device=device,
|
||||
)
|
||||
fn = _make_plan_callable(inputs)
|
||||
else:
|
||||
fn = naive_cumsum_fn(bs=bs, device=device)
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=_X_NAMES_PC,
|
||||
x_vals=_X_VALS_PC,
|
||||
line_arg="provider",
|
||||
line_vals=["canary", "naive"],
|
||||
line_names=["canary_plan_step", "naive torch.cumsum"],
|
||||
styles=[("blue", "-"), ("red", "--")],
|
||||
ylabel="us",
|
||||
plot_name="kv-canary-plan-pool-capacity-perf",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark_pool_capacity(
|
||||
bs: int,
|
||||
bs_padded: int,
|
||||
prefix_len: int,
|
||||
verify_capacity: int,
|
||||
pool_kind: str,
|
||||
provider: str,
|
||||
) -> Tuple[float, float, float]:
|
||||
device = torch.device(DEFAULT_DEVICE)
|
||||
if provider == "canary":
|
||||
inputs = _build_plan_inputs(
|
||||
bs=bs,
|
||||
prefix_len=prefix_len,
|
||||
extend_len=1,
|
||||
pool_kind=pool_kind,
|
||||
device=device,
|
||||
verify_capacity_override=verify_capacity,
|
||||
bs_padded=bs_padded,
|
||||
)
|
||||
fn = _make_plan_callable(inputs)
|
||||
else:
|
||||
fn = naive_cumsum_fn(bs=bs, device=device)
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark_matrix.run(print_data=True)
|
||||
benchmark_total_tokens.run(print_data=True)
|
||||
benchmark_pool_capacity.run(print_data=True)
|
||||
@@ -0,0 +1,95 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import (
|
||||
DEFAULT_DEVICE,
|
||||
get_benchmark_range,
|
||||
run_benchmark_no_cudagraph,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary.scatter_req_token_ids import (
|
||||
launch_scatter_req_token_ids_kernel,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=180, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
_BS_AXIS_FULL: list[int] = [1, 8, 64, 256]
|
||||
_SEQ_LEN_AXIS_FULL: list[int] = [128, 512, 2048, 8192]
|
||||
_BS_AXIS_CI: list[int] = [1, 64]
|
||||
_SEQ_LEN_AXIS_CI: list[int] = [512, 2048]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True, kw_only=True)
|
||||
class _BenchCase:
|
||||
bs: int
|
||||
seq_len: int
|
||||
|
||||
|
||||
def _build_cases() -> list[_BenchCase]:
|
||||
bs_axis = get_benchmark_range(full_range=_BS_AXIS_FULL, ci_range=_BS_AXIS_CI)
|
||||
seq_axis = get_benchmark_range(
|
||||
full_range=_SEQ_LEN_AXIS_FULL, ci_range=_SEQ_LEN_AXIS_CI
|
||||
)
|
||||
return [
|
||||
_BenchCase(bs=bs, seq_len=seq_len) for bs in bs_axis for seq_len in seq_axis
|
||||
]
|
||||
|
||||
|
||||
_X_NAMES = ["bs", "seq_len"]
|
||||
_X_VALS = [(c.bs, c.seq_len) for c in _build_cases()]
|
||||
|
||||
|
||||
def _build_inputs(*, bs: int, seq_len: int, device: torch.device) -> dict:
|
||||
max_reqs = max(bs + 1, 4)
|
||||
max_context_len = max(seq_len + 1, 1)
|
||||
total_tokens = bs * seq_len
|
||||
|
||||
flat = torch.randint(
|
||||
low=0,
|
||||
high=1 << 30,
|
||||
size=(total_tokens,),
|
||||
dtype=torch.int64,
|
||||
device=device,
|
||||
)
|
||||
lens = torch.full((bs,), seq_len, dtype=torch.int64, device=device)
|
||||
offsets = torch.zeros(bs + 1, dtype=torch.int64, device=device)
|
||||
offsets[1:] = torch.cumsum(lens, dim=0)
|
||||
|
||||
req_pool_indices = torch.arange(1, bs + 1, dtype=torch.int64, device=device)
|
||||
pool = torch.zeros((max_reqs, max_context_len), dtype=torch.int32, device=device)
|
||||
return dict(
|
||||
flat_in=flat,
|
||||
offsets=offsets,
|
||||
req_pool_indices=req_pool_indices,
|
||||
pool_out=pool,
|
||||
)
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=_X_NAMES,
|
||||
x_vals=_X_VALS,
|
||||
line_arg="provider",
|
||||
line_vals=["triton"],
|
||||
line_names=["Triton"],
|
||||
styles=[("blue", "-")],
|
||||
ylabel="time (us)",
|
||||
plot_name="kv-canary-scatter-req-token-ids",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(bs: int, seq_len: int, provider: str) -> tuple[float, float, float]:
|
||||
inputs = _build_inputs(bs=bs, seq_len=seq_len, device=torch.device(DEFAULT_DEVICE))
|
||||
return run_benchmark_no_cudagraph(
|
||||
lambda: launch_scatter_req_token_ids_kernel(**inputs)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run(print_data=True)
|
||||
@@ -0,0 +1,314 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.kv_canary.utils import (
|
||||
RING_CAPACITY,
|
||||
SWA_WINDOW,
|
||||
BenchCase,
|
||||
build_fast_matrix_cases,
|
||||
build_full_matrix_cases,
|
||||
cases_to_x_vals,
|
||||
make_real_kv_sources,
|
||||
naive_slot_copy_fn,
|
||||
)
|
||||
from sglang.jit_kernel.benchmark.utils import (
|
||||
DEFAULT_DEVICE,
|
||||
get_benchmark_range,
|
||||
run_benchmark,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary import consts
|
||||
from sglang.jit_kernel.kv_canary.verify import (
|
||||
CANARY_SLOT_BYTES,
|
||||
CanaryLaunchTag,
|
||||
RealKvSource,
|
||||
VerifyOrWriteContext,
|
||||
VerifyPlan,
|
||||
launch_canary_verify_kernel,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=900, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
_X_NAMES = [
|
||||
"scenario",
|
||||
"bs",
|
||||
"prefix_len",
|
||||
"mode",
|
||||
"extend_len",
|
||||
"pool_kind",
|
||||
"real_kv_kind",
|
||||
"hash_mode",
|
||||
]
|
||||
_X_VALS = cases_to_x_vals(
|
||||
get_benchmark_range(
|
||||
full_range=build_full_matrix_cases(),
|
||||
ci_range=build_fast_matrix_cases(),
|
||||
)
|
||||
)
|
||||
|
||||
_KERNEL_KIND_X_NAMES = ["kernel_kind_name"]
|
||||
_KERNEL_KIND_X_VALS = [(tag.name,) for tag in CanaryLaunchTag]
|
||||
|
||||
|
||||
def _verify_entry_count(case: BenchCase) -> int:
|
||||
if case.pool_kind == "swa_window_128":
|
||||
per_req = min(case.prefix_len, SWA_WINDOW)
|
||||
else:
|
||||
per_req = case.prefix_len
|
||||
return case.bs * per_req
|
||||
|
||||
|
||||
def _verify_num_slots(case: BenchCase) -> int:
|
||||
if case.pool_kind == "swa_window_128":
|
||||
per_req_slots = SWA_WINDOW
|
||||
else:
|
||||
per_req_slots = max(1, case.prefix_len)
|
||||
return max(2, case.bs * per_req_slots + 1)
|
||||
|
||||
|
||||
def _build_verify_inputs(case: BenchCase, *, device: torch.device) -> Tuple[
|
||||
torch.Tensor,
|
||||
VerifyPlan,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
tuple[RealKvSource, ...],
|
||||
]:
|
||||
total_entries = _verify_entry_count(case)
|
||||
capacity = max(1, total_entries)
|
||||
num_slots = _verify_num_slots(case)
|
||||
|
||||
canary_buf = torch.zeros(
|
||||
num_slots, CANARY_SLOT_BYTES, dtype=torch.uint8, device=device
|
||||
)
|
||||
|
||||
slot_indices = torch.empty(capacity, dtype=torch.int64, device=device)
|
||||
positions = torch.empty(capacity, dtype=torch.int64, device=device)
|
||||
prev_slots = torch.empty(capacity, dtype=torch.int64, device=device)
|
||||
if total_entries > 0:
|
||||
flat_idx = torch.arange(total_entries, device=device, dtype=torch.int64)
|
||||
per_req = total_entries // case.bs if case.bs > 0 else 0
|
||||
slot_indices[:total_entries] = (flat_idx % max(num_slots - 1, 1)).to(
|
||||
torch.int64
|
||||
)
|
||||
positions[:total_entries] = (flat_idx % max(per_req, 1)).to(torch.int64)
|
||||
is_head = (flat_idx % max(per_req, 1)) == 0
|
||||
prev_seq = (flat_idx - 1) % max(num_slots - 1, 1)
|
||||
prev_slots[:total_entries] = torch.where(
|
||||
is_head, torch.full_like(flat_idx, -1), prev_seq
|
||||
).to(torch.int64)
|
||||
if capacity > total_entries:
|
||||
slot_indices[total_entries:] = 0
|
||||
positions[total_entries:] = 0
|
||||
prev_slots[total_entries:] = -1
|
||||
|
||||
num_valid = torch.tensor([total_entries], dtype=torch.int32, device=device)
|
||||
enable = torch.ones(1, dtype=torch.int32, device=device)
|
||||
expected_input_ids = torch.full((capacity,), -1, dtype=torch.int64, device=device)
|
||||
plan = VerifyPlan(
|
||||
verify_slot_indices=slot_indices,
|
||||
verify_expected_tokens=expected_input_ids,
|
||||
verify_expected_positions=positions,
|
||||
verify_prev_slot_indices=prev_slots,
|
||||
verify_num_valid=num_valid,
|
||||
enable=enable,
|
||||
)
|
||||
|
||||
violation_ring = torch.zeros(
|
||||
RING_CAPACITY, consts.VIOLATION_FIELDS, dtype=torch.int64, device=device
|
||||
)
|
||||
violation_write_index = torch.zeros(1, dtype=torch.int32, device=device)
|
||||
slot_run_counter = torch.zeros(1, dtype=torch.int64, device=device)
|
||||
kernel_run_counter = torch.zeros(1, dtype=torch.int64, device=device)
|
||||
enable_chain_position_assert = torch.ones(1, dtype=torch.int32, device=device)
|
||||
|
||||
real_kv_sources = make_real_kv_sources(
|
||||
kind=case.real_kv_kind, num_slots=num_slots, device=device
|
||||
)
|
||||
|
||||
return (
|
||||
canary_buf,
|
||||
plan,
|
||||
violation_ring,
|
||||
violation_write_index,
|
||||
slot_run_counter,
|
||||
kernel_run_counter,
|
||||
enable_chain_position_assert,
|
||||
real_kv_sources,
|
||||
)
|
||||
|
||||
|
||||
def _build_context(
|
||||
*,
|
||||
canary_buf: torch.Tensor,
|
||||
violation_ring: torch.Tensor,
|
||||
violation_write_index: torch.Tensor,
|
||||
slot_run_counter: torch.Tensor,
|
||||
kernel_run_counter: torch.Tensor,
|
||||
enable_chain_position_assert: torch.Tensor,
|
||||
real_kv_sources: tuple[RealKvSource, ...],
|
||||
kernel_kind: CanaryLaunchTag,
|
||||
hash_mode: consts.RealKvHashMode,
|
||||
) -> VerifyOrWriteContext:
|
||||
return VerifyOrWriteContext(
|
||||
canary_buf=canary_buf,
|
||||
kernel_kind=kernel_kind,
|
||||
violation_ring=violation_ring,
|
||||
violation_write_index=violation_write_index,
|
||||
slot_run_counter=slot_run_counter,
|
||||
kernel_run_counter=kernel_run_counter,
|
||||
real_kv_sources=real_kv_sources,
|
||||
real_kv_hash_mode=hash_mode,
|
||||
enable_chain_position_assert=enable_chain_position_assert,
|
||||
)
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=_X_NAMES,
|
||||
x_vals=_X_VALS,
|
||||
line_arg="provider",
|
||||
line_vals=["canary", "naive"],
|
||||
line_names=["canary_verify_step", "naive index_copy_"],
|
||||
styles=[("blue", "-"), ("red", "--")],
|
||||
ylabel="us",
|
||||
plot_name="kv-canary-verify-perf",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(
|
||||
scenario: str,
|
||||
bs: int,
|
||||
prefix_len: int,
|
||||
mode: str,
|
||||
extend_len: int,
|
||||
pool_kind: str,
|
||||
real_kv_kind: str,
|
||||
hash_mode: str,
|
||||
provider: str,
|
||||
) -> Tuple[float, float, float]:
|
||||
case = BenchCase(
|
||||
scenario=scenario,
|
||||
bs=bs,
|
||||
prefix_len=prefix_len,
|
||||
mode=mode,
|
||||
extend_len=extend_len,
|
||||
pool_kind=pool_kind,
|
||||
real_kv_kind=real_kv_kind,
|
||||
hash_mode=hash_mode,
|
||||
)
|
||||
device = torch.device(DEFAULT_DEVICE)
|
||||
|
||||
if provider == "canary":
|
||||
(
|
||||
canary_buf,
|
||||
plan,
|
||||
violation_ring,
|
||||
violation_write_index,
|
||||
slot_run_counter,
|
||||
kernel_run_counter,
|
||||
enable_chain_position_assert,
|
||||
real_kv_sources,
|
||||
) = _build_verify_inputs(case, device=device)
|
||||
hash_mode_enum = consts.RealKvHashMode[case.hash_mode.upper()]
|
||||
context = _build_context(
|
||||
canary_buf=canary_buf,
|
||||
violation_ring=violation_ring,
|
||||
violation_write_index=violation_write_index,
|
||||
slot_run_counter=slot_run_counter,
|
||||
kernel_run_counter=kernel_run_counter,
|
||||
enable_chain_position_assert=enable_chain_position_assert,
|
||||
real_kv_sources=real_kv_sources,
|
||||
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
|
||||
hash_mode=hash_mode_enum,
|
||||
)
|
||||
|
||||
def fn() -> None:
|
||||
violation_write_index.zero_()
|
||||
launch_canary_verify_kernel(
|
||||
context=context,
|
||||
plan=plan,
|
||||
check_verify_expected_token=True,
|
||||
)
|
||||
|
||||
else:
|
||||
fn = naive_slot_copy_fn(total=_verify_entry_count(case), device=device)
|
||||
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=_KERNEL_KIND_X_NAMES,
|
||||
x_vals=_KERNEL_KIND_X_VALS,
|
||||
line_arg="provider",
|
||||
line_vals=["canary"],
|
||||
line_names=["canary_verify_step"],
|
||||
styles=[("blue", "-")],
|
||||
ylabel="us",
|
||||
plot_name="kv-canary-verify-kernel-kind-perf",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark_kernel_kind(
|
||||
kernel_kind_name: str,
|
||||
provider: str,
|
||||
) -> Tuple[float, float, float]:
|
||||
case = BenchCase(
|
||||
scenario="kernel_kind",
|
||||
bs=32,
|
||||
prefix_len=4096,
|
||||
mode="extend",
|
||||
extend_len=128,
|
||||
pool_kind="full",
|
||||
real_kv_kind="none",
|
||||
hash_mode="none",
|
||||
)
|
||||
device = torch.device(DEFAULT_DEVICE)
|
||||
|
||||
(
|
||||
canary_buf,
|
||||
plan,
|
||||
violation_ring,
|
||||
violation_write_index,
|
||||
slot_run_counter,
|
||||
kernel_run_counter,
|
||||
enable_chain_position_assert,
|
||||
real_kv_sources,
|
||||
) = _build_verify_inputs(case, device=device)
|
||||
kernel_kind = CanaryLaunchTag[kernel_kind_name]
|
||||
hash_mode_enum = consts.RealKvHashMode[case.hash_mode.upper()]
|
||||
context = _build_context(
|
||||
canary_buf=canary_buf,
|
||||
violation_ring=violation_ring,
|
||||
violation_write_index=violation_write_index,
|
||||
slot_run_counter=slot_run_counter,
|
||||
kernel_run_counter=kernel_run_counter,
|
||||
enable_chain_position_assert=enable_chain_position_assert,
|
||||
real_kv_sources=real_kv_sources,
|
||||
kernel_kind=kernel_kind,
|
||||
hash_mode=hash_mode_enum,
|
||||
)
|
||||
|
||||
def fn() -> None:
|
||||
violation_write_index.zero_()
|
||||
launch_canary_verify_kernel(
|
||||
context=context,
|
||||
plan=plan,
|
||||
check_verify_expected_token=True,
|
||||
)
|
||||
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run(print_data=True)
|
||||
benchmark_kernel_kind.run(print_data=True)
|
||||
@@ -0,0 +1,310 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.kv_canary.utils import (
|
||||
RING_CAPACITY,
|
||||
SWA_WINDOW,
|
||||
BenchCase,
|
||||
build_fast_matrix_cases,
|
||||
build_full_matrix_cases,
|
||||
cases_to_x_vals,
|
||||
make_real_kv_sources,
|
||||
naive_slot_copy_fn,
|
||||
)
|
||||
from sglang.jit_kernel.benchmark.utils import (
|
||||
DEFAULT_DEVICE,
|
||||
get_benchmark_range,
|
||||
run_benchmark,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary import consts
|
||||
from sglang.jit_kernel.kv_canary.verify import (
|
||||
CANARY_SLOT_BYTES,
|
||||
CanaryLaunchTag,
|
||||
VerifyOrWriteContext,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary.write import WritePlan, launch_canary_write_kernel
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=900, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
_X_NAMES = [
|
||||
"scenario",
|
||||
"bs",
|
||||
"prefix_len",
|
||||
"mode",
|
||||
"extend_len",
|
||||
"pool_kind",
|
||||
"real_kv_kind",
|
||||
"hash_mode",
|
||||
]
|
||||
_X_VALS = cases_to_x_vals(
|
||||
get_benchmark_range(
|
||||
full_range=build_full_matrix_cases(),
|
||||
ci_range=build_fast_matrix_cases(),
|
||||
)
|
||||
)
|
||||
|
||||
_KERNEL_KIND_X_NAMES = ["kernel_kind_name", "enable_write_verify_inputs_name"]
|
||||
_KERNEL_KIND_X_VALS = [
|
||||
(tag.name, str(enable)) for tag in CanaryLaunchTag for enable in (False, True)
|
||||
]
|
||||
|
||||
|
||||
def _write_entry_count(case: BenchCase) -> int:
|
||||
return case.bs * case.extend_len
|
||||
|
||||
|
||||
def _write_num_slots(case: BenchCase) -> int:
|
||||
per_req_slots = max(
|
||||
SWA_WINDOW if case.pool_kind == "swa_window_128" else 1,
|
||||
case.prefix_len + case.extend_len,
|
||||
)
|
||||
return max(2, case.bs * per_req_slots + 1)
|
||||
|
||||
|
||||
def _build_write_inputs(
|
||||
case: BenchCase, *, device: torch.device, mirror_expected_inputs: bool = False
|
||||
) -> dict:
|
||||
total_entries = _write_entry_count(case)
|
||||
num_tokens_padded = max(1, total_entries)
|
||||
|
||||
per_req_slots = max(
|
||||
SWA_WINDOW if case.pool_kind == "swa_window_128" else 1,
|
||||
case.prefix_len + case.extend_len,
|
||||
)
|
||||
num_slots = _write_num_slots(case)
|
||||
|
||||
canary_buf = torch.zeros(
|
||||
num_slots, CANARY_SLOT_BYTES, dtype=torch.uint8, device=device
|
||||
)
|
||||
|
||||
write_offsets = torch.zeros(case.bs + 1, dtype=torch.int64, device=device)
|
||||
if case.bs > 0:
|
||||
offsets_host = torch.arange(0, case.bs + 1, dtype=torch.int64) * case.extend_len
|
||||
write_offsets.copy_(offsets_host.to(device))
|
||||
|
||||
write_seed_slots = torch.empty(case.bs, dtype=torch.int64, device=device)
|
||||
if case.bs > 0:
|
||||
if case.prefix_len == 0:
|
||||
write_seed_slots.fill_(-1)
|
||||
else:
|
||||
per_req_stride = per_req_slots
|
||||
seeds = (
|
||||
torch.arange(case.bs, dtype=torch.int32, device=device) * per_req_stride
|
||||
+ case.prefix_len
|
||||
- 1
|
||||
)
|
||||
write_seed_slots.copy_(seeds.to(torch.int64))
|
||||
|
||||
write_num_valid_reqs = torch.tensor([case.bs], dtype=torch.int32, device=device)
|
||||
|
||||
plan = WritePlan(
|
||||
write_offsets=write_offsets,
|
||||
write_seed_slot_indices=write_seed_slots,
|
||||
write_num_valid_reqs=write_num_valid_reqs,
|
||||
)
|
||||
|
||||
input_ids = torch.zeros(num_tokens_padded, dtype=torch.int64, device=device)
|
||||
positions = torch.zeros(num_tokens_padded, dtype=torch.int64, device=device)
|
||||
out_cache_loc = torch.zeros(num_tokens_padded, dtype=torch.int64, device=device)
|
||||
if total_entries > 0:
|
||||
flat_idx = torch.arange(total_entries, device=device, dtype=torch.int64)
|
||||
per_req_idx = flat_idx % max(case.extend_len, 1)
|
||||
req_idx = flat_idx // max(case.extend_len, 1)
|
||||
per_req_stride = per_req_slots
|
||||
slots = (req_idx * per_req_stride + case.prefix_len + per_req_idx) % max(
|
||||
num_slots, 1
|
||||
)
|
||||
input_ids[:total_entries] = (flat_idx % 32768).to(torch.int64)
|
||||
positions[:total_entries] = (case.prefix_len + per_req_idx).to(torch.int64)
|
||||
out_cache_loc[:total_entries] = slots.to(torch.int64)
|
||||
|
||||
if case.pool_kind == "swa_window_128":
|
||||
full_to_swa = torch.arange(num_slots + 1, dtype=torch.int64, device=device)
|
||||
full_to_swa[-1] = -1
|
||||
out_cache_loc = full_to_swa[out_cache_loc]
|
||||
|
||||
if mirror_expected_inputs:
|
||||
expected_input_tokens = input_ids.clone()
|
||||
expected_input_positions = positions.clone()
|
||||
else:
|
||||
expected_input_tokens = None
|
||||
expected_input_positions = None
|
||||
|
||||
violation_ring = torch.zeros(
|
||||
RING_CAPACITY, consts.VIOLATION_FIELDS, dtype=torch.int64, device=device
|
||||
)
|
||||
violation_write_index = torch.zeros(1, dtype=torch.int32, device=device)
|
||||
slot_run_counter = torch.zeros(1, dtype=torch.int64, device=device)
|
||||
kernel_run_counter = torch.zeros(1, dtype=torch.int64, device=device)
|
||||
enable_chain_position_assert = torch.ones(1, dtype=torch.int32, device=device)
|
||||
|
||||
real_kv_sources = make_real_kv_sources(
|
||||
kind=case.real_kv_kind, num_slots=num_slots, device=device
|
||||
)
|
||||
|
||||
return dict(
|
||||
canary_buf=canary_buf,
|
||||
plan=plan,
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
expected_input_tokens=expected_input_tokens,
|
||||
expected_input_positions=expected_input_positions,
|
||||
violation_ring=violation_ring,
|
||||
violation_write_index=violation_write_index,
|
||||
slot_run_counter=slot_run_counter,
|
||||
kernel_run_counter=kernel_run_counter,
|
||||
enable_chain_position_assert=enable_chain_position_assert,
|
||||
real_kv_sources=real_kv_sources,
|
||||
)
|
||||
|
||||
|
||||
def _build_context(
|
||||
*,
|
||||
inputs: dict,
|
||||
kernel_kind: CanaryLaunchTag,
|
||||
hash_mode: consts.RealKvHashMode,
|
||||
) -> VerifyOrWriteContext:
|
||||
return VerifyOrWriteContext(
|
||||
canary_buf=inputs["canary_buf"],
|
||||
kernel_kind=kernel_kind,
|
||||
violation_ring=inputs["violation_ring"],
|
||||
violation_write_index=inputs["violation_write_index"],
|
||||
slot_run_counter=inputs["slot_run_counter"],
|
||||
kernel_run_counter=inputs["kernel_run_counter"],
|
||||
real_kv_sources=inputs["real_kv_sources"],
|
||||
real_kv_hash_mode=hash_mode,
|
||||
enable_chain_position_assert=inputs["enable_chain_position_assert"],
|
||||
)
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=_X_NAMES,
|
||||
x_vals=_X_VALS,
|
||||
line_arg="provider",
|
||||
line_vals=["canary", "naive"],
|
||||
line_names=["canary_write_step", "naive index_copy_"],
|
||||
styles=[("blue", "-"), ("red", "--")],
|
||||
ylabel="us",
|
||||
plot_name="kv-canary-write-perf",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(
|
||||
scenario: str,
|
||||
bs: int,
|
||||
prefix_len: int,
|
||||
mode: str,
|
||||
extend_len: int,
|
||||
pool_kind: str,
|
||||
real_kv_kind: str,
|
||||
hash_mode: str,
|
||||
provider: str,
|
||||
) -> Tuple[float, float, float]:
|
||||
case = BenchCase(
|
||||
scenario=scenario,
|
||||
bs=bs,
|
||||
prefix_len=prefix_len,
|
||||
mode=mode,
|
||||
extend_len=extend_len,
|
||||
pool_kind=pool_kind,
|
||||
real_kv_kind=real_kv_kind,
|
||||
hash_mode=hash_mode,
|
||||
)
|
||||
device = torch.device(DEFAULT_DEVICE)
|
||||
|
||||
if provider == "canary":
|
||||
inputs = _build_write_inputs(case, device=device)
|
||||
hash_mode_enum = consts.RealKvHashMode[case.hash_mode.upper()]
|
||||
context = _build_context(
|
||||
inputs=inputs,
|
||||
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
|
||||
hash_mode=hash_mode_enum,
|
||||
)
|
||||
|
||||
def fn() -> None:
|
||||
launch_canary_write_kernel(
|
||||
context=context,
|
||||
plan=inputs["plan"],
|
||||
input_ids=inputs["input_ids"],
|
||||
positions=inputs["positions"],
|
||||
out_cache_loc=inputs["out_cache_loc"],
|
||||
enable_write_input_assert=False,
|
||||
expected_input_tokens=inputs["expected_input_tokens"],
|
||||
expected_input_positions=inputs["expected_input_positions"],
|
||||
)
|
||||
|
||||
else:
|
||||
fn = naive_slot_copy_fn(total=_write_entry_count(case), device=device)
|
||||
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=_KERNEL_KIND_X_NAMES,
|
||||
x_vals=_KERNEL_KIND_X_VALS,
|
||||
line_arg="provider",
|
||||
line_vals=["canary"],
|
||||
line_names=["canary_write_step"],
|
||||
styles=[("blue", "-")],
|
||||
ylabel="us",
|
||||
plot_name="kv-canary-write-kernel-kind-perf",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark_kernel_kind(
|
||||
kernel_kind_name: str,
|
||||
enable_write_verify_inputs_name: str,
|
||||
provider: str,
|
||||
) -> Tuple[float, float, float]:
|
||||
case = BenchCase(
|
||||
scenario="kernel_kind",
|
||||
bs=32,
|
||||
prefix_len=4096,
|
||||
mode="extend",
|
||||
extend_len=128,
|
||||
pool_kind="full",
|
||||
real_kv_kind="none",
|
||||
hash_mode="none",
|
||||
)
|
||||
device = torch.device(DEFAULT_DEVICE)
|
||||
|
||||
enable_write_verify_inputs = enable_write_verify_inputs_name == "True"
|
||||
inputs = _build_write_inputs(
|
||||
case, device=device, mirror_expected_inputs=enable_write_verify_inputs
|
||||
)
|
||||
kernel_kind = CanaryLaunchTag[kernel_kind_name]
|
||||
hash_mode_enum = consts.RealKvHashMode[case.hash_mode.upper()]
|
||||
context = _build_context(
|
||||
inputs=inputs,
|
||||
kernel_kind=kernel_kind,
|
||||
hash_mode=hash_mode_enum,
|
||||
)
|
||||
|
||||
def fn() -> None:
|
||||
launch_canary_write_kernel(
|
||||
context=context,
|
||||
plan=inputs["plan"],
|
||||
input_ids=inputs["input_ids"],
|
||||
positions=inputs["positions"],
|
||||
out_cache_loc=inputs["out_cache_loc"],
|
||||
enable_write_input_assert=enable_write_verify_inputs,
|
||||
expected_input_tokens=inputs["expected_input_tokens"],
|
||||
expected_input_positions=inputs["expected_input_positions"],
|
||||
)
|
||||
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run(print_data=True)
|
||||
benchmark_kernel_kind.run(print_data=True)
|
||||
@@ -0,0 +1,260 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from typing import Tuple, Union
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.jit_kernel.dsv4 import compress_forward
|
||||
from sglang.jit_kernel.tests.deepseek_v4.common import (
|
||||
LegacyContext,
|
||||
PagedContext,
|
||||
make_legacy_context,
|
||||
make_paged_context,
|
||||
make_state_pool,
|
||||
to_seq_extend,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=30, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
Context = Union[LegacyContext, PagedContext]
|
||||
|
||||
# c128 input row layout: | kv | score | each [head_dim]
|
||||
HEAD_DIM = 512
|
||||
RATIO = 128
|
||||
ATOL = 5e-3
|
||||
RTOL = 5e-3
|
||||
|
||||
|
||||
def _gt_compress(
|
||||
kv_score_input_cpu: torch.Tensor, # [num_q, head_dim*2]
|
||||
ape_cpu: torch.Tensor, # [128, head_dim]
|
||||
P: int,
|
||||
head_dim: int,
|
||||
) -> torch.Tensor:
|
||||
"""fp64 reference for compress event at ragged position ``P`` (P % 128 == 127)."""
|
||||
lo = P - (RATIO - 1)
|
||||
kv = kv_score_input_cpu[lo : P + 1, :head_dim].double()
|
||||
sc = kv_score_input_cpu[lo : P + 1, head_dim:].double()
|
||||
return ((kv * (sc + ape_cpu.double()).softmax(dim=0)).sum(dim=0)).float()
|
||||
|
||||
|
||||
def _make_inputs(
|
||||
num_q: int, head_dim: int, seed: int
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
g = torch.Generator(device="cpu").manual_seed(seed)
|
||||
kv_score_input_cpu = torch.randn(
|
||||
num_q, head_dim * 2, generator=g, dtype=torch.float32
|
||||
)
|
||||
ape_cpu = torch.randn(RATIO, head_dim, generator=g, dtype=torch.float32)
|
||||
return kv_score_input_cpu, ape_cpu
|
||||
|
||||
|
||||
def _run_prefill(
|
||||
ctx: Context,
|
||||
pool: torch.Tensor,
|
||||
kv_score_input: torch.Tensor,
|
||||
ape: torch.Tensor,
|
||||
seq_lens_cpu: torch.Tensor,
|
||||
extend_lens_cpu: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
num_q = int(extend_lens_cpu.sum().item())
|
||||
plan = ctx.make_prefill_plan(seq_lens_cpu, extend_lens_cpu, num_q)
|
||||
return compress_forward(
|
||||
pool,
|
||||
kv_score_input,
|
||||
ape,
|
||||
plan,
|
||||
head_dim=ctx.head_dim,
|
||||
compress_ratio=RATIO,
|
||||
)
|
||||
|
||||
|
||||
def _run_decode(
|
||||
ctx: Context,
|
||||
pool: torch.Tensor,
|
||||
kv_score_input: torch.Tensor,
|
||||
ape: torch.Tensor,
|
||||
seq_lens_gpu: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
plan = ctx.make_decode_plan(seq_lens_gpu)
|
||||
return compress_forward(
|
||||
pool,
|
||||
kv_score_input,
|
||||
ape,
|
||||
plan,
|
||||
head_dim=ctx.head_dim,
|
||||
compress_ratio=RATIO,
|
||||
)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Tests
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["legacy", "paged"])
|
||||
@pytest.mark.parametrize("seq_len", [128, 256, 512])
|
||||
def test_prefill_no_context(mode: str, seq_len: int) -> None:
|
||||
"""Single-shot prefill, no prefix. Every compress event must match fp64 GT."""
|
||||
if mode == "legacy":
|
||||
ctx: Context = make_legacy_context(
|
||||
bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM
|
||||
)
|
||||
else:
|
||||
ctx = make_paged_context(bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM)
|
||||
|
||||
seq_lens_cpu, extend_lens_cpu, num_q = to_seq_extend([(seq_len, seq_len)])
|
||||
kv_in_cpu, ape_cpu = _make_inputs(num_q, ctx.head_dim, seed=seq_len)
|
||||
|
||||
pool = make_state_pool(ctx.num_pages, RATIO, ctx.head_dim)
|
||||
out = _run_prefill(
|
||||
ctx, pool, kv_in_cpu.cuda(), ape_cpu.cuda(), seq_lens_cpu, extend_lens_cpu
|
||||
)
|
||||
|
||||
# Compact prefill output: row per compress plan, in CPU-planner order.
|
||||
for plan_id, P in enumerate(range(RATIO - 1, seq_len, RATIO)):
|
||||
gt = _gt_compress(kv_in_cpu, ape_cpu, P=P, head_dim=ctx.head_dim)
|
||||
triton.testing.assert_close(out[plan_id].cpu(), gt, atol=ATOL, rtol=RTOL)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["legacy", "paged"])
|
||||
@pytest.mark.parametrize("prefix_len", [0, 128, 256])
|
||||
def test_prefill_then_decode(mode: str, prefix_len: int) -> None:
|
||||
"""Prefill ``prefix_len`` tokens, then decode through to the next 128 boundary."""
|
||||
seq_len = prefix_len + RATIO # one full compress chunk after prefix
|
||||
|
||||
if mode == "legacy":
|
||||
ctx: Context = make_legacy_context(
|
||||
bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM
|
||||
)
|
||||
else:
|
||||
ctx = make_paged_context(bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM)
|
||||
|
||||
kv_full_cpu, ape_cpu = _make_inputs(
|
||||
seq_len, ctx.head_dim, seed=seq_len + prefix_len
|
||||
)
|
||||
pool = make_state_pool(ctx.num_pages, RATIO, ctx.head_dim)
|
||||
|
||||
if prefix_len > 0:
|
||||
seq_lens_cpu, extend_lens_cpu, _ = to_seq_extend([(prefix_len, prefix_len)])
|
||||
_run_prefill(
|
||||
ctx,
|
||||
pool,
|
||||
kv_full_cpu[:prefix_len].cuda(),
|
||||
ape_cpu.cuda(),
|
||||
seq_lens_cpu,
|
||||
extend_lens_cpu,
|
||||
)
|
||||
|
||||
final_out = None
|
||||
for k in range(RATIO):
|
||||
cur_seq_len = prefix_len + k + 1
|
||||
seq_lens_gpu = torch.tensor([cur_seq_len], dtype=torch.int64, device="cuda")
|
||||
kv_step = kv_full_cpu[prefix_len + k : prefix_len + k + 1].cuda()
|
||||
out = _run_decode(ctx, pool, kv_step, ape_cpu.cuda(), seq_lens_gpu)
|
||||
if cur_seq_len % RATIO == 0:
|
||||
final_out = out
|
||||
|
||||
P = seq_len - 1
|
||||
gt = _gt_compress(kv_full_cpu, ape_cpu, P=P, head_dim=ctx.head_dim)
|
||||
assert final_out is not None
|
||||
triton.testing.assert_close(final_out[0].cpu(), gt, atol=ATOL, rtol=RTOL)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["legacy", "paged"])
|
||||
@pytest.mark.parametrize("prefix_len", [128, 256])
|
||||
def test_prefill_then_extend(mode: str, prefix_len: int) -> None:
|
||||
"""Prefill once, then a second prefill that extends across one compress event.
|
||||
|
||||
First prefill ends at a 128-boundary so the second prefill starts fresh.
|
||||
"""
|
||||
extend_len = RATIO
|
||||
seq_len = prefix_len + extend_len
|
||||
|
||||
if mode == "legacy":
|
||||
ctx: Context = make_legacy_context(
|
||||
bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM
|
||||
)
|
||||
else:
|
||||
ctx = make_paged_context(bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM)
|
||||
|
||||
kv_full_cpu, ape_cpu = _make_inputs(seq_len, ctx.head_dim, seed=prefix_len)
|
||||
pool = make_state_pool(ctx.num_pages, RATIO, ctx.head_dim)
|
||||
|
||||
seq_lens_cpu, extend_lens_cpu, _ = to_seq_extend([(prefix_len, prefix_len)])
|
||||
_run_prefill(
|
||||
ctx,
|
||||
pool,
|
||||
kv_full_cpu[:prefix_len].cuda(),
|
||||
ape_cpu.cuda(),
|
||||
seq_lens_cpu,
|
||||
extend_lens_cpu,
|
||||
)
|
||||
|
||||
seq_lens_cpu, extend_lens_cpu, _ = to_seq_extend([(seq_len, extend_len)])
|
||||
out = _run_prefill(
|
||||
ctx,
|
||||
pool,
|
||||
kv_full_cpu[prefix_len:].cuda(),
|
||||
ape_cpu.cuda(),
|
||||
seq_lens_cpu,
|
||||
extend_lens_cpu,
|
||||
)
|
||||
|
||||
P = seq_len - 1
|
||||
gt = _gt_compress(kv_full_cpu, ape_cpu, P=P, head_dim=ctx.head_dim)
|
||||
# Single compress event in this extend; compact plan_id 0.
|
||||
triton.testing.assert_close(out[0].cpu(), gt, atol=ATOL, rtol=RTOL)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["legacy", "paged"])
|
||||
def test_prefill_multibatch(mode: str) -> None:
|
||||
"""Multi-batch prefill, each batch ending at a different chunk count."""
|
||||
seq_extend = [(128, 128), (256, 256), (384, 384)]
|
||||
bs = len(seq_extend)
|
||||
if mode == "legacy":
|
||||
ctx: Context = make_legacy_context(
|
||||
bs=bs, compress_ratio=RATIO, head_dim=HEAD_DIM
|
||||
)
|
||||
else:
|
||||
ctx = make_paged_context(bs=bs, compress_ratio=RATIO, head_dim=HEAD_DIM)
|
||||
|
||||
seq_lens_cpu, extend_lens_cpu, num_q = to_seq_extend(seq_extend)
|
||||
kv_in_cpu, ape_cpu = _make_inputs(num_q, ctx.head_dim, seed=99)
|
||||
pool = make_state_pool(ctx.num_pages, RATIO, ctx.head_dim)
|
||||
out = _run_prefill(
|
||||
ctx, pool, kv_in_cpu.cuda(), ape_cpu.cuda(), seq_lens_cpu, extend_lens_cpu
|
||||
)
|
||||
|
||||
# Compact: walk batches in order, then positions in order; matches the
|
||||
# CPU planner's emit order for plan_c.
|
||||
base = 0
|
||||
plan_id = 0
|
||||
for b, (seq, ext) in enumerate(seq_extend):
|
||||
for j in range(ext):
|
||||
P = j # prefix=0
|
||||
if (P + 1) % RATIO != 0:
|
||||
continue
|
||||
gt = _gt_compress(
|
||||
kv_in_cpu[base : base + ext],
|
||||
ape_cpu,
|
||||
P=P,
|
||||
head_dim=ctx.head_dim,
|
||||
)
|
||||
triton.testing.assert_close(
|
||||
out[plan_id].cpu(),
|
||||
gt,
|
||||
atol=ATOL,
|
||||
rtol=RTOL,
|
||||
)
|
||||
plan_id += 1
|
||||
base += ext
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -0,0 +1,337 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from typing import Tuple, Union
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.jit_kernel.dsv4 import compress_forward
|
||||
from sglang.jit_kernel.tests.deepseek_v4.common import (
|
||||
LegacyContext,
|
||||
PagedContext,
|
||||
make_legacy_context,
|
||||
make_paged_context,
|
||||
make_state_pool,
|
||||
to_seq_extend,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=30, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
Context = Union[LegacyContext, PagedContext]
|
||||
|
||||
# c4 input row layout: | kv_overlap | kv | score_overlap | score |
|
||||
HEAD_DIM = 512
|
||||
RATIO = 4
|
||||
WINDOW = 8 # = 2 * RATIO (overlap + current)
|
||||
ATOL = 5e-3
|
||||
RTOL = 5e-3
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# fp64 ground truth (single compress event over a 8-token window).
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _gt_compress(
|
||||
kv_score_input_cpu: torch.Tensor, # [num_q, head_dim*4]
|
||||
ape_cpu: torch.Tensor, # [8, head_dim]
|
||||
P: int,
|
||||
head_dim: int,
|
||||
) -> torch.Tensor:
|
||||
"""fp64 reference for compress event at ragged position ``P``.
|
||||
|
||||
Tokens at positions [P-7..P-4] contribute their *overlap* halves, tokens
|
||||
at [P-3..P] contribute their *fresh* halves. Bias[0..3] for overlap,
|
||||
bias[4..7] for fresh. When P < 7, the overlap is masked (kv=0, score=-inf)
|
||||
so the softmax sees only the 4 fresh tokens.
|
||||
"""
|
||||
if P < 7:
|
||||
kv_ov = torch.zeros(4, head_dim, dtype=torch.float64)
|
||||
sc_ov = torch.full((4, head_dim), float("-inf"), dtype=torch.float64)
|
||||
else:
|
||||
kv_ov = kv_score_input_cpu[P - 7 : P - 3, :head_dim].double()
|
||||
sc_ov = kv_score_input_cpu[P - 7 : P - 3, 2 * head_dim : 3 * head_dim].double()
|
||||
kv_fr = kv_score_input_cpu[P - 3 : P + 1, head_dim : 2 * head_dim].double()
|
||||
sc_fr = kv_score_input_cpu[P - 3 : P + 1, 3 * head_dim :].double()
|
||||
kv = torch.cat([kv_ov, kv_fr], dim=0)
|
||||
sc = torch.cat([sc_ov, sc_fr], dim=0) + ape_cpu.double()
|
||||
return ((kv * sc.softmax(dim=0)).sum(dim=0)).float()
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Driver
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _run_prefill(
|
||||
ctx: Context,
|
||||
pool: torch.Tensor,
|
||||
kv_score_input: torch.Tensor,
|
||||
ape: torch.Tensor,
|
||||
seq_lens_cpu: torch.Tensor,
|
||||
extend_lens_cpu: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
num_q = int(extend_lens_cpu.sum().item())
|
||||
plan = ctx.make_prefill_plan(seq_lens_cpu, extend_lens_cpu, num_q)
|
||||
return compress_forward(
|
||||
pool,
|
||||
kv_score_input,
|
||||
ape,
|
||||
plan,
|
||||
head_dim=ctx.head_dim,
|
||||
compress_ratio=RATIO,
|
||||
)
|
||||
|
||||
|
||||
def _run_decode(
|
||||
ctx: Context,
|
||||
pool: torch.Tensor,
|
||||
kv_score_input: torch.Tensor,
|
||||
ape: torch.Tensor,
|
||||
seq_lens_gpu: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
plan = ctx.make_decode_plan(seq_lens_gpu)
|
||||
return compress_forward(
|
||||
pool,
|
||||
kv_score_input,
|
||||
ape,
|
||||
plan,
|
||||
head_dim=ctx.head_dim,
|
||||
compress_ratio=RATIO,
|
||||
)
|
||||
|
||||
|
||||
def _make_inputs(
|
||||
num_q: int, head_dim: int, seed: int
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
g = torch.Generator(device="cpu").manual_seed(seed)
|
||||
kv_score_input_cpu = torch.randn(
|
||||
num_q, head_dim * 4, generator=g, dtype=torch.float32
|
||||
)
|
||||
ape_cpu = torch.randn(WINDOW, head_dim, generator=g, dtype=torch.float32)
|
||||
return kv_score_input_cpu, ape_cpu
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Tests
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["legacy", "paged"])
|
||||
@pytest.mark.parametrize("seq_len", [4, 8, 32, 256, 1024])
|
||||
def test_prefill_no_context(mode: str, seq_len: int) -> None:
|
||||
"""Prefill once, no prefix. Every compress event must match fp64 GT."""
|
||||
if mode == "legacy":
|
||||
ctx: Context = make_legacy_context(
|
||||
bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM
|
||||
)
|
||||
else:
|
||||
ctx = make_paged_context(bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM)
|
||||
|
||||
seq_lens_cpu, extend_lens_cpu, num_q = to_seq_extend([(seq_len, seq_len)])
|
||||
kv_in_cpu, ape_cpu = _make_inputs(num_q, ctx.head_dim, seed=seq_len)
|
||||
|
||||
pool = make_state_pool(ctx.num_pages, RATIO, ctx.head_dim)
|
||||
out = _run_prefill(
|
||||
ctx, pool, kv_in_cpu.cuda(), ape_cpu.cuda(), seq_lens_cpu, extend_lens_cpu
|
||||
)
|
||||
|
||||
# Compact prefill output: row per compress plan, in CPU-planner order
|
||||
# (batch-major, position-ascending).
|
||||
for plan_id, P in enumerate(range(RATIO - 1, seq_len, RATIO)):
|
||||
gt = _gt_compress(kv_in_cpu, ape_cpu, P=P, head_dim=ctx.head_dim)
|
||||
triton.testing.assert_close(out[plan_id].cpu(), gt, atol=ATOL, rtol=RTOL)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["legacy", "paged"])
|
||||
@pytest.mark.parametrize("prefix_len", [4, 256])
|
||||
def test_prefill_then_decode(mode: str, prefix_len: int) -> None:
|
||||
"""Prefill once, then decode 4 more tokens through one compress boundary."""
|
||||
extend_decode = 4
|
||||
seq_len = prefix_len + extend_decode
|
||||
|
||||
if mode == "legacy":
|
||||
ctx: Context = make_legacy_context(
|
||||
bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM
|
||||
)
|
||||
else:
|
||||
ctx = make_paged_context(bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM)
|
||||
|
||||
kv_full_cpu, ape_cpu = _make_inputs(
|
||||
seq_len, ctx.head_dim, seed=seq_len + prefix_len
|
||||
)
|
||||
pool = make_state_pool(ctx.num_pages, RATIO, ctx.head_dim)
|
||||
|
||||
# Prefill the prefix.
|
||||
seq_lens_cpu, extend_lens_cpu, _ = to_seq_extend([(prefix_len, prefix_len)])
|
||||
_run_prefill(
|
||||
ctx,
|
||||
pool,
|
||||
kv_full_cpu[:prefix_len].cuda(),
|
||||
ape_cpu.cuda(),
|
||||
seq_lens_cpu,
|
||||
extend_lens_cpu,
|
||||
)
|
||||
|
||||
# Decode `extend_decode` tokens one at a time.
|
||||
final_out = None
|
||||
for k in range(extend_decode):
|
||||
cur_seq_len = prefix_len + k + 1
|
||||
seq_lens_gpu = torch.tensor([cur_seq_len], dtype=torch.int64, device="cuda")
|
||||
kv_step = kv_full_cpu[prefix_len + k : prefix_len + k + 1].cuda()
|
||||
out = _run_decode(ctx, pool, kv_step, ape_cpu.cuda(), seq_lens_gpu)
|
||||
if cur_seq_len % RATIO == 0:
|
||||
final_out = out
|
||||
|
||||
# Check the trailing compress: position P = seq_len - 1 = prefix + 3.
|
||||
P = seq_len - 1
|
||||
gt = _gt_compress(kv_full_cpu, ape_cpu, P=P, head_dim=ctx.head_dim)
|
||||
assert final_out is not None
|
||||
triton.testing.assert_close(final_out[0].cpu(), gt, atol=ATOL, rtol=RTOL)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["legacy", "paged"])
|
||||
@pytest.mark.parametrize("prefix_len", [256, 512, 768])
|
||||
def test_prefill_then_extend(mode: str, prefix_len: int) -> None:
|
||||
"""Prefill once, then prefill an extend that crosses one compress event.
|
||||
|
||||
The first prefill ends at a swa_page boundary (only relevant for paged),
|
||||
so the second prefill's overlap must be read out of the buffer.
|
||||
"""
|
||||
extend_len = 4
|
||||
|
||||
if mode == "legacy":
|
||||
ctx: Context = make_legacy_context(
|
||||
bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM
|
||||
)
|
||||
else:
|
||||
ctx = make_paged_context(bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM)
|
||||
|
||||
seq_len = prefix_len + extend_len
|
||||
kv_full_cpu, ape_cpu = _make_inputs(seq_len, ctx.head_dim, seed=prefix_len)
|
||||
pool = make_state_pool(ctx.num_pages, RATIO, ctx.head_dim)
|
||||
|
||||
# First prefill: seq=prefix, ext=prefix.
|
||||
seq_lens_cpu, extend_lens_cpu, _ = to_seq_extend([(prefix_len, prefix_len)])
|
||||
_run_prefill(
|
||||
ctx,
|
||||
pool,
|
||||
kv_full_cpu[:prefix_len].cuda(),
|
||||
ape_cpu.cuda(),
|
||||
seq_lens_cpu,
|
||||
extend_lens_cpu,
|
||||
)
|
||||
|
||||
# Second prefill: seq=prefix+extend, ext=extend, prefix=prefix_len.
|
||||
seq_lens_cpu, extend_lens_cpu, num_q = to_seq_extend([(seq_len, extend_len)])
|
||||
out = _run_prefill(
|
||||
ctx,
|
||||
pool,
|
||||
kv_full_cpu[prefix_len:].cuda(),
|
||||
ape_cpu.cuda(),
|
||||
seq_lens_cpu,
|
||||
extend_lens_cpu,
|
||||
)
|
||||
|
||||
P = seq_len - 1
|
||||
gt = _gt_compress(kv_full_cpu, ape_cpu, P=P, head_dim=ctx.head_dim)
|
||||
# Single compress event in this extend; compact plan_id 0.
|
||||
triton.testing.assert_close(out[0].cpu(), gt, atol=ATOL, rtol=RTOL)
|
||||
|
||||
|
||||
def test_paged_buffer_intermediate() -> None:
|
||||
"""Paged-only: after a multi-page prefill, verify the trailing 4 tokens of
|
||||
every swa_page sit in the correct state-pool slots.
|
||||
|
||||
These slots are what radix-cache resume reads when prefix-matching from a
|
||||
swa_page boundary, so they MUST match the original token data.
|
||||
"""
|
||||
ctx = make_paged_context(
|
||||
bs=1,
|
||||
compress_ratio=RATIO,
|
||||
head_dim=HEAD_DIM,
|
||||
swa_page_size=256,
|
||||
ring_size=8,
|
||||
num_swa_pages_per_req=8,
|
||||
)
|
||||
seq_len = 1024 # 4 swa_pages
|
||||
seq_lens_cpu, extend_lens_cpu, num_q = to_seq_extend([(seq_len, seq_len)])
|
||||
kv_in_cpu, ape_cpu = _make_inputs(num_q, ctx.head_dim, seed=42)
|
||||
|
||||
pool = make_state_pool(ctx.num_pages, RATIO, ctx.head_dim)
|
||||
_run_prefill(
|
||||
ctx, pool, kv_in_cpu.cuda(), ape_cpu.cuda(), seq_lens_cpu, extend_lens_cpu
|
||||
)
|
||||
|
||||
pool_cpu = pool.cpu()
|
||||
# For each swa_page boundary, the trailing `RATIO` tokens must have been
|
||||
# written. The state slot for token at position p is
|
||||
# `state_loc(0, p) = (p // swa_page_size) * ring_size + p % ring_size`.
|
||||
for swa_page_end in range(ctx.swa_page_size, seq_len + 1, ctx.swa_page_size):
|
||||
for offset in range(RATIO):
|
||||
p = swa_page_end - RATIO + offset
|
||||
sl = ctx.state_loc(0, p)
|
||||
page_idx = sl // RATIO
|
||||
slot_idx = sl % RATIO
|
||||
actual = pool_cpu[page_idx, slot_idx]
|
||||
# Token-row layout: the c4 prefill write copies the full
|
||||
# head_dim*4 row from kv_input verbatim into the state pool.
|
||||
expected = kv_in_cpu[p]
|
||||
triton.testing.assert_close(
|
||||
actual,
|
||||
expected,
|
||||
atol=ATOL,
|
||||
rtol=RTOL,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["legacy", "paged"])
|
||||
def test_prefill_multibatch(mode: str) -> None:
|
||||
"""Multi-batch prefill, both modes."""
|
||||
seq_extend = [(8, 8), (256, 256), (260, 260), (1023, 1023)]
|
||||
bs = len(seq_extend)
|
||||
if mode == "legacy":
|
||||
ctx: Context = make_legacy_context(
|
||||
bs=bs, compress_ratio=RATIO, head_dim=HEAD_DIM
|
||||
)
|
||||
else:
|
||||
ctx = make_paged_context(bs=bs, compress_ratio=RATIO, head_dim=HEAD_DIM)
|
||||
|
||||
seq_lens_cpu, extend_lens_cpu, num_q = to_seq_extend(seq_extend)
|
||||
kv_in_cpu, ape_cpu = _make_inputs(num_q, ctx.head_dim, seed=99)
|
||||
pool = make_state_pool(ctx.num_pages, RATIO, ctx.head_dim)
|
||||
out = _run_prefill(
|
||||
ctx, pool, kv_in_cpu.cuda(), ape_cpu.cuda(), seq_lens_cpu, extend_lens_cpu
|
||||
)
|
||||
|
||||
# Compact: walk batches in order, then positions in order; matches the
|
||||
# CPU planner's emit order for plan_c.
|
||||
base = 0
|
||||
plan_id = 0
|
||||
for b, (seq, ext) in enumerate(seq_extend):
|
||||
for j in range(ext):
|
||||
P = j # prefix=0 here
|
||||
if (P + 1) % RATIO != 0:
|
||||
continue
|
||||
gt = _gt_compress(
|
||||
kv_in_cpu[base : base + ext],
|
||||
ape_cpu,
|
||||
P=P,
|
||||
head_dim=ctx.head_dim,
|
||||
)
|
||||
triton.testing.assert_close(
|
||||
out[plan_id].cpu(),
|
||||
gt,
|
||||
atol=ATOL,
|
||||
rtol=RTOL,
|
||||
)
|
||||
plan_id += 1
|
||||
base += ext
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -0,0 +1,228 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.dsv4 import (
|
||||
CompressorDecodePlan,
|
||||
compress_norm_rope_store,
|
||||
fused_q_indexer_rope_hadamard_fp4_quant,
|
||||
)
|
||||
from sglang.jit_kernel.hadamard import hadamard_transform
|
||||
from sglang.srt.layers.attention.dsv4.fp4_indexer import (
|
||||
quantize_fp4_indexer_tensor,
|
||||
store_fp4_index_k_cache,
|
||||
)
|
||||
from sglang.srt.layers.deepseek_v4_rope import (
|
||||
apply_rotary_emb_triton,
|
||||
precompute_freqs_cis,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=60, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=60, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
HEAD_DIM = 128
|
||||
FP4_DIM = HEAD_DIM // 2
|
||||
GROUP_SIZE = 32
|
||||
SCALE_GROUPS = HEAD_DIM // GROUP_SIZE
|
||||
SCALE_BYTES = 4
|
||||
PAGE_SIZE = 64
|
||||
E2M1_MAX = 6.0
|
||||
|
||||
|
||||
def _ceil_ue8m0_exp_ref(x: torch.Tensor) -> torch.Tensor:
|
||||
bits = x.to(torch.float32).contiguous().view(torch.int32)
|
||||
exp = (bits >> 23) & 0xFF
|
||||
mantissa = bits & 0x7FFFFF
|
||||
exp = exp + (mantissa != 0).to(torch.int32)
|
||||
return exp.clamp(1, 254)
|
||||
|
||||
|
||||
def _fp4_e2m1_code_ref(x: torch.Tensor) -> torch.Tensor:
|
||||
ax = torch.minimum(x.abs(), torch.tensor(E2M1_MAX, device=x.device))
|
||||
idx = torch.zeros_like(ax, dtype=torch.uint8)
|
||||
for threshold in (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0):
|
||||
idx += (ax > threshold).to(torch.uint8)
|
||||
sign = ((x < 0) & (idx != 0)).to(torch.uint8) * 8
|
||||
return idx | sign
|
||||
|
||||
|
||||
def _ref_quantize_fp4_indexer(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
x = x.contiguous().view(-1, HEAD_DIM).float()
|
||||
groups = x.view(-1, SCALE_GROUPS, GROUP_SIZE)
|
||||
scale_raw = (groups.abs().amax(dim=-1) / E2M1_MAX).clamp_min(1.0e-4)
|
||||
scale_exp = _ceil_ue8m0_exp_ref(scale_raw)
|
||||
scale = (scale_exp << 23).contiguous().view(torch.float32)
|
||||
|
||||
scaled = (groups / scale.unsqueeze(-1)).view(-1, HEAD_DIM)
|
||||
code = _fp4_e2m1_code_ref(scaled)
|
||||
packed = (code[:, 0::2].to(torch.int16) | (code[:, 1::2].to(torch.int16) << 4)).to(
|
||||
torch.uint8
|
||||
)
|
||||
|
||||
packed_sf = scale_exp[:, 0].clone()
|
||||
for group_id in range(1, SCALE_GROUPS):
|
||||
packed_sf |= scale_exp[:, group_id] << (8 * group_id)
|
||||
return packed, packed_sf
|
||||
|
||||
|
||||
def _ref_store_fp4_index_cache(
|
||||
x_fp4: torch.Tensor,
|
||||
x_sf: torch.Tensor,
|
||||
loc: torch.Tensor,
|
||||
num_pages: int,
|
||||
) -> torch.Tensor:
|
||||
expected = torch.zeros(
|
||||
num_pages,
|
||||
PAGE_SIZE * (FP4_DIM + SCALE_BYTES),
|
||||
device=x_fp4.device,
|
||||
dtype=torch.uint8,
|
||||
)
|
||||
sf_shifts = torch.arange(0, 32, 8, device=x_fp4.device, dtype=torch.int32)
|
||||
for token_id in range(x_fp4.shape[0]):
|
||||
cache_loc = int(loc[token_id].item())
|
||||
page = cache_loc // PAGE_SIZE
|
||||
offset = cache_loc % PAGE_SIZE
|
||||
expected[page, offset * FP4_DIM : (offset + 1) * FP4_DIM] = x_fp4[token_id]
|
||||
sf_start = PAGE_SIZE * FP4_DIM + offset * SCALE_BYTES
|
||||
expected[page, sf_start : sf_start + SCALE_BYTES] = (
|
||||
(x_sf[token_id] >> sf_shifts) & 0xFF
|
||||
).to(torch.uint8)
|
||||
return expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [1, 7, 96])
|
||||
def test_quantize_fp4_indexer_tensor(num_tokens: int) -> None:
|
||||
torch.manual_seed(num_tokens)
|
||||
x = torch.randn(num_tokens, HEAD_DIM, device="cuda", dtype=torch.bfloat16)
|
||||
x[0, :8] = torch.tensor(
|
||||
[-8.0, -6.0, -3.0, -1.5, 0.0, 0.5, 2.0, 8.0],
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
x_fp4, x_sf = quantize_fp4_indexer_tensor(x)
|
||||
ref_fp4, ref_sf = _ref_quantize_fp4_indexer(x)
|
||||
|
||||
torch.testing.assert_close(x_fp4.view(torch.uint8), ref_fp4)
|
||||
torch.testing.assert_close(x_sf, ref_sf)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [1, 16, 96])
|
||||
def test_fp4_index_cache_store_layout(num_tokens: int) -> None:
|
||||
torch.manual_seed(num_tokens)
|
||||
num_pages = max(1, (num_tokens + PAGE_SIZE - 1) // PAGE_SIZE)
|
||||
x = torch.randn(num_tokens, HEAD_DIM, device="cuda", dtype=torch.bfloat16)
|
||||
loc = torch.randperm(num_pages * PAGE_SIZE, device="cuda")[:num_tokens].to(
|
||||
torch.int64
|
||||
)
|
||||
cache = torch.zeros(
|
||||
num_pages,
|
||||
PAGE_SIZE * (FP4_DIM + SCALE_BYTES),
|
||||
device="cuda",
|
||||
dtype=torch.uint8,
|
||||
)
|
||||
|
||||
store_fp4_index_k_cache(x, cache, loc, page_size=PAGE_SIZE)
|
||||
|
||||
ref_fp4, ref_sf = _ref_quantize_fp4_indexer(x)
|
||||
expected = _ref_store_fp4_index_cache(ref_fp4, ref_sf, loc, num_pages)
|
||||
torch.testing.assert_close(cache, expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [1, 16, 96])
|
||||
def test_fp4_fused_norm_rope_store_layout(num_tokens: int) -> None:
|
||||
torch.manual_seed(num_tokens + 100)
|
||||
num_pages = max(1, (num_tokens + PAGE_SIZE - 1) // PAGE_SIZE)
|
||||
compress_ratio = 4
|
||||
kv = torch.randn(num_tokens, HEAD_DIM, device="cuda", dtype=torch.bfloat16)
|
||||
norm_weight = torch.randn(HEAD_DIM, device="cuda", dtype=torch.bfloat16)
|
||||
seq_lens = (
|
||||
torch.arange(1, num_tokens + 1, device="cuda", dtype=torch.int64)
|
||||
* compress_ratio
|
||||
)
|
||||
req_pool_indices = torch.arange(num_tokens, device="cuda", dtype=torch.int64)
|
||||
plan = CompressorDecodePlan.generate_legacy(
|
||||
compress_ratio, req_pool_indices, seq_lens
|
||||
)
|
||||
loc = torch.arange(num_tokens, device="cuda", dtype=torch.int32)
|
||||
freqs_cis = precompute_freqs_cis(
|
||||
64, int(seq_lens.max().item()) + 1, 0, 10000, 1, 32, 1
|
||||
).to("cuda")
|
||||
cache = torch.zeros(
|
||||
num_pages,
|
||||
PAGE_SIZE * (FP4_DIM + SCALE_BYTES),
|
||||
device="cuda",
|
||||
dtype=torch.uint8,
|
||||
)
|
||||
|
||||
compress_norm_rope_store(
|
||||
kv.clone(),
|
||||
plan,
|
||||
norm_weight=norm_weight,
|
||||
norm_eps=1.0e-6,
|
||||
freq_cis=freqs_cis,
|
||||
out_loc=loc,
|
||||
kvcache=cache,
|
||||
page_size=PAGE_SIZE,
|
||||
use_fp4=True,
|
||||
)
|
||||
|
||||
ref = kv.float()
|
||||
ref = ref * torch.rsqrt((ref * ref).sum(dim=-1, keepdim=True) / HEAD_DIM + 1.0e-6)
|
||||
ref = ref * norm_weight.float()
|
||||
freqs = torch.view_as_real(freqs_cis).flatten(-2)[
|
||||
(seq_lens - compress_ratio).long()
|
||||
]
|
||||
rope = ref[:, 64:].reshape(num_tokens, 32, 2)
|
||||
freqs = freqs.reshape(num_tokens, 32, 2)
|
||||
rope_out = torch.empty_like(rope)
|
||||
rope_out[..., 0] = rope[..., 0] * freqs[..., 0] - rope[..., 1] * freqs[..., 1]
|
||||
rope_out[..., 1] = rope[..., 0] * freqs[..., 1] + rope[..., 1] * freqs[..., 0]
|
||||
ref[:, 64:] = rope_out.reshape(num_tokens, 64)
|
||||
ref = hadamard_transform(ref.contiguous(), scale=HEAD_DIM**-0.5)
|
||||
ref_fp4, ref_sf = _ref_quantize_fp4_indexer(ref)
|
||||
|
||||
expected = _ref_store_fp4_index_cache(
|
||||
ref_fp4,
|
||||
ref_sf,
|
||||
loc.to(torch.int64),
|
||||
num_pages,
|
||||
)
|
||||
torch.testing.assert_close(cache, expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", [1, 5, 17])
|
||||
def test_fp4_fused_q_indexer_rope_hadamard_quant(batch_size: int) -> None:
|
||||
torch.manual_seed(batch_size + 200)
|
||||
num_heads = 8
|
||||
rope_dim = 64
|
||||
weight_scale = HEAD_DIM**-0.5 * num_heads**-0.5
|
||||
q = torch.randn(
|
||||
batch_size, num_heads, HEAD_DIM, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
weight = torch.randn(batch_size, num_heads, device="cuda", dtype=torch.bfloat16)
|
||||
positions = (torch.arange(batch_size, device="cuda", dtype=torch.int32) * 7) % 63
|
||||
freqs_cis = precompute_freqs_cis(rope_dim, 64, 0, 10000, 1, 32, 1).to("cuda")
|
||||
|
||||
(q_fp4, q_sf), weights_out = fused_q_indexer_rope_hadamard_fp4_quant(
|
||||
q, weight, weight_scale, freqs_cis, positions
|
||||
)
|
||||
|
||||
ref = q.clone()
|
||||
apply_rotary_emb_triton(ref[..., -rope_dim:], freqs_cis, positions=positions)
|
||||
ref = hadamard_transform(ref.contiguous(), scale=HEAD_DIM**-0.5)
|
||||
ref_fp4, ref_sf = _ref_quantize_fp4_indexer(ref.view(-1, HEAD_DIM))
|
||||
ref_fp4 = ref_fp4.view(batch_size, num_heads, FP4_DIM)
|
||||
ref_sf = ref_sf.view(batch_size, num_heads)
|
||||
|
||||
torch.testing.assert_close(q_fp4.view(torch.uint8), ref_fp4)
|
||||
torch.testing.assert_close(q_sf, ref_sf)
|
||||
torch.testing.assert_close(weights_out.squeeze(-1), weight.float() * weight_scale)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -0,0 +1,142 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
|
||||
ModelOptFp8Config,
|
||||
ModelOptFp8LinearMethod,
|
||||
)
|
||||
from sglang.srt.layers.quantization.fp8_kernel import static_quant_fp8
|
||||
from sglang.srt.layers.quantization.fp8_utils import (
|
||||
cutlass_fp8_supported,
|
||||
input_to_float8,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=20, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=80, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
DEVICE = "cuda"
|
||||
DTYPE = torch.bfloat16
|
||||
MAX_FP8_DIFF = 5e-4
|
||||
TEST_CASES = [
|
||||
pytest.param(19, 150, 80, id="misaligned_projection_shape"),
|
||||
pytest.param(512, 3072, 4096, id="flux2_added_kv_projection_shape"),
|
||||
]
|
||||
|
||||
|
||||
def _modelopt_fp8_supported() -> bool:
|
||||
return torch.cuda.is_available() and cutlass_fp8_supported()
|
||||
|
||||
|
||||
def _calc_diff(x: torch.Tensor, y: torch.Tensor) -> float:
|
||||
x, y = x.double(), y.double()
|
||||
denominator = (x * x + y * y).sum()
|
||||
if denominator == 0:
|
||||
return 0.0
|
||||
sim = 2 * (x * y).sum() / denominator
|
||||
return (1 - sim).item()
|
||||
|
||||
|
||||
def _dequantize_fp8_input(qinput: torch.Tensor, x_scale: torch.Tensor) -> torch.Tensor:
|
||||
return qinput.to(torch.float32) * x_scale.to(torch.float32)
|
||||
|
||||
|
||||
def _dequantize_fp8_weight(
|
||||
weight: torch.Tensor, weight_scale: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
if weight_scale.ndim == 0 or weight_scale.numel() == 1:
|
||||
scale = weight_scale.to(torch.float32)
|
||||
else:
|
||||
scale = weight_scale.to(torch.float32).reshape(-1, 1).t()
|
||||
return weight.to(torch.float32) * scale
|
||||
|
||||
|
||||
def _build_layer(
|
||||
weight_q: torch.Tensor,
|
||||
weight_scale: torch.Tensor,
|
||||
input_scale: torch.Tensor,
|
||||
) -> tuple[torch.nn.Module, ModelOptFp8LinearMethod]:
|
||||
output_size, input_size = weight_q.shape
|
||||
method = ModelOptFp8LinearMethod(
|
||||
ModelOptFp8Config(is_checkpoint_fp8_serialized=True)
|
||||
)
|
||||
layer = torch.nn.Module()
|
||||
method.create_weights(
|
||||
layer=layer,
|
||||
input_size_per_partition=input_size,
|
||||
output_partition_sizes=[output_size],
|
||||
input_size=input_size,
|
||||
output_size=output_size,
|
||||
params_dtype=DTYPE,
|
||||
weight_loader=lambda *args, **kwargs: None,
|
||||
)
|
||||
layer = layer.to(device=DEVICE)
|
||||
|
||||
layer.weight.data.copy_(weight_q)
|
||||
layer.weight_scale.data.copy_(weight_scale.reshape_as(layer.weight_scale))
|
||||
layer.input_scale.data.copy_(input_scale.reshape_as(layer.input_scale))
|
||||
method.process_weights_after_loading(layer)
|
||||
return layer, method
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _modelopt_fp8_supported(),
|
||||
reason="Diffusion ModelOpt FP8 scaled mm correctness requires CUDA FP8 support",
|
||||
)
|
||||
@pytest.mark.parametrize("m,n,k", TEST_CASES)
|
||||
def test_checkpoint_processing(m: int, n: int, k: int) -> None:
|
||||
generator = torch.Generator(device=DEVICE)
|
||||
generator.manual_seed(20260410 + m + n + k)
|
||||
|
||||
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
weight_q, weight_scale = input_to_float8(weight)
|
||||
input_scale = torch.tensor(1.0, device=DEVICE, dtype=torch.float32)
|
||||
|
||||
layer, _ = _build_layer(weight_q, weight_scale, input_scale)
|
||||
|
||||
assert tuple(layer.weight.shape) == (k, n)
|
||||
assert tuple(layer.weight.stride()) == (1, k)
|
||||
assert layer.weight.dtype == torch.float8_e4m3fn
|
||||
assert layer.input_scale.ndim == 0
|
||||
assert tuple(layer.weight_scale.shape) == (n, 1)
|
||||
|
||||
expected_weight = weight_q.t().to(torch.float32) * weight_scale.to(torch.float32)
|
||||
actual_weight = _dequantize_fp8_weight(layer.weight, layer.weight_scale)
|
||||
torch.testing.assert_close(actual_weight, expected_weight, atol=0.0, rtol=0.0)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _modelopt_fp8_supported(),
|
||||
reason="Diffusion ModelOpt FP8 scaled mm correctness requires CUDA FP8 support",
|
||||
)
|
||||
@pytest.mark.parametrize("m,n,k", TEST_CASES)
|
||||
def test_shape_correctness(m: int, n: int, k: int) -> None:
|
||||
generator = torch.Generator(device=DEVICE)
|
||||
generator.manual_seed(20260410 + m + n + k)
|
||||
|
||||
x = torch.randn((m, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
weight_q, weight_scale = input_to_float8(weight)
|
||||
_, input_scale = input_to_float8(x)
|
||||
|
||||
layer, method = _build_layer(weight_q, weight_scale, input_scale)
|
||||
|
||||
qinput, x_scale = static_quant_fp8(
|
||||
x.contiguous(),
|
||||
layer.input_scale,
|
||||
repeat_scale=method.cutlass_fp8_supported,
|
||||
)
|
||||
expected = torch.matmul(
|
||||
_dequantize_fp8_input(qinput, x_scale),
|
||||
_dequantize_fp8_weight(layer.weight, layer.weight_scale),
|
||||
)
|
||||
|
||||
actual = method.apply(layer, x)
|
||||
diff = _calc_diff(actual, expected.to(dtype=DTYPE))
|
||||
assert diff < MAX_FP8_DIFF, f"{m=}, {n=}, {k=}, {diff=:.6f}"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,460 @@
|
||||
import sys
|
||||
|
||||
import flashinfer
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant
|
||||
from sglang.multimodal_gen.runtime.layers.quantization import (
|
||||
modelopt_quant as diffusion_modelopt_quant,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
|
||||
ModelOptFp4Config,
|
||||
ModelOptFp4LinearMethod,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
from sglang.srt.layers.quantization.modelopt_quant import pad_nvfp4_weight
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
# B200-only correctness coverage for diffusion NVFP4 scaled mm.
|
||||
register_cuda_ci(est_time=15, suite="base-b-kernel-unit-1-gpu-b200")
|
||||
|
||||
DEVICE = "cuda"
|
||||
DTYPE = torch.bfloat16
|
||||
BLOCK_SIZE = 16
|
||||
FLOAT4_E2M1_MAX = 6.0
|
||||
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
|
||||
FP4_VALUE_LUT = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
|
||||
DEEPGEMM_FP4_MAX_DIFF = 0.02
|
||||
TEST_CASES = [
|
||||
pytest.param(19, 150, 80, id="padding_regression"),
|
||||
pytest.param(512, 6144, 128, id="flux2_projection_shape"),
|
||||
]
|
||||
FLUX2_PROJECTION_SHAPE = (512, 6144, 128)
|
||||
|
||||
|
||||
def _nvfp4_supported() -> bool:
|
||||
return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0)
|
||||
|
||||
|
||||
def _make_global_scale(x: torch.Tensor) -> torch.Tensor:
|
||||
max_abs = torch.amax(x.abs()).clamp_min_(1e-6)
|
||||
return (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / max_abs).to(torch.float32)
|
||||
|
||||
|
||||
def _calc_diff(x: torch.Tensor, y: torch.Tensor) -> float:
|
||||
x, y = x.double(), y.double()
|
||||
denominator = (x * x + y * y).sum()
|
||||
if denominator == 0:
|
||||
return 0.0
|
||||
sim = 2 * (x * y).sum() / denominator
|
||||
return (1 - sim).item()
|
||||
|
||||
|
||||
def _swap_fp4_nibbles(packed: torch.Tensor) -> torch.Tensor:
|
||||
return ((packed >> 4) | (packed << 4)).contiguous()
|
||||
|
||||
|
||||
def _fp4_lut(device: torch.device) -> torch.Tensor:
|
||||
return torch.tensor(FP4_VALUE_LUT, dtype=torch.float32, device=device)
|
||||
|
||||
|
||||
def _unpack_fp4_bytes(packed: torch.Tensor) -> torch.Tensor:
|
||||
assert packed.dtype == torch.uint8
|
||||
lut = _fp4_lut(packed.device)
|
||||
|
||||
def _decode(nibbles: torch.Tensor) -> torch.Tensor:
|
||||
values = lut[(nibbles & 0x7).to(torch.long)]
|
||||
return torch.where((nibbles & 0x8) != 0, -values, values)
|
||||
|
||||
low = _decode(packed & 0x0F)
|
||||
high = _decode((packed & 0xF0) >> 4)
|
||||
return torch.stack((low, high), dim=-1).reshape(
|
||||
packed.shape[0], packed.shape[1] * 2
|
||||
)
|
||||
|
||||
|
||||
def _swizzled_to_linear(
|
||||
scales_swizzled: torch.Tensor,
|
||||
rows: int,
|
||||
cols: int,
|
||||
) -> torch.Tensor:
|
||||
scales_swizzled = scales_swizzled.view(torch.float8_e4m3fn)
|
||||
row_tiles = (rows + 128 - 1) // 128
|
||||
tile_cols = BLOCK_SIZE * 4
|
||||
col_tiles = (cols + tile_cols - 1) // tile_cols
|
||||
tmp = scales_swizzled.reshape(1, row_tiles, col_tiles, 32, 4, 4)
|
||||
tmp = tmp.permute(0, 1, 4, 3, 2, 5)
|
||||
linear = tmp.reshape(row_tiles * 128, col_tiles * tile_cols // BLOCK_SIZE)
|
||||
return linear[:rows, : cols // BLOCK_SIZE]
|
||||
|
||||
|
||||
def _dequantize_nvfp4(
|
||||
packed: torch.Tensor,
|
||||
scales_swizzled: torch.Tensor,
|
||||
global_scale: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
rows, packed_cols = packed.shape
|
||||
cols = packed_cols * 2
|
||||
unpacked = _unpack_fp4_bytes(packed).reshape(rows, cols // BLOCK_SIZE, BLOCK_SIZE)
|
||||
scales_linear = _swizzled_to_linear(scales_swizzled, rows, cols).to(torch.float32)
|
||||
return (unpacked * (scales_linear / global_scale).unsqueeze(-1)).reshape(rows, cols)
|
||||
|
||||
|
||||
def _quantize_weight_for_checkpoint(
|
||||
weight: torch.Tensor, weight_global_scale: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
weight_fp4, weight_scale_linear = flashinfer.fp4_quantize(
|
||||
weight,
|
||||
weight_global_scale,
|
||||
is_sf_swizzled_layout=False,
|
||||
)
|
||||
if weight_scale_linear.dtype == torch.uint8:
|
||||
weight_scale_linear = weight_scale_linear.view(torch.float8_e4m3fn)
|
||||
return weight_fp4, weight_scale_linear.contiguous()
|
||||
|
||||
|
||||
def _set_diffusion_fp4_backend(
|
||||
monkeypatch: pytest.MonkeyPatch, backend: str | None
|
||||
) -> None:
|
||||
if backend is None:
|
||||
monkeypatch.delenv(
|
||||
"SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND", raising=False
|
||||
)
|
||||
else:
|
||||
monkeypatch.setenv("SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND", backend)
|
||||
|
||||
current_platform.__class__.get_modelopt_flashinfer_fp4_backend.cache_clear()
|
||||
current_platform.__class__.get_modelopt_fp4_gemm_op.cache_clear()
|
||||
diffusion_modelopt_quant._get_fp4_gemm_op.cache_clear()
|
||||
|
||||
|
||||
def _build_layer(
|
||||
weight_fp4: torch.Tensor,
|
||||
weight_scale_linear: torch.Tensor,
|
||||
input_global_scale: torch.Tensor,
|
||||
weight_global_scale: torch.Tensor,
|
||||
*,
|
||||
weight_scale_device: torch.device | str | None = None,
|
||||
checkpoint_weight_scale_layout: str = "linear",
|
||||
) -> tuple[ModelOptFp4LinearMethod, torch.nn.Module]:
|
||||
output_size, input_size_half = weight_fp4.shape
|
||||
input_size = input_size_half * 2
|
||||
method = ModelOptFp4LinearMethod(
|
||||
ModelOptFp4Config(
|
||||
is_checkpoint_nvfp4_serialized=True,
|
||||
group_size=BLOCK_SIZE,
|
||||
swap_weight_nibbles=True,
|
||||
checkpoint_weight_scale_layout=checkpoint_weight_scale_layout,
|
||||
)
|
||||
)
|
||||
layer = torch.nn.Module()
|
||||
method.create_weights(
|
||||
layer,
|
||||
input_size_per_partition=input_size,
|
||||
output_partition_sizes=[output_size],
|
||||
input_size=input_size,
|
||||
output_size=output_size,
|
||||
params_dtype=DTYPE,
|
||||
weight_loader=lambda *args, **kwargs: None,
|
||||
)
|
||||
layer = layer.to(device=DEVICE)
|
||||
|
||||
checkpoint_weight = _swap_fp4_nibbles(weight_fp4)
|
||||
layer.weight.data.copy_(checkpoint_weight)
|
||||
layer.input_scale.data.copy_(
|
||||
(1.0 / input_global_scale).reshape_as(layer.input_scale)
|
||||
)
|
||||
layer.weight_scale_2.data.copy_(
|
||||
(1.0 / weight_global_scale).reshape_as(layer.weight_scale_2)
|
||||
)
|
||||
layer.weight_scale.data.copy_(weight_scale_linear)
|
||||
if weight_scale_device is not None:
|
||||
layer.weight_scale = torch.nn.Parameter(
|
||||
layer.weight_scale.detach().to(weight_scale_device), requires_grad=False
|
||||
)
|
||||
|
||||
method.process_weights_after_loading(layer)
|
||||
|
||||
_, flashinfer_backend = current_platform.get_modelopt_fp4_gemm_op()
|
||||
if flashinfer_backend == "trtllm":
|
||||
expected_weight, _ = pad_nvfp4_weight(
|
||||
weight_fp4, n_alignment=128, k_alignment=0
|
||||
)
|
||||
expected_scale = (
|
||||
_swizzled_to_linear(weight_scale_linear, output_size, input_size)
|
||||
if checkpoint_weight_scale_layout == "swizzled"
|
||||
else weight_scale_linear
|
||||
)
|
||||
if expected_scale.shape[0] != expected_weight.shape[0]:
|
||||
pad_n = expected_weight.shape[0] - expected_scale.shape[0]
|
||||
expected_scale = torch.nn.functional.pad(expected_scale, (0, 0, 0, pad_n))
|
||||
|
||||
expected_padding_cols = 0
|
||||
if expected_scale.shape[1] % 4 != 0:
|
||||
padded_scale_k = ((expected_scale.shape[1] + 4 - 1) // 4) * 4
|
||||
pad_scale_k = padded_scale_k - expected_scale.shape[1]
|
||||
expected_scale = torch.nn.functional.pad(
|
||||
expected_scale, (0, pad_scale_k, 0, 0)
|
||||
)
|
||||
pad_weight_k = pad_scale_k * 8
|
||||
expected_weight = torch.nn.functional.pad(
|
||||
expected_weight, (0, pad_weight_k, 0, 0)
|
||||
)
|
||||
expected_padding_cols = pad_weight_k
|
||||
|
||||
expected_weight = flashinfer.shuffle_matrix_a(
|
||||
expected_weight.view(torch.uint8), 128
|
||||
)
|
||||
expected_scale = (
|
||||
flashinfer.shuffle_matrix_sf_a(expected_scale.view(torch.uint8), 128)
|
||||
.reshape(expected_scale.shape)
|
||||
.view(torch.float8_e4m3fn)
|
||||
)
|
||||
|
||||
assert torch.equal(layer.weight, expected_weight)
|
||||
assert torch.equal(
|
||||
layer.weight_scale_interleaved.view(torch.uint8),
|
||||
expected_scale.view(torch.uint8),
|
||||
)
|
||||
assert layer.weights_padding_cols == expected_padding_cols
|
||||
else:
|
||||
expected_weight, expected_padding_cols = pad_nvfp4_weight(weight_fp4)
|
||||
expected_scale_shape = (
|
||||
((output_size + 128 - 1) // 128) * 128,
|
||||
(((input_size // BLOCK_SIZE) + 4 - 1) // 4) * 4,
|
||||
)
|
||||
|
||||
assert torch.equal(layer.weight, expected_weight)
|
||||
assert layer.weight_scale_interleaved.shape == expected_scale_shape
|
||||
assert layer.weight_scale_interleaved.dtype == torch.float8_e4m3fn
|
||||
assert layer.weights_padding_cols == expected_padding_cols
|
||||
torch.testing.assert_close(
|
||||
layer.alpha,
|
||||
(1.0 / (input_global_scale * weight_global_scale)).to(torch.float32),
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
layer.input_scale_inv,
|
||||
input_global_scale.to(torch.float32),
|
||||
)
|
||||
return method, layer
|
||||
|
||||
|
||||
def _resolve_mode(mode: str):
|
||||
if mode == "jit_cutlass":
|
||||
return scaled_fp4_quant, cutlass_scaled_fp4_mm, None
|
||||
if mode == "flashinfer2":
|
||||
return flashinfer.fp4_quantize, flashinfer.mm_fp4, "cudnn"
|
||||
if mode == "flashinfer_trtllm":
|
||||
return flashinfer.fp4_quantize, flashinfer.mm_fp4, "trtllm"
|
||||
raise ValueError(f"Unknown mode: {mode}")
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _nvfp4_supported(),
|
||||
reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"backend", [None, "flashinfer_trtllm"], ids=["default", "flashinfer_trtllm"]
|
||||
)
|
||||
@pytest.mark.parametrize("m,n,k", TEST_CASES)
|
||||
def test_checkpoint_processing(
|
||||
monkeypatch: pytest.MonkeyPatch, backend: str | None, m: int, n: int, k: int
|
||||
) -> None:
|
||||
_set_diffusion_fp4_backend(monkeypatch, backend)
|
||||
generator = torch.Generator(device=DEVICE)
|
||||
generator.manual_seed(20260404 + m + n + k)
|
||||
|
||||
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
input_global_scale = torch.tensor(512.0, device=DEVICE, dtype=torch.float32)
|
||||
weight_global_scale = _make_global_scale(weight)
|
||||
weight_fp4, weight_scale_linear = _quantize_weight_for_checkpoint(
|
||||
weight, weight_global_scale
|
||||
)
|
||||
|
||||
_build_layer(
|
||||
weight_fp4, weight_scale_linear, input_global_scale, weight_global_scale
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _nvfp4_supported(),
|
||||
reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
|
||||
)
|
||||
@pytest.mark.parametrize("mode", ["jit_cutlass", "flashinfer2"])
|
||||
def test_flux2_shape_correctness(mode: str) -> None:
|
||||
m, n, k = FLUX2_PROJECTION_SHAPE
|
||||
quantize_op, gemm_op, gemm_backend = _resolve_mode(mode)
|
||||
generator = torch.Generator(device=DEVICE)
|
||||
generator.manual_seed(20260404 + m + n + k)
|
||||
|
||||
x = torch.randn((m, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
input_global_scale = _make_global_scale(x)
|
||||
weight_global_scale = _make_global_scale(weight)
|
||||
alpha = (1.0 / (input_global_scale * weight_global_scale)).to(torch.float32)
|
||||
|
||||
x_fp4, x_scale_swizzled = quantize_op(x, input_global_scale)
|
||||
weight_fp4, weight_scale_swizzled = quantize_op(weight, weight_global_scale)
|
||||
if x_scale_swizzled.dtype == torch.uint8:
|
||||
x_scale_swizzled = x_scale_swizzled.view(torch.float8_e4m3fn)
|
||||
if weight_scale_swizzled.dtype == torch.uint8:
|
||||
weight_scale_swizzled = weight_scale_swizzled.view(torch.float8_e4m3fn)
|
||||
|
||||
expected = torch.matmul(
|
||||
_dequantize_nvfp4(x_fp4, x_scale_swizzled, input_global_scale),
|
||||
_dequantize_nvfp4(weight_fp4, weight_scale_swizzled, weight_global_scale).t(),
|
||||
)
|
||||
|
||||
if gemm_backend is None:
|
||||
actual = gemm_op(
|
||||
x_fp4,
|
||||
weight_fp4,
|
||||
x_scale_swizzled,
|
||||
weight_scale_swizzled,
|
||||
alpha,
|
||||
DTYPE,
|
||||
)
|
||||
else:
|
||||
actual = gemm_op(
|
||||
x_fp4,
|
||||
weight_fp4.t(),
|
||||
x_scale_swizzled,
|
||||
weight_scale_swizzled.t(),
|
||||
alpha,
|
||||
DTYPE,
|
||||
backend=gemm_backend,
|
||||
)
|
||||
|
||||
diff = _calc_diff(actual, expected.to(dtype=DTYPE))
|
||||
assert diff < DEEPGEMM_FP4_MAX_DIFF, f"{mode=}, {m=}, {n=}, {k=}, {diff=:.6f}"
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _nvfp4_supported(),
|
||||
reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
|
||||
)
|
||||
def test_flux2_shape_correctness_flashinfer_trtllm(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_set_diffusion_fp4_backend(monkeypatch, "flashinfer_trtllm")
|
||||
|
||||
m, n, k = FLUX2_PROJECTION_SHAPE
|
||||
generator = torch.Generator(device=DEVICE)
|
||||
generator.manual_seed(20260404 + m + n + k + 17)
|
||||
|
||||
x = torch.randn((m, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
input_global_scale = _make_global_scale(x)
|
||||
weight_global_scale = _make_global_scale(weight)
|
||||
weight_fp4, weight_scale_linear = _quantize_weight_for_checkpoint(
|
||||
weight, weight_global_scale
|
||||
)
|
||||
|
||||
method, layer = _build_layer(
|
||||
weight_fp4, weight_scale_linear, input_global_scale, weight_global_scale
|
||||
)
|
||||
actual = method.apply(layer, x)
|
||||
|
||||
x_fp4, x_scale_swizzled = flashinfer.fp4_quantize(x, input_global_scale)
|
||||
weight_fp4_ref, weight_scale_swizzled = flashinfer.fp4_quantize(
|
||||
weight, weight_global_scale
|
||||
)
|
||||
if x_scale_swizzled.dtype == torch.uint8:
|
||||
x_scale_swizzled = x_scale_swizzled.view(torch.float8_e4m3fn)
|
||||
if weight_scale_swizzled.dtype == torch.uint8:
|
||||
weight_scale_swizzled = weight_scale_swizzled.view(torch.float8_e4m3fn)
|
||||
|
||||
expected = torch.matmul(
|
||||
_dequantize_nvfp4(x_fp4, x_scale_swizzled, input_global_scale),
|
||||
_dequantize_nvfp4(
|
||||
weight_fp4_ref, weight_scale_swizzled, weight_global_scale
|
||||
).t(),
|
||||
)
|
||||
|
||||
diff = _calc_diff(actual, expected.to(dtype=DTYPE))
|
||||
assert diff < DEEPGEMM_FP4_MAX_DIFF, f"{m=}, {n=}, {k=}, {diff=:.6f}"
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _nvfp4_supported(),
|
||||
reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
|
||||
)
|
||||
def test_flux2_swizzled_scale_checkpoint_flashinfer_trtllm_matches_cudnn(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_set_diffusion_fp4_backend(monkeypatch, "flashinfer_trtllm")
|
||||
|
||||
m, n, k = FLUX2_PROJECTION_SHAPE
|
||||
generator = torch.Generator(device=DEVICE)
|
||||
generator.manual_seed(20260517 + m + n + k)
|
||||
|
||||
x = torch.randn((m, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
input_global_scale = _make_global_scale(x)
|
||||
weight_global_scale = _make_global_scale(weight)
|
||||
alpha = (1.0 / (input_global_scale * weight_global_scale)).to(torch.float32)
|
||||
|
||||
x_fp4, x_scale_swizzled = flashinfer.fp4_quantize(x, input_global_scale)
|
||||
weight_fp4, weight_scale_swizzled = flashinfer.fp4_quantize(
|
||||
weight, weight_global_scale
|
||||
)
|
||||
if x_scale_swizzled.dtype == torch.uint8:
|
||||
x_scale_swizzled = x_scale_swizzled.view(torch.float8_e4m3fn)
|
||||
if weight_scale_swizzled.dtype == torch.uint8:
|
||||
weight_scale_swizzled = weight_scale_swizzled.view(torch.float8_e4m3fn)
|
||||
|
||||
method, layer = _build_layer(
|
||||
weight_fp4,
|
||||
weight_scale_swizzled,
|
||||
input_global_scale,
|
||||
weight_global_scale,
|
||||
checkpoint_weight_scale_layout="swizzled",
|
||||
)
|
||||
actual = method.apply(layer, x)
|
||||
|
||||
expected = flashinfer.mm_fp4(
|
||||
x_fp4,
|
||||
weight_fp4.t(),
|
||||
x_scale_swizzled,
|
||||
weight_scale_swizzled.t(),
|
||||
alpha,
|
||||
DTYPE,
|
||||
backend="cudnn",
|
||||
)
|
||||
|
||||
diff = _calc_diff(actual, expected)
|
||||
assert diff < DEEPGEMM_FP4_MAX_DIFF, f"{m=}, {n=}, {k=}, {diff=:.6f}"
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _nvfp4_supported(),
|
||||
reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
|
||||
)
|
||||
def test_checkpoint_processing_flashinfer_trtllm_cpu_weight_scale(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_set_diffusion_fp4_backend(monkeypatch, "flashinfer_trtllm")
|
||||
|
||||
m, n, k = FLUX2_PROJECTION_SHAPE
|
||||
generator = torch.Generator(device=DEVICE)
|
||||
generator.manual_seed(20260413 + m + n + k)
|
||||
|
||||
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
input_global_scale = torch.tensor(512.0, device=DEVICE, dtype=torch.float32)
|
||||
weight_global_scale = _make_global_scale(weight)
|
||||
weight_fp4, weight_scale_linear = _quantize_weight_for_checkpoint(
|
||||
weight, weight_global_scale
|
||||
)
|
||||
|
||||
_build_layer(
|
||||
weight_fp4,
|
||||
weight_scale_linear,
|
||||
input_global_scale,
|
||||
weight_global_scale,
|
||||
weight_scale_device="cpu",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,133 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from sglang.test.ci.ci_register import register_amd_ci
|
||||
|
||||
register_amd_ci(est_time=30, suite="jit-kernel-unit-test-amd")
|
||||
|
||||
DEVICE = "cuda"
|
||||
D = 5120
|
||||
EPS = 1e-6
|
||||
|
||||
|
||||
def _ref_rms_norm(x_f32, weight, eps):
|
||||
var = x_f32.pow(2).mean(-1, keepdim=True)
|
||||
return x_f32 * torch.rsqrt(var + eps)
|
||||
|
||||
|
||||
def _ref_fused_residual_norm_ss(
|
||||
residual, x, gate, weight, bias, scale, shift, norm_type, eps
|
||||
):
|
||||
ref_res = residual.float() + x.float() * (gate.float() if gate is not None else 1)
|
||||
ref_res_bf16 = ref_res.to(torch.bfloat16)
|
||||
if norm_type == "layer":
|
||||
normed = F.layer_norm(ref_res_bf16.float(), (D,), weight, bias, eps)
|
||||
else:
|
||||
normed = _ref_rms_norm(ref_res_bf16.float(), weight, eps) * weight.float()
|
||||
y = (normed * (1.0 + scale.float()) + shift.float()).to(torch.bfloat16)
|
||||
return y, ref_res_bf16
|
||||
|
||||
|
||||
def _ref_norm_ss(x, weight, bias, scale, shift, norm_type, eps):
|
||||
if norm_type == "layer":
|
||||
normed = F.layer_norm(x.float(), (D,), weight, bias, eps)
|
||||
else:
|
||||
normed = _ref_rms_norm(x.float(), weight, eps) * weight.float()
|
||||
return (normed * (1.0 + scale.float()) + shift.float()).to(torch.bfloat16)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def cuda_setup():
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA required")
|
||||
if not hasattr(torch.version, "hip") or not torch.version.hip:
|
||||
pytest.skip("ROCm/HIP required for FlyDSL kernels")
|
||||
torch.manual_seed(42)
|
||||
|
||||
|
||||
FUSED_CASES = [
|
||||
("rms", 1, 16),
|
||||
("rms", 2, 16),
|
||||
("layer", 2, 16),
|
||||
("rms", 1, 90000),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("norm_type,B,L", FUSED_CASES)
|
||||
def test_fused_residual_norm_scale_shift(norm_type, B, L):
|
||||
from sglang.jit_kernel.diffusion.flydsl.fused_residual_norm import (
|
||||
flydsl_fused_residual_norm_scale_shift,
|
||||
)
|
||||
|
||||
residual = torch.randn(B, L, D, device=DEVICE, dtype=torch.bfloat16)
|
||||
x = torch.randn(B, L, D, device=DEVICE, dtype=torch.bfloat16)
|
||||
gate = torch.randn(B, 1, D, device=DEVICE, dtype=torch.bfloat16)
|
||||
weight = torch.randn(D, device=DEVICE, dtype=torch.float32)
|
||||
bias = (
|
||||
torch.randn(D, device=DEVICE, dtype=torch.float32)
|
||||
if norm_type == "layer"
|
||||
else None
|
||||
)
|
||||
scale = torch.randn(B, 1, D, device=DEVICE, dtype=torch.bfloat16)
|
||||
shift = torch.randn(B, 1, D, device=DEVICE, dtype=torch.bfloat16)
|
||||
|
||||
y, res_out = flydsl_fused_residual_norm_scale_shift(
|
||||
residual,
|
||||
x,
|
||||
gate,
|
||||
weight,
|
||||
bias,
|
||||
scale,
|
||||
shift,
|
||||
norm_type,
|
||||
EPS,
|
||||
)
|
||||
y_ref, res_ref = _ref_fused_residual_norm_ss(
|
||||
residual,
|
||||
x,
|
||||
gate,
|
||||
weight,
|
||||
bias,
|
||||
scale,
|
||||
shift,
|
||||
norm_type,
|
||||
EPS,
|
||||
)
|
||||
torch.testing.assert_close(res_out, res_ref, atol=5e-2, rtol=5e-2)
|
||||
torch.testing.assert_close(y, y_ref, atol=1.0, rtol=5e-2)
|
||||
|
||||
|
||||
NSS_CASES = [
|
||||
("rms", 2, 16),
|
||||
("layer", 2, 16),
|
||||
("rms", 1, 90000),
|
||||
("layer", 1, 90000),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("norm_type,B,L", NSS_CASES)
|
||||
def test_norm_scale_shift(norm_type, B, L):
|
||||
from sglang.jit_kernel.diffusion.flydsl.fused_residual_norm import (
|
||||
flydsl_norm_scale_shift,
|
||||
)
|
||||
|
||||
x = torch.randn(B, L, D, device=DEVICE, dtype=torch.bfloat16)
|
||||
weight = torch.randn(D, device=DEVICE, dtype=torch.float32)
|
||||
bias = (
|
||||
torch.randn(D, device=DEVICE, dtype=torch.float32)
|
||||
if norm_type == "layer"
|
||||
else None
|
||||
)
|
||||
scale = torch.randn(B, 1, D, device=DEVICE, dtype=torch.bfloat16)
|
||||
shift = torch.randn(B, 1, D, device=DEVICE, dtype=torch.bfloat16)
|
||||
|
||||
y = flydsl_norm_scale_shift(x, weight, bias, scale, shift, norm_type, EPS)
|
||||
y_ref = _ref_norm_ss(x, weight, bias, scale, shift, norm_type, EPS)
|
||||
torch.testing.assert_close(y, y_ref, atol=1.0, rtol=5e-2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,252 @@
|
||||
import sys
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import Tensor
|
||||
|
||||
from sglang.jit_kernel.diffusion.cutedsl.scale_residual_norm_scale_shift import (
|
||||
fused_norm_scale_shift,
|
||||
fused_scale_residual_norm_scale_shift,
|
||||
validate_scale_shift,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=28, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
DEVICE = "cuda"
|
||||
SHAPE_MAP = {
|
||||
"1": lambda B, S, F, D: (1,),
|
||||
"D": lambda B, S, F, D: (D,),
|
||||
"1D": lambda B, S, F, D: (1, D),
|
||||
"BD": lambda B, S, F, D: (B, D),
|
||||
"11D": lambda B, S, F, D: (1, 1, D),
|
||||
"B1D": lambda B, S, F, D: (B, 1, D),
|
||||
"1SD": lambda B, S, F, D: (1, S, D),
|
||||
"BSD": lambda B, S, F, D: (B, S, D),
|
||||
"BF1D": lambda B, S, F, D: (B, F, 1, D),
|
||||
}
|
||||
SHAPES = [
|
||||
# (B, S, F, D)
|
||||
(1, 115200, 1, 3072), # Hunyuan
|
||||
(1, 32760, 1, 1536), # Wan
|
||||
(1, 6, 1, 3072), # Qwen
|
||||
(1, 1024, 8, 3072),
|
||||
(4, 512, 16, 3072),
|
||||
]
|
||||
DTYPES = [torch.float16, torch.bfloat16, torch.float32]
|
||||
NORM_TYPES = ["layer", "rms"]
|
||||
AFFINE_MODES = ["D", "NAT"]
|
||||
INDEX_MODES = ["BSD", "1", "1SD", "BD", "B1D", "D", "1D", "11D", "BF1D"]
|
||||
|
||||
|
||||
def _tol(dtype: torch.dtype):
|
||||
return 1e-5 if dtype == torch.float32 else 5e-2
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def cuda_setup():
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA required")
|
||||
torch.cuda.manual_seed(0)
|
||||
|
||||
|
||||
def _apply_scale_shift(y: Tensor, scale: Tensor, shift: Tensor) -> Tensor:
|
||||
if scale.ndim == 4:
|
||||
num_frame = scale.shape[1]
|
||||
return rearrange(
|
||||
rearrange(y, "b (f l) d -> b f l d", f=num_frame) * (1 + scale) + shift,
|
||||
"b f l d -> b (f l) d",
|
||||
)
|
||||
else:
|
||||
scale = rearrange(scale, "b d -> b 1 d") if scale.ndim == 2 else scale
|
||||
shift = rearrange(shift, "b d -> b 1 d") if shift.ndim == 2 else shift
|
||||
return y * (1 + scale) + shift
|
||||
|
||||
|
||||
def fused_norm_scale_shift_ref(
|
||||
x: Tensor,
|
||||
weight: Optional[Tensor],
|
||||
bias: Optional[Tensor],
|
||||
scale: Tensor,
|
||||
shift: Tensor,
|
||||
norm_type: str,
|
||||
eps: float,
|
||||
) -> Tensor:
|
||||
original_dtype = x.dtype
|
||||
x, weight, bias, scale, shift = (
|
||||
v.float() if v is not None else v for v in [x, weight, bias, scale, shift]
|
||||
)
|
||||
if norm_type == "layer":
|
||||
norm = torch.layer_norm(x, x.shape[-1:], eps=eps, weight=weight, bias=bias)
|
||||
else:
|
||||
norm = torch.rms_norm(x, x.shape[-1:], eps=eps, weight=weight)
|
||||
return _apply_scale_shift(norm, scale, shift).to(original_dtype)
|
||||
|
||||
|
||||
def fused_scale_residual_norm_scale_shift_ref(
|
||||
residual: Tensor,
|
||||
x: Tensor,
|
||||
gate: Optional[Tensor] | int,
|
||||
weight: Optional[Tensor],
|
||||
bias: Optional[Tensor],
|
||||
scale: Tensor,
|
||||
shift: Tensor,
|
||||
norm_type: str,
|
||||
eps: float,
|
||||
):
|
||||
original_dtype = x.dtype
|
||||
residual, x, gate, weight, bias, scale, shift = (
|
||||
v.float() if isinstance(v, Tensor) else v
|
||||
for v in [residual, x, gate, weight, bias, scale, shift]
|
||||
)
|
||||
if isinstance(gate, int):
|
||||
x = residual + gate * x
|
||||
else:
|
||||
if gate.ndim == 4:
|
||||
num_frame = gate.shape[1]
|
||||
x_fld = rearrange(x, "b (f l) d -> b f l d", f=num_frame)
|
||||
x = residual + rearrange(x_fld * gate, "b f l d -> b (f l) d")
|
||||
else:
|
||||
gate = rearrange(gate, "b d -> b 1 d") if gate.ndim == 2 else gate
|
||||
x = residual + gate * x
|
||||
if norm_type == "layer":
|
||||
norm = torch.layer_norm(x, x.shape[-1:], eps=eps, weight=weight, bias=bias)
|
||||
else:
|
||||
norm = torch.rms_norm(x, x.shape[-1:], eps=eps, weight=weight)
|
||||
y_ref = _apply_scale_shift(norm, scale, shift)
|
||||
return y_ref.to(original_dtype), x.to(original_dtype)
|
||||
|
||||
|
||||
def _make_tensor(index_mode: str, shape: Tuple, dtype: torch.dtype):
|
||||
if index_mode == "NAT":
|
||||
return None
|
||||
return torch.randn(*SHAPE_MAP[index_mode](*shape), device=DEVICE, dtype=dtype)
|
||||
|
||||
|
||||
def test_validate_scale_shift_rejects_non_divisible_frames():
|
||||
with pytest.raises(ValueError, match=r"S\(10\) must be divisible by F\(4\)"):
|
||||
validate_scale_shift(
|
||||
torch.empty((1, 4, 1, 256), device=DEVICE, dtype=torch.float16),
|
||||
1,
|
||||
10,
|
||||
256,
|
||||
)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def run_norm_scale_shift(
|
||||
shape=SHAPES[0],
|
||||
dtype=DTYPES[0],
|
||||
affine_dtype=DTYPES[0],
|
||||
scale_dtype=DTYPES[0],
|
||||
shift_dtype=DTYPES[0],
|
||||
norm_type=NORM_TYPES[0],
|
||||
affine_mode=AFFINE_MODES[0],
|
||||
scale_mode="BSD",
|
||||
shift_mode="BSD",
|
||||
eps=1e-5,
|
||||
):
|
||||
x = _make_tensor("BSD", shape, dtype)
|
||||
weight = _make_tensor(affine_mode, shape, affine_dtype)
|
||||
bias = _make_tensor(affine_mode, shape, affine_dtype)
|
||||
scale = _make_tensor(scale_mode, shape, scale_dtype)
|
||||
shift = _make_tensor(shift_mode, shape, shift_dtype)
|
||||
y_dev = fused_norm_scale_shift(x, weight, bias, scale, shift, norm_type, eps)
|
||||
y_ref = fused_norm_scale_shift_ref(x, weight, bias, scale, shift, norm_type, eps)
|
||||
torch.testing.assert_close(y_dev, y_ref, atol=_tol(dtype), rtol=_tol(dtype))
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def run_scale_resi_norm_scale_shift(
|
||||
shape=SHAPES[0],
|
||||
dtype=DTYPES[0],
|
||||
affine_dtype=DTYPES[0],
|
||||
scale_dtype=DTYPES[0],
|
||||
shift_dtype=DTYPES[0],
|
||||
norm_type=NORM_TYPES[0],
|
||||
affine_mode=AFFINE_MODES[0],
|
||||
gate_mode="B1D",
|
||||
scale_mode="BSD",
|
||||
shift_mode="BSD",
|
||||
eps=1e-5,
|
||||
):
|
||||
residual = _make_tensor("BSD", shape, dtype)
|
||||
x = _make_tensor("BSD", shape, dtype)
|
||||
gate = _make_tensor(gate_mode, shape, dtype)
|
||||
weight = _make_tensor(affine_mode, shape, affine_dtype)
|
||||
bias = _make_tensor(affine_mode, shape, affine_dtype)
|
||||
scale = _make_tensor(scale_mode, shape, scale_dtype)
|
||||
shift = _make_tensor(shift_mode, shape, shift_dtype)
|
||||
y_dev, res_dev = fused_scale_residual_norm_scale_shift(
|
||||
residual, x, gate, weight, bias, scale, shift, norm_type, eps
|
||||
)
|
||||
y_ref, res_ref = fused_scale_residual_norm_scale_shift_ref(
|
||||
residual, x, gate, weight, bias, scale, shift, norm_type, eps
|
||||
)
|
||||
torch.testing.assert_close(y_dev, y_ref, atol=_tol(dtype), rtol=_tol(dtype))
|
||||
torch.testing.assert_close(res_dev, res_ref, atol=_tol(dtype), rtol=_tol(dtype))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("norm_type", NORM_TYPES)
|
||||
class TestFusedNormScaleShift:
|
||||
@pytest.mark.parametrize("shape", SHAPES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
def test_shape_dtype(self, shape, dtype, norm_type):
|
||||
run_norm_scale_shift(shape=shape, dtype=dtype, norm_type=norm_type)
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
def test_dtype_0(self, dtype, norm_type):
|
||||
run_norm_scale_shift(affine_dtype=dtype, norm_type=norm_type)
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
def test_dtype_1(self, dtype, norm_type):
|
||||
run_norm_scale_shift(scale_dtype=dtype, shift_dtype=dtype, norm_type=norm_type)
|
||||
|
||||
@pytest.mark.parametrize("affine_mode", AFFINE_MODES)
|
||||
def test_normtype_affine(self, affine_mode, norm_type):
|
||||
run_norm_scale_shift(affine_mode=affine_mode, norm_type=norm_type)
|
||||
|
||||
@pytest.mark.parametrize("index_mode", INDEX_MODES)
|
||||
def test_index_mode(self, index_mode, norm_type):
|
||||
run_norm_scale_shift(
|
||||
scale_mode=index_mode, shift_mode=index_mode, norm_type=norm_type
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("norm_type", NORM_TYPES)
|
||||
class TestFusedScaleResidualNormScaleShift:
|
||||
@pytest.mark.parametrize("shape", SHAPES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
def test_shape_dtype(self, shape, dtype, norm_type):
|
||||
run_scale_resi_norm_scale_shift(shape=shape, dtype=dtype, norm_type=norm_type)
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
def test_dtype_0(self, dtype, norm_type):
|
||||
run_scale_resi_norm_scale_shift(affine_dtype=dtype, norm_type=norm_type)
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
def test_dtype_1(self, dtype, norm_type):
|
||||
run_scale_resi_norm_scale_shift(
|
||||
scale_dtype=dtype, shift_dtype=dtype, norm_type=norm_type
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("affine_mode", AFFINE_MODES)
|
||||
def test_normtype_affine(self, affine_mode, norm_type):
|
||||
run_scale_resi_norm_scale_shift(affine_mode=affine_mode, norm_type=norm_type)
|
||||
|
||||
@pytest.mark.parametrize("index_mode", INDEX_MODES)
|
||||
def test_scale_shift_index_mode(self, index_mode, norm_type):
|
||||
run_scale_resi_norm_scale_shift(
|
||||
scale_mode=index_mode, shift_mode=index_mode, norm_type=norm_type
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("index_mode", INDEX_MODES)
|
||||
def test_gate_index_mode(self, index_mode, norm_type):
|
||||
run_scale_resi_norm_scale_shift(gate_mode=index_mode, norm_type=norm_type)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,104 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from sglang.jit_kernel.diffusion.group_norm_silu import apply_group_norm_silu
|
||||
from sglang.jit_kernel.diffusion.triton.group_norm_silu import triton_group_norm_silu
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=8, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
DEVICE = "cuda"
|
||||
DTYPES = [torch.float16, torch.bfloat16, torch.float32]
|
||||
TEST_CASES = [
|
||||
pytest.param((2, 64, 32, 32), 32, id="image_2d"),
|
||||
pytest.param((1, 64, 4, 16, 16), 32, id="video_3d"),
|
||||
pytest.param((4, 128), 32, id="token_2d"),
|
||||
]
|
||||
LARGE_TILE_CASE = ((1, 128, 20, 256, 256), 32)
|
||||
|
||||
|
||||
def _tol(dtype: torch.dtype) -> tuple[float, float]:
|
||||
if dtype == torch.float32:
|
||||
return 1e-5, 1e-5
|
||||
if dtype == torch.bfloat16:
|
||||
return 7e-2, 2e-2
|
||||
return 3e-3, 3e-3
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def cuda_setup():
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA required")
|
||||
torch.cuda.manual_seed(0)
|
||||
|
||||
|
||||
def _reference(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor,
|
||||
num_groups: int,
|
||||
eps: float = 1e-5,
|
||||
) -> torch.Tensor:
|
||||
return F.silu(F.group_norm(x, num_groups, weight=weight, bias=bias, eps=eps))
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
@pytest.mark.parametrize("shape,num_groups", TEST_CASES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
def test_triton_group_norm_silu(
|
||||
shape: tuple[int, ...], num_groups: int, dtype: torch.dtype
|
||||
) -> None:
|
||||
channels = shape[1]
|
||||
x = torch.randn(shape, device=DEVICE, dtype=dtype)
|
||||
weight = torch.randn(channels, device=DEVICE, dtype=dtype)
|
||||
bias = torch.randn(channels, device=DEVICE, dtype=dtype)
|
||||
|
||||
actual = triton_group_norm_silu(x, weight, bias, num_groups=num_groups)
|
||||
expected = _reference(x, weight, bias, num_groups)
|
||||
|
||||
atol, rtol = _tol(dtype)
|
||||
torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
@pytest.mark.parametrize("shape,num_groups", TEST_CASES[:2])
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
def test_apply_group_norm_silu(
|
||||
shape: tuple[int, ...],
|
||||
num_groups: int,
|
||||
dtype: torch.dtype,
|
||||
) -> None:
|
||||
norm = nn.GroupNorm(num_groups, shape[1], eps=1e-5, affine=True).to(
|
||||
device=DEVICE, dtype=dtype
|
||||
)
|
||||
activation = nn.SiLU()
|
||||
hidden_states = torch.randn(shape, device=DEVICE, dtype=dtype)
|
||||
|
||||
actual = apply_group_norm_silu(hidden_states, norm, activation)
|
||||
expected = activation(norm(hidden_states))
|
||||
|
||||
atol, rtol = _tol(dtype)
|
||||
torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def test_triton_group_norm_silu_large_tile_bf16() -> None:
|
||||
shape, num_groups = LARGE_TILE_CASE
|
||||
x = torch.randn(shape, device=DEVICE, dtype=torch.bfloat16)
|
||||
weight = torch.randn(shape[1], device=DEVICE, dtype=torch.bfloat16)
|
||||
bias = torch.randn(shape[1], device=DEVICE, dtype=torch.bfloat16)
|
||||
|
||||
actual = triton_group_norm_silu(x, weight, bias, num_groups=num_groups)
|
||||
expected = _reference(x, weight, bias, num_groups)
|
||||
|
||||
atol, rtol = _tol(torch.bfloat16)
|
||||
torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,153 @@
|
||||
import itertools
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.jit_kernel.utils import get_ci_test_range
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=44, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=176, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
DEVICE = "cuda"
|
||||
DTYPE = torch.bfloat16
|
||||
MAX_SEQ_LEN = 131072
|
||||
ROPE_BASE = 10000.0
|
||||
ATOL = 8e-2
|
||||
RTOL = 1e-2
|
||||
|
||||
|
||||
def create_cos_sin_cache(
|
||||
rotary_dim: int,
|
||||
max_position: int = MAX_SEQ_LEN,
|
||||
base: float = ROPE_BASE,
|
||||
) -> torch.Tensor:
|
||||
inv_freq = 1.0 / (
|
||||
base
|
||||
** (
|
||||
torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=DEVICE)
|
||||
/ rotary_dim
|
||||
)
|
||||
)
|
||||
t = torch.arange(max_position, dtype=torch.float32, device=DEVICE)
|
||||
freqs = torch.einsum("i,j->ij", t, inv_freq)
|
||||
return torch.cat((freqs.cos(), freqs.sin()), dim=-1)
|
||||
|
||||
|
||||
def split_qknorm_rope(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
cos_sin_cache: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
is_neox: bool,
|
||||
) -> None:
|
||||
from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace
|
||||
|
||||
from sglang.jit_kernel.norm import fused_inplace_qknorm
|
||||
|
||||
fused_inplace_qknorm(q, k, q_weight, k_weight)
|
||||
apply_rope_with_cos_sin_cache_inplace(
|
||||
positions=positions.long(),
|
||||
query=q.view(q.shape[0], -1),
|
||||
key=k.view(k.shape[0], -1),
|
||||
head_size=q.shape[-1],
|
||||
cos_sin_cache=cos_sin_cache,
|
||||
is_neox=is_neox,
|
||||
)
|
||||
|
||||
|
||||
def fused_qknorm_rope(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
cos_sin_cache: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
is_neox: bool,
|
||||
) -> None:
|
||||
from sglang.jit_kernel.diffusion.qknorm_rope import fused_inplace_qknorm_rope
|
||||
|
||||
fused_inplace_qknorm_rope(
|
||||
q,
|
||||
k,
|
||||
q_weight,
|
||||
k_weight,
|
||||
cos_sin_cache,
|
||||
positions,
|
||||
is_neox=is_neox,
|
||||
rope_dim=cos_sin_cache.shape[-1],
|
||||
)
|
||||
|
||||
|
||||
BS_LIST = [2**n for n in range(13)]
|
||||
BS_LIST += [x + 1 for x in BS_LIST]
|
||||
BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 129, 257, 2049, 4097])
|
||||
HEADS_LIST = get_ci_test_range([8, 16, 24, 32], [8, 24])
|
||||
HEAD_DIM_LIST = get_ci_test_range([64, 128, 256], [64, 128, 256])
|
||||
IS_NEOX_LIST = [False, True]
|
||||
POSITION_DTYPES = [torch.int32, torch.int64]
|
||||
ROPE_DIM_CHOICES = {
|
||||
64: [64],
|
||||
128: [64, 128],
|
||||
256: [64, 128, 256],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_size,num_heads,head_dim,is_neox,position_dtype",
|
||||
list(
|
||||
itertools.product(
|
||||
BS_LIST,
|
||||
HEADS_LIST,
|
||||
HEAD_DIM_LIST,
|
||||
IS_NEOX_LIST,
|
||||
POSITION_DTYPES,
|
||||
)
|
||||
),
|
||||
)
|
||||
def test_qknorm_rope(
|
||||
batch_size: int,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
is_neox: bool,
|
||||
position_dtype: torch.dtype,
|
||||
) -> None:
|
||||
rope_dims = ROPE_DIM_CHOICES[head_dim]
|
||||
for rope_dim in rope_dims:
|
||||
if is_neox:
|
||||
elems_per_thread = head_dim // 32
|
||||
rotary_lanes = rope_dim // elems_per_thread
|
||||
if rotary_lanes < 2 or rotary_lanes & (rotary_lanes - 1):
|
||||
continue
|
||||
|
||||
q = torch.randn(batch_size, num_heads, head_dim, device=DEVICE, dtype=DTYPE)
|
||||
k = torch.randn(batch_size, num_heads, head_dim, device=DEVICE, dtype=DTYPE)
|
||||
q_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
|
||||
k_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
|
||||
positions = torch.randint(
|
||||
0, MAX_SEQ_LEN, (batch_size,), device=DEVICE, dtype=position_dtype
|
||||
)
|
||||
cos_sin_cache = create_cos_sin_cache(rope_dim)
|
||||
|
||||
q_ref, k_ref = q.clone(), k.clone()
|
||||
q_fused, k_fused = q.clone(), k.clone()
|
||||
|
||||
split_qknorm_rope(
|
||||
q_ref, k_ref, q_weight, k_weight, cos_sin_cache, positions, is_neox
|
||||
)
|
||||
fused_qknorm_rope(
|
||||
q_fused, k_fused, q_weight, k_weight, cos_sin_cache, positions, is_neox
|
||||
)
|
||||
|
||||
# The split baseline mixes a separate BF16 qknorm kernel with FlashInfer RoPE,
|
||||
# which differs from the fused path by about one BF16 rounding step on H200.
|
||||
triton.testing.assert_close(q_ref, q_fused, atol=ATOL, rtol=RTOL)
|
||||
triton.testing.assert_close(k_ref, k_fused, atol=ATOL, rtol=RTOL)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,226 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.jit_kernel.diffusion.triton.norm import norm_infer
|
||||
from sglang.jit_kernel.diffusion.triton.scale_shift import (
|
||||
fuse_layernorm_scale_shift_gate_select01_kernel,
|
||||
fuse_residual_layernorm_scale_shift_gate_select01_kernel,
|
||||
)
|
||||
from sglang.jit_kernel.utils import get_ci_test_range
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=15, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
DEVICE = "cuda"
|
||||
DTYPES = get_ci_test_range(
|
||||
[torch.float16, torch.bfloat16, torch.float32], [torch.float16, torch.bfloat16]
|
||||
)
|
||||
BATCH_SIZES = get_ci_test_range([1, 2, 4], [1, 2])
|
||||
SEQ_LENS = get_ci_test_range([6, 33, 128, 257], [6, 128])
|
||||
HIDDEN_SIZES = get_ci_test_range([512, 1024, 1536, 3072], [512, 3072])
|
||||
EPS = 1e-6
|
||||
|
||||
|
||||
def _tol(dtype: torch.dtype) -> tuple[float, float]:
|
||||
if dtype == torch.float32:
|
||||
return 1e-5, 1e-5
|
||||
return 5e-2, 5e-2
|
||||
|
||||
|
||||
def _make_modulation_tensors(batch_size: int, hidden_size: int, dtype: torch.dtype):
|
||||
scale0 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
|
||||
shift0 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
|
||||
gate0 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
|
||||
scale1 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
|
||||
shift1 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
|
||||
gate1 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
|
||||
return scale0, shift0, gate0, scale1, shift1, gate1
|
||||
|
||||
|
||||
def _baseline_select01_modulation(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor | None,
|
||||
bias: torch.Tensor | None,
|
||||
scale0: torch.Tensor,
|
||||
shift0: torch.Tensor,
|
||||
gate0: torch.Tensor,
|
||||
scale1: torch.Tensor,
|
||||
shift1: torch.Tensor,
|
||||
gate1: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
eps: float,
|
||||
):
|
||||
normalized = norm_infer(
|
||||
x.view(-1, x.shape[-1]),
|
||||
weight,
|
||||
bias,
|
||||
eps=eps,
|
||||
is_rms_norm=False,
|
||||
).view_as(x)
|
||||
return _apply_select01_modulation(
|
||||
normalized, scale0, shift0, gate0, scale1, shift1, gate1, index
|
||||
)
|
||||
|
||||
|
||||
def _baseline_residual_select01_modulation(
|
||||
x: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
residual_gate: torch.Tensor,
|
||||
weight: torch.Tensor | None,
|
||||
bias: torch.Tensor | None,
|
||||
scale0: torch.Tensor,
|
||||
shift0: torch.Tensor,
|
||||
gate0: torch.Tensor,
|
||||
scale1: torch.Tensor,
|
||||
shift1: torch.Tensor,
|
||||
gate1: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
eps: float,
|
||||
):
|
||||
residual_out = residual + residual_gate * x
|
||||
normalized = norm_infer(
|
||||
residual_out.view(-1, residual_out.shape[-1]),
|
||||
weight,
|
||||
bias,
|
||||
eps=eps,
|
||||
is_rms_norm=False,
|
||||
).view_as(residual_out)
|
||||
output, gate_out = _apply_select01_modulation(
|
||||
normalized, scale0, shift0, gate0, scale1, shift1, gate1, index
|
||||
)
|
||||
return output, residual_out, gate_out
|
||||
|
||||
|
||||
def _apply_select01_modulation(
|
||||
x: torch.Tensor,
|
||||
scale0: torch.Tensor,
|
||||
shift0: torch.Tensor,
|
||||
gate0: torch.Tensor,
|
||||
scale1: torch.Tensor,
|
||||
shift1: torch.Tensor,
|
||||
gate1: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
):
|
||||
idx = index.bool().unsqueeze(-1)
|
||||
scale = torch.where(idx, scale1.unsqueeze(1), scale0.unsqueeze(1))
|
||||
shift = torch.where(idx, shift1.unsqueeze(1), shift0.unsqueeze(1))
|
||||
gate = torch.where(idx, gate1.unsqueeze(1), gate0.unsqueeze(1))
|
||||
return x * (1 + scale) + shift, gate
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def cuda_setup():
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA required")
|
||||
torch.cuda.manual_seed(0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("batch_size", BATCH_SIZES)
|
||||
@pytest.mark.parametrize("seq_len", SEQ_LENS)
|
||||
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
|
||||
def test_fused_layernorm_scale_shift_gate_select01(
|
||||
dtype, batch_size, seq_len, hidden_size
|
||||
):
|
||||
x = torch.randn(batch_size, seq_len, hidden_size, device=DEVICE, dtype=dtype)
|
||||
weight = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
|
||||
bias = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
|
||||
index = torch.randint(0, 2, (batch_size, seq_len), device=DEVICE, dtype=torch.int32)
|
||||
scale0, shift0, gate0, scale1, shift1, gate1 = _make_modulation_tensors(
|
||||
batch_size, hidden_size, dtype
|
||||
)
|
||||
|
||||
out_ref, gate_ref = _baseline_select01_modulation(
|
||||
x,
|
||||
weight,
|
||||
bias,
|
||||
scale0,
|
||||
shift0,
|
||||
gate0,
|
||||
scale1,
|
||||
shift1,
|
||||
gate1,
|
||||
index,
|
||||
EPS,
|
||||
)
|
||||
out_fused, gate_fused = fuse_layernorm_scale_shift_gate_select01_kernel(
|
||||
x.contiguous(),
|
||||
weight=weight,
|
||||
bias=bias,
|
||||
scale0=scale0,
|
||||
shift0=shift0,
|
||||
gate0=gate0,
|
||||
scale1=scale1,
|
||||
shift1=shift1,
|
||||
gate1=gate1,
|
||||
index=index,
|
||||
eps=EPS,
|
||||
)
|
||||
|
||||
atol, rtol = _tol(dtype)
|
||||
triton.testing.assert_close(out_ref, out_fused, atol=atol, rtol=rtol)
|
||||
triton.testing.assert_close(gate_ref, gate_fused, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("batch_size", BATCH_SIZES)
|
||||
@pytest.mark.parametrize("seq_len", SEQ_LENS)
|
||||
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
|
||||
def test_fused_residual_layernorm_scale_shift_gate_select01(
|
||||
dtype, batch_size, seq_len, hidden_size
|
||||
):
|
||||
x = torch.randn(batch_size, seq_len, hidden_size, device=DEVICE, dtype=dtype)
|
||||
residual = torch.randn_like(x)
|
||||
residual_gate = torch.randn_like(x)
|
||||
weight = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
|
||||
bias = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
|
||||
index = torch.randint(0, 2, (batch_size, seq_len), device=DEVICE, dtype=torch.int32)
|
||||
scale0, shift0, gate0, scale1, shift1, gate1 = _make_modulation_tensors(
|
||||
batch_size, hidden_size, dtype
|
||||
)
|
||||
|
||||
out_ref, residual_ref, gate_ref = _baseline_residual_select01_modulation(
|
||||
x,
|
||||
residual,
|
||||
residual_gate,
|
||||
weight,
|
||||
bias,
|
||||
scale0,
|
||||
shift0,
|
||||
gate0,
|
||||
scale1,
|
||||
shift1,
|
||||
gate1,
|
||||
index,
|
||||
EPS,
|
||||
)
|
||||
out_fused, residual_fused, gate_fused = (
|
||||
fuse_residual_layernorm_scale_shift_gate_select01_kernel(
|
||||
x.contiguous(),
|
||||
residual=residual.contiguous(),
|
||||
residual_gate=residual_gate.contiguous(),
|
||||
weight=weight,
|
||||
bias=bias,
|
||||
scale0=scale0,
|
||||
shift0=shift0,
|
||||
gate0=gate0,
|
||||
scale1=scale1,
|
||||
shift1=shift1,
|
||||
gate1=gate1,
|
||||
index=index,
|
||||
eps=EPS,
|
||||
)
|
||||
)
|
||||
|
||||
atol, rtol = _tol(dtype)
|
||||
triton.testing.assert_close(out_ref, out_fused, atol=atol, rtol=rtol)
|
||||
triton.testing.assert_close(residual_ref, residual_fused, atol=atol, rtol=rtol)
|
||||
triton.testing.assert_close(gate_ref, gate_fused, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,193 @@
|
||||
"""Numerical correctness for fused varlen pack/scatter Triton kernels.
|
||||
|
||||
Bit-exact comparison against the equivalent PyTorch ops (index_select,
|
||||
zeros + index_copy_) across bf16/fp16 and several shape/mask cases.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.diffusion.triton.varlen_pack_pad import (
|
||||
build_inv_indices,
|
||||
fused_pack_qkv,
|
||||
fused_scatter_to_padded,
|
||||
)
|
||||
from sglang.jit_kernel.utils import get_ci_test_range
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=60, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
DEVICE = "cuda"
|
||||
DTYPES = get_ci_test_range([torch.bfloat16, torch.float16], [torch.bfloat16])
|
||||
# (bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens) tuples
|
||||
SHAPES = get_ci_test_range(
|
||||
[
|
||||
# name, bs, s_txt, s_img, H, D, valid_txt_lens
|
||||
("small_c2", 2, 64, 128, 4, 64, [32, 48]),
|
||||
("prod_c2", 2, 256, 1024, 24, 128, [128, 200]),
|
||||
("all_valid_b1", 1, 64, 128, 4, 64, [64]),
|
||||
("all_valid_b4", 4, 64, 128, 4, 64, [64, 64, 64, 64]),
|
||||
("c8_prod", 8, 256, 4096, 24, 128, [128, 200, 256, 100, 50, 256, 256, 50]),
|
||||
# one batch with zero valid text tokens (image side still valid)
|
||||
("zero_txt_one_batch", 2, 64, 128, 4, 64, [0, 32]),
|
||||
# bs=1 with no text validity (only image rows packed)
|
||||
("bs1_zero_txt", 1, 64, 128, 4, 64, [0]),
|
||||
],
|
||||
[
|
||||
("small_c2", 2, 64, 128, 4, 64, [32, 48]),
|
||||
("prod_c2", 2, 256, 1024, 24, 128, [128, 200]),
|
||||
("all_valid_b4", 4, 64, 128, 4, 64, [64, 64, 64, 64]),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _build_mask(bs, s_txt, s_img, valid_txt_lens):
|
||||
s = s_txt + s_img
|
||||
mask = torch.zeros(bs, s, dtype=torch.bool, device=DEVICE)
|
||||
for b, vt in enumerate(valid_txt_lens):
|
||||
mask[b, :vt] = True
|
||||
mask[b, s_txt:] = True
|
||||
return mask
|
||||
|
||||
|
||||
def _ref_pack(q, k, v, indices):
|
||||
bs, seq = q.shape[:2]
|
||||
flat = lambda t: t.reshape(bs * seq, *t.shape[2:])
|
||||
return (
|
||||
flat(q).index_select(0, indices),
|
||||
flat(k).index_select(0, indices),
|
||||
flat(v).index_select(0, indices),
|
||||
)
|
||||
|
||||
|
||||
def _ref_scatter(out_unpad, indices, bs, seq):
|
||||
n_valid = indices.shape[0]
|
||||
_, num_heads, head_dim = out_unpad.shape
|
||||
flat = torch.zeros(
|
||||
bs * seq, num_heads, head_dim, dtype=out_unpad.dtype, device=DEVICE
|
||||
)
|
||||
flat.index_copy_(0, indices, out_unpad)
|
||||
return flat.view(bs, seq, num_heads, head_dim)
|
||||
|
||||
|
||||
def _build_meta(mask):
|
||||
bs, seq = mask.shape
|
||||
indices = mask.reshape(-1).nonzero(as_tuple=False).flatten()
|
||||
inv_indices = build_inv_indices(indices, bs * seq)
|
||||
return indices, inv_indices
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize(
|
||||
"shape", SHAPES, ids=lambda s: s[0] if isinstance(s, tuple) else str(s)
|
||||
)
|
||||
def test_pack_matches_index_select(dtype, shape):
|
||||
_, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens = shape
|
||||
torch.manual_seed(0)
|
||||
s = s_txt + s_img
|
||||
mask = _build_mask(bs, s_txt, s_img, valid_txt_lens)
|
||||
indices, _ = _build_meta(mask)
|
||||
|
||||
q = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
|
||||
k = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
|
||||
v = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
|
||||
|
||||
q_ref, k_ref, v_ref = _ref_pack(q, k, v, indices)
|
||||
q_fused, k_fused, v_fused = fused_pack_qkv(q, k, v, indices)
|
||||
|
||||
# bit-exact: pack is pure gather, no math
|
||||
assert torch.equal(q_ref, q_fused)
|
||||
assert torch.equal(k_ref, k_fused)
|
||||
assert torch.equal(v_ref, v_fused)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize(
|
||||
"shape", SHAPES, ids=lambda s: s[0] if isinstance(s, tuple) else str(s)
|
||||
)
|
||||
def test_scatter_matches_index_copy(dtype, shape):
|
||||
_, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens = shape
|
||||
torch.manual_seed(1)
|
||||
s = s_txt + s_img
|
||||
mask = _build_mask(bs, s_txt, s_img, valid_txt_lens)
|
||||
indices, inv_indices = _build_meta(mask)
|
||||
n_valid = indices.shape[0]
|
||||
|
||||
out_unpad = torch.randn(n_valid, num_heads, head_dim, dtype=dtype, device=DEVICE)
|
||||
out_ref = _ref_scatter(out_unpad, indices, bs, s)
|
||||
out_fused = fused_scatter_to_padded(out_unpad, inv_indices, bs, s)
|
||||
|
||||
# bit-exact: scatter is pure copy + zero-fill
|
||||
assert torch.equal(out_ref, out_fused)
|
||||
# Padding rows must be exactly zero
|
||||
invalid = ~mask
|
||||
if invalid.any():
|
||||
assert out_fused[invalid].abs().max().item() == 0.0
|
||||
|
||||
|
||||
def test_pack_handles_non_contiguous_input():
|
||||
"""Helper must accept non-contiguous Q/K/V (auto .contiguous() inside)."""
|
||||
torch.manual_seed(2)
|
||||
bs, s_txt, s_img, num_heads, head_dim = 2, 64, 128, 4, 64
|
||||
s = s_txt + s_img
|
||||
mask = _build_mask(bs, s_txt, s_img, [32, 48])
|
||||
indices, _ = _build_meta(mask)
|
||||
|
||||
# Build non-contiguous tensors via permute
|
||||
qkv_pre = torch.randn(
|
||||
bs, num_heads, s, head_dim, dtype=torch.bfloat16, device=DEVICE
|
||||
)
|
||||
q = qkv_pre.permute(0, 2, 1, 3)
|
||||
k = torch.randn_like(qkv_pre).permute(0, 2, 1, 3)
|
||||
v = torch.randn_like(qkv_pre).permute(0, 2, 1, 3)
|
||||
assert not q.is_contiguous()
|
||||
|
||||
q_ref, k_ref, v_ref = _ref_pack(
|
||||
q.contiguous(), k.contiguous(), v.contiguous(), indices
|
||||
)
|
||||
q_fused, k_fused, v_fused = fused_pack_qkv(q, k, v, indices)
|
||||
assert torch.equal(q_ref, q_fused)
|
||||
assert torch.equal(k_ref, k_fused)
|
||||
assert torch.equal(v_ref, v_fused)
|
||||
|
||||
|
||||
def test_build_inv_indices_matches_manual():
|
||||
"""build_inv_indices output should match the manual full+scatter form."""
|
||||
torch.manual_seed(3)
|
||||
bs, s = 2, 32
|
||||
mask = torch.bernoulli(torch.full((bs, s), 0.6, device=DEVICE)).to(torch.bool)
|
||||
indices = mask.reshape(-1).nonzero(as_tuple=False).flatten()
|
||||
n_valid = indices.shape[0]
|
||||
|
||||
manual = torch.full((bs * s,), -1, dtype=torch.int32, device=DEVICE)
|
||||
if n_valid > 0:
|
||||
manual[indices.long()] = torch.arange(n_valid, dtype=torch.int32, device=DEVICE)
|
||||
|
||||
built = build_inv_indices(indices, bs * s)
|
||||
assert torch.equal(built, manual)
|
||||
|
||||
|
||||
def test_empty_valid_set_handled():
|
||||
"""All-False mask: pack returns empty tensors; scatter writes all zeros."""
|
||||
bs, s, num_heads, head_dim = 2, 16, 4, 64
|
||||
mask = torch.zeros(bs, s, dtype=torch.bool, device=DEVICE)
|
||||
indices = mask.reshape(-1).nonzero(as_tuple=False).flatten()
|
||||
inv_indices = build_inv_indices(indices, bs * s)
|
||||
assert indices.numel() == 0
|
||||
|
||||
q = torch.randn(bs, s, num_heads, head_dim, dtype=torch.bfloat16, device=DEVICE)
|
||||
q_unpad, k_unpad, v_unpad = fused_pack_qkv(q, q.clone(), q.clone(), indices)
|
||||
assert q_unpad.shape == (0, num_heads, head_dim)
|
||||
assert k_unpad.shape == (0, num_heads, head_dim)
|
||||
assert v_unpad.shape == (0, num_heads, head_dim)
|
||||
|
||||
out_padded = fused_scatter_to_padded(q_unpad, inv_indices, bs, s)
|
||||
assert out_padded.shape == (bs, s, num_heads, head_dim)
|
||||
assert out_padded.abs().max().item() == 0.0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,156 @@
|
||||
"""End-to-end equivalence between USPAttention varlen path and SDPA reference.
|
||||
|
||||
Compares the production varlen path (``build_varlen_mask_meta`` +
|
||||
``fused_pack_qkv`` + ``flash_attn_varlen_func`` + ``fused_scatter_to_padded``)
|
||||
against ``torch.nn.functional.scaled_dot_product_attention`` with a broadcast
|
||||
key mask, for inputs the gating in ``USPAttention.forward`` would accept.
|
||||
|
||||
Verifies the documented contract:
|
||||
* Valid (non-masked) query rows match SDPA within FA-vs-SDPA tolerance.
|
||||
* Masked query rows are exactly zero in the varlen path (differs from
|
||||
SDPA, which produces deterministic attention output at those rows).
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from sglang.jit_kernel.diffusion.triton.varlen_pack_pad import (
|
||||
fused_pack_qkv,
|
||||
fused_scatter_to_padded,
|
||||
)
|
||||
from sglang.jit_kernel.flash_attention import flash_attn_varlen_func
|
||||
from sglang.jit_kernel.utils import get_ci_test_range
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends import (
|
||||
flash_attn as _fa_backend,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.attention.layer import (
|
||||
build_varlen_mask_meta,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=15, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=60, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
DEVICE = "cuda"
|
||||
DTYPES = get_ci_test_range([torch.bfloat16, torch.float16], [torch.bfloat16])
|
||||
# (name, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens)
|
||||
SHAPES = get_ci_test_range(
|
||||
[
|
||||
("small_c2", 2, 64, 128, 4, 64, [32, 48]),
|
||||
("prod_c2", 2, 256, 1024, 24, 128, [128, 200]),
|
||||
("all_valid_b1", 1, 64, 128, 4, 64, [64]),
|
||||
("zero_txt_one_batch", 2, 64, 128, 4, 64, [0, 32]),
|
||||
],
|
||||
[
|
||||
("small_c2", 2, 64, 128, 4, 64, [32, 48]),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _build_mask(bs, s_txt, s_img, valid_txt_lens):
|
||||
s = s_txt + s_img
|
||||
mask = torch.zeros(bs, s, dtype=torch.bool, device=DEVICE)
|
||||
for b, vt in enumerate(valid_txt_lens):
|
||||
mask[b, :vt] = True
|
||||
mask[b, s_txt:] = True
|
||||
return mask
|
||||
|
||||
|
||||
def _sdpa_with_key_mask(q, k, v, key_mask, softmax_scale):
|
||||
"""Reference: SDPA with a ``[B, S]`` key mask broadcast to ``[B, 1, 1, S]``."""
|
||||
q_ = q.transpose(1, 2)
|
||||
k_ = k.transpose(1, 2)
|
||||
v_ = v.transpose(1, 2)
|
||||
mask = key_mask.to(dtype=q.dtype)[:, None, None, :]
|
||||
mask = (mask - 1.0) * torch.finfo(q.dtype).max
|
||||
out = F.scaled_dot_product_attention(
|
||||
q_,
|
||||
k_,
|
||||
v_,
|
||||
attn_mask=mask,
|
||||
dropout_p=0.0,
|
||||
is_causal=False,
|
||||
scale=softmax_scale,
|
||||
)
|
||||
return out.transpose(1, 2)
|
||||
|
||||
|
||||
def _varlen_path(q, k, v, key_mask, softmax_scale):
|
||||
"""Production varlen path matching USPAttention.forward masked branch."""
|
||||
bs, seq = q.shape[0], q.shape[1]
|
||||
meta = build_varlen_mask_meta(key_mask)
|
||||
indices = meta["indices"]
|
||||
if indices.shape[0] == 0:
|
||||
return torch.zeros_like(q)
|
||||
q_unpad, k_unpad, v_unpad = fused_pack_qkv(q, k, v, indices)
|
||||
out_unpad = flash_attn_varlen_func(
|
||||
q=q_unpad,
|
||||
k=k_unpad,
|
||||
v=v_unpad,
|
||||
cu_seqlens_q=meta["cu_seqlens"],
|
||||
cu_seqlens_k=meta["cu_seqlens"],
|
||||
max_seqlen_q=meta["max_seqlen"],
|
||||
max_seqlen_k=meta["max_seqlen"],
|
||||
softmax_scale=softmax_scale,
|
||||
causal=False,
|
||||
ver=_fa_backend.fa_ver,
|
||||
)
|
||||
return fused_scatter_to_padded(out_unpad, meta["inv_indices"], bs, seq)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize(
|
||||
"shape", SHAPES, ids=lambda s: s[0] if isinstance(s, tuple) else str(s)
|
||||
)
|
||||
def test_varlen_path_matches_sdpa_on_valid_rows(dtype, shape):
|
||||
"""Valid rows: varlen output ≈ SDPA output within FA tolerance."""
|
||||
_, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens = shape
|
||||
torch.manual_seed(0)
|
||||
s = s_txt + s_img
|
||||
softmax_scale = head_dim**-0.5
|
||||
mask = _build_mask(bs, s_txt, s_img, valid_txt_lens)
|
||||
q = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
|
||||
k = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
|
||||
v = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
|
||||
|
||||
out_sdpa = _sdpa_with_key_mask(q, k, v, mask, softmax_scale)
|
||||
out_varlen = _varlen_path(q, k, v, mask, softmax_scale)
|
||||
|
||||
valid = mask[..., None, None].expand_as(out_sdpa)
|
||||
rtol = 1e-2 if dtype == torch.bfloat16 else 5e-3
|
||||
atol = 5e-2 if dtype == torch.bfloat16 else 1e-2
|
||||
torch.testing.assert_close(
|
||||
out_sdpa[valid],
|
||||
out_varlen[valid],
|
||||
rtol=rtol,
|
||||
atol=atol,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize(
|
||||
"shape", SHAPES, ids=lambda s: s[0] if isinstance(s, tuple) else str(s)
|
||||
)
|
||||
def test_varlen_path_zeros_masked_rows(dtype, shape):
|
||||
"""Masked rows: varlen path produces exact zeros (documented contract)."""
|
||||
_, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens = shape
|
||||
torch.manual_seed(1)
|
||||
s = s_txt + s_img
|
||||
softmax_scale = head_dim**-0.5
|
||||
mask = _build_mask(bs, s_txt, s_img, valid_txt_lens)
|
||||
q = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
|
||||
k = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
|
||||
v = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
|
||||
|
||||
out_varlen = _varlen_path(q, k, v, mask, softmax_scale)
|
||||
|
||||
invalid = ~mask
|
||||
if invalid.any():
|
||||
assert (out_varlen[invalid] == 0).all(), "masked rows must be zero-filled"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,76 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import sglang.jit_kernel
|
||||
from sglang.jit_kernel.kv_canary import consts
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=5, suite="base-b-kernel-unit-1-gpu-large")
|
||||
|
||||
|
||||
# Resolve the kernel source against the installed jit_kernel package rather
|
||||
# than this file's location, so the test stays correct wherever it lives.
|
||||
_CONSTS_CUH: Path = (
|
||||
Path(sglang.jit_kernel.__file__).resolve().parent
|
||||
/ "csrc"
|
||||
/ "kv_canary"
|
||||
/ "consts.cuh"
|
||||
)
|
||||
|
||||
|
||||
def _camel_to_upper_snake(name: str) -> str:
|
||||
return re.sub(r"([A-Z])", r"_\1", name).lstrip("_").upper()
|
||||
|
||||
|
||||
def _decode(expr: str) -> int:
|
||||
expr = expr.strip().rstrip("UuLl")
|
||||
if "<<" in expr:
|
||||
return 1 << int(expr.split("<<")[1].strip())
|
||||
return int(expr, 0)
|
||||
|
||||
|
||||
def _parse_constexpr_ints(source: str) -> dict[str, int]:
|
||||
pattern = re.compile(r"constexpr\s+(?:[\w:]+)\s+(k[A-Za-z]\w*)\s*=\s*([^;]+);")
|
||||
return {name: _decode(rhs) for name, rhs in pattern.findall(source)}
|
||||
|
||||
|
||||
def _parse_enum_class(source: str, enum_name: str) -> dict[str, int]:
|
||||
pattern = re.compile(
|
||||
r"enum\s+class\s+" + re.escape(enum_name) + r"\s*:\s*[^\{]+\{([^}]+)\}"
|
||||
)
|
||||
body = pattern.search(source).group(1)
|
||||
member_re = re.compile(r"(k[A-Za-z]\w*)\s*=\s*([^,]+)")
|
||||
return {name: _decode(rhs) for name, rhs in member_re.findall(body)}
|
||||
|
||||
|
||||
def test_int_consts_sync() -> None:
|
||||
cpp = _parse_constexpr_ints(_CONSTS_CUH.read_text(encoding="utf-8"))
|
||||
cpp_normalized = {_camel_to_upper_snake(n[1:]): v for n, v in cpp.items()}
|
||||
py = {
|
||||
n: v
|
||||
for n, v in vars(consts).items()
|
||||
if isinstance(v, int) and not isinstance(v, bool) and not n.startswith("_")
|
||||
}
|
||||
assert cpp_normalized == py
|
||||
|
||||
|
||||
def test_enums_sync() -> None:
|
||||
cuh = _CONSTS_CUH.read_text(encoding="utf-8")
|
||||
for enum_name in ("RealKvHashMode", "FailReason"):
|
||||
cpp_members = _parse_enum_class(cuh, enum_name)
|
||||
py_enum = getattr(consts, enum_name)
|
||||
cpp_normalized = {
|
||||
_camel_to_upper_snake(n[1:]): v for n, v in cpp_members.items()
|
||||
}
|
||||
py_normalized = {m.name: int(m.value) for m in py_enum}
|
||||
assert cpp_normalized == py_normalized
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -0,0 +1,440 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary import consts
|
||||
from sglang.jit_kernel.kv_canary.verify import (
|
||||
CanaryLaunchTag,
|
||||
VerifyOrWriteContext,
|
||||
VerifyPlan,
|
||||
launch_canary_verify_kernel,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary.write import WritePlan
|
||||
from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
|
||||
FakeViolationLog,
|
||||
assert_canary_buf_equal,
|
||||
assert_canary_state_equal,
|
||||
make_canary_buf,
|
||||
make_canary_buf_pair,
|
||||
make_log_pair,
|
||||
make_verify_plan,
|
||||
make_verify_plan_pair,
|
||||
make_write_plan_pair,
|
||||
stamp_clean_chain,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._differential import (
|
||||
_assert_plans_byte_equal,
|
||||
_run_both_plan,
|
||||
_run_both_verify,
|
||||
_run_both_write,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._fixtures import (
|
||||
dummy_pseudo_tensors,
|
||||
empty_extras,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=60, suite="base-b-kernel-unit-1-gpu-large")
|
||||
|
||||
_DEVICE = torch.device("cuda")
|
||||
|
||||
|
||||
def _build_verify_plan_5_entries(
|
||||
*, device: torch.device
|
||||
) -> tuple[VerifyPlan, VerifyPlan]:
|
||||
num_slots = 16
|
||||
cuda_buf, ref_buf = make_canary_buf_pair(
|
||||
num_slots=num_slots, slot_stride_bytes=32, device=_DEVICE
|
||||
)
|
||||
|
||||
plan_cuda, plan_ref = make_verify_plan_pair(
|
||||
slot_indices=[0, 1, 2, 3, 4],
|
||||
positions=[0, 1, 2, 3, 4],
|
||||
prev_slot_indices=[-1, 0, 1, 2, 3],
|
||||
capacity=8,
|
||||
device=device,
|
||||
)
|
||||
return plan_cuda, plan_ref
|
||||
|
||||
|
||||
def _build_write_fixtures(
|
||||
*, device: torch.device
|
||||
) -> tuple[WritePlan, WritePlan, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
num_tokens = 5
|
||||
plan_cuda, plan_ref = make_write_plan_pair(
|
||||
write_offsets=[0, num_tokens],
|
||||
seed_slot_indices=[-1],
|
||||
num_valid_reqs=1,
|
||||
req_capacity=4,
|
||||
device=device,
|
||||
)
|
||||
input_ids = torch.tensor([10, 20, 30, 40, 50], dtype=torch.int64, device=device)
|
||||
positions = torch.tensor([0, 1, 2, 3, 4], dtype=torch.int64, device=device)
|
||||
out_cache_loc = torch.tensor([0, 1, 2, 3, 4], dtype=torch.int64, device=device)
|
||||
return plan_cuda, plan_ref, input_ids, positions, out_cache_loc
|
||||
|
||||
|
||||
def _build_plan_fixtures(
|
||||
*, device: torch.device, int64_req_to_token: bool = False
|
||||
) -> tuple[
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
]:
|
||||
bs = 3
|
||||
max_reqs = 4
|
||||
max_seq_len = 16
|
||||
req_pool_indices = torch.tensor([1, 2, 3], dtype=torch.int64, device=device)
|
||||
prefix_lens = torch.tensor([0, 4, 8], dtype=torch.int64, device=device)
|
||||
extend_seq_lens = torch.tensor([5, 1, 1], dtype=torch.int64, device=device)
|
||||
rp_axis = torch.arange(max_reqs, device=device, dtype=torch.int32).unsqueeze(1)
|
||||
pos_axis = torch.arange(max_seq_len, device=device, dtype=torch.int32).unsqueeze(0)
|
||||
req_to_token_int32 = (rp_axis * max_seq_len + pos_axis).contiguous()
|
||||
if int64_req_to_token:
|
||||
req_to_token = req_to_token_int32.to(torch.int64)
|
||||
else:
|
||||
req_to_token = req_to_token_int32
|
||||
return req_pool_indices, prefix_lens, extend_seq_lens, req_to_token
|
||||
|
||||
|
||||
def test_verify_byte_equal_across_repeated_launches_10x() -> None:
|
||||
num_launches = 10
|
||||
plan_cuda, plan_ref = _build_verify_plan_5_entries(device=_DEVICE)
|
||||
|
||||
snapshot_rings: list[torch.Tensor] = []
|
||||
snapshot_write_indices: list[torch.Tensor] = []
|
||||
snapshot_bufs: list[torch.Tensor] = []
|
||||
|
||||
for _ in range(num_launches):
|
||||
cuda_buf, ref_buf = make_canary_buf_pair(
|
||||
num_slots=16, slot_stride_bytes=32, device=_DEVICE
|
||||
)
|
||||
cuda_log, ref_log = make_log_pair(capacity=64, device=_DEVICE)
|
||||
|
||||
_run_both_verify(
|
||||
cuda_canary_buf=cuda_buf,
|
||||
ref_canary_buf=ref_buf,
|
||||
plan_cuda=plan_cuda,
|
||||
plan_ref=plan_ref,
|
||||
cuda_log=cuda_log,
|
||||
ref_log=ref_log,
|
||||
real_kv_sources_cuda=(),
|
||||
real_kv_sources_ref=(),
|
||||
real_kv_hash_mode=consts.RealKvHashMode.NONE,
|
||||
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
|
||||
)
|
||||
|
||||
assert_canary_buf_equal(buf_a=cuda_buf, buf_b=ref_buf)
|
||||
assert_canary_state_equal(log_a=cuda_log, log_b=ref_log)
|
||||
|
||||
snapshot_rings.append(cuda_log.ring.clone())
|
||||
snapshot_write_indices.append(cuda_log.write_index.clone())
|
||||
snapshot_bufs.append(cuda_buf.clone())
|
||||
|
||||
for i in range(1, num_launches):
|
||||
assert torch.equal(
|
||||
snapshot_rings[0], snapshot_rings[i]
|
||||
), f"violation_ring differs between launch 0 and {i}"
|
||||
assert torch.equal(
|
||||
snapshot_write_indices[0], snapshot_write_indices[i]
|
||||
), f"violation_write_index differs between launch 0 and {i}"
|
||||
assert torch.equal(
|
||||
snapshot_bufs[0], snapshot_bufs[i]
|
||||
), f"canary_buf differs between launch 0 and {i}"
|
||||
|
||||
|
||||
def test_write_byte_equal_across_repeated_launches_10x() -> None:
|
||||
num_launches = 10
|
||||
plan_cuda, plan_ref, input_ids, positions, out_cache_loc = _build_write_fixtures(
|
||||
device=_DEVICE
|
||||
)
|
||||
|
||||
snapshot_bufs: list[torch.Tensor] = []
|
||||
snapshot_rings: list[torch.Tensor] = []
|
||||
snapshot_counters: list[torch.Tensor] = []
|
||||
|
||||
for _ in range(num_launches):
|
||||
cuda_buf, ref_buf = make_canary_buf_pair(
|
||||
num_slots=16, slot_stride_bytes=32, device=_DEVICE
|
||||
)
|
||||
cuda_log, ref_log = make_log_pair(capacity=64, device=_DEVICE)
|
||||
pseudo_tok, pseudo_pos = dummy_pseudo_tensors(input_ids.shape[0])
|
||||
|
||||
_run_both_write(
|
||||
cuda_canary_buf=cuda_buf,
|
||||
ref_canary_buf=ref_buf,
|
||||
plan_cuda=plan_cuda,
|
||||
plan_ref=plan_ref,
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
enable_write_verify_inputs=False,
|
||||
expected_input_tokens=pseudo_tok,
|
||||
expected_input_positions=pseudo_pos,
|
||||
cuda_log=cuda_log,
|
||||
ref_log=ref_log,
|
||||
real_kv_sources_cuda=(),
|
||||
real_kv_sources_ref=(),
|
||||
real_kv_hash_mode=consts.RealKvHashMode.NONE,
|
||||
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
|
||||
)
|
||||
|
||||
assert_canary_buf_equal(buf_a=cuda_buf, buf_b=ref_buf)
|
||||
assert_canary_state_equal(log_a=cuda_log, log_b=ref_log)
|
||||
|
||||
snapshot_bufs.append(cuda_buf.clone())
|
||||
snapshot_rings.append(cuda_log.ring.clone())
|
||||
snapshot_counters.append(cuda_log.slot_run_counter.clone())
|
||||
|
||||
for i in range(1, num_launches):
|
||||
assert torch.equal(
|
||||
snapshot_bufs[0], snapshot_bufs[i]
|
||||
), f"canary_buf differs between launch 0 and {i}"
|
||||
assert torch.equal(
|
||||
snapshot_rings[0], snapshot_rings[i]
|
||||
), f"violation_ring differs between launch 0 and {i}"
|
||||
assert torch.equal(
|
||||
snapshot_counters[0], snapshot_counters[i]
|
||||
), f"slot_run_counter differs between launch 0 and {i}"
|
||||
|
||||
|
||||
def test_plan_byte_equal_across_repeated_launches_10x() -> None:
|
||||
num_launches = 10
|
||||
req_pool_indices, prefix_lens, extend_seq_lens, req_to_token = _build_plan_fixtures(
|
||||
device=_DEVICE
|
||||
)
|
||||
|
||||
snapshot_slots: list[torch.Tensor] = []
|
||||
snapshot_positions: list[torch.Tensor] = []
|
||||
snapshot_prevs: list[torch.Tensor] = []
|
||||
snapshot_write_offsets: list[torch.Tensor] = []
|
||||
|
||||
for _ in range(num_launches):
|
||||
triton_v = VerifyPlan.allocate(
|
||||
verify_capacity=64, device=_DEVICE
|
||||
).zero_for_testing_()
|
||||
triton_w = WritePlan.allocate(
|
||||
write_req_capacity=8, device=_DEVICE
|
||||
).zero_for_testing_()
|
||||
ref_v = VerifyPlan.allocate(
|
||||
verify_capacity=64, device=_DEVICE
|
||||
).zero_for_testing_()
|
||||
ref_w = WritePlan.allocate(
|
||||
write_req_capacity=8, device=_DEVICE
|
||||
).zero_for_testing_()
|
||||
|
||||
_run_both_plan(
|
||||
triton_verify=triton_v,
|
||||
triton_write=triton_w,
|
||||
ref_verify=ref_v,
|
||||
ref_write=ref_w,
|
||||
req_pool_indices=req_pool_indices,
|
||||
prefix_lens=prefix_lens,
|
||||
extend_seq_lens=extend_seq_lens,
|
||||
req_to_token=req_to_token,
|
||||
extras=empty_extras(),
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
)
|
||||
|
||||
_assert_plans_byte_equal(
|
||||
triton_verify=triton_v,
|
||||
triton_write=triton_w,
|
||||
ref_verify=ref_v,
|
||||
ref_write=ref_w,
|
||||
)
|
||||
|
||||
n_verify = int(triton_v.verify_num_valid[0].item())
|
||||
snapshot_slots.append(triton_v.verify_slot_indices[:n_verify].clone())
|
||||
snapshot_positions.append(triton_v.verify_expected_positions[:n_verify].clone())
|
||||
snapshot_prevs.append(triton_v.verify_prev_slot_indices[:n_verify].clone())
|
||||
snapshot_write_offsets.append(triton_w.write_offsets.clone())
|
||||
|
||||
for i in range(1, num_launches):
|
||||
assert torch.equal(
|
||||
snapshot_slots[0], snapshot_slots[i]
|
||||
), f"verify_slot_indices differs between launch 0 and {i}"
|
||||
assert torch.equal(
|
||||
snapshot_positions[0], snapshot_positions[i]
|
||||
), f"verify_expected_positions differs between launch 0 and {i}"
|
||||
assert torch.equal(
|
||||
snapshot_prevs[0], snapshot_prevs[i]
|
||||
), f"verify_prev_slot_indices differs between launch 0 and {i}"
|
||||
assert torch.equal(
|
||||
snapshot_write_offsets[0], snapshot_write_offsets[i]
|
||||
), f"write_offsets differs between launch 0 and {i}"
|
||||
|
||||
|
||||
def test_verify_multi_launch_100x_counter_linear() -> None:
|
||||
num_launches = 100
|
||||
|
||||
plan_cuda = make_verify_plan(
|
||||
slot_indices=[0],
|
||||
positions=[0],
|
||||
prev_slot_indices=[-1],
|
||||
capacity=4,
|
||||
device=_DEVICE,
|
||||
)
|
||||
|
||||
cuda_log = FakeViolationLog.allocate(capacity=64, device=_DEVICE)
|
||||
|
||||
for _ in range(num_launches):
|
||||
cuda_buf = make_canary_buf(num_slots=16, slot_stride_bytes=32, device=_DEVICE)
|
||||
launch_canary_verify_kernel(
|
||||
context=VerifyOrWriteContext(
|
||||
canary_buf=cuda_buf,
|
||||
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
|
||||
violation_ring=cuda_log.ring,
|
||||
violation_write_index=cuda_log.write_index,
|
||||
slot_run_counter=cuda_log.slot_run_counter,
|
||||
kernel_run_counter=cuda_log.kernel_run_counter,
|
||||
enable_chain_position_assert=cuda_log.enable_chain_position_assert,
|
||||
real_kv_sources=(),
|
||||
real_kv_hash_mode=consts.RealKvHashMode.NONE,
|
||||
),
|
||||
plan=plan_cuda,
|
||||
check_verify_expected_token=True,
|
||||
)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
assert (
|
||||
int(cuda_log.kernel_run_counter[0].item()) == num_launches
|
||||
), f"kernel_run_counter expected {num_launches}, got {cuda_log.kernel_run_counter[0].item()}"
|
||||
assert int(cuda_log.slot_run_counter[0].item()) == num_launches, (
|
||||
f"slot_run_counter expected {num_launches} (1 active entry x 100 launches), "
|
||||
f"got {cuda_log.slot_run_counter[0].item()}"
|
||||
)
|
||||
|
||||
|
||||
def test_verify_check_disabled_byte_equal() -> None:
|
||||
"""check_verify_expected_token True vs False produce equivalent violation logs on a clean plan."""
|
||||
plan_true_cuda, plan_true_ref = _build_verify_plan_5_entries(device=_DEVICE)
|
||||
plan_false_cuda, plan_false_ref = _build_verify_plan_5_entries(device=_DEVICE)
|
||||
|
||||
chain_slot_indices = [0, 1, 2, 3, 4]
|
||||
chain_tokens = [10, 20, 30, 40, 50]
|
||||
chain_positions = [0, 1, 2, 3, 4]
|
||||
|
||||
cuda_buf_true, ref_buf_true = make_canary_buf_pair(
|
||||
num_slots=16, slot_stride_bytes=32, device=_DEVICE
|
||||
)
|
||||
stamp_clean_chain(
|
||||
cuda_buf=cuda_buf_true,
|
||||
ref_buf=ref_buf_true,
|
||||
slot_indices=chain_slot_indices,
|
||||
tokens=chain_tokens,
|
||||
positions=chain_positions,
|
||||
)
|
||||
cuda_buf_false, ref_buf_false = make_canary_buf_pair(
|
||||
num_slots=16, slot_stride_bytes=32, device=_DEVICE
|
||||
)
|
||||
stamp_clean_chain(
|
||||
cuda_buf=cuda_buf_false,
|
||||
ref_buf=ref_buf_false,
|
||||
slot_indices=chain_slot_indices,
|
||||
tokens=chain_tokens,
|
||||
positions=chain_positions,
|
||||
)
|
||||
cuda_log_true, ref_log_true = make_log_pair(capacity=64, device=_DEVICE)
|
||||
cuda_log_false, ref_log_false = make_log_pair(capacity=64, device=_DEVICE)
|
||||
|
||||
_run_both_verify(
|
||||
cuda_canary_buf=cuda_buf_true,
|
||||
ref_canary_buf=ref_buf_true,
|
||||
plan_cuda=plan_true_cuda,
|
||||
plan_ref=plan_true_ref,
|
||||
cuda_log=cuda_log_true,
|
||||
ref_log=ref_log_true,
|
||||
real_kv_sources_cuda=(),
|
||||
real_kv_sources_ref=(),
|
||||
real_kv_hash_mode=consts.RealKvHashMode.NONE,
|
||||
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
|
||||
check_verify_expected_token=True,
|
||||
)
|
||||
_run_both_verify(
|
||||
cuda_canary_buf=cuda_buf_false,
|
||||
ref_canary_buf=ref_buf_false,
|
||||
plan_cuda=plan_false_cuda,
|
||||
plan_ref=plan_false_ref,
|
||||
cuda_log=cuda_log_false,
|
||||
ref_log=ref_log_false,
|
||||
real_kv_sources_cuda=(),
|
||||
real_kv_sources_ref=(),
|
||||
real_kv_hash_mode=consts.RealKvHashMode.NONE,
|
||||
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
|
||||
check_verify_expected_token=False,
|
||||
)
|
||||
|
||||
assert int(cuda_log_true.write_index[0].item()) == 0
|
||||
assert int(cuda_log_false.write_index[0].item()) == 0
|
||||
assert torch.equal(cuda_log_true.ring, cuda_log_false.ring)
|
||||
assert torch.equal(cuda_log_true.write_index, cuda_log_false.write_index)
|
||||
assert torch.equal(cuda_log_true.slot_run_counter, cuda_log_false.slot_run_counter)
|
||||
assert torch.equal(
|
||||
cuda_log_true.kernel_run_counter, cuda_log_false.kernel_run_counter
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("per_req_present", [False, True])
|
||||
def test_plan_per_req_present_or_absent(per_req_present: bool) -> None:
|
||||
max_reqs = 4
|
||||
max_seq_len = 16
|
||||
rp_axis = torch.arange(max_reqs, device=_DEVICE, dtype=torch.int32).unsqueeze(1)
|
||||
pos_axis = torch.arange(max_seq_len, device=_DEVICE, dtype=torch.int32).unsqueeze(0)
|
||||
req_to_token = (rp_axis * max_seq_len + pos_axis).contiguous()
|
||||
|
||||
if per_req_present:
|
||||
req_pool_indices = torch.tensor([1, 2], dtype=torch.int64, device=_DEVICE)
|
||||
prefix_lens = torch.tensor([3, 5], dtype=torch.int64, device=_DEVICE)
|
||||
extend_seq_lens = torch.tensor([1, 1], dtype=torch.int64, device=_DEVICE)
|
||||
else:
|
||||
req_pool_indices = torch.tensor([0], dtype=torch.int64, device=_DEVICE)
|
||||
prefix_lens = torch.tensor([0], dtype=torch.int64, device=_DEVICE)
|
||||
extend_seq_lens = torch.tensor([0], dtype=torch.int64, device=_DEVICE)
|
||||
|
||||
triton_v = VerifyPlan.allocate(
|
||||
verify_capacity=64, device=_DEVICE
|
||||
).zero_for_testing_()
|
||||
triton_w = WritePlan.allocate(
|
||||
write_req_capacity=8, device=_DEVICE
|
||||
).zero_for_testing_()
|
||||
ref_v = VerifyPlan.allocate(verify_capacity=64, device=_DEVICE).zero_for_testing_()
|
||||
ref_w = WritePlan.allocate(write_req_capacity=8, device=_DEVICE).zero_for_testing_()
|
||||
|
||||
_run_both_plan(
|
||||
triton_verify=triton_v,
|
||||
triton_write=triton_w,
|
||||
ref_verify=ref_v,
|
||||
ref_write=ref_w,
|
||||
req_pool_indices=req_pool_indices,
|
||||
prefix_lens=prefix_lens,
|
||||
extend_seq_lens=extend_seq_lens,
|
||||
req_to_token=req_to_token,
|
||||
extras=empty_extras(),
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
)
|
||||
|
||||
_assert_plans_byte_equal(
|
||||
triton_verify=triton_v,
|
||||
triton_write=triton_w,
|
||||
ref_verify=ref_v,
|
||||
ref_write=ref_w,
|
||||
)
|
||||
|
||||
if not per_req_present:
|
||||
assert int(triton_v.verify_num_valid[0].item()) == 0
|
||||
bs = int(req_pool_indices.shape[0])
|
||||
assert int(triton_w.write_offsets[bs].item()) == 0
|
||||
|
||||
if per_req_present:
|
||||
assert int(triton_v.verify_num_valid[0].item()) == 8
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -0,0 +1,842 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Optional
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary import consts
|
||||
from sglang.jit_kernel.kv_canary.plan import launch_canary_plan_kernels
|
||||
from sglang.jit_kernel.kv_canary.plan_ref import (
|
||||
launch_canary_plan_kernels_torch_reference,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary.verify import (
|
||||
CanaryLaunchTag,
|
||||
RealKvSource,
|
||||
VerifyOrWriteContext,
|
||||
VerifyPlan,
|
||||
launch_canary_verify_kernel,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary.verify_ref import (
|
||||
launch_canary_verify_kernel_torch_reference,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary.write import WritePlan, launch_canary_write_kernel
|
||||
from sglang.jit_kernel.kv_canary.write_ref import (
|
||||
launch_canary_write_kernel_torch_reference,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
|
||||
FakeViolationLog,
|
||||
assert_canary_buf_equal,
|
||||
assert_canary_state_equal,
|
||||
make_canary_buf,
|
||||
make_real_kv_sources,
|
||||
stamp_clean_chain,
|
||||
write_slot_fields,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._fixtures import (
|
||||
clone_real_kv_sources,
|
||||
empty_extras,
|
||||
make_req_to_token,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
|
||||
|
||||
|
||||
_DEVICE = torch.device("cuda")
|
||||
|
||||
|
||||
def _run_pipeline(
|
||||
*,
|
||||
real: bool,
|
||||
req_pool_indices: torch.Tensor,
|
||||
prefix_lens: torch.Tensor,
|
||||
extend_seq_lens: torch.Tensor,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
out_cache_loc: torch.Tensor,
|
||||
req_to_token: torch.Tensor,
|
||||
canary_buf: torch.Tensor,
|
||||
log: FakeViolationLog,
|
||||
extras: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
|
||||
swa_window_size: int,
|
||||
full_to_swa_index_mapping: Optional[torch.Tensor],
|
||||
kernel_kind: CanaryLaunchTag,
|
||||
enable_write_verify_inputs: bool,
|
||||
expected_input_tokens: torch.Tensor,
|
||||
expected_input_positions: torch.Tensor,
|
||||
real_kv_sources: tuple[RealKvSource, ...],
|
||||
real_kv_hash_mode: consts.RealKvHashMode,
|
||||
verify_capacity: int,
|
||||
write_req_capacity: int,
|
||||
req_to_verify_expected_tokens: Optional[torch.Tensor] = None,
|
||||
req_to_verify_expected_tokens_valid_lens: Optional[torch.Tensor] = None,
|
||||
kv_token_id_vs_position_offset: int = 0,
|
||||
check_verify_expected_token: bool = True,
|
||||
) -> tuple[VerifyPlan, WritePlan]:
|
||||
_ = extras
|
||||
plan_v = VerifyPlan.allocate(verify_capacity=verify_capacity, device=_DEVICE)
|
||||
plan_w = WritePlan.allocate(write_req_capacity=write_req_capacity, device=_DEVICE)
|
||||
|
||||
# Existing pipeline tests that supply a pool but no per-req lens want "bound by full
|
||||
# row width" semantics. Synthesise that bound here so callers don't have to.
|
||||
if (
|
||||
req_to_verify_expected_tokens is not None
|
||||
and req_to_verify_expected_tokens_valid_lens is None
|
||||
):
|
||||
req_to_verify_expected_tokens_valid_lens = torch.full(
|
||||
(int(req_pool_indices.shape[0]),),
|
||||
int(req_to_verify_expected_tokens.shape[1]),
|
||||
dtype=torch.int64,
|
||||
device=req_pool_indices.device,
|
||||
)
|
||||
|
||||
plan_fn = (
|
||||
launch_canary_plan_kernels
|
||||
if real
|
||||
else launch_canary_plan_kernels_torch_reference
|
||||
)
|
||||
plan_fn(
|
||||
verify_plan_out=plan_v,
|
||||
write_plan_out=plan_w,
|
||||
req_pool_indices=req_pool_indices,
|
||||
prefix_lens=prefix_lens,
|
||||
extend_seq_lens=extend_seq_lens,
|
||||
req_to_token=req_to_token,
|
||||
swa_window_size=swa_window_size,
|
||||
full_to_swa_index_mapping=full_to_swa_index_mapping,
|
||||
verify_capacity=verify_capacity,
|
||||
req_to_verify_expected_tokens=req_to_verify_expected_tokens,
|
||||
req_to_verify_expected_tokens_valid_lens=req_to_verify_expected_tokens_valid_lens,
|
||||
kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
|
||||
)
|
||||
|
||||
if real:
|
||||
context = VerifyOrWriteContext(
|
||||
canary_buf=canary_buf,
|
||||
kernel_kind=kernel_kind,
|
||||
violation_ring=log.ring,
|
||||
violation_write_index=log.write_index,
|
||||
slot_run_counter=log.slot_run_counter,
|
||||
kernel_run_counter=log.kernel_run_counter,
|
||||
enable_chain_position_assert=log.enable_chain_position_assert,
|
||||
real_kv_sources=real_kv_sources,
|
||||
real_kv_hash_mode=real_kv_hash_mode,
|
||||
)
|
||||
launch_canary_write_kernel(
|
||||
context=context,
|
||||
plan=plan_w,
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
enable_write_input_assert=enable_write_verify_inputs,
|
||||
expected_input_tokens=expected_input_tokens,
|
||||
expected_input_positions=expected_input_positions,
|
||||
)
|
||||
launch_canary_verify_kernel(
|
||||
context=context,
|
||||
plan=plan_v,
|
||||
check_verify_expected_token=check_verify_expected_token,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
else:
|
||||
launch_canary_write_kernel_torch_reference(
|
||||
context=VerifyOrWriteContext(
|
||||
canary_buf=canary_buf,
|
||||
kernel_kind=kernel_kind,
|
||||
violation_ring=log.ring,
|
||||
violation_write_index=log.write_index,
|
||||
slot_run_counter=log.slot_run_counter,
|
||||
kernel_run_counter=log.kernel_run_counter,
|
||||
enable_chain_position_assert=log.enable_chain_position_assert,
|
||||
real_kv_sources=real_kv_sources,
|
||||
real_kv_hash_mode=real_kv_hash_mode,
|
||||
),
|
||||
plan=plan_w,
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
enable_write_input_assert=enable_write_verify_inputs,
|
||||
expected_input_tokens=expected_input_tokens,
|
||||
expected_input_positions=expected_input_positions,
|
||||
)
|
||||
launch_canary_verify_kernel_torch_reference(
|
||||
context=VerifyOrWriteContext(
|
||||
canary_buf=canary_buf,
|
||||
kernel_kind=kernel_kind,
|
||||
violation_ring=log.ring,
|
||||
violation_write_index=log.write_index,
|
||||
slot_run_counter=log.slot_run_counter,
|
||||
kernel_run_counter=log.kernel_run_counter,
|
||||
enable_chain_position_assert=log.enable_chain_position_assert,
|
||||
real_kv_sources=real_kv_sources,
|
||||
real_kv_hash_mode=real_kv_hash_mode,
|
||||
),
|
||||
plan=plan_v,
|
||||
check_verify_expected_token=check_verify_expected_token,
|
||||
)
|
||||
|
||||
return plan_v, plan_w
|
||||
|
||||
|
||||
def _run_both_and_assert_pipeline_equal(
|
||||
*,
|
||||
req_pool_indices: torch.Tensor,
|
||||
prefix_lens: torch.Tensor,
|
||||
extend_seq_lens: torch.Tensor,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
out_cache_loc: torch.Tensor,
|
||||
req_to_token: torch.Tensor,
|
||||
num_slots: int,
|
||||
extras: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
|
||||
swa_window_size: int = 0,
|
||||
full_to_swa_index_mapping: Optional[torch.Tensor] = None,
|
||||
kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL,
|
||||
enable_write_verify_inputs: bool = False,
|
||||
expected_input_tokens: Optional[torch.Tensor] = None,
|
||||
expected_input_positions: Optional[torch.Tensor] = None,
|
||||
real_kv_sources_real: tuple[RealKvSource, ...] = (),
|
||||
real_kv_sources_ref: tuple[RealKvSource, ...] = (),
|
||||
real_kv_hash_mode: consts.RealKvHashMode = consts.RealKvHashMode.NONE,
|
||||
ring_capacity: int = 64,
|
||||
verify_capacity: int = 256,
|
||||
write_req_capacity: int = 16,
|
||||
assert_ring_equal: bool = True,
|
||||
initial_canary_buf: Optional[torch.Tensor] = None,
|
||||
req_to_verify_expected_tokens: Optional[torch.Tensor] = None,
|
||||
kv_token_id_vs_position_offset: int = 0,
|
||||
check_verify_expected_token: bool = True,
|
||||
) -> tuple[
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
FakeViolationLog,
|
||||
FakeViolationLog,
|
||||
VerifyPlan,
|
||||
WritePlan,
|
||||
VerifyPlan,
|
||||
WritePlan,
|
||||
]:
|
||||
# The kernel rejects non-None expected_* tensors when enable_write_verify_inputs=False
|
||||
# (sanity check to catch caller bugs), so only synthesise zero placeholders in the
|
||||
# branch that will actually assert against them.
|
||||
if enable_write_verify_inputs:
|
||||
total_tokens = int(input_ids.shape[0])
|
||||
if expected_input_tokens is None:
|
||||
expected_input_tokens = torch.zeros(
|
||||
total_tokens, dtype=torch.int64, device=_DEVICE
|
||||
)
|
||||
if expected_input_positions is None:
|
||||
expected_input_positions = torch.zeros(
|
||||
total_tokens, dtype=torch.int64, device=_DEVICE
|
||||
)
|
||||
|
||||
if initial_canary_buf is None:
|
||||
buf_real = make_canary_buf(num_slots=num_slots, device=_DEVICE)
|
||||
else:
|
||||
buf_real = initial_canary_buf.clone()
|
||||
buf_ref = buf_real.clone()
|
||||
log_real = FakeViolationLog.allocate(capacity=ring_capacity, device=_DEVICE)
|
||||
log_ref = FakeViolationLog.allocate(capacity=ring_capacity, device=_DEVICE)
|
||||
|
||||
shared: dict[str, Any] = dict(
|
||||
req_pool_indices=req_pool_indices,
|
||||
prefix_lens=prefix_lens,
|
||||
extend_seq_lens=extend_seq_lens,
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
req_to_token=req_to_token,
|
||||
extras=extras,
|
||||
swa_window_size=swa_window_size,
|
||||
full_to_swa_index_mapping=full_to_swa_index_mapping,
|
||||
kernel_kind=kernel_kind,
|
||||
enable_write_verify_inputs=enable_write_verify_inputs,
|
||||
expected_input_tokens=expected_input_tokens,
|
||||
expected_input_positions=expected_input_positions,
|
||||
real_kv_hash_mode=real_kv_hash_mode,
|
||||
verify_capacity=verify_capacity,
|
||||
write_req_capacity=write_req_capacity,
|
||||
req_to_verify_expected_tokens=req_to_verify_expected_tokens,
|
||||
kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
|
||||
check_verify_expected_token=check_verify_expected_token,
|
||||
)
|
||||
|
||||
plan_v_real, plan_w_real = _run_pipeline(
|
||||
real=True,
|
||||
canary_buf=buf_real,
|
||||
log=log_real,
|
||||
real_kv_sources=real_kv_sources_real,
|
||||
**shared,
|
||||
)
|
||||
plan_v_ref, plan_w_ref = _run_pipeline(
|
||||
real=False,
|
||||
canary_buf=buf_ref,
|
||||
log=log_ref,
|
||||
real_kv_sources=real_kv_sources_ref,
|
||||
**shared,
|
||||
)
|
||||
|
||||
assert_canary_buf_equal(buf_a=buf_real, buf_b=buf_ref)
|
||||
if assert_ring_equal:
|
||||
assert_canary_state_equal(log_a=log_real, log_b=log_ref)
|
||||
else:
|
||||
assert torch.equal(log_real.write_index, log_ref.write_index)
|
||||
assert torch.equal(log_real.slot_run_counter, log_ref.slot_run_counter)
|
||||
assert torch.equal(log_real.kernel_run_counter, log_ref.kernel_run_counter)
|
||||
|
||||
return (
|
||||
buf_real,
|
||||
buf_ref,
|
||||
log_real,
|
||||
log_ref,
|
||||
plan_v_real,
|
||||
plan_w_real,
|
||||
plan_v_ref,
|
||||
plan_w_ref,
|
||||
)
|
||||
|
||||
|
||||
def _t(values: list[int]) -> torch.Tensor:
|
||||
return torch.tensor(values, dtype=torch.int64, device=_DEVICE)
|
||||
|
||||
|
||||
def _linear_r2t(*, max_reqs: int = 4, max_seq_len: int = 16) -> torch.Tensor:
|
||||
return make_req_to_token(
|
||||
kind="linear", max_reqs=max_reqs, max_seq_len=max_seq_len, device=_DEVICE
|
||||
)
|
||||
|
||||
|
||||
def _zero_no_write_inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""``(input_ids, positions, out_cache_loc)`` zero placeholders for extend_seq_lens=0 tests."""
|
||||
zeros = torch.zeros(1, dtype=torch.int64, device=_DEVICE)
|
||||
return zeros.clone(), zeros.clone(), zeros.clone()
|
||||
|
||||
|
||||
def _contiguous_out_cache_loc(
|
||||
*, req_pool_idx: int, start: int, count: int, max_seq_len: int = 16
|
||||
) -> torch.Tensor:
|
||||
return _t([req_pool_idx * max_seq_len + start + i for i in range(count)])
|
||||
|
||||
|
||||
def _stamp_linear_prefix(
|
||||
*,
|
||||
initial_buf: torch.Tensor,
|
||||
initial_ref: torch.Tensor,
|
||||
req_pool_idx: int,
|
||||
prefix_len: int,
|
||||
tokens: list[int],
|
||||
max_seq_len: int = 16,
|
||||
) -> None:
|
||||
"""Stamp clean chain for slots ``[rp*max_seq_len + 0 .. + prefix_len)`` at positions ``0..prefix_len``."""
|
||||
stamp_clean_chain(
|
||||
cuda_buf=initial_buf,
|
||||
ref_buf=initial_ref,
|
||||
slot_indices=[req_pool_idx * max_seq_len + pos for pos in range(prefix_len)],
|
||||
tokens=tokens,
|
||||
positions=list(range(prefix_len)),
|
||||
)
|
||||
|
||||
|
||||
def test_pipeline_basic_5_step_single_req() -> None:
|
||||
"""Single req, prefix_len=0, extend_seq_len=5: basic plan→write→verify byte-equal."""
|
||||
_run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([1]),
|
||||
prefix_lens=_t([0]),
|
||||
extend_seq_lens=_t([5]),
|
||||
input_ids=_t([10, 20, 30, 40, 50]),
|
||||
positions=_t([0, 1, 2, 3, 4]),
|
||||
out_cache_loc=_contiguous_out_cache_loc(req_pool_idx=1, start=0, count=5),
|
||||
req_to_token=_linear_r2t(),
|
||||
num_slots=64,
|
||||
extras=empty_extras(),
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
)
|
||||
|
||||
|
||||
def test_pipeline_multi_req_mixed_extend_decode() -> None:
|
||||
"""bs=3: pure extend req, decode req (prefix+1 extend), and padding sentinel row."""
|
||||
max_seq_len = 16
|
||||
_run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([1, 2, 0]),
|
||||
prefix_lens=_t([0, 5, 0]),
|
||||
extend_seq_lens=_t([4, 1, 0]),
|
||||
input_ids=_t([11, 12, 13, 14, 21]),
|
||||
positions=_t([0, 1, 2, 3, 5]),
|
||||
out_cache_loc=_t(
|
||||
[
|
||||
1 * max_seq_len + 0,
|
||||
1 * max_seq_len + 1,
|
||||
1 * max_seq_len + 2,
|
||||
1 * max_seq_len + 3,
|
||||
2 * max_seq_len + 5,
|
||||
]
|
||||
),
|
||||
req_to_token=_linear_r2t(max_reqs=8, max_seq_len=max_seq_len),
|
||||
num_slots=128,
|
||||
extras=empty_extras(),
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
write_req_capacity=4,
|
||||
)
|
||||
|
||||
|
||||
def test_pipeline_swa_window() -> None:
|
||||
"""SWA window=4, prefix_len=6: verify covers window [2,6), write covers extend tokens."""
|
||||
max_seq_len = 16
|
||||
max_reqs = 4
|
||||
full_to_swa_index_mapping = torch.arange(
|
||||
max_reqs * max_seq_len + 1, dtype=torch.int64, device=_DEVICE
|
||||
)
|
||||
|
||||
_run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([1]),
|
||||
prefix_lens=_t([6]),
|
||||
extend_seq_lens=_t([2]),
|
||||
input_ids=_t([100, 101]),
|
||||
positions=_t([6, 7]),
|
||||
out_cache_loc=_t(
|
||||
[
|
||||
full_to_swa_index_mapping[1 * max_seq_len + 6].item(),
|
||||
full_to_swa_index_mapping[1 * max_seq_len + 7].item(),
|
||||
]
|
||||
),
|
||||
req_to_token=_linear_r2t(max_reqs=max_reqs, max_seq_len=max_seq_len),
|
||||
num_slots=64,
|
||||
extras=empty_extras(),
|
||||
swa_window_size=4,
|
||||
full_to_swa_index_mapping=full_to_swa_index_mapping,
|
||||
)
|
||||
|
||||
|
||||
def test_pipeline_sweep_no_write() -> None:
|
||||
"""All extend_seq_lens=0: write_step is no-op, verify sweeps prefix, buf unchanged."""
|
||||
prefix_len = 4
|
||||
input_ids, positions, out_cache_loc = _zero_no_write_inputs()
|
||||
|
||||
initial_buf = make_canary_buf(num_slots=64, device=_DEVICE)
|
||||
initial_ref = initial_buf.clone()
|
||||
_stamp_linear_prefix(
|
||||
initial_buf=initial_buf,
|
||||
initial_ref=initial_ref,
|
||||
req_pool_idx=1,
|
||||
prefix_len=prefix_len,
|
||||
tokens=[100 + pos for pos in range(prefix_len)],
|
||||
)
|
||||
|
||||
buf_real, buf_ref, log_real, log_ref, plan_v_real, plan_w_real, _, _ = (
|
||||
_run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([1]),
|
||||
prefix_lens=_t([prefix_len]),
|
||||
extend_seq_lens=_t([0]),
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
req_to_token=_linear_r2t(),
|
||||
num_slots=64,
|
||||
extras=empty_extras(),
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
initial_canary_buf=initial_buf,
|
||||
)
|
||||
)
|
||||
|
||||
assert int(plan_v_real.verify_num_valid[0].item()) == prefix_len
|
||||
assert int(plan_w_real.write_num_valid_reqs[0].item()) == 1
|
||||
assert int(plan_w_real.write_offsets[1].item()) == 0
|
||||
assert torch.equal(buf_real, initial_buf)
|
||||
assert torch.equal(buf_ref, initial_buf)
|
||||
assert int(log_real.write_index[0].item()) == 0
|
||||
assert int(log_ref.write_index[0].item()) == 0
|
||||
assert int(log_real.slot_run_counter[0].item()) == prefix_len
|
||||
assert int(log_ref.slot_run_counter[0].item()) == prefix_len
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"real_kv_hash_mode",
|
||||
[
|
||||
consts.RealKvHashMode.NONE,
|
||||
consts.RealKvHashMode.PARTIAL,
|
||||
consts.RealKvHashMode.ALL,
|
||||
],
|
||||
)
|
||||
def test_pipeline_real_kv_mode(real_kv_hash_mode: consts.RealKvHashMode) -> None:
|
||||
"""real_kv_hash_mode OFF/PARTIAL/ALL: real and ref use cloned sources to prevent ALL-mode hash aliasing."""
|
||||
sources_real = make_real_kv_sources(count=2, num_slots=64, device=_DEVICE)
|
||||
sources_ref = clone_real_kv_sources(sources_real)
|
||||
|
||||
_run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([1]),
|
||||
prefix_lens=_t([0]),
|
||||
extend_seq_lens=_t([3]),
|
||||
input_ids=_t([5, 6, 7]),
|
||||
positions=_t([0, 1, 2]),
|
||||
out_cache_loc=_contiguous_out_cache_loc(req_pool_idx=1, start=0, count=3),
|
||||
req_to_token=_linear_r2t(),
|
||||
num_slots=64,
|
||||
extras=empty_extras(),
|
||||
real_kv_sources_real=sources_real,
|
||||
real_kv_sources_ref=sources_ref,
|
||||
real_kv_hash_mode=real_kv_hash_mode,
|
||||
)
|
||||
|
||||
|
||||
def test_pipeline_pseudo_mode_on_match() -> None:
|
||||
"""enable_write_verify_inputs=ON, expected==actual: zero violations, buf byte-equal."""
|
||||
input_ids = _t([1, 2, 3, 4])
|
||||
positions = _t([0, 1, 2, 3])
|
||||
|
||||
_, _, log_real, log_ref, _, _, _, _ = _run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([1]),
|
||||
prefix_lens=_t([0]),
|
||||
extend_seq_lens=_t([4]),
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=_contiguous_out_cache_loc(req_pool_idx=1, start=0, count=4),
|
||||
req_to_token=_linear_r2t(),
|
||||
num_slots=64,
|
||||
extras=empty_extras(),
|
||||
enable_write_verify_inputs=True,
|
||||
expected_input_tokens=input_ids.clone(),
|
||||
expected_input_positions=positions.clone(),
|
||||
)
|
||||
|
||||
assert int(log_real.write_index[0].item()) == 0
|
||||
assert int(log_ref.write_index[0].item()) == 0
|
||||
|
||||
|
||||
def test_pipeline_pseudo_mode_on_token_mismatch_then_verify_clean() -> None:
|
||||
"""enable_write_verify_inputs=ON, expected tokens all wrong: write records N violations."""
|
||||
n_tokens = 3
|
||||
positions = _t([0, 1, 2])
|
||||
|
||||
_, _, log_real, log_ref, _, _, _, _ = _run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([1]),
|
||||
prefix_lens=_t([0]),
|
||||
extend_seq_lens=_t([3]),
|
||||
input_ids=_t([10, 20, 30]),
|
||||
positions=positions,
|
||||
out_cache_loc=_contiguous_out_cache_loc(req_pool_idx=1, start=0, count=3),
|
||||
req_to_token=_linear_r2t(),
|
||||
num_slots=64,
|
||||
extras=empty_extras(),
|
||||
enable_write_verify_inputs=True,
|
||||
expected_input_tokens=_t([99, 99, 99]),
|
||||
expected_input_positions=positions.clone(),
|
||||
ring_capacity=64,
|
||||
)
|
||||
|
||||
write_violations = int(log_real.write_index[0].item())
|
||||
assert (
|
||||
write_violations == n_tokens
|
||||
), f"expected {n_tokens} write violations, got {write_violations}"
|
||||
|
||||
|
||||
def test_pipeline_empty_batch() -> None:
|
||||
"""bs=1 with req_pool_idx=0 (padding): write and verify are no-op, kernel_run_counter == 2 (write+verify)."""
|
||||
input_ids, positions, out_cache_loc = _zero_no_write_inputs()
|
||||
|
||||
_, _, log_real, log_ref, _, _, _, _ = _run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([0]),
|
||||
prefix_lens=_t([0]),
|
||||
extend_seq_lens=_t([0]),
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
req_to_token=_linear_r2t(),
|
||||
num_slots=64,
|
||||
extras=empty_extras(),
|
||||
)
|
||||
|
||||
assert int(log_real.kernel_run_counter[0].item()) == 2
|
||||
assert int(log_ref.kernel_run_counter[0].item()) == 2
|
||||
assert int(log_real.write_index[0].item()) == 0
|
||||
|
||||
|
||||
def test_pipeline_negative_slot_swa_out_of_window() -> None:
|
||||
"""SWA: some out_cache_loc entries map to -1 (out-of-window); write_step skips them, buf unchanged."""
|
||||
max_seq_len = 16
|
||||
max_reqs = 4
|
||||
|
||||
full_to_swa_index_mapping = torch.arange(
|
||||
max_reqs * max_seq_len + 1, dtype=torch.int64, device=_DEVICE
|
||||
)
|
||||
full_to_swa_index_mapping[1 * max_seq_len + 6] = -1
|
||||
full_to_swa_index_mapping[1 * max_seq_len + 7] = -1
|
||||
|
||||
_run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([1]),
|
||||
prefix_lens=_t([6]),
|
||||
extend_seq_lens=_t([4]),
|
||||
input_ids=_t([100, 101, 102, 103]),
|
||||
positions=_t([6, 7, 8, 9]),
|
||||
out_cache_loc=_t([-1, -1, 1 * max_seq_len + 8, 1 * max_seq_len + 9]),
|
||||
req_to_token=_linear_r2t(max_reqs=max_reqs, max_seq_len=max_seq_len),
|
||||
num_slots=128,
|
||||
extras=empty_extras(),
|
||||
swa_window_size=4,
|
||||
full_to_swa_index_mapping=full_to_swa_index_mapping,
|
||||
)
|
||||
|
||||
|
||||
def test_pipeline_ring_overflow_via_real_plan() -> None:
|
||||
"""Verify detects >capacity violations when prev_hash is pre-corrupted; write_index byte-equal, ring relaxed."""
|
||||
max_seq_len = 16
|
||||
max_reqs = 4
|
||||
req_to_token = _linear_r2t(max_reqs=max_reqs, max_seq_len=max_seq_len)
|
||||
n_slots = 8
|
||||
|
||||
req_pool_indices = _t([1])
|
||||
prefix_lens = _t([n_slots])
|
||||
extend_seq_lens = _t([0])
|
||||
input_ids, positions, out_cache_loc = _zero_no_write_inputs()
|
||||
|
||||
num_slots = max_reqs * max_seq_len
|
||||
|
||||
# Step 1: pre-pollute canary_buf slots [0..n_slots) with wrong prev_hash so verify fires n_slots violations.
|
||||
buf_real = make_canary_buf(num_slots=num_slots, device=_DEVICE)
|
||||
buf_ref = make_canary_buf(num_slots=num_slots, device=_DEVICE)
|
||||
for slot_idx in range(n_slots):
|
||||
full_slot = 1 * max_seq_len + slot_idx
|
||||
for buf in (buf_real, buf_ref):
|
||||
write_slot_fields(
|
||||
canary_buf=buf,
|
||||
slot_idx=full_slot,
|
||||
token=slot_idx + 1,
|
||||
position=slot_idx,
|
||||
prev_hash=0x1234_DEAD_BEEF_0000 + slot_idx,
|
||||
real_kv_hash=0,
|
||||
)
|
||||
|
||||
# Step 2: run real pipeline (plan + no write + verify); overflow ring capacity=4 with all n_slots violations.
|
||||
ring_capacity = 4
|
||||
log_real = FakeViolationLog.allocate(capacity=ring_capacity, device=_DEVICE)
|
||||
log_ref = FakeViolationLog.allocate(capacity=ring_capacity, device=_DEVICE)
|
||||
plan_v_real = VerifyPlan.allocate(verify_capacity=256, device=_DEVICE)
|
||||
plan_w_real = WritePlan.allocate(write_req_capacity=4, device=_DEVICE)
|
||||
plan_v_ref = VerifyPlan.allocate(verify_capacity=256, device=_DEVICE)
|
||||
plan_w_ref = WritePlan.allocate(write_req_capacity=4, device=_DEVICE)
|
||||
|
||||
launch_canary_plan_kernels(
|
||||
verify_plan_out=plan_v_real,
|
||||
write_plan_out=plan_w_real,
|
||||
req_pool_indices=req_pool_indices,
|
||||
prefix_lens=prefix_lens,
|
||||
extend_seq_lens=extend_seq_lens,
|
||||
req_to_token=req_to_token,
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
verify_capacity=int(plan_v_real.verify_slot_indices.shape[0]),
|
||||
req_to_verify_expected_tokens=None,
|
||||
req_to_verify_expected_tokens_valid_lens=None,
|
||||
kv_token_id_vs_position_offset=0,
|
||||
)
|
||||
launch_canary_plan_kernels_torch_reference(
|
||||
verify_plan_out=plan_v_ref,
|
||||
write_plan_out=plan_w_ref,
|
||||
req_pool_indices=req_pool_indices,
|
||||
prefix_lens=prefix_lens,
|
||||
extend_seq_lens=extend_seq_lens,
|
||||
req_to_token=req_to_token,
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
verify_capacity=int(plan_v_ref.verify_slot_indices.shape[0]),
|
||||
req_to_verify_expected_tokens=None,
|
||||
req_to_verify_expected_tokens_valid_lens=None,
|
||||
kv_token_id_vs_position_offset=0,
|
||||
)
|
||||
|
||||
launch_canary_verify_kernel(
|
||||
context=VerifyOrWriteContext(
|
||||
canary_buf=buf_real,
|
||||
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
|
||||
violation_ring=log_real.ring,
|
||||
violation_write_index=log_real.write_index,
|
||||
slot_run_counter=log_real.slot_run_counter,
|
||||
kernel_run_counter=log_real.kernel_run_counter,
|
||||
enable_chain_position_assert=log_real.enable_chain_position_assert,
|
||||
real_kv_sources=(),
|
||||
real_kv_hash_mode=consts.RealKvHashMode.NONE,
|
||||
),
|
||||
plan=plan_v_real,
|
||||
check_verify_expected_token=True,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
launch_canary_verify_kernel_torch_reference(
|
||||
context=VerifyOrWriteContext(
|
||||
canary_buf=buf_ref,
|
||||
kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
|
||||
violation_ring=log_ref.ring,
|
||||
violation_write_index=log_ref.write_index,
|
||||
slot_run_counter=log_ref.slot_run_counter,
|
||||
kernel_run_counter=log_ref.kernel_run_counter,
|
||||
enable_chain_position_assert=log_ref.enable_chain_position_assert,
|
||||
real_kv_sources=(),
|
||||
real_kv_hash_mode=consts.RealKvHashMode.NONE,
|
||||
),
|
||||
plan=plan_v_ref,
|
||||
check_verify_expected_token=True,
|
||||
)
|
||||
|
||||
# Step 3: write_index byte-equal; ring contents relaxed (atomic order not guaranteed under overflow).
|
||||
assert torch.equal(log_real.write_index, log_ref.write_index)
|
||||
assert int(log_real.write_index[0].item()) == n_slots
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"kernel_kind", [CanaryLaunchTag.HEAD_K_FULL, CanaryLaunchTag.SWEEP_V_SWA]
|
||||
)
|
||||
def test_pipeline_kernel_kind_propagates(kernel_kind: CanaryLaunchTag) -> None:
|
||||
"""Different CanaryLaunchTag values: violation ring's kernel_kind field matches on both sides."""
|
||||
max_seq_len = 16
|
||||
input_ids, positions, out_cache_loc = _zero_no_write_inputs()
|
||||
|
||||
initial_buf = make_canary_buf(num_slots=64, device=_DEVICE)
|
||||
write_slot_fields(
|
||||
canary_buf=initial_buf,
|
||||
slot_idx=1 * max_seq_len,
|
||||
token=7,
|
||||
position=99,
|
||||
prev_hash=0,
|
||||
real_kv_hash=0,
|
||||
)
|
||||
|
||||
_, _, log_real, log_ref, _, _, _, _ = _run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([1]),
|
||||
prefix_lens=_t([1]),
|
||||
extend_seq_lens=_t([0]),
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
req_to_token=_linear_r2t(max_seq_len=max_seq_len),
|
||||
num_slots=64,
|
||||
extras=empty_extras(),
|
||||
kernel_kind=kernel_kind,
|
||||
initial_canary_buf=initial_buf,
|
||||
)
|
||||
|
||||
assert int(log_real.write_index[0].item()) == 1
|
||||
assert int(log_ref.write_index[0].item()) == 1
|
||||
assert int(log_real.ring[0, consts.VIOLATION_FIELD_KERNEL_KIND].item()) == int(
|
||||
kernel_kind
|
||||
)
|
||||
assert int(log_ref.ring[0, consts.VIOLATION_FIELD_KERNEL_KIND].item()) == int(
|
||||
kernel_kind
|
||||
)
|
||||
|
||||
|
||||
def test_pipeline_token_mismatch_detected_via_pool() -> None:
|
||||
"""plan-pool gather + verify-token check: stamped wrong token id raises VERIFY_TOKEN_MISMATCH."""
|
||||
max_seq_len = 16
|
||||
max_reqs = 4
|
||||
prefix_len = 4
|
||||
input_ids, positions, out_cache_loc = _zero_no_write_inputs()
|
||||
|
||||
expected_tokens = [1000 + pos for pos in range(prefix_len)]
|
||||
pool = torch.full((max_reqs, max_seq_len), -999, dtype=torch.int32, device=_DEVICE)
|
||||
for pos, token in enumerate(expected_tokens):
|
||||
pool[1, pos] = token
|
||||
|
||||
stored_tokens = [token + 1 for token in expected_tokens]
|
||||
|
||||
initial_buf = make_canary_buf(num_slots=64, device=_DEVICE)
|
||||
initial_ref = initial_buf.clone()
|
||||
_stamp_linear_prefix(
|
||||
initial_buf=initial_buf,
|
||||
initial_ref=initial_ref,
|
||||
req_pool_idx=1,
|
||||
prefix_len=prefix_len,
|
||||
tokens=stored_tokens,
|
||||
max_seq_len=max_seq_len,
|
||||
)
|
||||
|
||||
_, _, log_real, log_ref, _, _, _, _ = _run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([1]),
|
||||
prefix_lens=_t([prefix_len]),
|
||||
extend_seq_lens=_t([0]),
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
req_to_token=_linear_r2t(max_reqs=max_reqs, max_seq_len=max_seq_len),
|
||||
num_slots=64,
|
||||
extras=empty_extras(),
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
initial_canary_buf=initial_buf,
|
||||
req_to_verify_expected_tokens=pool,
|
||||
kv_token_id_vs_position_offset=0,
|
||||
check_verify_expected_token=True,
|
||||
)
|
||||
|
||||
assert int(log_real.write_index[0].item()) == prefix_len
|
||||
assert int(log_ref.write_index[0].item()) == prefix_len
|
||||
# Ring rows may land in any order; collect stored/expected pairs and compare as sets.
|
||||
observed_pairs: set[tuple[int, int]] = set()
|
||||
for row_idx in range(prefix_len):
|
||||
fail_bits = int(
|
||||
log_real.ring[row_idx, consts.VIOLATION_FIELD_FAIL_REASON_BITS].item()
|
||||
)
|
||||
assert fail_bits & int(
|
||||
consts.FailReason.VERIFY_TOKEN_MISMATCH
|
||||
), f"row {row_idx}: VERIFY_TOKEN_MISMATCH bit missing in {fail_bits:#b}"
|
||||
stored = int(log_real.ring[row_idx, consts.VIOLATION_FIELD_STORED_TOKEN].item())
|
||||
expected = int(
|
||||
log_real.ring[row_idx, consts.VIOLATION_FIELD_EXPECTED_TOKEN].item()
|
||||
)
|
||||
observed_pairs.add((stored, expected))
|
||||
expected_pairs = {(stored_tokens[i], expected_tokens[i]) for i in range(prefix_len)}
|
||||
assert observed_pairs == expected_pairs
|
||||
|
||||
|
||||
def test_pipeline_eagle_offset_plus_1_byte_equal() -> None:
|
||||
"""plan-pool + offset=+1 full pipeline: stamped tokens match pool[rp, pos+1], no violations CUDA vs ref byte-equal."""
|
||||
max_seq_len = 16
|
||||
max_reqs = 4
|
||||
prefix_len = 4
|
||||
input_ids, positions, out_cache_loc = _zero_no_write_inputs()
|
||||
|
||||
stored_tokens = [2000 + pos for pos in range(prefix_len)]
|
||||
pool = torch.full((max_reqs, max_seq_len), -999, dtype=torch.int32, device=_DEVICE)
|
||||
for pos in range(prefix_len):
|
||||
# offset=+1 means kernel gathers from pool[rp, pos + 1], so place stored_tokens[pos] there.
|
||||
pool[1, pos + 1] = stored_tokens[pos]
|
||||
|
||||
initial_buf = make_canary_buf(num_slots=64, device=_DEVICE)
|
||||
initial_ref = initial_buf.clone()
|
||||
_stamp_linear_prefix(
|
||||
initial_buf=initial_buf,
|
||||
initial_ref=initial_ref,
|
||||
req_pool_idx=1,
|
||||
prefix_len=prefix_len,
|
||||
tokens=stored_tokens,
|
||||
max_seq_len=max_seq_len,
|
||||
)
|
||||
|
||||
_, _, log_real, log_ref, _, _, _, _ = _run_both_and_assert_pipeline_equal(
|
||||
req_pool_indices=_t([1]),
|
||||
prefix_lens=_t([prefix_len]),
|
||||
extend_seq_lens=_t([0]),
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
req_to_token=_linear_r2t(max_reqs=max_reqs, max_seq_len=max_seq_len),
|
||||
num_slots=64,
|
||||
extras=empty_extras(),
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
initial_canary_buf=initial_buf,
|
||||
req_to_verify_expected_tokens=pool,
|
||||
kv_token_id_vs_position_offset=1,
|
||||
check_verify_expected_token=True,
|
||||
)
|
||||
|
||||
assert int(log_real.write_index[0].item()) == 0
|
||||
assert int(log_ref.write_index[0].item()) == 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -0,0 +1,222 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.tests.kv_canary._differential import _run_both_plan
|
||||
from sglang.jit_kernel.tests.kv_canary._fixtures import (
|
||||
allocate_plan_pair,
|
||||
derive_plan_capacity,
|
||||
make_lut,
|
||||
make_padding_mask,
|
||||
make_req_to_token,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._fuzz_driver import (
|
||||
FUZZ_SEEDS_PR,
|
||||
run_fuzz_combo,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._invariants import PlanInvariants
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
|
||||
|
||||
|
||||
_DEVICE = torch.device("cuda")
|
||||
|
||||
_FUZZ_ITER_PER_SEED = 50
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True, kw_only=True)
|
||||
class PlanFuzzInputs:
|
||||
req_pool_indices: torch.Tensor
|
||||
prefix_lens: torch.Tensor
|
||||
extend_seq_lens: torch.Tensor
|
||||
req_to_token: torch.Tensor
|
||||
swa_window_size: int
|
||||
full_to_swa_index_mapping: Optional[torch.Tensor]
|
||||
verify_capacity: int
|
||||
write_req_capacity: int
|
||||
req_to_verify_expected_tokens: Optional[torch.Tensor]
|
||||
kv_token_id_vs_position_offset: int
|
||||
|
||||
|
||||
def _draw_random_plan_inputs(rng: random.Random) -> PlanFuzzInputs:
|
||||
bs = rng.randint(1, 16)
|
||||
max_seq_len = rng.choice([8, 16, 64, 128, 256])
|
||||
swa_enabled = rng.random() < 0.5
|
||||
swa_window_size = (
|
||||
rng.choice([4, 16, 64, max_seq_len, max(2, max_seq_len // 3)])
|
||||
if swa_enabled
|
||||
else 0
|
||||
)
|
||||
swa_window_size = min(swa_window_size, max_seq_len)
|
||||
lut_kind = (
|
||||
rng.choice(["identity", "shift", "permutation", "with_oob"])
|
||||
if swa_enabled
|
||||
else None
|
||||
)
|
||||
rtt_kind = rng.choice(["linear", "sparse_permuted"])
|
||||
padding_kind = rng.choice(["none", "trailing", "interleaved"])
|
||||
capacity_kind = rng.choice(["loose", "tight_match", "under_by_one"])
|
||||
|
||||
max_reqs = max(bs + 2, 4)
|
||||
pool_size = max_reqs * max_seq_len
|
||||
rtt = make_req_to_token(
|
||||
kind=rtt_kind,
|
||||
max_reqs=max_reqs,
|
||||
max_seq_len=max_seq_len,
|
||||
device=_DEVICE,
|
||||
rng=rng,
|
||||
)
|
||||
padding_mask = make_padding_mask(bs=bs, kind=padding_kind, rng=rng)
|
||||
req_pool_indices_list: list[int] = []
|
||||
prefix_lens_list: list[int] = []
|
||||
extend_seq_lens_list: list[int] = []
|
||||
for r in range(bs):
|
||||
if padding_mask[r]:
|
||||
req_pool_indices_list.append(0)
|
||||
prefix_lens_list.append(0)
|
||||
extend_seq_lens_list.append(0)
|
||||
else:
|
||||
req_pool_indices_list.append(rng.randint(1, max_reqs - 1))
|
||||
prefix_lens_list.append(rng.randint(0, max_seq_len - 1))
|
||||
extend_seq_lens_list.append(rng.randint(1, max(1, max_seq_len // 4)))
|
||||
req_pool_indices = torch.tensor(
|
||||
req_pool_indices_list, dtype=torch.int64, device=_DEVICE
|
||||
)
|
||||
prefix_lens = torch.tensor(prefix_lens_list, dtype=torch.int64, device=_DEVICE)
|
||||
extend_seq_lens = torch.tensor(
|
||||
extend_seq_lens_list, dtype=torch.int64, device=_DEVICE
|
||||
)
|
||||
|
||||
total_verify = 0
|
||||
for rpi, pfx in zip(req_pool_indices_list, prefix_lens_list):
|
||||
if rpi == 0:
|
||||
continue
|
||||
if swa_window_size > 0:
|
||||
window_start = max(0, pfx - swa_window_size)
|
||||
total_verify += max(0, pfx - window_start)
|
||||
else:
|
||||
total_verify += pfx
|
||||
|
||||
verify_capacity, write_req_capacity = derive_plan_capacity(
|
||||
kind=capacity_kind,
|
||||
total_verify=total_verify,
|
||||
extras_count=0,
|
||||
bs=bs,
|
||||
)
|
||||
|
||||
full_to_swa: Optional[torch.Tensor]
|
||||
if swa_window_size > 0 and lut_kind is not None:
|
||||
full_to_swa = make_lut(
|
||||
kind=lut_kind, pool_size=pool_size, device=_DEVICE, rng=rng
|
||||
)
|
||||
else:
|
||||
full_to_swa = None
|
||||
|
||||
expected_pool_present = rng.random() < 0.5
|
||||
kv_token_id_vs_position_offset = rng.choice([0, 1])
|
||||
expected_pool: Optional[torch.Tensor]
|
||||
if expected_pool_present:
|
||||
pool_max_context_len = rng.choice(
|
||||
[
|
||||
max(1, max_seq_len // 4),
|
||||
max(1, max_seq_len // 2),
|
||||
max_seq_len,
|
||||
]
|
||||
)
|
||||
expected_pool = torch.randint(
|
||||
low=0,
|
||||
high=50000,
|
||||
size=(max_reqs, pool_max_context_len),
|
||||
dtype=torch.int32,
|
||||
device=_DEVICE,
|
||||
)
|
||||
else:
|
||||
expected_pool = None
|
||||
|
||||
return PlanFuzzInputs(
|
||||
req_pool_indices=req_pool_indices,
|
||||
prefix_lens=prefix_lens,
|
||||
extend_seq_lens=extend_seq_lens,
|
||||
req_to_token=rtt,
|
||||
swa_window_size=swa_window_size,
|
||||
full_to_swa_index_mapping=full_to_swa,
|
||||
verify_capacity=verify_capacity,
|
||||
write_req_capacity=write_req_capacity,
|
||||
req_to_verify_expected_tokens=expected_pool,
|
||||
kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
|
||||
)
|
||||
|
||||
|
||||
def _run_one(inputs: PlanFuzzInputs) -> tuple:
|
||||
triton_v, triton_w, ref_v, ref_w = allocate_plan_pair(
|
||||
verify_capacity=inputs.verify_capacity,
|
||||
write_req_capacity=inputs.write_req_capacity,
|
||||
)
|
||||
_run_both_plan(
|
||||
triton_verify=triton_v,
|
||||
triton_write=triton_w,
|
||||
ref_verify=ref_v,
|
||||
ref_write=ref_w,
|
||||
req_pool_indices=inputs.req_pool_indices,
|
||||
prefix_lens=inputs.prefix_lens,
|
||||
extend_seq_lens=inputs.extend_seq_lens,
|
||||
req_to_token=inputs.req_to_token,
|
||||
extras=(
|
||||
torch.empty(0, dtype=torch.int64, device=_DEVICE),
|
||||
torch.empty(0, dtype=torch.int64, device=_DEVICE),
|
||||
torch.empty(0, dtype=torch.int64, device=_DEVICE),
|
||||
torch.zeros(1, dtype=torch.int32, device=_DEVICE),
|
||||
),
|
||||
swa_window_size=inputs.swa_window_size,
|
||||
full_to_swa_index_mapping=inputs.full_to_swa_index_mapping,
|
||||
req_to_verify_expected_tokens=inputs.req_to_verify_expected_tokens,
|
||||
kv_token_id_vs_position_offset=inputs.kv_token_id_vs_position_offset,
|
||||
)
|
||||
PlanInvariants.assert_all(
|
||||
verify_plan=triton_v,
|
||||
write_plan=triton_w,
|
||||
req_pool_indices=inputs.req_pool_indices,
|
||||
prefix_lens=inputs.prefix_lens,
|
||||
extend_seq_lens=inputs.extend_seq_lens,
|
||||
swa_window_size=inputs.swa_window_size,
|
||||
extras_slot_indices=torch.empty(0, dtype=torch.int64, device=_DEVICE),
|
||||
extras_positions=torch.empty(0, dtype=torch.int64, device=_DEVICE),
|
||||
extras_prev_slot_indices=torch.empty(0, dtype=torch.int64, device=_DEVICE),
|
||||
extras_count=0,
|
||||
)
|
||||
return triton_v, triton_w
|
||||
|
||||
|
||||
def _summarize(inputs: PlanFuzzInputs) -> str:
|
||||
return (
|
||||
f"bs={int(inputs.req_pool_indices.shape[0])} "
|
||||
f"swa={inputs.swa_window_size} "
|
||||
f"verify_cap={inputs.verify_capacity} write_cap={inputs.write_req_capacity} "
|
||||
f"has_lut={inputs.full_to_swa_index_mapping is not None} "
|
||||
f"has_pool={inputs.req_to_verify_expected_tokens is not None} "
|
||||
f"offset={inputs.kv_token_id_vs_position_offset}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seed", FUZZ_SEEDS_PR)
|
||||
def test_plan_fuzz_full_combo(seed: int) -> None:
|
||||
"""Multi-dim plan fuzzer: random LUT/rtt/padding/capacity/swa × N iters, byte-equal."""
|
||||
run_fuzz_combo(
|
||||
seed,
|
||||
draw_fn=_draw_random_plan_inputs,
|
||||
run_one_fn=_run_one,
|
||||
summarize_fn=_summarize,
|
||||
n_iter=_FUZZ_ITER_PER_SEED,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,278 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary.scatter_req_token_ids import (
|
||||
_SCATTER_BATCH_BLOCK,
|
||||
launch_scatter_req_token_ids_kernel,
|
||||
scatter_req_token_ids_torch_reference,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
|
||||
|
||||
|
||||
_DEVICE = torch.device("cuda")
|
||||
|
||||
|
||||
def _build_pool(*, max_reqs: int, max_context_len: int) -> torch.Tensor:
|
||||
return torch.zeros((max_reqs, max_context_len), dtype=torch.int32, device=_DEVICE)
|
||||
|
||||
|
||||
def _build_offsets(lens: list[int]) -> torch.Tensor:
|
||||
cumsum = [0]
|
||||
for n in lens:
|
||||
cumsum.append(cumsum[-1] + n)
|
||||
return torch.tensor(cumsum, dtype=torch.int64, device=_DEVICE)
|
||||
|
||||
|
||||
def _build_flat(seqs: list[list[int]]) -> torch.Tensor:
|
||||
flat: list[int] = []
|
||||
for s in seqs:
|
||||
flat.extend(s)
|
||||
return torch.tensor(flat, dtype=torch.int64, device=_DEVICE)
|
||||
|
||||
|
||||
class TestScatterReqTokenIds(CustomTestCase):
|
||||
def test_scatter_byte_equal_basic(self) -> None:
|
||||
"""Triton output matches the PyTorch reference for a small mixed batch."""
|
||||
seqs = [[10, 20, 30], [40, 50], [60, 70, 80, 90]]
|
||||
lens = [len(s) for s in seqs]
|
||||
rp = [3, 1, 5]
|
||||
|
||||
flat = _build_flat(seqs)
|
||||
offsets = _build_offsets(lens)
|
||||
req_pool_indices = torch.tensor(rp, dtype=torch.int64, device=_DEVICE)
|
||||
triton_pool = _build_pool(max_reqs=8, max_context_len=16)
|
||||
ref_pool = _build_pool(max_reqs=8, max_context_len=16)
|
||||
|
||||
launch_scatter_req_token_ids_kernel(
|
||||
flat_in=flat,
|
||||
offsets=offsets,
|
||||
req_pool_indices=req_pool_indices,
|
||||
pool_out=triton_pool,
|
||||
)
|
||||
scatter_req_token_ids_torch_reference(
|
||||
flat_in=flat,
|
||||
offsets=offsets,
|
||||
req_pool_indices=req_pool_indices,
|
||||
pool_out=ref_pool,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
self.assertTrue(torch.equal(triton_pool, ref_pool))
|
||||
# Spot-check: req in slot 3 holds [10,20,30,0,0,...] etc.
|
||||
self.assertEqual(triton_pool[3, :3].tolist(), [10, 20, 30])
|
||||
self.assertEqual(triton_pool[1, :2].tolist(), [40, 50])
|
||||
self.assertEqual(triton_pool[5, :4].tolist(), [60, 70, 80, 90])
|
||||
|
||||
def test_scatter_empty_batch_no_op(self) -> None:
|
||||
"""Empty input (num_tokens == 0) returns without touching the pool."""
|
||||
flat = torch.empty(0, dtype=torch.int64, device=_DEVICE)
|
||||
offsets = torch.zeros(1, dtype=torch.int64, device=_DEVICE)
|
||||
req_pool_indices = torch.empty(0, dtype=torch.int64, device=_DEVICE)
|
||||
pool = _build_pool(max_reqs=8, max_context_len=16)
|
||||
pool_before = pool.clone()
|
||||
|
||||
launch_scatter_req_token_ids_kernel(
|
||||
flat_in=flat,
|
||||
offsets=offsets,
|
||||
req_pool_indices=req_pool_indices,
|
||||
pool_out=pool,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
self.assertTrue(torch.equal(pool, pool_before))
|
||||
|
||||
def test_scatter_single_req(self) -> None:
|
||||
"""One req with a full row of tokens writes byte-equal to the reference."""
|
||||
seqs = [list(range(10))]
|
||||
lens = [10]
|
||||
rp = [2]
|
||||
|
||||
flat = _build_flat(seqs)
|
||||
offsets = _build_offsets(lens)
|
||||
req_pool_indices = torch.tensor(rp, dtype=torch.int64, device=_DEVICE)
|
||||
triton_pool = _build_pool(max_reqs=4, max_context_len=16)
|
||||
ref_pool = _build_pool(max_reqs=4, max_context_len=16)
|
||||
|
||||
launch_scatter_req_token_ids_kernel(
|
||||
flat_in=flat,
|
||||
offsets=offsets,
|
||||
req_pool_indices=req_pool_indices,
|
||||
pool_out=triton_pool,
|
||||
)
|
||||
scatter_req_token_ids_torch_reference(
|
||||
flat_in=flat,
|
||||
offsets=offsets,
|
||||
req_pool_indices=req_pool_indices,
|
||||
pool_out=ref_pool,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
self.assertTrue(torch.equal(triton_pool, ref_pool))
|
||||
|
||||
def test_scatter_truncates_at_max_context_len(self) -> None:
|
||||
"""Tokens past the pool's max_context_len are silently dropped (no row spill)."""
|
||||
# Two reqs; first req longer than max_context_len. Second req must remain
|
||||
# uncorrupted.
|
||||
seqs = [list(range(20)), [777, 888, 999]]
|
||||
lens = [len(s) for s in seqs]
|
||||
rp = [1, 2]
|
||||
max_context_len = 8
|
||||
|
||||
flat = _build_flat(seqs)
|
||||
offsets = _build_offsets(lens)
|
||||
req_pool_indices = torch.tensor(rp, dtype=torch.int64, device=_DEVICE)
|
||||
pool = _build_pool(max_reqs=4, max_context_len=max_context_len)
|
||||
|
||||
launch_scatter_req_token_ids_kernel(
|
||||
flat_in=flat,
|
||||
offsets=offsets,
|
||||
req_pool_indices=req_pool_indices,
|
||||
pool_out=pool,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
self.assertEqual(
|
||||
pool[1, :max_context_len].tolist(), list(range(max_context_len))
|
||||
)
|
||||
self.assertEqual(pool[2, :3].tolist(), [777, 888, 999])
|
||||
|
||||
def test_scatter_mixed_empty_and_nonempty_reqs(self) -> None:
|
||||
"""Middle req has length 0 between two non-empty reqs: pool rows are byte-equal and untouched rows stay zero."""
|
||||
seqs = [[1, 2], [], [3, 4, 5]]
|
||||
lens = [len(s) for s in seqs]
|
||||
rp = [2, 4, 6]
|
||||
|
||||
flat = _build_flat(seqs)
|
||||
offsets = _build_offsets(lens)
|
||||
req_pool_indices = torch.tensor(rp, dtype=torch.int64, device=_DEVICE)
|
||||
triton_pool = _build_pool(max_reqs=8, max_context_len=8)
|
||||
ref_pool = _build_pool(max_reqs=8, max_context_len=8)
|
||||
|
||||
launch_scatter_req_token_ids_kernel(
|
||||
flat_in=flat,
|
||||
offsets=offsets,
|
||||
req_pool_indices=req_pool_indices,
|
||||
pool_out=triton_pool,
|
||||
)
|
||||
scatter_req_token_ids_torch_reference(
|
||||
flat_in=flat,
|
||||
offsets=offsets,
|
||||
req_pool_indices=req_pool_indices,
|
||||
pool_out=ref_pool,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
self.assertTrue(torch.equal(triton_pool, ref_pool))
|
||||
# Middle req contributes nothing; its pool row stays zero.
|
||||
zero_row = torch.zeros(8, dtype=torch.int32, device=_DEVICE)
|
||||
self.assertTrue(torch.equal(triton_pool[4], zero_row))
|
||||
# First and third reqs are written to their respective rows.
|
||||
self.assertEqual(triton_pool[2, :2].tolist(), [1, 2])
|
||||
self.assertEqual(triton_pool[6, :3].tolist(), [3, 4, 5])
|
||||
|
||||
def test_scatter_random_byte_equal(self) -> None:
|
||||
"""Randomized fuzz across bs, seq lengths, and req pool indices."""
|
||||
rng = random.Random(0)
|
||||
max_reqs = 64
|
||||
max_context_len = 32
|
||||
|
||||
for _ in range(8):
|
||||
bs = rng.randint(1, 16)
|
||||
lens = [rng.randint(0, max_context_len) for _ in range(bs)]
|
||||
# All distinct req pool indices in [1, max_reqs)
|
||||
rp = rng.sample(range(1, max_reqs), k=bs)
|
||||
seqs = [[rng.randint(0, 1 << 30) for _ in range(n)] for n in lens]
|
||||
|
||||
flat = _build_flat(seqs)
|
||||
offsets = _build_offsets(lens)
|
||||
req_pool_indices = torch.tensor(rp, dtype=torch.int64, device=_DEVICE)
|
||||
triton_pool = _build_pool(
|
||||
max_reqs=max_reqs, max_context_len=max_context_len
|
||||
)
|
||||
ref_pool = _build_pool(max_reqs=max_reqs, max_context_len=max_context_len)
|
||||
|
||||
launch_scatter_req_token_ids_kernel(
|
||||
flat_in=flat,
|
||||
offsets=offsets,
|
||||
req_pool_indices=req_pool_indices,
|
||||
pool_out=triton_pool,
|
||||
)
|
||||
scatter_req_token_ids_torch_reference(
|
||||
flat_in=flat,
|
||||
offsets=offsets,
|
||||
req_pool_indices=req_pool_indices,
|
||||
pool_out=ref_pool,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
self.assertTrue(torch.equal(triton_pool, ref_pool))
|
||||
|
||||
|
||||
class TestScatterInputValidation(CustomTestCase):
|
||||
"""Cover the strict input checks in launch_scatter_req_token_ids_kernel."""
|
||||
|
||||
def test_raises_on_2d_flat_in(self) -> None:
|
||||
"""A 2-D flat_in tensor triggers a ValueError before any kernel launch."""
|
||||
flat = torch.zeros((2, 2), dtype=torch.int64, device=_DEVICE)
|
||||
offsets = torch.tensor([0, 1, 2], dtype=torch.int64, device=_DEVICE)
|
||||
req_pool_indices = torch.tensor([1, 2], dtype=torch.int64, device=_DEVICE)
|
||||
pool = _build_pool(max_reqs=4, max_context_len=4)
|
||||
with self.assertRaises(ValueError):
|
||||
launch_scatter_req_token_ids_kernel(
|
||||
flat_in=flat,
|
||||
offsets=offsets,
|
||||
req_pool_indices=req_pool_indices,
|
||||
pool_out=pool,
|
||||
)
|
||||
|
||||
def test_raises_on_wrong_dtype_pool(self) -> None:
|
||||
"""A pool_out with non-int32 dtype triggers a TypeError."""
|
||||
flat = torch.tensor([10, 20], dtype=torch.int64, device=_DEVICE)
|
||||
offsets = torch.tensor([0, 2], dtype=torch.int64, device=_DEVICE)
|
||||
req_pool_indices = torch.tensor([1], dtype=torch.int64, device=_DEVICE)
|
||||
pool = torch.zeros((4, 4), dtype=torch.int64, device=_DEVICE)
|
||||
with self.assertRaises(TypeError):
|
||||
launch_scatter_req_token_ids_kernel(
|
||||
flat_in=flat,
|
||||
offsets=offsets,
|
||||
req_pool_indices=req_pool_indices,
|
||||
pool_out=pool,
|
||||
)
|
||||
|
||||
def test_raises_on_offsets_len_mismatch(self) -> None:
|
||||
"""offsets.shape[0] must equal bs + 1; mismatch triggers a ValueError."""
|
||||
flat = torch.tensor([10, 20], dtype=torch.int64, device=_DEVICE)
|
||||
# bs = 2 but offsets has length 2 instead of 3.
|
||||
offsets = torch.tensor([0, 2], dtype=torch.int64, device=_DEVICE)
|
||||
req_pool_indices = torch.tensor([1, 2], dtype=torch.int64, device=_DEVICE)
|
||||
pool = _build_pool(max_reqs=4, max_context_len=4)
|
||||
with self.assertRaises(ValueError):
|
||||
launch_scatter_req_token_ids_kernel(
|
||||
flat_in=flat,
|
||||
offsets=offsets,
|
||||
req_pool_indices=req_pool_indices,
|
||||
pool_out=pool,
|
||||
)
|
||||
|
||||
def test_raises_on_bs_plus_one_exceeds_batch_block(self) -> None:
|
||||
"""bs+1 must fit in _SCATTER_BATCH_BLOCK; exceeding it triggers a ValueError."""
|
||||
bs = _SCATTER_BATCH_BLOCK
|
||||
flat = torch.empty(0, dtype=torch.int64, device=_DEVICE)
|
||||
offsets = torch.zeros(bs + 1, dtype=torch.int64, device=_DEVICE)
|
||||
req_pool_indices = torch.zeros(bs, dtype=torch.int64, device=_DEVICE)
|
||||
pool = _build_pool(max_reqs=4, max_context_len=4)
|
||||
with self.assertRaises(ValueError):
|
||||
launch_scatter_req_token_ids_kernel(
|
||||
flat_in=flat,
|
||||
offsets=offsets,
|
||||
req_pool_indices=req_pool_indices,
|
||||
pool_out=pool,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,64 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from sglang.jit_kernel.benchmark.kv_canary.utils import (
|
||||
MAX_EXTEND_TOKENS_PER_FORWARD,
|
||||
build_fast_matrix_cases,
|
||||
build_full_matrix_cases,
|
||||
cases_to_x_vals,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=3, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
def test_fast_matrix_cases_include_e2e_decode_and_chunked_prefill_scenarios() -> None:
|
||||
cases = build_fast_matrix_cases()
|
||||
scenarios = {case.scenario for case in cases}
|
||||
|
||||
assert {
|
||||
"e2e_decode_steady",
|
||||
"e2e_decode_tail",
|
||||
"e2e_prefill_chunk_first",
|
||||
"e2e_prefill_chunk_second",
|
||||
"e2e_prefill_chunk_mid",
|
||||
"e2e_prefill_chunk_last",
|
||||
} <= scenarios
|
||||
|
||||
|
||||
def test_extend_cases_are_bounded_to_scheduler_chunk_size() -> None:
|
||||
cases = build_full_matrix_cases()
|
||||
bad_cases = [
|
||||
case
|
||||
for case in cases
|
||||
if case.mode == "extend"
|
||||
and case.bs * case.extend_len > MAX_EXTEND_TOKENS_PER_FORWARD
|
||||
]
|
||||
|
||||
assert bad_cases == []
|
||||
|
||||
|
||||
def test_cases_to_x_vals_includes_scenario_axis() -> None:
|
||||
case = build_fast_matrix_cases()[0]
|
||||
|
||||
x_vals = cases_to_x_vals([case])
|
||||
|
||||
assert x_vals == [
|
||||
(
|
||||
case.scenario,
|
||||
case.bs,
|
||||
case.prefix_len,
|
||||
case.mode,
|
||||
case.extend_len,
|
||||
case.pool_kind,
|
||||
case.real_kv_kind,
|
||||
case.hash_mode,
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -0,0 +1,214 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary import consts
|
||||
from sglang.jit_kernel.kv_canary.verify import (
|
||||
CanaryLaunchTag,
|
||||
RealKvSource,
|
||||
VerifyPlan,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
|
||||
FakeViolationLog,
|
||||
make_canary_buf,
|
||||
make_log_pair,
|
||||
make_verify_plan_pair,
|
||||
stamp_clean_chain,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._differential import _run_both_verify
|
||||
from sglang.jit_kernel.tests.kv_canary._fixtures import (
|
||||
clone_real_kv_sources,
|
||||
make_real_kv_sources,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._fuzz_driver import (
|
||||
FUZZ_SEEDS_PR,
|
||||
run_fuzz_combo,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._invariants import VerifyInvariants
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
|
||||
|
||||
|
||||
_DEVICE = torch.device("cuda")
|
||||
|
||||
_FUZZ_ITER_PER_SEED = 30
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True, kw_only=True)
|
||||
class VerifyFuzzInputs:
|
||||
cuda_canary_buf: torch.Tensor
|
||||
ref_canary_buf: torch.Tensor
|
||||
plan_cuda: VerifyPlan
|
||||
plan_ref: VerifyPlan
|
||||
kernel_kind: CanaryLaunchTag
|
||||
real_kv_sources_cuda: tuple[RealKvSource, ...]
|
||||
real_kv_sources_ref: tuple[RealKvSource, ...]
|
||||
real_kv_hash_mode: consts.RealKvHashMode
|
||||
ring_capacity: int
|
||||
check_verify_expected_token: bool
|
||||
|
||||
|
||||
def _draw_random_verify_inputs(rng: random.Random) -> VerifyFuzzInputs:
|
||||
hash_mode = rng.choice(
|
||||
[
|
||||
consts.RealKvHashMode.NONE,
|
||||
consts.RealKvHashMode.PARTIAL,
|
||||
consts.RealKvHashMode.ALL,
|
||||
]
|
||||
)
|
||||
src_count = rng.choice([1, 2, 4])
|
||||
page_size = rng.choice([1, 16])
|
||||
bytes_per = rng.choice([16, 64, 128])
|
||||
kernel_kind = rng.choice(list(CanaryLaunchTag))
|
||||
plan_size = rng.randint(0, 32)
|
||||
num_slots = max(plan_size + 8, 16)
|
||||
ring_capacity = rng.choice([16, 64, 256])
|
||||
|
||||
sources_cuda = make_real_kv_sources(
|
||||
count=src_count,
|
||||
num_bytes_per_token=bytes_per,
|
||||
page_size=page_size,
|
||||
num_slots=num_slots,
|
||||
device=_DEVICE,
|
||||
rng=rng,
|
||||
)
|
||||
sources_ref = clone_real_kv_sources(sources_cuda)
|
||||
|
||||
cuda_buf = make_canary_buf(
|
||||
num_slots=num_slots, slot_stride_bytes=32, device=_DEVICE
|
||||
)
|
||||
ref_buf = cuda_buf.clone()
|
||||
|
||||
slot_universe = list(range(1, num_slots))
|
||||
rng.shuffle(slot_universe)
|
||||
slot_indices = slot_universe[:plan_size]
|
||||
tokens = [rng.randint(0, 0xFFFFFFFF) for _ in range(plan_size)]
|
||||
positions = list(range(plan_size))
|
||||
prev_slot_indices: list[int] = []
|
||||
for i in range(plan_size):
|
||||
if i == 0:
|
||||
prev_slot_indices.append(-1)
|
||||
else:
|
||||
prev_slot_indices.append(slot_indices[i - 1])
|
||||
|
||||
if hash_mode == consts.RealKvHashMode.NONE and plan_size > 0:
|
||||
stamp_clean_chain(
|
||||
cuda_buf=cuda_buf,
|
||||
ref_buf=ref_buf,
|
||||
slot_indices=slot_indices,
|
||||
tokens=tokens,
|
||||
positions=positions,
|
||||
)
|
||||
|
||||
# Inject prev_slot == TOKEN_TO_KV_SLOT_PADDING into ~15% of entries so the differential
|
||||
# harness exercises the chain-check-skip branch (added for SWA-evicted ancestor handling).
|
||||
# Done AFTER stamp_clean_chain so the stored prev_hash on those slots is still chain-clean;
|
||||
# the kernel must rely on prev_slot==padding (not on stored hash) to decide whether to skip.
|
||||
for i in range(plan_size):
|
||||
if rng.random() < 0.15:
|
||||
prev_slot_indices[i] = consts.TOKEN_TO_KV_SLOT_PADDING
|
||||
|
||||
check_verify_expected_token = rng.random() < 0.5
|
||||
expected_input_ids: list[int] = []
|
||||
for i in range(plan_size):
|
||||
# Always pick a value; with check=False the kernel must not deref this column.
|
||||
if rng.random() < 0.3:
|
||||
expected_input_ids.append(-1)
|
||||
elif rng.random() < 0.5:
|
||||
expected_input_ids.append(int(tokens[i]))
|
||||
else:
|
||||
mutated = (int(tokens[i]) ^ 0x1) & 0xFFFFFFFF
|
||||
expected_input_ids.append(mutated)
|
||||
|
||||
plan_cuda, plan_ref = make_verify_plan_pair(
|
||||
slot_indices=slot_indices,
|
||||
positions=positions,
|
||||
prev_slot_indices=prev_slot_indices,
|
||||
expected_input_ids=expected_input_ids if plan_size > 0 else None,
|
||||
capacity=max(plan_size, 1),
|
||||
device=_DEVICE,
|
||||
)
|
||||
|
||||
return VerifyFuzzInputs(
|
||||
cuda_canary_buf=cuda_buf,
|
||||
ref_canary_buf=ref_buf,
|
||||
plan_cuda=plan_cuda,
|
||||
plan_ref=plan_ref,
|
||||
kernel_kind=kernel_kind,
|
||||
real_kv_sources_cuda=sources_cuda,
|
||||
real_kv_sources_ref=sources_ref,
|
||||
real_kv_hash_mode=hash_mode,
|
||||
ring_capacity=ring_capacity,
|
||||
check_verify_expected_token=check_verify_expected_token,
|
||||
)
|
||||
|
||||
|
||||
def _run_one(inputs: VerifyFuzzInputs) -> None:
|
||||
cuda_buf_before = inputs.cuda_canary_buf.clone()
|
||||
cuda_log, ref_log = make_log_pair(capacity=inputs.ring_capacity, device=_DEVICE)
|
||||
log_before = FakeViolationLog.allocate(
|
||||
capacity=inputs.ring_capacity, device=_DEVICE
|
||||
)
|
||||
_run_both_verify(
|
||||
cuda_canary_buf=inputs.cuda_canary_buf,
|
||||
ref_canary_buf=inputs.ref_canary_buf,
|
||||
plan_cuda=inputs.plan_cuda,
|
||||
plan_ref=inputs.plan_ref,
|
||||
cuda_log=cuda_log,
|
||||
ref_log=ref_log,
|
||||
real_kv_sources_cuda=inputs.real_kv_sources_cuda,
|
||||
real_kv_sources_ref=inputs.real_kv_sources_ref,
|
||||
real_kv_hash_mode=inputs.real_kv_hash_mode,
|
||||
kernel_kind=inputs.kernel_kind,
|
||||
assert_equal=False,
|
||||
check_verify_expected_token=inputs.check_verify_expected_token,
|
||||
)
|
||||
assert int(cuda_log.kernel_run_counter[0].item()) == int(
|
||||
ref_log.kernel_run_counter[0].item()
|
||||
)
|
||||
assert int(cuda_log.slot_run_counter[0].item()) == int(
|
||||
ref_log.slot_run_counter[0].item()
|
||||
)
|
||||
assert int(cuda_log.write_index[0].item()) == int(ref_log.write_index[0].item())
|
||||
VerifyInvariants.assert_all(
|
||||
canary_buf_before=cuda_buf_before,
|
||||
canary_buf_after=inputs.cuda_canary_buf,
|
||||
log_before=log_before,
|
||||
log_after=cuda_log,
|
||||
plan=inputs.plan_cuda,
|
||||
kernel_kind=inputs.kernel_kind,
|
||||
)
|
||||
|
||||
|
||||
def _summarize(inputs: VerifyFuzzInputs) -> str:
|
||||
n_active = int(inputs.plan_cuda.verify_num_valid[0].item())
|
||||
return (
|
||||
f"plan_size={n_active} kind={inputs.kernel_kind.name} "
|
||||
f"hash_mode={inputs.real_kv_hash_mode.name} "
|
||||
f"sources={len(inputs.real_kv_sources_cuda)} "
|
||||
f"ring={inputs.ring_capacity} "
|
||||
f"check_token={inputs.check_verify_expected_token}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seed", FUZZ_SEEDS_PR)
|
||||
def test_verify_fuzz_full_combo(seed: int) -> None:
|
||||
"""Multi-dim verify fuzzer: random hash mode × kernel kind × page × bytes × N iters, byte-equal."""
|
||||
run_fuzz_combo(
|
||||
seed,
|
||||
draw_fn=_draw_random_verify_inputs,
|
||||
run_one_fn=_run_one,
|
||||
summarize_fn=_summarize,
|
||||
n_iter=_FUZZ_ITER_PER_SEED,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,256 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary import consts
|
||||
from sglang.jit_kernel.kv_canary.verify import (
|
||||
CanaryLaunchTag,
|
||||
RealKvSource,
|
||||
)
|
||||
from sglang.jit_kernel.kv_canary.write import WritePlan
|
||||
from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
|
||||
FakeViolationLog,
|
||||
make_canary_buf,
|
||||
make_log_pair,
|
||||
make_write_plan_pair,
|
||||
stamp_pair,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._differential import _run_both_write
|
||||
from sglang.jit_kernel.tests.kv_canary._fixtures import (
|
||||
clone_real_kv_sources,
|
||||
make_real_kv_sources,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._fuzz_driver import (
|
||||
FUZZ_SEEDS_PR,
|
||||
run_fuzz_combo,
|
||||
)
|
||||
from sglang.jit_kernel.tests.kv_canary._invariants import WriteInvariants
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
|
||||
|
||||
|
||||
_DEVICE = torch.device("cuda")
|
||||
|
||||
_FUZZ_ITER_PER_SEED = 30
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True, kw_only=True)
|
||||
class WriteFuzzInputs:
|
||||
cuda_canary_buf: torch.Tensor
|
||||
ref_canary_buf: torch.Tensor
|
||||
plan_cuda: WritePlan
|
||||
plan_ref: WritePlan
|
||||
input_ids: torch.Tensor
|
||||
positions: torch.Tensor
|
||||
out_cache_loc: torch.Tensor
|
||||
kernel_kind: CanaryLaunchTag
|
||||
enable_write_verify_inputs: bool
|
||||
expected_input_tokens: torch.Tensor
|
||||
expected_input_positions: torch.Tensor
|
||||
real_kv_sources_cuda: tuple[RealKvSource, ...]
|
||||
real_kv_sources_ref: tuple[RealKvSource, ...]
|
||||
real_kv_hash_mode: consts.RealKvHashMode
|
||||
ring_capacity: int
|
||||
|
||||
|
||||
def _draw_random_write_inputs(rng: random.Random) -> WriteFuzzInputs:
|
||||
enable_write_verify_inputs = rng.choice([False, True])
|
||||
hash_mode = rng.choice(
|
||||
[
|
||||
consts.RealKvHashMode.NONE,
|
||||
consts.RealKvHashMode.PARTIAL,
|
||||
consts.RealKvHashMode.ALL,
|
||||
]
|
||||
)
|
||||
src_count = rng.choice([1, 2, 4])
|
||||
page_size = rng.choice([1, 16])
|
||||
bytes_per = rng.choice([16, 64, 128])
|
||||
kernel_kind = rng.choice(list(CanaryLaunchTag))
|
||||
ring_capacity = rng.choice([16, 64, 256])
|
||||
|
||||
n_reqs = rng.randint(1, 4)
|
||||
per_req_tokens: list[int] = [rng.randint(1, 5) for _ in range(n_reqs)]
|
||||
total_tokens = sum(per_req_tokens)
|
||||
num_slots = max(total_tokens + 8, 16)
|
||||
|
||||
sources_cuda = make_real_kv_sources(
|
||||
count=src_count,
|
||||
num_bytes_per_token=bytes_per,
|
||||
page_size=page_size,
|
||||
num_slots=num_slots,
|
||||
device=_DEVICE,
|
||||
rng=rng,
|
||||
)
|
||||
sources_ref = clone_real_kv_sources(sources_cuda)
|
||||
|
||||
cuda_buf = make_canary_buf(
|
||||
num_slots=num_slots, slot_stride_bytes=32, device=_DEVICE
|
||||
)
|
||||
ref_buf = cuda_buf.clone()
|
||||
|
||||
write_offsets: list[int] = [0]
|
||||
running = 0
|
||||
for t in per_req_tokens:
|
||||
running += t
|
||||
write_offsets.append(running)
|
||||
slot_pool = list(range(1, num_slots))
|
||||
rng.shuffle(slot_pool)
|
||||
seed_slot_indices: list[int] = []
|
||||
for _ in range(n_reqs):
|
||||
if rng.random() < 0.4 or len(slot_pool) <= total_tokens:
|
||||
seed_slot_indices.append(-1)
|
||||
else:
|
||||
seed_slot_indices.append(slot_pool.pop())
|
||||
|
||||
out_cache_loc_list: list[int] = []
|
||||
for _ in range(total_tokens):
|
||||
if not slot_pool:
|
||||
out_cache_loc_list.append(-1)
|
||||
else:
|
||||
out_cache_loc_list.append(slot_pool.pop())
|
||||
|
||||
plan_cuda, plan_ref = make_write_plan_pair(
|
||||
write_offsets=write_offsets,
|
||||
seed_slot_indices=seed_slot_indices,
|
||||
num_valid_reqs=n_reqs,
|
||||
device=_DEVICE,
|
||||
)
|
||||
|
||||
input_ids = torch.tensor(
|
||||
[rng.randint(-(1 << 31), (1 << 31) - 1) for _ in range(total_tokens)],
|
||||
dtype=torch.int64,
|
||||
device=_DEVICE,
|
||||
)
|
||||
# Per-chain sequential positions so the write kernel's chain-step position assert holds.
|
||||
# For chains with a real seed slot, stamp the seed with (chain_start_position - 1) so the
|
||||
# first chain entry's position == seed.position + 1.
|
||||
chain_start_positions: list[int] = [rng.randint(0, 1024) for _ in range(n_reqs)]
|
||||
positions_list: list[int] = []
|
||||
for r in range(n_reqs):
|
||||
start = chain_start_positions[r]
|
||||
positions_list.extend(start + i for i in range(per_req_tokens[r]))
|
||||
seed_slot = seed_slot_indices[r]
|
||||
if seed_slot >= 0:
|
||||
stamp_pair(
|
||||
(cuda_buf, ref_buf),
|
||||
slot_idx=seed_slot,
|
||||
token=0,
|
||||
position=start - 1,
|
||||
prev_hash=0,
|
||||
)
|
||||
positions = torch.tensor(positions_list, dtype=torch.int64, device=_DEVICE)
|
||||
out_cache_loc = torch.tensor(out_cache_loc_list, dtype=torch.int64, device=_DEVICE)
|
||||
expected_input_tokens = input_ids.clone()
|
||||
expected_input_positions = positions.clone()
|
||||
if enable_write_verify_inputs:
|
||||
candidate_indices = [
|
||||
idx for idx, slot in enumerate(out_cache_loc_list) if slot >= 0
|
||||
]
|
||||
rng.shuffle(candidate_indices)
|
||||
mismatch_count = rng.randint(0, len(candidate_indices))
|
||||
for idx in candidate_indices[:mismatch_count]:
|
||||
if rng.choice([False, True]):
|
||||
expected_input_tokens[idx] = expected_input_tokens[idx] + 1
|
||||
else:
|
||||
expected_input_positions[idx] = expected_input_positions[idx] + 1
|
||||
|
||||
return WriteFuzzInputs(
|
||||
cuda_canary_buf=cuda_buf,
|
||||
ref_canary_buf=ref_buf,
|
||||
plan_cuda=plan_cuda,
|
||||
plan_ref=plan_ref,
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
out_cache_loc=out_cache_loc,
|
||||
kernel_kind=kernel_kind,
|
||||
enable_write_verify_inputs=enable_write_verify_inputs,
|
||||
expected_input_tokens=expected_input_tokens,
|
||||
expected_input_positions=expected_input_positions,
|
||||
real_kv_sources_cuda=sources_cuda,
|
||||
real_kv_sources_ref=sources_ref,
|
||||
real_kv_hash_mode=hash_mode,
|
||||
ring_capacity=ring_capacity,
|
||||
)
|
||||
|
||||
|
||||
def _run_one(inputs: WriteFuzzInputs) -> None:
|
||||
cuda_buf_before = inputs.cuda_canary_buf.clone()
|
||||
cuda_log, ref_log = make_log_pair(capacity=inputs.ring_capacity, device=_DEVICE)
|
||||
log_before = FakeViolationLog.allocate(
|
||||
capacity=inputs.ring_capacity, device=_DEVICE
|
||||
)
|
||||
_run_both_write(
|
||||
cuda_canary_buf=inputs.cuda_canary_buf,
|
||||
ref_canary_buf=inputs.ref_canary_buf,
|
||||
plan_cuda=inputs.plan_cuda,
|
||||
plan_ref=inputs.plan_ref,
|
||||
input_ids=inputs.input_ids,
|
||||
positions=inputs.positions,
|
||||
out_cache_loc=inputs.out_cache_loc,
|
||||
enable_write_verify_inputs=inputs.enable_write_verify_inputs,
|
||||
expected_input_tokens=inputs.expected_input_tokens,
|
||||
expected_input_positions=inputs.expected_input_positions,
|
||||
cuda_log=cuda_log,
|
||||
ref_log=ref_log,
|
||||
real_kv_sources_cuda=inputs.real_kv_sources_cuda,
|
||||
real_kv_sources_ref=inputs.real_kv_sources_ref,
|
||||
real_kv_hash_mode=inputs.real_kv_hash_mode,
|
||||
kernel_kind=inputs.kernel_kind,
|
||||
assert_equal=False,
|
||||
)
|
||||
assert torch.equal(
|
||||
inputs.cuda_canary_buf, inputs.ref_canary_buf
|
||||
), "CUDA vs ref canary_buf diverged"
|
||||
assert int(cuda_log.write_index[0].item()) == int(ref_log.write_index[0].item())
|
||||
assert int(cuda_log.slot_run_counter[0].item()) == int(
|
||||
ref_log.slot_run_counter[0].item()
|
||||
)
|
||||
assert int(cuda_log.kernel_run_counter[0].item()) == int(
|
||||
ref_log.kernel_run_counter[0].item()
|
||||
)
|
||||
WriteInvariants.assert_all(
|
||||
canary_buf_before=cuda_buf_before,
|
||||
canary_buf_after=inputs.cuda_canary_buf,
|
||||
plan=inputs.plan_cuda,
|
||||
input_ids=inputs.input_ids,
|
||||
positions=inputs.positions,
|
||||
out_cache_loc=inputs.out_cache_loc,
|
||||
enable_write_verify_inputs=inputs.enable_write_verify_inputs,
|
||||
expected_input_tokens=inputs.expected_input_tokens,
|
||||
expected_input_positions=inputs.expected_input_positions,
|
||||
log_before=log_before,
|
||||
log_after=cuda_log,
|
||||
)
|
||||
|
||||
|
||||
def _summarize(inputs: WriteFuzzInputs) -> str:
|
||||
n_active = int(inputs.plan_cuda.write_num_valid_reqs[0].item())
|
||||
total = int(inputs.plan_cuda.write_offsets[n_active].item())
|
||||
return (
|
||||
f"n_reqs={n_active} total_tokens={total} kind={inputs.kernel_kind.name} "
|
||||
f"pseudo={inputs.enable_write_verify_inputs} hash_mode={inputs.real_kv_hash_mode.name} "
|
||||
f"sources={len(inputs.real_kv_sources_cuda)}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seed", FUZZ_SEEDS_PR)
|
||||
def test_write_fuzz_full_combo(seed: int) -> None:
|
||||
"""Multi-dim write fuzzer: random pseudo/hash/kernel/page/source × N iters, byte-equal."""
|
||||
run_fuzz_combo(
|
||||
seed,
|
||||
draw_fn=_draw_random_write_inputs,
|
||||
run_one_fn=_run_one,
|
||||
summarize_fn=_summarize,
|
||||
n_iter=_FUZZ_ITER_PER_SEED,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,205 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from sglang.jit_kernel.activation import (
|
||||
SUPPORTED_ACTIVATIONS,
|
||||
relu2,
|
||||
run_activation,
|
||||
)
|
||||
from sglang.jit_kernel.utils import get_ci_test_range
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=20, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=30, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
OPS = SUPPORTED_ACTIVATIONS
|
||||
DTYPES = [torch.float16, torch.bfloat16, torch.float32]
|
||||
SHAPES = get_ci_test_range(
|
||||
full_range=[
|
||||
(7, 16),
|
||||
(83, 1024),
|
||||
(3, 5, 16),
|
||||
(2, 3, 512),
|
||||
(1, 17, 4096),
|
||||
*[(2**x, 2048) for x in range(0, 15, 2)],
|
||||
*[(2**x, 65536) for x in range(0, 5, 2)],
|
||||
],
|
||||
ci_range=[(7, 16), (2, 3, 512)],
|
||||
)
|
||||
|
||||
|
||||
def _reference(op_name: str, x: torch.Tensor) -> torch.Tensor:
|
||||
d = x.shape[-1] // 2
|
||||
lhs = x[..., :d].float()
|
||||
rhs = x[..., d:]
|
||||
if op_name == "silu":
|
||||
act = F.silu(lhs)
|
||||
elif op_name == "gelu":
|
||||
act = F.gelu(lhs, approximate="none")
|
||||
else:
|
||||
act = F.gelu(lhs, approximate="tanh")
|
||||
return act.to(dtype=x.dtype) * rhs
|
||||
|
||||
|
||||
def _tolerances(dtype: torch.dtype) -> tuple[float, float]:
|
||||
if dtype == torch.float32:
|
||||
return 1e-4, 1e-4
|
||||
return 1e-2, 1e-2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("op_name", OPS)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("shape", SHAPES)
|
||||
def test_activation_correctness(
|
||||
op_name: str, dtype: torch.dtype, shape: tuple[int, ...]
|
||||
) -> None:
|
||||
x = torch.randn(shape, dtype=dtype, device="cuda")
|
||||
out = run_activation(op_name, x, None)
|
||||
expected = _reference(op_name, x)
|
||||
atol, rtol = _tolerances(dtype)
|
||||
torch.testing.assert_close(out, expected, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("op_name", OPS)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("shape", SHAPES)
|
||||
def test_activation_out_param(
|
||||
op_name: str, dtype: torch.dtype, shape: tuple[int, ...]
|
||||
) -> None:
|
||||
x = torch.randn(shape, dtype=dtype, device="cuda")
|
||||
out = torch.empty(shape[:-1] + (shape[-1] // 2,), dtype=dtype, device="cuda")
|
||||
result = run_activation(op_name, x, out)
|
||||
assert result is out
|
||||
expected = _reference(op_name, x)
|
||||
atol, rtol = _tolerances(dtype)
|
||||
torch.testing.assert_close(out, expected, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
FILTER_SHAPES = get_ci_test_range(
|
||||
full_range=[(83, 1024), (256, 2048), (1024, 4096)],
|
||||
ci_range=[(83, 1024)],
|
||||
)
|
||||
EXPERT_STEPS = [1, 16]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("op_name", OPS)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("shape", FILTER_SHAPES)
|
||||
@pytest.mark.parametrize("expert_step", EXPERT_STEPS)
|
||||
def test_activation_filter_expert(
|
||||
op_name: str,
|
||||
dtype: torch.dtype,
|
||||
shape: tuple[int, int],
|
||||
expert_step: int,
|
||||
) -> None:
|
||||
"""expert_ids[token // expert_step] == -1 must leave the output row untouched."""
|
||||
num_tokens = shape[0]
|
||||
x = torch.randn(shape, dtype=dtype, device="cuda")
|
||||
# Pre-fill out with a sentinel so we can detect untouched rows.
|
||||
sentinel = float("nan")
|
||||
out = torch.full(
|
||||
shape[:-1] + (shape[-1] // 2,),
|
||||
sentinel,
|
||||
dtype=dtype,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
num_groups = (num_tokens + expert_step - 1) // expert_step
|
||||
expert_ids = torch.randint(
|
||||
low=0, high=8, size=(num_groups,), dtype=torch.int32, device="cuda"
|
||||
)
|
||||
skip_mask = torch.rand(num_groups, device="cuda") < 0.4
|
||||
expert_ids[skip_mask] = -1
|
||||
|
||||
result = run_activation(op_name, x, out, expert_ids, expert_step)
|
||||
assert result is out
|
||||
|
||||
token_skip = skip_mask[torch.arange(num_tokens, device="cuda") // expert_step]
|
||||
expected = _reference(op_name, x)
|
||||
atol, rtol = _tolerances(dtype)
|
||||
|
||||
kept = ~token_skip
|
||||
if kept.any():
|
||||
torch.testing.assert_close(out[kept], expected[kept], atol=atol, rtol=rtol)
|
||||
if token_skip.any():
|
||||
assert torch.isnan(
|
||||
out[token_skip]
|
||||
).all(), "filter_expert kernel touched rows whose expert_id is -1"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("op_name", OPS)
|
||||
def test_activation_filter_expert_all_skipped(op_name: str) -> None:
|
||||
"""If every expert id is -1, the output must be left entirely untouched."""
|
||||
shape = (32, 512)
|
||||
x = torch.randn(shape, dtype=torch.bfloat16, device="cuda")
|
||||
out = torch.full(
|
||||
shape[:-1] + (shape[-1] // 2,),
|
||||
float("nan"),
|
||||
dtype=torch.bfloat16,
|
||||
device="cuda",
|
||||
)
|
||||
expert_ids = torch.full((shape[0],), -1, dtype=torch.int32, device="cuda")
|
||||
run_activation(op_name, x, out, expert_ids, 1)
|
||||
assert torch.isnan(out).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("op_name", OPS)
|
||||
def test_activation_filter_expert_none_skipped(op_name: str) -> None:
|
||||
"""No -1 in expert_ids must yield bit-identical output to the unfiltered path."""
|
||||
shape = (64, 512)
|
||||
dtype = torch.bfloat16
|
||||
x = torch.randn(shape, dtype=dtype, device="cuda")
|
||||
expert_ids = torch.zeros((shape[0],), dtype=torch.int32, device="cuda")
|
||||
out_filtered = run_activation(op_name, x, None, expert_ids, 1)
|
||||
out_unfiltered = run_activation(op_name, x, None)
|
||||
torch.testing.assert_close(out_filtered, out_unfiltered, atol=0.0, rtol=0.0)
|
||||
|
||||
|
||||
UNARY_SHAPES = get_ci_test_range(
|
||||
full_range=[
|
||||
(7, 16),
|
||||
(83, 1024),
|
||||
(3, 5, 16),
|
||||
(2, 3, 512),
|
||||
(1, 17, 4096),
|
||||
*[(2**x, 2048) for x in range(0, 15, 2)],
|
||||
],
|
||||
ci_range=[(7, 16), (2, 3, 512)],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("shape", UNARY_SHAPES)
|
||||
def test_relu2_correctness(dtype: torch.dtype, shape: tuple[int, ...]) -> None:
|
||||
x = torch.randn(shape, dtype=dtype, device="cuda")
|
||||
out = relu2(x)
|
||||
expected = F.relu(x.float()).pow(2).to(dtype=dtype)
|
||||
atol, rtol = _tolerances(dtype)
|
||||
torch.testing.assert_close(out, expected, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("shape", UNARY_SHAPES)
|
||||
def test_relu2_out_param(dtype: torch.dtype, shape: tuple[int, ...]) -> None:
|
||||
x = torch.randn(shape, dtype=dtype, device="cuda")
|
||||
out = torch.empty(shape, dtype=dtype, device="cuda")
|
||||
result = relu2(x, out)
|
||||
assert result is out
|
||||
expected = F.relu(x.float()).pow(2).to(dtype=dtype)
|
||||
atol, rtol = _tolerances(dtype)
|
||||
torch.testing.assert_close(out, expected, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
def test_relu2_negative_inputs_zeroed() -> None:
|
||||
"""All-negative input must produce an all-zero output."""
|
||||
x = -torch.rand((64, 512), dtype=torch.bfloat16, device="cuda") - 1e-3
|
||||
out = relu2(x)
|
||||
assert torch.count_nonzero(out) == 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,41 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.add_constant import add_constant
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=45, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=180, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("size", [1, 2, 127, 128, 1024, 1025, 4096, 4097])
|
||||
@pytest.mark.parametrize("constant", [0, 1, 7, 1024, -3])
|
||||
def test_add_constant(size: int, constant: int) -> None:
|
||||
src = torch.arange(0, size, dtype=torch.int32, device="cuda")
|
||||
dst = add_constant(src, constant)
|
||||
assert torch.all(dst == src + constant)
|
||||
|
||||
|
||||
def test_add_constant_unaligned_input() -> None:
|
||||
src = torch.arange(0, 4098, dtype=torch.int32, device="cuda")[1:]
|
||||
dst = add_constant(src, 7)
|
||||
assert torch.all(dst == src + 7)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("size", [2**20, 2**20 + 3])
|
||||
def test_add_constant_large_aligned_input(size: int) -> None:
|
||||
src = torch.arange(0, size, dtype=torch.int32, device="cuda")
|
||||
dst = add_constant(src, -3)
|
||||
assert torch.all(dst == src - 3)
|
||||
|
||||
|
||||
def test_add_constant_large_unaligned_input() -> None:
|
||||
src = torch.arange(0, 2**20 + 4, dtype=torch.int32, device="cuda")[1:]
|
||||
dst = add_constant(src, 7)
|
||||
assert torch.all(dst == src + 7)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,169 @@
|
||||
import itertools
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.awq_dequantize import awq_dequantize as jit_awq_dequantize
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=9, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
try:
|
||||
from sgl_kernel import awq_dequantize as aot_awq_dequantize
|
||||
|
||||
AOT_AVAILABLE = True
|
||||
except ImportError:
|
||||
AOT_AVAILABLE = False
|
||||
|
||||
|
||||
def reverse_awq_order(t: torch.Tensor):
|
||||
bits = 4
|
||||
AWQ_REVERSE_ORDER = [0, 4, 1, 5, 2, 6, 3, 7]
|
||||
reverse_order_tensor = torch.arange(
|
||||
t.shape[-1],
|
||||
dtype=torch.int32,
|
||||
device=t.device,
|
||||
)
|
||||
reverse_order_tensor = reverse_order_tensor.view(-1, 32 // bits)
|
||||
reverse_order_tensor = reverse_order_tensor[:, AWQ_REVERSE_ORDER]
|
||||
reverse_order_tensor = reverse_order_tensor.view(-1)
|
||||
|
||||
t = t[:, reverse_order_tensor] & 0xF
|
||||
return t
|
||||
|
||||
|
||||
# qweights - [R , C // 8], int32
|
||||
# scales - [R // G, C ], float16
|
||||
# zeros - [R // G, C // 8], int32
|
||||
def awq_dequantize_torch(
|
||||
qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor, group_size: int
|
||||
) -> torch.Tensor:
|
||||
if group_size == -1:
|
||||
group_size = qweight.shape[0]
|
||||
|
||||
bits = 4
|
||||
shifts = torch.arange(0, 32, bits, device=qzeros.device)
|
||||
|
||||
iweights = torch.bitwise_right_shift(qweight[:, :, None], shifts[None, None, :]).to(
|
||||
torch.int8
|
||||
)
|
||||
|
||||
iweights = iweights.view(iweights.shape[0], -1)
|
||||
|
||||
zeros = torch.bitwise_right_shift(qzeros[:, :, None], shifts[None, None, :]).to(
|
||||
torch.int8
|
||||
)
|
||||
zeros = zeros.view(qzeros.shape[0], -1)
|
||||
zeros = reverse_awq_order(zeros)
|
||||
|
||||
iweights = reverse_awq_order(iweights)
|
||||
|
||||
iweights = torch.bitwise_and(iweights, (2**bits) - 1)
|
||||
zeros = torch.bitwise_and(zeros, (2**bits) - 1)
|
||||
|
||||
scales = scales.repeat_interleave(group_size, dim=0)
|
||||
zeros = zeros.repeat_interleave(group_size, dim=0)
|
||||
return (iweights - zeros) * scales
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"qweight_row,qweight_col,is_bf16_act",
|
||||
list(
|
||||
itertools.product(
|
||||
[128, 256, 512, 1024, 3584],
|
||||
[16, 32, 64, 128, 448],
|
||||
[True, False],
|
||||
)
|
||||
),
|
||||
)
|
||||
def test_awq_dequantize_jit_vs_torch(
|
||||
qweight_row: int, qweight_col: int, is_bf16_act: bool
|
||||
):
|
||||
device = torch.device("cuda")
|
||||
qweight = torch.randint(
|
||||
0,
|
||||
torch.iinfo(torch.int32).max,
|
||||
(qweight_row, qweight_col),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
group_size = qweight_row
|
||||
scales_row = qweight_row // group_size
|
||||
scales_col = qweight_col * 8
|
||||
|
||||
if is_bf16_act:
|
||||
scales = torch.rand(scales_row, scales_col, dtype=torch.bfloat16, device=device)
|
||||
else:
|
||||
scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device)
|
||||
|
||||
qzeros = torch.randint(
|
||||
0,
|
||||
torch.iinfo(torch.int32).max,
|
||||
(scales_row, qweight_col),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# Run both implementations
|
||||
torch_out = awq_dequantize_torch(qweight, scales, qzeros, group_size)
|
||||
jit_out = jit_awq_dequantize(qweight, scales, qzeros)
|
||||
|
||||
# Compare results (approximate due to different computation paths)
|
||||
torch.testing.assert_close(
|
||||
torch_out.to(torch.float32), jit_out.to(torch.float32), rtol=1e-3, atol=1e-5
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"qweight_row,qweight_col,is_bf16_act",
|
||||
list(
|
||||
itertools.product(
|
||||
[128, 256, 512, 1024, 3584],
|
||||
[16, 32, 64, 128, 448],
|
||||
[True, False],
|
||||
)
|
||||
),
|
||||
)
|
||||
def test_awq_dequantize_jit_vs_aot(
|
||||
qweight_row: int, qweight_col: int, is_bf16_act: bool
|
||||
):
|
||||
if not AOT_AVAILABLE:
|
||||
pytest.skip("sgl_kernel AOT not available")
|
||||
|
||||
device = torch.device("cuda")
|
||||
qweight = torch.randint(
|
||||
0,
|
||||
torch.iinfo(torch.int32).max,
|
||||
(qweight_row, qweight_col),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
group_size = qweight_row
|
||||
scales_row = qweight_row // group_size
|
||||
scales_col = qweight_col * 8
|
||||
|
||||
if is_bf16_act:
|
||||
scales = torch.rand(scales_row, scales_col, dtype=torch.bfloat16, device=device)
|
||||
else:
|
||||
scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device)
|
||||
|
||||
qzeros = torch.randint(
|
||||
0,
|
||||
torch.iinfo(torch.int32).max,
|
||||
(scales_row, qweight_col),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# Run both implementations
|
||||
aot_out = aot_awq_dequantize(qweight, scales, qzeros)
|
||||
jit_out = jit_awq_dequantize(qweight, scales, qzeros)
|
||||
|
||||
# Bitwise equality
|
||||
torch.testing.assert_close(jit_out, aot_out, rtol=0, atol=0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,125 @@
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from sgl_kernel.scalar_type import scalar_types
|
||||
|
||||
from sglang.jit_kernel.awq_marlin_repack import (
|
||||
awq_marlin_moe_repack as jit_awq_marlin_moe_repack,
|
||||
)
|
||||
from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
def _has_aot_awq_marlin_moe_repack() -> bool:
|
||||
return hasattr(torch.ops.sgl_kernel, "awq_marlin_moe_repack") and hasattr(
|
||||
torch.ops.sgl_kernel.awq_marlin_moe_repack, "default"
|
||||
)
|
||||
|
||||
|
||||
AOT_AVAILABLE = _has_aot_awq_marlin_moe_repack()
|
||||
|
||||
|
||||
def awq_pack(
|
||||
q_w: torch.Tensor,
|
||||
num_bits: int,
|
||||
size_k: int,
|
||||
size_n: int,
|
||||
):
|
||||
assert q_w.shape == (size_k, size_n)
|
||||
|
||||
if num_bits == 4:
|
||||
interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7])
|
||||
elif num_bits == 8:
|
||||
interleave = np.array([0, 2, 1, 3])
|
||||
else:
|
||||
raise Exception("num_bits must be 4 or 8, got {}".format(num_bits))
|
||||
|
||||
q_w = q_w.reshape((-1, len(interleave)))[:, interleave].ravel()
|
||||
q_w = q_w.reshape((-1, size_n)).contiguous()
|
||||
|
||||
return pack_cols(q_w, num_bits, size_k, size_n)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_bits", [4])
|
||||
@pytest.mark.parametrize("num_experts", [2, 4, 8])
|
||||
@pytest.mark.parametrize("k_tiles,n_tiles", [(1, 1), (2, 2), (4, 4)])
|
||||
@pytest.mark.parametrize("group_size", [16, 32])
|
||||
def test_awq_marlin_moe_repack_jit_vs_aot(
|
||||
num_bits, num_experts, k_tiles, n_tiles, group_size
|
||||
):
|
||||
if not AOT_AVAILABLE:
|
||||
pytest.skip("sgl_kernel AOT not available")
|
||||
|
||||
tile_k, tile_n = 16, 64
|
||||
size_k = k_tiles * tile_k
|
||||
size_n = n_tiles * tile_n
|
||||
pack_factor = 32 // num_bits
|
||||
|
||||
# Create per-expert AWQ-packed weights
|
||||
b_q_weight = torch.empty(
|
||||
(num_experts, size_k, size_n // pack_factor),
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
for e in range(num_experts):
|
||||
b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda")
|
||||
w_ref, q_w, s, zp = quantize_weights(
|
||||
b_weight, scalar_types.uint4, group_size, zero_points=True
|
||||
)
|
||||
b_q_weight[e] = awq_pack(q_w, num_bits, size_k, size_n)
|
||||
|
||||
perm = torch.empty((num_experts, 0), dtype=torch.int32, device="cuda")
|
||||
|
||||
out_jit = jit_awq_marlin_moe_repack(b_q_weight, perm, size_k, size_n, num_bits)
|
||||
out_aot = torch.ops.sgl_kernel.awq_marlin_moe_repack.default(
|
||||
b_q_weight, perm, size_k, size_n, num_bits
|
||||
)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Bitwise equality
|
||||
torch.testing.assert_close(out_jit, out_aot, rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_bits", [4])
|
||||
@pytest.mark.parametrize("num_experts", [2, 4])
|
||||
@pytest.mark.parametrize("k_tiles,n_tiles", [(1, 1), (2, 2)])
|
||||
@pytest.mark.parametrize("group_size", [16, 32])
|
||||
def test_awq_marlin_moe_repack_shape(
|
||||
num_bits, num_experts, k_tiles, n_tiles, group_size
|
||||
):
|
||||
tile_k, tile_n = 16, 64
|
||||
size_k = k_tiles * tile_k
|
||||
size_n = n_tiles * tile_n
|
||||
pack_factor = 32 // num_bits
|
||||
|
||||
# Create per-expert AWQ-packed weights
|
||||
b_q_weight = torch.empty(
|
||||
(num_experts, size_k, size_n // pack_factor),
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
for e in range(num_experts):
|
||||
b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda")
|
||||
w_ref, q_w, s, zp = quantize_weights(
|
||||
b_weight, scalar_types.uint4, group_size, zero_points=True
|
||||
)
|
||||
b_q_weight[e] = awq_pack(q_w, num_bits, size_k, size_n)
|
||||
|
||||
perm = torch.empty((num_experts, 0), dtype=torch.int32, device="cuda")
|
||||
|
||||
out = jit_awq_marlin_moe_repack(b_q_weight, perm, size_k, size_n, num_bits)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
assert out.is_cuda and out.dtype == torch.int32
|
||||
expected_shape = (num_experts, size_k // 16, size_n * (num_bits // 2))
|
||||
assert list(out.shape) == list(expected_shape)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,111 @@
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from sgl_kernel.scalar_type import scalar_types
|
||||
|
||||
from sglang.jit_kernel.awq_marlin_repack import (
|
||||
awq_marlin_repack as jit_awq_marlin_repack,
|
||||
)
|
||||
from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_marlin_utils import get_weight_perm, marlin_weights
|
||||
|
||||
register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
def _has_aot_awq_marlin_repack() -> bool:
|
||||
return hasattr(torch.ops.sgl_kernel, "awq_marlin_repack") and hasattr(
|
||||
torch.ops.sgl_kernel.awq_marlin_repack, "default"
|
||||
)
|
||||
|
||||
|
||||
AOT_AVAILABLE = _has_aot_awq_marlin_repack()
|
||||
|
||||
|
||||
def awq_pack(
|
||||
q_w: torch.Tensor,
|
||||
num_bits: int,
|
||||
size_k: int,
|
||||
size_n: int,
|
||||
):
|
||||
assert q_w.shape == (size_k, size_n)
|
||||
|
||||
if num_bits == 4:
|
||||
interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7])
|
||||
elif num_bits == 8:
|
||||
interleave = np.array([0, 2, 1, 3])
|
||||
else:
|
||||
raise Exception("num_bits must be 4 or 8, got {}".format(num_bits))
|
||||
|
||||
q_w = q_w.reshape((-1, len(interleave)))[:, interleave].ravel()
|
||||
q_w = q_w.reshape((-1, size_n)).contiguous()
|
||||
|
||||
return pack_cols(q_w, num_bits, size_k, size_n)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_bits", [4, 8])
|
||||
@pytest.mark.parametrize("k_tiles,n_tiles", [(1, 1), (2, 2), (4, 4)])
|
||||
@pytest.mark.parametrize("group_size", [16, 32])
|
||||
def test_awq_marlin_repack_jit_vs_aot(num_bits, k_tiles, n_tiles, group_size):
|
||||
if not AOT_AVAILABLE:
|
||||
pytest.skip("sgl_kernel AOT not available")
|
||||
|
||||
tile_k, tile_n = 16, 64
|
||||
size_k = k_tiles * tile_k
|
||||
size_n = n_tiles * tile_n
|
||||
|
||||
b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda")
|
||||
|
||||
w_ref, q_w, s, zp = quantize_weights(
|
||||
b_weight, scalar_types.uint4, group_size, zero_points=True
|
||||
)
|
||||
|
||||
q_w_awq = awq_pack(q_w, num_bits, size_k, size_n)
|
||||
|
||||
out_jit = jit_awq_marlin_repack(q_w_awq, size_k, size_n, num_bits)
|
||||
out_aot = torch.ops.sgl_kernel.awq_marlin_repack.default(
|
||||
q_w_awq, size_k, size_n, num_bits
|
||||
)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Bitwise equality
|
||||
torch.testing.assert_close(out_jit, out_aot, rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_bits", [4, 8])
|
||||
@pytest.mark.parametrize("k_tiles,n_tiles", [(1, 1), (2, 2)])
|
||||
@pytest.mark.parametrize("group_size", [16, 32])
|
||||
def test_awq_marlin_repack_correct(num_bits, k_tiles, n_tiles, group_size):
|
||||
tile_k, tile_n = 16, 64
|
||||
size_k = k_tiles * tile_k
|
||||
size_n = n_tiles * tile_n
|
||||
pack_factor = 32 // num_bits
|
||||
|
||||
b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda")
|
||||
|
||||
w_ref, q_w, s, zp = quantize_weights(
|
||||
b_weight, scalar_types.uint4, group_size, zero_points=True
|
||||
)
|
||||
|
||||
q_w_awq = awq_pack(q_w, num_bits, size_k, size_n)
|
||||
|
||||
weight_perm = get_weight_perm(num_bits)
|
||||
q_w_marlin = marlin_weights(q_w, size_k, size_n, num_bits, weight_perm)
|
||||
|
||||
out_gpu = jit_awq_marlin_repack(q_w_awq, size_k, size_n, num_bits)
|
||||
assert out_gpu.is_cuda and out_gpu.dtype == torch.int32
|
||||
|
||||
expected_cols = size_n * tile_k // pack_factor
|
||||
assert list(out_gpu.shape) == [size_k // tile_k, expected_cols]
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(out_gpu, q_w_marlin)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,46 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.clamp_position import clamp_position_cuda
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=12, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
def _reference_clamp_position(seq_lens):
|
||||
return torch.clamp(seq_lens - 1, min=0).to(seq_lens.dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("size", [1, 2, 127, 128, 255, 256, 1024, 4097])
|
||||
@pytest.mark.parametrize("dtype", [torch.int32, torch.int64])
|
||||
class TestClampPosition:
|
||||
def test_normal(self, size: int, dtype: torch.dtype) -> None:
|
||||
seq_lens = torch.randint(1, 10000, (size,), dtype=dtype, device="cuda")
|
||||
expected = _reference_clamp_position(seq_lens)
|
||||
result = clamp_position_cuda(seq_lens)
|
||||
assert torch.equal(result, expected)
|
||||
|
||||
def test_zeros(self, size: int, dtype: torch.dtype) -> None:
|
||||
seq_lens = torch.zeros(size, dtype=dtype, device="cuda")
|
||||
expected = _reference_clamp_position(seq_lens)
|
||||
result = clamp_position_cuda(seq_lens)
|
||||
assert torch.equal(result, expected)
|
||||
|
||||
def test_ones(self, size: int, dtype: torch.dtype) -> None:
|
||||
seq_lens = torch.ones(size, dtype=dtype, device="cuda")
|
||||
expected = _reference_clamp_position(seq_lens)
|
||||
result = clamp_position_cuda(seq_lens)
|
||||
assert torch.equal(result, expected)
|
||||
|
||||
def test_mixed(self, size: int, dtype: torch.dtype) -> None:
|
||||
seq_lens = torch.randint(0, 10000, (size,), dtype=dtype, device="cuda")
|
||||
expected = _reference_clamp_position(seq_lens)
|
||||
result = clamp_position_cuda(seq_lens)
|
||||
assert torch.equal(result, expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,175 @@
|
||||
import itertools
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=17, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
def torch_concat_mla_k(
|
||||
k: torch.Tensor, k_nope: torch.Tensor, k_rope: torch.Tensor
|
||||
) -> None:
|
||||
"""Reference PyTorch implementation for concat_mla_k."""
|
||||
# k_nope: [num_tokens, num_heads, nope_head_dim]
|
||||
# k_rope: [num_tokens, 1, rope_head_dim]
|
||||
# k: [num_tokens, num_heads, nope_head_dim + rope_head_dim]
|
||||
nope_head_dim = k_nope.shape[-1]
|
||||
k[:, :, :nope_head_dim] = k_nope
|
||||
# Broadcast k_rope across all heads
|
||||
k[:, :, nope_head_dim:] = k_rope.expand(-1, k.shape[1], -1)
|
||||
|
||||
|
||||
def torch_concat_mla_absorb_q(
|
||||
a: torch.Tensor, b: torch.Tensor, out: torch.Tensor
|
||||
) -> None:
|
||||
"""Reference PyTorch implementation for concat_mla_absorb_q."""
|
||||
# a: [dim_0, dim_1, a_last_dim]
|
||||
# b: [dim_0, dim_1, b_last_dim]
|
||||
# out: [dim_0, dim_1, a_last_dim + b_last_dim]
|
||||
a_last_dim = a.shape[-1]
|
||||
out[:, :, :a_last_dim] = a
|
||||
out[:, :, a_last_dim:] = b
|
||||
|
||||
|
||||
def sgl_kernel_concat_mla_k(
|
||||
k: torch.Tensor, k_nope: torch.Tensor, k_rope: torch.Tensor
|
||||
) -> None:
|
||||
"""AOT compiled sgl_kernel implementation."""
|
||||
from sgl_kernel import concat_mla_k
|
||||
|
||||
concat_mla_k(k, k_nope, k_rope)
|
||||
|
||||
|
||||
def sgl_kernel_concat_mla_absorb_q(
|
||||
a: torch.Tensor, b: torch.Tensor, out: torch.Tensor
|
||||
) -> None:
|
||||
"""AOT compiled sgl_kernel implementation."""
|
||||
from sgl_kernel import concat_mla_absorb_q
|
||||
|
||||
result = concat_mla_absorb_q(a, b) # AOT returns output
|
||||
out.copy_(result) # Copy to provided tensor for comparison
|
||||
|
||||
|
||||
def jit_concat_mla_k(
|
||||
k: torch.Tensor, k_nope: torch.Tensor, k_rope: torch.Tensor
|
||||
) -> None:
|
||||
"""JIT compiled implementation."""
|
||||
from sglang.jit_kernel.concat_mla import concat_mla_k
|
||||
|
||||
concat_mla_k(k, k_nope, k_rope)
|
||||
|
||||
|
||||
def jit_concat_mla_absorb_q(
|
||||
a: torch.Tensor, b: torch.Tensor, out: torch.Tensor
|
||||
) -> None:
|
||||
"""JIT compiled implementation - wrapper for test compatibility."""
|
||||
from sglang.jit_kernel.concat_mla import concat_mla_absorb_q
|
||||
|
||||
result = concat_mla_absorb_q(a, b)
|
||||
out.copy_(result)
|
||||
|
||||
|
||||
# Constants matching the kernel
|
||||
NUM_LOCAL_HEADS = 128
|
||||
QK_NOPE_HEAD_DIM = 128
|
||||
QK_ROPE_HEAD_DIM = 64
|
||||
K_HEAD_DIM = QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM
|
||||
|
||||
A_LAST_DIM = 512
|
||||
B_LAST_DIM = 64
|
||||
OUT_LAST_DIM = A_LAST_DIM + B_LAST_DIM
|
||||
|
||||
DEVICE = "cuda"
|
||||
DTYPE = torch.bfloat16
|
||||
|
||||
# Test configurations
|
||||
NUM_TOKENS_LIST = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS_LIST)
|
||||
def test_concat_mla_k_jit_vs_torch(num_tokens: int) -> None:
|
||||
"""Test JIT kernel against PyTorch reference."""
|
||||
k_jit = torch.empty(
|
||||
num_tokens, NUM_LOCAL_HEADS, K_HEAD_DIM, device=DEVICE, dtype=DTYPE
|
||||
)
|
||||
k_torch = torch.empty(
|
||||
num_tokens, NUM_LOCAL_HEADS, K_HEAD_DIM, device=DEVICE, dtype=DTYPE
|
||||
)
|
||||
|
||||
k_nope = torch.randn(
|
||||
num_tokens, NUM_LOCAL_HEADS, QK_NOPE_HEAD_DIM, device=DEVICE, dtype=DTYPE
|
||||
)
|
||||
k_rope = torch.randn(num_tokens, 1, QK_ROPE_HEAD_DIM, device=DEVICE, dtype=DTYPE)
|
||||
|
||||
torch_concat_mla_k(k_torch, k_nope, k_rope)
|
||||
jit_concat_mla_k(k_jit, k_nope, k_rope)
|
||||
|
||||
triton.testing.assert_close(k_jit, k_torch, atol=0, rtol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS_LIST)
|
||||
def test_concat_mla_k_jit_vs_aot(num_tokens: int) -> None:
|
||||
"""Test JIT kernel against AOT kernel for bitwise equivalence."""
|
||||
k_jit = torch.empty(
|
||||
num_tokens, NUM_LOCAL_HEADS, K_HEAD_DIM, device=DEVICE, dtype=DTYPE
|
||||
)
|
||||
k_aot = torch.empty(
|
||||
num_tokens, NUM_LOCAL_HEADS, K_HEAD_DIM, device=DEVICE, dtype=DTYPE
|
||||
)
|
||||
|
||||
k_nope = torch.randn(
|
||||
num_tokens, NUM_LOCAL_HEADS, QK_NOPE_HEAD_DIM, device=DEVICE, dtype=DTYPE
|
||||
)
|
||||
k_rope = torch.randn(num_tokens, 1, QK_ROPE_HEAD_DIM, device=DEVICE, dtype=DTYPE)
|
||||
|
||||
sgl_kernel_concat_mla_k(k_aot, k_nope, k_rope)
|
||||
jit_concat_mla_k(k_jit, k_nope, k_rope)
|
||||
|
||||
triton.testing.assert_close(k_jit, k_aot, atol=0, rtol=0)
|
||||
|
||||
|
||||
DIM_0_LIST = [1, 2, 4, 8, 16, 32]
|
||||
DIM_1_LIST = [1, 2, 4, 8, 16, 128]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"dim_0,dim_1",
|
||||
list(itertools.product(DIM_0_LIST, DIM_1_LIST)),
|
||||
)
|
||||
def test_concat_mla_absorb_q_jit_vs_torch(dim_0: int, dim_1: int) -> None:
|
||||
"""Test JIT kernel against PyTorch reference."""
|
||||
a = torch.randn(dim_0, dim_1, A_LAST_DIM, device=DEVICE, dtype=DTYPE)
|
||||
b = torch.randn(dim_0, dim_1, B_LAST_DIM, device=DEVICE, dtype=DTYPE)
|
||||
out_jit = torch.empty(dim_0, dim_1, OUT_LAST_DIM, device=DEVICE, dtype=DTYPE)
|
||||
out_torch = torch.empty(dim_0, dim_1, OUT_LAST_DIM, device=DEVICE, dtype=DTYPE)
|
||||
|
||||
torch_concat_mla_absorb_q(a, b, out_torch)
|
||||
jit_concat_mla_absorb_q(a, b, out_jit)
|
||||
|
||||
triton.testing.assert_close(out_jit, out_torch, atol=0, rtol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"dim_0,dim_1",
|
||||
list(itertools.product(DIM_0_LIST, DIM_1_LIST)),
|
||||
)
|
||||
def test_concat_mla_absorb_q_jit_vs_aot(dim_0: int, dim_1: int) -> None:
|
||||
"""Test JIT kernel against AOT kernel for bitwise equivalence."""
|
||||
a = torch.randn(dim_0, dim_1, A_LAST_DIM, device=DEVICE, dtype=DTYPE)
|
||||
b = torch.randn(dim_0, dim_1, B_LAST_DIM, device=DEVICE, dtype=DTYPE)
|
||||
out_jit = torch.empty(dim_0, dim_1, OUT_LAST_DIM, device=DEVICE, dtype=DTYPE)
|
||||
out_aot = torch.empty(dim_0, dim_1, OUT_LAST_DIM, device=DEVICE, dtype=DTYPE)
|
||||
|
||||
sgl_kernel_concat_mla_absorb_q(a, b, out_aot)
|
||||
jit_concat_mla_absorb_q(a, b, out_jit)
|
||||
|
||||
triton.testing.assert_close(out_jit, out_aot, atol=0, rtol=0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,239 @@
|
||||
"""
|
||||
Correctness test for the JIT custom all-reduce (v2) kernel.
|
||||
|
||||
The test compares the JIT custom all-reduce output against NCCL all-reduce
|
||||
for various tensor sizes and dtypes, in both eager and CUDA-graph modes.
|
||||
|
||||
Usage:
|
||||
python -m pytest test_jit_custom_all_reduce.py -v
|
||||
|
||||
This file doubles as the torchrun worker script. The test class launches
|
||||
torchrun --nproc_per_node=N <this_file>
|
||||
and asserts that all worker processes exit successfully.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
import logging
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
import sglang.srt.distributed.parallel_state as ps
|
||||
from sglang.jit_kernel.all_reduce import (
|
||||
AllReduceAlgo,
|
||||
_jit_custom_all_reduce_pull_module,
|
||||
_jit_custom_all_reduce_push_module,
|
||||
)
|
||||
from sglang.jit_kernel.tests.utils import multiprocess_main, multiprocess_test
|
||||
from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
|
||||
CustomAllReduceV2,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=300,
|
||||
suite="base-b-kernel-unit-8-gpu-h200",
|
||||
)
|
||||
register_cuda_ci(
|
||||
est_time=300,
|
||||
suite="nightly-kernel-8-gpu-h200",
|
||||
nightly=True,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test parameters (shared between test class and worker)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
TEST_SIZES = [
|
||||
16,
|
||||
32,
|
||||
512,
|
||||
1024,
|
||||
1024 + 16, # weird case
|
||||
4 * 1024,
|
||||
32 * 1024,
|
||||
256 * 1024,
|
||||
2 * 1024 * 1024, # 2M elements
|
||||
4 * 1024 * 1024, # 4M elements
|
||||
]
|
||||
TEST_DTYPES = [torch.float16, torch.bfloat16, torch.float32]
|
||||
SHOTS = [
|
||||
AllReduceAlgo.ONE_SHOT_PULL,
|
||||
AllReduceAlgo.ONE_SHOT_PUSH,
|
||||
AllReduceAlgo.TWO_SHOT_PULL,
|
||||
]
|
||||
USE_GRAPH_OPTIONS = [True, False]
|
||||
TEST_CONFIG = itertools.product(TEST_SIZES, TEST_DTYPES, SHOTS, USE_GRAPH_OPTIONS)
|
||||
TEST_LAYERS = 4
|
||||
TEST_LOOP = 16
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test class (runs via pytest, launches torchrun subprocesses)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _compile_one(dtype: torch.dtype, world_size: int):
|
||||
_jit_custom_all_reduce_push_module(dtype, world_size)
|
||||
_jit_custom_all_reduce_pull_module(dtype, world_size)
|
||||
|
||||
|
||||
def _precompile_kernels() -> None:
|
||||
# NOTE: even when device count < 8, we should be able to compile all
|
||||
process_map: Dict[Tuple[torch.dtype, int], mp.Process] = {}
|
||||
COMPILE_SPACE = itertools.product(TEST_DTYPES, [2, 3, 4, 5, 6, 7, 8])
|
||||
mp.set_start_method("spawn")
|
||||
for config in COMPILE_SPACE:
|
||||
process_map[config] = mp.Process(target=_compile_one, args=config)
|
||||
for process in process_map.values():
|
||||
process.start()
|
||||
for (dtype, world_size), process in process_map.items():
|
||||
process.join()
|
||||
if process.exitcode != 0:
|
||||
raise RuntimeError(f"Custom All Reduce {world_size=} {dtype=} failed")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("nproc", [1, 2, 3, 4, 5, 6, 7, 8])
|
||||
def test_custom_allreduce(nproc: int) -> None:
|
||||
if nproc == 1: # NOTE: special case to speed up tests
|
||||
return _precompile_kernels()
|
||||
|
||||
device_count = torch.cuda.device_count()
|
||||
if device_count < nproc:
|
||||
pytest.skip(
|
||||
f"Requires at least {nproc} GPUs, but only {device_count} available"
|
||||
)
|
||||
multiprocess_test(__file__, nproc)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Worker logic (executed by each torchrun process)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def init_distributed():
|
||||
"""Initialize distributed groups via torchrun env vars.
|
||||
|
||||
Returns (rank, device, cpu_group, nccl_group, comm).
|
||||
"""
|
||||
local_rank = int(os.environ["LOCAL_RANK"])
|
||||
world_size = int(os.environ["WORLD_SIZE"])
|
||||
rank = local_rank
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.cuda.set_device(device)
|
||||
|
||||
dist.init_process_group(backend="gloo")
|
||||
ps._WORLD = coord = ps.init_world_group(
|
||||
ranks=list(range(world_size)),
|
||||
local_rank=local_rank,
|
||||
backend="nccl",
|
||||
)
|
||||
|
||||
cpu_group = coord.cpu_group
|
||||
nccl_group = coord.device_group
|
||||
assert nccl_group is not None
|
||||
|
||||
max_size = max(TEST_SIZES) * 4
|
||||
comm = CustomAllReduceV2(cpu_group, device, max_size, max_size)
|
||||
if comm.disabled:
|
||||
raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system")
|
||||
|
||||
return rank, device, cpu_group, nccl_group, comm
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def worker_test(
|
||||
device: torch.device,
|
||||
nccl_group: dist.ProcessGroup,
|
||||
comm: CustomAllReduceV2,
|
||||
size: int,
|
||||
dtype: torch.dtype,
|
||||
use_graph: bool,
|
||||
algo: AllReduceAlgo,
|
||||
) -> Optional[RuntimeError]:
|
||||
comm.override_algo = algo
|
||||
|
||||
def get_run_graph_fn():
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
graph_inp = torch.zeros((TEST_LAYERS, size), dtype=dtype, device=device)
|
||||
out_jits = []
|
||||
with comm.capture():
|
||||
with torch.cuda.graph(graph):
|
||||
for i in range(TEST_LAYERS):
|
||||
out_jits.append(comm.custom_all_reduce(graph_inp[i]))
|
||||
out_jit = torch.stack(out_jits)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
def run_graph(x: torch.Tensor) -> torch.Tensor:
|
||||
graph_inp.copy_(x)
|
||||
graph.replay()
|
||||
return out_jit.clone()
|
||||
|
||||
return run_graph
|
||||
|
||||
def get_run_eager_fn():
|
||||
def run_eager(x: torch.Tensor) -> torch.Tensor:
|
||||
eager_inp = x.clone()
|
||||
out_eagers = []
|
||||
for i in range(TEST_LAYERS):
|
||||
out_eagers.append(comm.custom_all_reduce(eager_inp[i]))
|
||||
torch.cuda.synchronize()
|
||||
return torch.stack(out_eagers)
|
||||
|
||||
return run_eager
|
||||
|
||||
run_fn = get_run_graph_fn() if use_graph else get_run_eager_fn()
|
||||
num_errors = 0
|
||||
for _ in range(TEST_LOOP):
|
||||
# NOTE: 15 * 8 < 128, which is the precision limit for bf16
|
||||
inp = torch.randint(0, 16, (TEST_LAYERS, size), dtype=dtype, device=device)
|
||||
assert comm.should_custom_ar(inp[0])
|
||||
out_ref = inp.clone()
|
||||
dist.all_reduce(out_ref, group=nccl_group)
|
||||
out_jit = run_fn(inp)
|
||||
num_errors += not torch.all(out_jit == out_ref)
|
||||
if num_errors > 0:
|
||||
return RuntimeError(
|
||||
f"Test failed for {size=}, {dtype=}, {algo=}, "
|
||||
f"{use_graph=} with {num_errors} errors. "
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def worker_main() -> None:
|
||||
"""Entry point for each torchrun worker process."""
|
||||
rank, device, cpu_group, nccl_group, comm = init_distributed()
|
||||
|
||||
torch.cuda.set_stream(torch.cuda.Stream())
|
||||
|
||||
logging.disable(logging.INFO) # Suppress internal logging for cleaner test output
|
||||
items = list(enumerate(TEST_CONFIG))
|
||||
for i, (size, dtype, algo, use_graph) in items:
|
||||
error = worker_test(device, nccl_group, comm, size, dtype, use_graph, algo)
|
||||
if error is not None:
|
||||
print(
|
||||
f"Worker {rank} failed for {size=}, {dtype=}, "
|
||||
f"{algo=}, {use_graph=}, iteration={i}\n"
|
||||
f"Error: {error}"
|
||||
)
|
||||
# communicate the result to rank 0 for logging
|
||||
result = torch.tensor([int(error is not None)])
|
||||
dist.all_reduce(result, group=cpu_group)
|
||||
failed = bool(result.item())
|
||||
if failed:
|
||||
raise RuntimeError(
|
||||
f"Test failed on rank {rank} for config: "
|
||||
f"{size=}, {dtype=}, {algo=}, {use_graph=}"
|
||||
)
|
||||
|
||||
comm.close()
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
multiprocess_main(__file__, worker_main)
|
||||
@@ -0,0 +1,312 @@
|
||||
"""Tests for CuTe DSL fused sigmoid gating delta rule kernel (GDN)."""
|
||||
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
try:
|
||||
import cuda.bindings.driver as cuda_driver
|
||||
import cutlass # noqa: F401
|
||||
from cutlass.cute.runtime import from_dlpack
|
||||
|
||||
from sglang.jit_kernel import cutedsl_gdn
|
||||
|
||||
CUTEDSL_AVAILABLE = True
|
||||
except ImportError:
|
||||
CUTEDSL_AVAILABLE = False
|
||||
cutedsl_gdn = None
|
||||
|
||||
try:
|
||||
from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
|
||||
fused_sigmoid_gating_delta_rule_update,
|
||||
)
|
||||
|
||||
TRITON_AVAILABLE = True
|
||||
except ImportError:
|
||||
TRITON_AVAILABLE = False
|
||||
|
||||
register_cuda_ci(est_time=5, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
def run_triton_kernel(A_log, dt_bias, q, k, v, a, b, initial_state, indices, scale):
|
||||
return fused_sigmoid_gating_delta_rule_update(
|
||||
A_log=A_log,
|
||||
a=a,
|
||||
dt_bias=dt_bias,
|
||||
softplus_beta=1.0,
|
||||
softplus_threshold=20.0,
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
b=b,
|
||||
initial_state_source=initial_state,
|
||||
initial_state_indices=indices,
|
||||
scale=scale,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
cu_seqlens=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not CUTEDSL_AVAILABLE, reason="CuTe DSL not available")
|
||||
@pytest.mark.skipif(not TRITON_AVAILABLE, reason="Triton kernel not available")
|
||||
@pytest.mark.skip(
|
||||
reason=(
|
||||
"Temporary CI workaround: CuTe DSL GDN precision is currently unstable "
|
||||
"against the Triton reference and needs follow-up investigation."
|
||||
)
|
||||
)
|
||||
@pytest.mark.parametrize("B", [16, 128])
|
||||
def test_cutedsl_gdn_precision(B: int):
|
||||
"""Test precision of CuTe DSL GDN kernel against Triton reference."""
|
||||
torch.manual_seed(2025)
|
||||
T, H, K, V, HV = 1, 16, 128, 128, 32
|
||||
scale = K**-0.5
|
||||
|
||||
A_log = torch.randn(HV, dtype=torch.float32, device="cuda")
|
||||
dt_bias = torch.randn(HV, dtype=torch.bfloat16, device="cuda")
|
||||
a = torch.randn(B, T, HV, dtype=torch.bfloat16, device="cuda")
|
||||
b = torch.randn(B, T, HV, dtype=torch.bfloat16, device="cuda")
|
||||
q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device="cuda")
|
||||
k = torch.randn(B, T, H, K, dtype=torch.bfloat16, device="cuda")
|
||||
v = torch.randn(B, T, HV, V, dtype=torch.bfloat16, device="cuda")
|
||||
indices = torch.arange(B, dtype=torch.int32, device="cuda")
|
||||
state_cutedsl = torch.randn(B, HV, K, V, dtype=torch.float32, device="cuda")
|
||||
state_triton = state_cutedsl.clone().reshape(-1).contiguous()
|
||||
|
||||
# Warmup compilation
|
||||
_ = cutedsl_gdn.cutedsl_fused_sigmoid_gating_delta_rule_update(
|
||||
A_log, dt_bias, q, k, v, a, b, state_cutedsl.clone(), indices, scale=scale
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Fresh state for actual test
|
||||
state_cutedsl = torch.randn(B, HV, K, V, dtype=torch.float32, device="cuda")
|
||||
state_triton = state_cutedsl.clone().reshape(-1).contiguous()
|
||||
|
||||
out_cutedsl = cutedsl_gdn.cutedsl_fused_sigmoid_gating_delta_rule_update(
|
||||
A_log, dt_bias, q, k, v, a, b, state_cutedsl, indices, scale=scale
|
||||
)
|
||||
out_triton = run_triton_kernel(
|
||||
A_log, dt_bias, q, k, v, a, b, state_triton, indices, scale
|
||||
)
|
||||
|
||||
# Check precision: diff > 0.1 must be < 1% of elements
|
||||
abs_diff = (out_triton.float() - out_cutedsl.float()).abs()
|
||||
max_diff = abs_diff.max().item()
|
||||
mean_diff = abs_diff.mean().item()
|
||||
fail_rate = (abs_diff > 0.1).float().mean().item() * 100
|
||||
has_nan = torch.isnan(out_cutedsl).any() or torch.isinf(out_cutedsl).any()
|
||||
|
||||
kernel_type = "SmallBatch" if B < 32 else "LargeBatch"
|
||||
print(
|
||||
f"\n B={B} ({kernel_type}): max_diff={max_diff:.2e}, mean_diff={mean_diff:.2e}, fail_rate={fail_rate:.2f}%"
|
||||
)
|
||||
|
||||
assert not has_nan, "Output contains NaN/Inf"
|
||||
assert fail_rate < 1.0, f"Fail rate {fail_rate:.2f}% >= 1%"
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
True,
|
||||
reason="Skip the performance test because the speedup ratio is highly unstable in the CI environment. ",
|
||||
)
|
||||
@pytest.mark.skipif(not CUTEDSL_AVAILABLE, reason="CuTe DSL not available")
|
||||
@pytest.mark.skipif(not TRITON_AVAILABLE, reason="Triton kernel not available")
|
||||
@pytest.mark.parametrize("B", [1, 128])
|
||||
def test_cutedsl_gdn_performance(B: int):
|
||||
"""Benchmark CuTe DSL GDN kernel against Triton reference."""
|
||||
torch.manual_seed(2025)
|
||||
T, H, K, V, HV = 1, 16, 128, 128, 32
|
||||
N = B
|
||||
scale = K**-0.5
|
||||
is_varlen = True
|
||||
warmup, bench_iters, run_iters = 10, 100, 10
|
||||
|
||||
A_log = torch.randn(HV, dtype=torch.float32, device="cuda")
|
||||
dt_bias = torch.randn(HV, dtype=torch.bfloat16, device="cuda")
|
||||
indices = torch.arange(N, dtype=torch.int32, device="cuda")
|
||||
state_cutedsl = torch.randn(N, HV, K, V, dtype=torch.float32, device="cuda")
|
||||
state_triton = state_cutedsl.reshape(-1).contiguous()
|
||||
cu_seqlens = torch.zeros(N + 1, dtype=torch.int32, device="cuda")
|
||||
o_cutedsl = torch.zeros(1, N, HV, V, dtype=torch.bfloat16, device="cuda")
|
||||
|
||||
# Prepare tensors for multiple runs
|
||||
q_list, k_list, v_list, a_list, b_list = [], [], [], [], []
|
||||
q_tensor_list, k_tensor_list, v_tensor_list, a_tensor_list, b_tensor_list = (
|
||||
[],
|
||||
[],
|
||||
[],
|
||||
[],
|
||||
[],
|
||||
)
|
||||
q_triton, k_triton, v_triton, a_triton, b_triton = [], [], [], [], []
|
||||
|
||||
for ri in range(run_iters):
|
||||
torch.manual_seed(2025 + ri)
|
||||
q_i = torch.randn(1, N, H, K, dtype=torch.bfloat16, device="cuda")
|
||||
k_i = torch.randn(1, N, H, K, dtype=torch.bfloat16, device="cuda")
|
||||
v_i = torch.randn(1, N, HV, V, dtype=torch.bfloat16, device="cuda")
|
||||
a_i = torch.randn(N, HV, dtype=torch.bfloat16, device="cuda")
|
||||
b_i = torch.randn(N, HV, dtype=torch.bfloat16, device="cuda")
|
||||
|
||||
q_list.append(q_i)
|
||||
k_list.append(k_i)
|
||||
v_list.append(v_i)
|
||||
a_list.append(a_i)
|
||||
b_list.append(b_i)
|
||||
q_tensor_list.append(from_dlpack(q_i, assumed_align=16))
|
||||
k_tensor_list.append(from_dlpack(k_i, assumed_align=16))
|
||||
v_tensor_list.append(from_dlpack(v_i, assumed_align=16))
|
||||
a_tensor_list.append(from_dlpack(a_i, assumed_align=16))
|
||||
b_tensor_list.append(from_dlpack(b_i, assumed_align=16))
|
||||
q_triton.append(q_i.transpose(0, 1).contiguous())
|
||||
k_triton.append(k_i.transpose(0, 1).contiguous())
|
||||
v_triton.append(v_i.transpose(0, 1).contiguous())
|
||||
a_triton.append(a_i.unsqueeze(1).contiguous())
|
||||
b_triton.append(b_i.unsqueeze(1).contiguous())
|
||||
|
||||
A_log_t = from_dlpack(A_log, assumed_align=16)
|
||||
dt_bias_t = from_dlpack(dt_bias, assumed_align=16)
|
||||
h0_t = from_dlpack(state_cutedsl, assumed_align=16)
|
||||
idx_t = from_dlpack(indices, assumed_align=16)
|
||||
o_t = from_dlpack(o_cutedsl, assumed_align=16)
|
||||
cu_t = from_dlpack(cu_seqlens, assumed_align=16)
|
||||
|
||||
torch_stream = torch.cuda.Stream()
|
||||
stream = cuda_driver.CUstream(torch_stream.cuda_stream)
|
||||
|
||||
# Compile kernels
|
||||
compiled = cutedsl_gdn._get_compiled_kernel(N, H, HV, K, V, N, N < 32, is_varlen)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
for ri in range(run_iters):
|
||||
_ = run_triton_kernel(
|
||||
A_log,
|
||||
dt_bias,
|
||||
q_triton[ri],
|
||||
k_triton[ri],
|
||||
v_triton[ri],
|
||||
a_triton[ri],
|
||||
b_triton[ri],
|
||||
state_triton,
|
||||
indices,
|
||||
scale,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
def run_cutedsl():
|
||||
for ri in range(run_iters):
|
||||
compiled(
|
||||
cu_t,
|
||||
q_tensor_list[ri],
|
||||
k_tensor_list[ri],
|
||||
v_tensor_list[ri],
|
||||
a_tensor_list[ri],
|
||||
b_tensor_list[ri],
|
||||
A_log_t,
|
||||
dt_bias_t,
|
||||
h0_t,
|
||||
idx_t,
|
||||
o_t,
|
||||
stream,
|
||||
)
|
||||
|
||||
def run_triton():
|
||||
for ri in range(run_iters):
|
||||
_ = run_triton_kernel(
|
||||
A_log,
|
||||
dt_bias,
|
||||
q_triton[ri],
|
||||
k_triton[ri],
|
||||
v_triton[ri],
|
||||
a_triton[ri],
|
||||
b_triton[ri],
|
||||
state_triton,
|
||||
indices,
|
||||
scale,
|
||||
)
|
||||
|
||||
# Warmup
|
||||
with torch.cuda.stream(torch_stream):
|
||||
run_cutedsl()
|
||||
torch.cuda.synchronize()
|
||||
run_triton()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Capture CUDA graphs
|
||||
graph_triton = torch.cuda.CUDAGraph()
|
||||
graph_cutedsl = torch.cuda.CUDAGraph()
|
||||
try:
|
||||
with torch.cuda.graph(graph_triton):
|
||||
run_triton()
|
||||
with torch.cuda.graph(graph_cutedsl, stream=torch_stream):
|
||||
run_cutedsl()
|
||||
torch.cuda.synchronize()
|
||||
except Exception:
|
||||
graph_triton = graph_cutedsl = None
|
||||
|
||||
# Warmup with graphs
|
||||
for _ in range(warmup):
|
||||
if graph_cutedsl:
|
||||
graph_cutedsl.replay()
|
||||
else:
|
||||
with torch.cuda.stream(torch_stream):
|
||||
run_cutedsl()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
if graph_triton:
|
||||
graph_triton.replay()
|
||||
else:
|
||||
run_triton()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Benchmark
|
||||
triton_times, cutedsl_times = [], []
|
||||
for _ in range(bench_iters):
|
||||
start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(
|
||||
enable_timing=True
|
||||
)
|
||||
start.record()
|
||||
if graph_triton:
|
||||
graph_triton.replay()
|
||||
else:
|
||||
run_triton()
|
||||
end.record()
|
||||
torch.cuda.synchronize()
|
||||
triton_times.append(start.elapsed_time(end))
|
||||
|
||||
start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(
|
||||
enable_timing=True
|
||||
)
|
||||
with torch.cuda.stream(torch_stream):
|
||||
start.record()
|
||||
if graph_cutedsl:
|
||||
graph_cutedsl.replay()
|
||||
else:
|
||||
run_cutedsl()
|
||||
end.record()
|
||||
torch.cuda.synchronize()
|
||||
cutedsl_times.append(start.elapsed_time(end))
|
||||
|
||||
triton_mean = np.mean(triton_times) / run_iters * 1000
|
||||
triton_std = np.std(triton_times) / run_iters * 1000
|
||||
cutedsl_mean = np.mean(cutedsl_times) / run_iters * 1000
|
||||
cutedsl_std = np.std(cutedsl_times) / run_iters * 1000
|
||||
speedup = triton_mean / cutedsl_mean
|
||||
|
||||
kernel_type = "SmallBatch" if B < 32 else "LargeBatch"
|
||||
print(
|
||||
f"\n B={B} ({kernel_type}): Triton={triton_mean:.2f}±{triton_std:.2f}μs, CuTeDSL={cutedsl_mean:.2f}±{cutedsl_std:.2f}μs, speedup={speedup:.2f}x"
|
||||
)
|
||||
|
||||
min_speedup = 1.0 if B < 32 else 1.15
|
||||
assert speedup >= min_speedup, f"Speedup {speedup:.2f}x < {min_speedup}x for B={B}"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,96 @@
|
||||
import itertools
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.utils import get_ci_test_range
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
def sglang_jit_fused_add_rmsnorm(
|
||||
input: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
eps: float,
|
||||
*,
|
||||
cast_x_before_out_mul: bool = False,
|
||||
) -> None:
|
||||
from sglang.jit_kernel.norm import fused_add_rmsnorm
|
||||
|
||||
fused_add_rmsnorm(
|
||||
input, residual, weight, eps, cast_x_before_out_mul=cast_x_before_out_mul
|
||||
)
|
||||
|
||||
|
||||
def flashinfer_fused_add_rmsnorm(
|
||||
input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float
|
||||
) -> None:
|
||||
from flashinfer.norm import fused_add_rmsnorm
|
||||
|
||||
fused_add_rmsnorm(input, residual, weight, eps=eps)
|
||||
|
||||
|
||||
def forward_native_hf_reference(
|
||||
x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor, eps: float
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
sum_fp32 = x.to(torch.float32) + residual.to(torch.float32)
|
||||
residual_out = sum_fp32.to(x.dtype)
|
||||
variance = sum_fp32.pow(2).mean(-1, keepdim=True)
|
||||
out = w * (sum_fp32 * torch.rsqrt(variance + eps)).to(x.dtype)
|
||||
return out, residual_out
|
||||
|
||||
|
||||
BS_LIST = [2**n for n in range(0, 14)]
|
||||
BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)]
|
||||
BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 256, 4109])
|
||||
HIDDEN_SIZE_LIST = get_ci_test_range(
|
||||
[512, 1024, 1536, 2048, 3072, 4096, 5120, 6144, 7168, 8192],
|
||||
[512, 2048, 8192],
|
||||
)
|
||||
DEVICE = "cuda"
|
||||
DTYPE = torch.bfloat16
|
||||
EPS = torch.finfo(torch.bfloat16).eps
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_size,hidden_size,cast_x_before_out_mul",
|
||||
list(itertools.product(BS_LIST, HIDDEN_SIZE_LIST, [False, True])),
|
||||
)
|
||||
def test_fused_add_rmsnorm(
|
||||
batch_size: int, hidden_size: int, cast_x_before_out_mul: bool
|
||||
) -> None:
|
||||
torch.manual_seed(0)
|
||||
input = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=DTYPE)
|
||||
residual = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=DTYPE)
|
||||
weight = torch.randn(hidden_size, device=DEVICE, dtype=DTYPE)
|
||||
|
||||
input_sglang = input.clone()
|
||||
residual_sglang = residual.clone()
|
||||
sglang_jit_fused_add_rmsnorm(
|
||||
input_sglang,
|
||||
residual_sglang,
|
||||
weight,
|
||||
EPS,
|
||||
cast_x_before_out_mul=cast_x_before_out_mul,
|
||||
)
|
||||
|
||||
if cast_x_before_out_mul:
|
||||
out_ref, residual_ref = forward_native_hf_reference(
|
||||
input, residual, weight, EPS
|
||||
)
|
||||
else:
|
||||
input_ref = input.clone()
|
||||
residual_ref_buf = residual.clone()
|
||||
flashinfer_fused_add_rmsnorm(input_ref, residual_ref_buf, weight, EPS)
|
||||
out_ref, residual_ref = input_ref, residual_ref_buf
|
||||
|
||||
torch.testing.assert_close(input_sglang, out_ref, atol=1e-2, rtol=1e-2)
|
||||
torch.testing.assert_close(residual_sglang, residual_ref, atol=1e-2, rtol=1e-2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,459 @@
|
||||
"""
|
||||
Test for fused_store_index_k_cache kernel.
|
||||
|
||||
Design Notes:
|
||||
1. torch.cuda.synchronize() needed after TVM FFI kernel call.
|
||||
2. _split_buffer used buf[:, :vb].reshape(-1) which COPIES data for
|
||||
non-contiguous slices → reference buffer stayed all-zeros.
|
||||
Fix: use flat byte-offset indexing.
|
||||
3. act_quant may use a different quantization scheme → generous tolerance.
|
||||
4. FP8 E4M3 1-ULP rounding differences between CUDA hardware cast
|
||||
(__nv_fp8_e4m3) and PyTorch .to(float8_e4m3fn) at tie-break points.
|
||||
Adjacent FP8 representable values at the high end differ by up to 32
|
||||
in float space (e.g. 288, 320, 352, ..., 448).
|
||||
Need to compare dequantized values with FP8-appropriate tolerance.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
try:
|
||||
from sglang.jit_kernel.fused_store_index_cache import (
|
||||
can_use_dsa_fused_store,
|
||||
fused_store_index_k_cache,
|
||||
)
|
||||
|
||||
HAS_FUSED = True
|
||||
except ImportError:
|
||||
HAS_FUSED = False
|
||||
|
||||
try:
|
||||
from sglang.srt.utils import is_hip
|
||||
|
||||
_is_hip = is_hip()
|
||||
except ImportError:
|
||||
_is_hip = False
|
||||
|
||||
try:
|
||||
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
|
||||
|
||||
_is_fp8_fnuz = is_fp8_fnuz()
|
||||
except ImportError:
|
||||
_is_fp8_fnuz = False
|
||||
|
||||
register_cuda_ci(est_time=24, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
PAGE_SIZE = 64
|
||||
HEAD_DIM = 128
|
||||
FP8_E4M3_MAX = 448.0
|
||||
FP8_DTYPE = torch.float8_e4m3fn
|
||||
BYTES_PER_TOKEN = 128 + 4 # 128 fp8 bytes + 4 scale bytes
|
||||
BYTES_PER_PAGE = PAGE_SIZE * BYTES_PER_TOKEN
|
||||
|
||||
|
||||
def _skip_if_unavailable(page_size: int = PAGE_SIZE):
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA required")
|
||||
if _is_hip:
|
||||
pytest.skip("Fused store kernel is CUDA-specific")
|
||||
if _is_fp8_fnuz:
|
||||
pytest.skip("Fused store path disabled for FP8 FNUZ")
|
||||
if not hasattr(torch, "float8_e4m3fn"):
|
||||
pytest.skip("torch.float8_e4m3fn not available")
|
||||
if not HAS_FUSED:
|
||||
pytest.skip("fused_store_index_cache not importable")
|
||||
if not can_use_dsa_fused_store(torch.bfloat16, torch.int64, page_size):
|
||||
pytest.skip("JIT kernel unavailable / failed to compile")
|
||||
|
||||
|
||||
def _num_pages(loc: torch.Tensor, page_size: int, extra: int = 1) -> int:
|
||||
return int(loc.max().item()) // page_size + 1 + extra
|
||||
|
||||
|
||||
def _make_buffer(num_pages: int, page_size: int = PAGE_SIZE) -> torch.Tensor:
|
||||
return torch.zeros(
|
||||
(num_pages, page_size * BYTES_PER_TOKEN),
|
||||
dtype=torch.uint8,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
|
||||
def _read_token_from_buffer(
|
||||
buf: torch.Tensor,
|
||||
token_idx: int,
|
||||
page_size: int = PAGE_SIZE,
|
||||
) -> Tuple[torch.Tensor, float]:
|
||||
"""
|
||||
Read a single token's fp8 values and scale from the paged buffer
|
||||
using flat byte offsets.
|
||||
"""
|
||||
page = token_idx // page_size
|
||||
offset = token_idx % page_size
|
||||
page_bytes = page_size * BYTES_PER_TOKEN
|
||||
|
||||
buf_flat = buf.reshape(-1)
|
||||
|
||||
val_start = page * page_bytes + offset * 128
|
||||
fp8_bytes = buf_flat[val_start : val_start + 128]
|
||||
fp8_vals = fp8_bytes.view(FP8_DTYPE).float()
|
||||
|
||||
scale_start = page * page_bytes + 128 * page_size + offset * 4
|
||||
scale_bytes = buf_flat[scale_start : scale_start + 4]
|
||||
scale = scale_bytes.view(torch.float32).item()
|
||||
|
||||
return fp8_vals, scale
|
||||
|
||||
|
||||
def _write_token_to_buffer(
|
||||
buf: torch.Tensor,
|
||||
token_idx: int,
|
||||
fp8_data: torch.Tensor,
|
||||
scale: float,
|
||||
page_size: int = PAGE_SIZE,
|
||||
) -> None:
|
||||
"""
|
||||
Write a single token's fp8 values and scale into the paged buffer
|
||||
using flat byte offsets on buf.reshape(-1) (which is a true view
|
||||
since buf is contiguous).
|
||||
"""
|
||||
page = token_idx // page_size
|
||||
offset = token_idx % page_size
|
||||
page_bytes = page_size * BYTES_PER_TOKEN
|
||||
|
||||
buf_flat = buf.reshape(-1)
|
||||
|
||||
val_start = page * page_bytes + offset * 128
|
||||
buf_flat[val_start : val_start + 128] = fp8_data.view(torch.uint8)
|
||||
|
||||
scale_start = page * page_bytes + 128 * page_size + offset * 4
|
||||
scale_t = torch.tensor([scale], dtype=torch.float32, device=buf.device)
|
||||
buf_flat[scale_start : scale_start + 4] = scale_t.view(torch.uint8)
|
||||
|
||||
|
||||
def _gather_tokens(
|
||||
buf: torch.Tensor,
|
||||
loc: torch.Tensor,
|
||||
page_size: int = PAGE_SIZE,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
N = loc.shape[0]
|
||||
fp8_f32 = torch.empty((N, HEAD_DIM), dtype=torch.float32, device=buf.device)
|
||||
scales = torch.empty((N,), dtype=torch.float32, device=buf.device)
|
||||
for i in range(N):
|
||||
idx = int(loc[i].item())
|
||||
vals, s = _read_token_from_buffer(buf, idx, page_size)
|
||||
fp8_f32[i] = vals
|
||||
scales[i] = s
|
||||
return fp8_f32, scales
|
||||
|
||||
|
||||
# Reference kernel
|
||||
def _reference_quantize_and_store(
|
||||
key_bf16: torch.Tensor,
|
||||
loc: torch.Tensor,
|
||||
num_pages: int,
|
||||
page_size: int = PAGE_SIZE,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Reference kernel of the fused kernel's quantization:
|
||||
abs_max = max(|row|)
|
||||
scale = max(1e-4, abs_max) / 448
|
||||
fp8_val = clip(val / scale, -448, 448) -> cast to fp8
|
||||
"""
|
||||
N = key_bf16.shape[0]
|
||||
key_f32 = key_bf16.float()
|
||||
buf = _make_buffer(num_pages, page_size)
|
||||
|
||||
for i in range(N):
|
||||
row = key_f32[i]
|
||||
abs_max = row.abs().max().item()
|
||||
scale = max(1e-4, abs_max) / FP8_E4M3_MAX
|
||||
inv_scale = 1.0 / scale
|
||||
quantized = (row * inv_scale).clamp(-FP8_E4M3_MAX, FP8_E4M3_MAX)
|
||||
quantized_fp8 = quantized.to(FP8_DTYPE)
|
||||
|
||||
idx = int(loc[i].item())
|
||||
_write_token_to_buffer(buf, idx, quantized_fp8, scale, page_size)
|
||||
|
||||
return buf
|
||||
|
||||
|
||||
def _import_act_quant():
|
||||
try:
|
||||
from sglang.srt.layers.attention.dsa.triton_kernel import act_quant
|
||||
|
||||
return act_quant
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _ref_store_via_act_quant(
|
||||
key_bf16: torch.Tensor,
|
||||
loc: torch.Tensor,
|
||||
num_pages: int,
|
||||
page_size: int = PAGE_SIZE,
|
||||
block_size: int = 128,
|
||||
scale_fmt: Optional[str] = None,
|
||||
) -> Optional[torch.Tensor]:
|
||||
act_quant = _import_act_quant()
|
||||
if act_quant is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
k_fp8, k_scale = act_quant(key_bf16, block_size, scale_fmt)
|
||||
except TypeError:
|
||||
k_fp8, k_scale = act_quant(key_bf16, block_size)
|
||||
|
||||
if k_fp8.dim() == 3 and k_fp8.shape[1] == 1:
|
||||
k_fp8 = k_fp8.squeeze(1)
|
||||
if k_scale is not None and k_scale.dim() == 3 and k_scale.shape[1] == 1:
|
||||
k_scale = k_scale.squeeze(1)
|
||||
k_scale = k_scale.view(-1).float()
|
||||
|
||||
buf = _make_buffer(num_pages, page_size)
|
||||
N = key_bf16.shape[0]
|
||||
for i in range(N):
|
||||
idx = int(loc[i].item())
|
||||
_write_token_to_buffer(
|
||||
buf, idx, k_fp8[i].to(FP8_DTYPE), k_scale[i].item(), page_size
|
||||
)
|
||||
return buf
|
||||
|
||||
|
||||
# TEST 1: Fused kernel vs. its own algorithm (pure-Python reference)
|
||||
#
|
||||
# NOTE on FP8 rounding:
|
||||
# CUDA hardware fp8 cast (__nv_fp8_e4m3) and PyTorch .to(float8_e4m3fn)
|
||||
# may round differently at tie-break points. This causes up to 1-ULP
|
||||
# differences in the FP8 codes. In FP8 E4M3, adjacent representable
|
||||
# values at the high end differ by up to 32 in float space (e.g.
|
||||
# 288 vs 320). After dequantization (fp8_float * scale), the error
|
||||
# from 1-ULP is: scale * ulp ≈ (abs_max/448) * 32 ≈ 0.07 * abs_max.
|
||||
# For randn inputs (abs_max ≈ 3-4), this is about 0.2-0.3.
|
||||
#
|
||||
# We therefore compare dequantized values with tolerances that
|
||||
# accommodate 1-ULP FP8 rounding, NOT byte-exact fp8 codes.
|
||||
@pytest.mark.parametrize(
|
||||
"num_tokens,base_index",
|
||||
[(1, 0), (32, 0), (64, 0), (128, 64), (257, 65), (512, 0)],
|
||||
)
|
||||
def test_fused_kernel_matches_own_algorithm(num_tokens: int, base_index: int):
|
||||
"""Compare fused CUDA kernel against a pure-Python implementation
|
||||
of the *same* quantization formula."""
|
||||
_skip_if_unavailable()
|
||||
device = torch.device("cuda")
|
||||
|
||||
key = torch.randn((num_tokens, HEAD_DIM), device=device, dtype=torch.bfloat16)
|
||||
loc = (
|
||||
base_index + torch.randperm(num_tokens, device=device, dtype=torch.int64)
|
||||
).contiguous()
|
||||
num_pages = _num_pages(loc, PAGE_SIZE)
|
||||
|
||||
# Reference kernel
|
||||
ref_buf = _reference_quantize_and_store(key, loc, num_pages)
|
||||
|
||||
# Fused kernel
|
||||
out_buf = _make_buffer(num_pages)
|
||||
fused_store_index_k_cache(key, out_buf, loc, page_size=PAGE_SIZE)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
out_f, out_s = _gather_tokens(out_buf, loc)
|
||||
ref_f, ref_s = _gather_tokens(ref_buf, loc)
|
||||
|
||||
# 1) Scales must match tightly (same f32 formula, no rounding ambiguity)
|
||||
torch.testing.assert_close(out_s, ref_s, rtol=1e-5, atol=1e-7)
|
||||
|
||||
# 2) Most FP8 codes should match; allow rare 1-ULP differences.
|
||||
# 1-ULP at FP8 E4M3 high end = 32 in float space.
|
||||
mismatch = out_f != ref_f
|
||||
mismatch_frac = mismatch.float().mean().item()
|
||||
assert mismatch_frac < 0.01, (
|
||||
f"Too many FP8 code mismatches: {mismatch_frac:.2%} "
|
||||
f"(expected < 1% from rounding tie-breaks)"
|
||||
)
|
||||
|
||||
# 3) Where codes differ, the difference should be exactly 1 ULP.
|
||||
# In FP8 E4M3: if the float-cast value is V, the adjacent value
|
||||
# differs by ~V * 0.1 (relative) at most.
|
||||
if mismatch.any():
|
||||
diff = (out_f[mismatch] - ref_f[mismatch]).abs()
|
||||
rel_diff = diff / ref_f[mismatch].abs().clamp(min=1e-6)
|
||||
# 1-ULP relative difference for E4M3 is at most ~12.5% (2^-3)
|
||||
assert rel_diff.max().item() <= 0.15, (
|
||||
f"FP8 code difference exceeds 1-ULP: max relative diff = "
|
||||
f"{rel_diff.max().item():.4f}"
|
||||
)
|
||||
|
||||
# 4) Dequantized values should be close.
|
||||
# Max error from 1-ULP: scale * fp8_ulp ≈ (abs_max/448) * 32
|
||||
# For randn abs_max ≈ 3-4: max_err ≈ 0.21 - 0.29
|
||||
out_deq = out_f * out_s.unsqueeze(-1)
|
||||
ref_deq = ref_f * ref_s.unsqueeze(-1)
|
||||
torch.testing.assert_close(out_deq, ref_deq, rtol=0.15, atol=0.5)
|
||||
|
||||
|
||||
# TEST 2: Cross-check against act_quant
|
||||
@pytest.mark.parametrize("scale_fmt", [None, "fp32"])
|
||||
def test_fused_kernel_vs_act_quant_semantic(scale_fmt: Optional[str]):
|
||||
"""Both fused kernel and act_quant should approximately reconstruct
|
||||
the original bf16 values."""
|
||||
_skip_if_unavailable()
|
||||
device = torch.device("cuda")
|
||||
|
||||
num_tokens = 257
|
||||
base_index = 65
|
||||
key = torch.randn((num_tokens, HEAD_DIM), device=device, dtype=torch.bfloat16)
|
||||
loc = (
|
||||
base_index + torch.randperm(num_tokens, device=device, dtype=torch.int64)
|
||||
).contiguous()
|
||||
num_pages = _num_pages(loc, PAGE_SIZE)
|
||||
|
||||
ref_buf = _ref_store_via_act_quant(key, loc, num_pages, scale_fmt=scale_fmt)
|
||||
if ref_buf is None:
|
||||
pytest.skip("act_quant not available")
|
||||
|
||||
out_buf = _make_buffer(num_pages)
|
||||
fused_store_index_k_cache(key, out_buf, loc, page_size=PAGE_SIZE)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
out_f, out_s = _gather_tokens(out_buf, loc)
|
||||
ref_f, ref_s = _gather_tokens(ref_buf, loc)
|
||||
|
||||
out_deq = out_f * out_s.unsqueeze(-1)
|
||||
ref_deq = ref_f * ref_s.unsqueeze(-1)
|
||||
orig_f32 = key.float()
|
||||
|
||||
# Fused kernel should reconstruct original within FP8 precision
|
||||
torch.testing.assert_close(
|
||||
out_deq,
|
||||
orig_f32,
|
||||
rtol=0.15,
|
||||
atol=5e-2,
|
||||
msg="Fused kernel dequantized values don't approximate original",
|
||||
)
|
||||
|
||||
# act_quant may use a very different scale policy.
|
||||
try:
|
||||
torch.testing.assert_close(
|
||||
ref_deq,
|
||||
orig_f32,
|
||||
rtol=0.25,
|
||||
atol=0.5,
|
||||
msg="act_quant dequantized values don't approximate original",
|
||||
)
|
||||
except AssertionError:
|
||||
nonzero_frac = (ref_deq.abs() > 1e-6).float().mean().item()
|
||||
if nonzero_frac < 0.5:
|
||||
pytest.fail(
|
||||
f"act_quant output looks mostly zero ({nonzero_frac:.1%} nonzero)."
|
||||
)
|
||||
else:
|
||||
pytest.skip(
|
||||
f"act_quant uses a very different quantization scheme "
|
||||
f"(scale_fmt={scale_fmt}). Fused kernel validated independently."
|
||||
)
|
||||
|
||||
torch.testing.assert_close(
|
||||
out_deq,
|
||||
ref_deq,
|
||||
rtol=0.3,
|
||||
atol=0.5,
|
||||
msg="Fused and act_quant dequantized values diverge too much",
|
||||
)
|
||||
|
||||
|
||||
# TEST 3: Roundtrip reconstruction
|
||||
@pytest.mark.parametrize("num_tokens", [1, 64, 257])
|
||||
def test_roundtrip_reconstruction(num_tokens: int):
|
||||
_skip_if_unavailable()
|
||||
device = torch.device("cuda")
|
||||
|
||||
key = torch.randn((num_tokens, HEAD_DIM), device=device, dtype=torch.bfloat16)
|
||||
loc = torch.arange(num_tokens, device=device, dtype=torch.int64)
|
||||
num_pages = _num_pages(loc, PAGE_SIZE)
|
||||
|
||||
buf = _make_buffer(num_pages)
|
||||
fused_store_index_k_cache(key, buf, loc, page_size=PAGE_SIZE)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
fp8_f32, scales = _gather_tokens(buf, loc)
|
||||
reconstructed = fp8_f32 * scales.unsqueeze(-1)
|
||||
original = key.float()
|
||||
|
||||
torch.testing.assert_close(reconstructed, original, rtol=0.15, atol=5e-2)
|
||||
|
||||
per_row_energy = reconstructed.abs().sum(dim=-1)
|
||||
orig_energy = original.abs().sum(dim=-1)
|
||||
mask = orig_energy > 0.1
|
||||
assert (
|
||||
per_row_energy[mask] > 0.01
|
||||
).all(), "Some tokens have zero reconstruction — kernel may not be writing output"
|
||||
|
||||
|
||||
# TEST 4: Boundary conditions
|
||||
def test_single_token():
|
||||
_skip_if_unavailable()
|
||||
device = torch.device("cuda")
|
||||
|
||||
key = torch.randn((1, HEAD_DIM), device=device, dtype=torch.bfloat16)
|
||||
loc = torch.tensor([0], device=device, dtype=torch.int64)
|
||||
|
||||
buf = _make_buffer(1)
|
||||
fused_store_index_k_cache(key, buf, loc, page_size=PAGE_SIZE)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
fp8_f32, scales = _gather_tokens(buf, loc)
|
||||
reconstructed = fp8_f32 * scales.unsqueeze(-1)
|
||||
torch.testing.assert_close(reconstructed, key.float(), rtol=0.15, atol=5e-2)
|
||||
|
||||
|
||||
# TEST 5: Zero input conditions
|
||||
def test_zero_input():
|
||||
_skip_if_unavailable()
|
||||
device = torch.device("cuda")
|
||||
|
||||
key = torch.zeros((4, HEAD_DIM), device=device, dtype=torch.bfloat16)
|
||||
loc = torch.arange(4, device=device, dtype=torch.int64)
|
||||
|
||||
buf = _make_buffer(1)
|
||||
fused_store_index_k_cache(key, buf, loc, page_size=PAGE_SIZE)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
fp8_f32, scales = _gather_tokens(buf, loc)
|
||||
|
||||
expected_scale = 1e-4 / FP8_E4M3_MAX
|
||||
torch.testing.assert_close(
|
||||
scales,
|
||||
torch.full_like(scales, expected_scale),
|
||||
rtol=1e-5,
|
||||
atol=1e-10,
|
||||
)
|
||||
assert (fp8_f32 == 0).all()
|
||||
|
||||
|
||||
# TEST 6: Sanity check — verify reference itself writes non-zero data
|
||||
def test_reference_writes_nonzero():
|
||||
_skip_if_unavailable()
|
||||
device = torch.device("cuda")
|
||||
|
||||
key = torch.randn((8, HEAD_DIM), device=device, dtype=torch.bfloat16)
|
||||
loc = torch.arange(8, device=device, dtype=torch.int64)
|
||||
|
||||
buf = _reference_quantize_and_store(key, loc, num_pages=1)
|
||||
|
||||
fp8_f32, scales = _gather_tokens(buf, loc)
|
||||
deq = fp8_f32 * scales.unsqueeze(-1)
|
||||
|
||||
assert deq.abs().sum().item() > 0, "Reference buffer is all zeros — error!"
|
||||
torch.testing.assert_close(deq, key.float(), rtol=0.15, atol=5e-2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,238 @@
|
||||
"""Tests for fused sigmoid gating delta rule MTP kernel (GDN target_verify).
|
||||
|
||||
Compares the fused kernel `fused_sigmoid_gating_delta_rule_update` against
|
||||
the reference two-step implementation:
|
||||
1. g, beta = fused_gdn_gating(A_log, a, b, dt_bias)
|
||||
2. o = fused_recurrent_gated_delta_rule_update(q, k, v, g, beta, ...)
|
||||
"""
|
||||
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
try:
|
||||
from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating
|
||||
from sglang.srt.layers.attention.fla.fused_recurrent import (
|
||||
fused_recurrent_gated_delta_rule_update,
|
||||
)
|
||||
from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
|
||||
fused_sigmoid_gating_delta_rule_update,
|
||||
)
|
||||
|
||||
KERNELS_AVAILABLE = True
|
||||
except ImportError:
|
||||
KERNELS_AVAILABLE = False
|
||||
|
||||
register_cuda_ci(est_time=6, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
def _make_tensors(N, T, H, HV, K, V, device="cuda", seed=2025):
|
||||
"""Create input tensors for GDN target_verify."""
|
||||
torch.manual_seed(seed)
|
||||
A_log = torch.randn(HV, dtype=torch.float32, device=device)
|
||||
dt_bias = torch.randn(HV, dtype=torch.bfloat16, device=device)
|
||||
a = torch.randn(1, N * T, HV, dtype=torch.bfloat16, device=device)
|
||||
b = torch.randn(1, N * T, HV, dtype=torch.bfloat16, device=device)
|
||||
q = torch.randn(1, N * T, H, K, dtype=torch.bfloat16, device=device)
|
||||
k = torch.randn(1, N * T, H, K, dtype=torch.bfloat16, device=device)
|
||||
v = torch.randn(1, N * T, HV, V, dtype=torch.bfloat16, device=device)
|
||||
indices = torch.arange(N, dtype=torch.int32, device=device)
|
||||
initial_state = torch.randn(N, HV, K, V, dtype=torch.float, device=device)
|
||||
cu_seqlens = torch.arange(0, N * T + 1, T, dtype=torch.int32, device=device)
|
||||
return A_log, dt_bias, a, b, q, k, v, initial_state, indices, cu_seqlens
|
||||
|
||||
|
||||
def run_reference(
|
||||
A_log,
|
||||
dt_bias,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
a,
|
||||
b,
|
||||
initial_state_source,
|
||||
initial_state_indices,
|
||||
cu_seqlens,
|
||||
disable_state_update=True,
|
||||
intermediate_states_buffer=None,
|
||||
intermediate_state_indices=None,
|
||||
cache_steps=None,
|
||||
retrieve_parent_token=None,
|
||||
):
|
||||
"""Reference: fused_gdn_gating + fused_recurrent_gated_delta_rule_update."""
|
||||
# fused_gdn_gating expects 2D [seq_len, HV]
|
||||
a_2d = a.view(-1, a.shape[-1])
|
||||
b_2d = b.view(-1, b.shape[-1])
|
||||
g, beta = fused_gdn_gating(A_log, a_2d, b_2d, dt_bias)
|
||||
# fused_recurrent expects 3D [B, T, HV]
|
||||
g = g.view(a.shape)
|
||||
beta = beta.view(b.shape)
|
||||
|
||||
# fused_recurrent requires intermediate_state_indices when cu_seqlens is used
|
||||
if cu_seqlens is not None and intermediate_state_indices is None:
|
||||
N = len(cu_seqlens) - 1
|
||||
intermediate_state_indices = torch.arange(N, dtype=torch.int32, device=q.device)
|
||||
|
||||
return fused_recurrent_gated_delta_rule_update(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
g=g,
|
||||
beta=beta,
|
||||
initial_state_source=initial_state_source,
|
||||
initial_state_indices=initial_state_indices,
|
||||
cu_seqlens=cu_seqlens,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
disable_state_update=disable_state_update,
|
||||
intermediate_states_buffer=intermediate_states_buffer,
|
||||
intermediate_state_indices=intermediate_state_indices,
|
||||
cache_steps=cache_steps,
|
||||
retrieve_parent_token=retrieve_parent_token,
|
||||
)
|
||||
|
||||
|
||||
def run_fused_mtp(
|
||||
A_log,
|
||||
dt_bias,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
a,
|
||||
b,
|
||||
initial_state_source,
|
||||
initial_state_indices,
|
||||
cu_seqlens,
|
||||
disable_state_update=True,
|
||||
intermediate_states_buffer=None,
|
||||
intermediate_state_indices=None,
|
||||
cache_steps=None,
|
||||
retrieve_parent_token=None,
|
||||
):
|
||||
"""Fused: fused_sigmoid_gating_delta_rule_update."""
|
||||
return fused_sigmoid_gating_delta_rule_update(
|
||||
A_log=A_log,
|
||||
dt_bias=dt_bias,
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
a=a,
|
||||
b=b,
|
||||
initial_state_source=initial_state_source,
|
||||
initial_state_indices=initial_state_indices,
|
||||
cu_seqlens=cu_seqlens,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
softplus_beta=1.0,
|
||||
softplus_threshold=20.0,
|
||||
is_kda=False,
|
||||
disable_state_update=disable_state_update,
|
||||
intermediate_states_buffer=intermediate_states_buffer,
|
||||
intermediate_state_indices=intermediate_state_indices,
|
||||
cache_steps=cache_steps,
|
||||
retrieve_parent_token=retrieve_parent_token,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not KERNELS_AVAILABLE, reason="Kernel not available")
|
||||
@pytest.mark.parametrize("N", [1, 8, 16])
|
||||
@pytest.mark.parametrize("T", [1, 4, 8])
|
||||
def test_fused_gdn_mtp_precision(N: int, T: int):
|
||||
"""Compare fused MTP output against reference."""
|
||||
H, HV, K, V = 16, 32, 128, 128
|
||||
|
||||
A_log, dt_bias, a, b, q, k, v, state, indices, cu_seqlens = _make_tensors(
|
||||
N, T, H, HV, K, V
|
||||
)
|
||||
|
||||
state_ref = state.clone()
|
||||
state_fused = state.clone()
|
||||
|
||||
out_ref = run_reference(
|
||||
A_log,
|
||||
dt_bias,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
a,
|
||||
b,
|
||||
state_ref,
|
||||
indices,
|
||||
cu_seqlens,
|
||||
disable_state_update=True,
|
||||
)
|
||||
out_fused = run_fused_mtp(
|
||||
A_log,
|
||||
dt_bias,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
a,
|
||||
b,
|
||||
state_fused,
|
||||
indices,
|
||||
cu_seqlens,
|
||||
disable_state_update=True,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(out_ref, out_fused, rtol=1e-2, atol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not KERNELS_AVAILABLE, reason="Kernels not available")
|
||||
@pytest.mark.parametrize("N", [1, 16, 128])
|
||||
def test_mtp_single_step_decode(N: int):
|
||||
"""Verify MTP kernel matches reference for T=1 (decode scenario)."""
|
||||
T = 1
|
||||
H, HV, K, V = 16, 32, 128, 128
|
||||
|
||||
A_log, dt_bias, a, b, q, k, v, state, indices, cu_seqlens = _make_tensors(
|
||||
N, T, H, HV, K, V
|
||||
)
|
||||
|
||||
state_ref = state.clone()
|
||||
state_fused = state.clone()
|
||||
|
||||
out_ref = run_reference(
|
||||
A_log,
|
||||
dt_bias,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
a,
|
||||
b,
|
||||
state_ref,
|
||||
indices,
|
||||
cu_seqlens,
|
||||
disable_state_update=False,
|
||||
)
|
||||
out_fused = run_fused_mtp(
|
||||
A_log,
|
||||
dt_bias,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
a,
|
||||
b,
|
||||
state_fused,
|
||||
indices,
|
||||
cu_seqlens,
|
||||
disable_state_update=False,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(out_ref, out_fused, rtol=1e-2, atol=1e-2)
|
||||
|
||||
# Also verify states match after update
|
||||
state_diff = (state_ref.float() - state_fused.float()).abs()
|
||||
state_max_diff = state_diff.max().item()
|
||||
state_fail_rate = (state_diff > 0.1).float().mean().item() * 100
|
||||
print(
|
||||
f" single_step state N={N}: max_diff={state_max_diff:.2e}, "
|
||||
f"fail_rate={state_fail_rate:.2f}%"
|
||||
)
|
||||
assert state_fail_rate < 0.01, f"State mismatch: fail_rate={state_fail_rate:.2f}%"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,194 @@
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from sgl_kernel.scalar_type import scalar_types
|
||||
|
||||
from sglang.jit_kernel.gptq_marlin import gptq_marlin_gemm
|
||||
from sglang.srt.layers.quantization.marlin_utils import (
|
||||
check_marlin_supported,
|
||||
marlin_make_workspace,
|
||||
)
|
||||
from sglang.srt.layers.quantization.marlin_utils_fp4 import (
|
||||
apply_fp4_marlin_linear,
|
||||
nvfp4_marlin_process_global_scale,
|
||||
prepare_nvfp4_layer_for_marlin,
|
||||
)
|
||||
from sglang.srt.utils.common import is_sm80_supported, is_sm90_supported
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_marlin_utils import (
|
||||
awq_marlin_quantize,
|
||||
make_nvfp4_weight_and_ref,
|
||||
marlin_quantize,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=13, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
MNK_FACTORS = [
|
||||
(1, 1, 1),
|
||||
(1, 4, 8),
|
||||
(13, 17, 67),
|
||||
(257, 13, 11),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("k_chunk", [128])
|
||||
@pytest.mark.parametrize("n_chunk", [64, 256])
|
||||
@pytest.mark.parametrize("quant_type", [scalar_types.uint4, scalar_types.uint4b8])
|
||||
@pytest.mark.parametrize("group_size", [-1, 128])
|
||||
@pytest.mark.parametrize("mnk_factors", MNK_FACTORS)
|
||||
@pytest.mark.parametrize("act_order", [False, True])
|
||||
def test_gptq_marlin_gemm(
|
||||
k_chunk,
|
||||
n_chunk,
|
||||
quant_type,
|
||||
group_size,
|
||||
mnk_factors,
|
||||
act_order,
|
||||
):
|
||||
m_factor, n_factor, k_factor = mnk_factors
|
||||
has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8]
|
||||
|
||||
size_m = m_factor
|
||||
size_k = k_chunk * k_factor
|
||||
size_n = n_chunk * n_factor
|
||||
|
||||
if act_order:
|
||||
if group_size == -1:
|
||||
return
|
||||
if group_size == size_k:
|
||||
return
|
||||
if has_zp:
|
||||
return
|
||||
|
||||
if size_k % group_size != 0:
|
||||
return
|
||||
|
||||
a_input = torch.randn((size_m, size_k), dtype=torch.float16, device="cuda")
|
||||
b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda")
|
||||
|
||||
if has_zp:
|
||||
w_ref, marlin_q_w, marlin_s, marlin_zp = awq_marlin_quantize(
|
||||
b_weight, quant_type, group_size
|
||||
)
|
||||
g_idx = None
|
||||
sort_indices = None
|
||||
marlin_s2 = None
|
||||
else:
|
||||
w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize(
|
||||
b_weight, quant_type, group_size, act_order
|
||||
)
|
||||
marlin_zp = None
|
||||
marlin_s2 = None
|
||||
|
||||
workspace = marlin_make_workspace(w_ref.device)
|
||||
|
||||
output = gptq_marlin_gemm(
|
||||
a_input,
|
||||
None,
|
||||
marlin_q_w,
|
||||
marlin_s,
|
||||
marlin_s2,
|
||||
marlin_zp,
|
||||
g_idx,
|
||||
sort_indices,
|
||||
workspace,
|
||||
quant_type,
|
||||
a_input.shape[0],
|
||||
b_weight.shape[1],
|
||||
a_input.shape[1],
|
||||
is_k_full=True,
|
||||
use_atomic_add=False,
|
||||
use_fp32_reduce=False,
|
||||
is_zp_float=False,
|
||||
)
|
||||
|
||||
output_ref = torch.matmul(a_input, w_ref)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# JIT kernel should produce approximately correct results vs torch.matmul
|
||||
max_diff = torch.mean(torch.abs(output - output_ref)) / torch.mean(
|
||||
torch.abs(output_ref)
|
||||
)
|
||||
assert max_diff < 0.04
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Skip, test pass locally but compiling takes too long in CI")
|
||||
@pytest.mark.skipif(
|
||||
not (is_sm80_supported() or is_sm90_supported()),
|
||||
reason="NVFP4 Marlin fallback tests require CUDA SM8X/SM9X",
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
def test_nvfp4_marlin_support_and_scale_transforms_sm80_sm90(dtype):
|
||||
major, minor = torch.cuda.get_device_capability()
|
||||
capability = major * 10 + minor
|
||||
assert check_marlin_supported(
|
||||
scalar_types.float4_e2m1f,
|
||||
group_size=16,
|
||||
has_zp=False,
|
||||
device_capability=capability,
|
||||
)
|
||||
|
||||
global_scale = torch.tensor(1.0, dtype=dtype, device="cuda")
|
||||
actual_global_scale = nvfp4_marlin_process_global_scale(global_scale)
|
||||
assert actual_global_scale.is_cuda
|
||||
assert actual_global_scale.ndim == 1
|
||||
assert actual_global_scale.numel() == 1
|
||||
if dtype == torch.float16:
|
||||
assert actual_global_scale.item() == 128.0
|
||||
else:
|
||||
assert actual_global_scale.item() == 2.0**119
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Skip, test pass locally but compiling takes too long in CI")
|
||||
@pytest.mark.skipif(
|
||||
not (is_sm80_supported() or is_sm90_supported()),
|
||||
reason="NVFP4 Marlin dense numeric test requires CUDA SM80, SM86, or SM90",
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
def test_nvfp4_marlin_dense_matches_dequant_reference(dtype):
|
||||
torch.manual_seed(0)
|
||||
|
||||
size_m = 17
|
||||
size_k = 256
|
||||
size_n = 192
|
||||
group_size = 16
|
||||
|
||||
a_input = torch.randn((size_m, size_k), dtype=dtype, device="cuda") / 10
|
||||
fp4_weight, scales, global_scale, weight_ref = make_nvfp4_weight_and_ref(
|
||||
size_n, size_k, dtype, group_size=group_size
|
||||
)
|
||||
|
||||
layer = torch.nn.Module()
|
||||
layer.quant_config = SimpleNamespace(group_size=group_size)
|
||||
layer.output_size_per_partition = size_n
|
||||
layer.input_size_per_partition = size_k
|
||||
layer.params_dtype = dtype
|
||||
layer.weight = torch.nn.Parameter(fp4_weight, requires_grad=False)
|
||||
layer.weight_scale = torch.nn.Parameter(scales, requires_grad=False)
|
||||
layer.weight_global_scale = torch.nn.Parameter(
|
||||
global_scale.reshape(1), requires_grad=False
|
||||
)
|
||||
prepare_nvfp4_layer_for_marlin(layer)
|
||||
|
||||
output = apply_fp4_marlin_linear(
|
||||
a_input,
|
||||
layer.weight,
|
||||
layer.weight_scale,
|
||||
layer.weight_global_scale,
|
||||
layer.workspace,
|
||||
size_n,
|
||||
size_k,
|
||||
use_fp32_reduce=True,
|
||||
)
|
||||
|
||||
output_ref = torch.matmul(a_input, weight_ref.T)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(output, output_ref, rtol=0.04, atol=0.04)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,96 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from sgl_kernel.scalar_type import scalar_types
|
||||
|
||||
from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack
|
||||
from sglang.srt.layers.quantization.utils import (
|
||||
gptq_quantize_weights,
|
||||
pack_rows,
|
||||
sort_weights,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_marlin_utils import get_weight_perm, marlin_weights
|
||||
|
||||
register_cuda_ci(est_time=16, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
MARLIN_K_CHUNKS = [128]
|
||||
MARLIN_N_CHUNKS = [64, 256]
|
||||
|
||||
MNK_FACTORS = [
|
||||
(1, 1, 1),
|
||||
(1, 4, 8),
|
||||
(1, 7, 5),
|
||||
(13, 17, 67),
|
||||
(26, 37, 13),
|
||||
(67, 13, 11),
|
||||
(257, 13, 11),
|
||||
(658, 13, 11),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS)
|
||||
@pytest.mark.parametrize("n_chunk", MARLIN_N_CHUNKS)
|
||||
@pytest.mark.parametrize("quant_type", [scalar_types.uint4b8])
|
||||
@pytest.mark.parametrize("group_size", [-1, 32, 64, 128])
|
||||
@pytest.mark.parametrize("act_order", [False, True])
|
||||
@pytest.mark.parametrize("mnk_factors", MNK_FACTORS)
|
||||
def test_gptq_marlin_repack(
|
||||
k_chunk, n_chunk, quant_type, group_size, act_order, mnk_factors
|
||||
):
|
||||
m_factor, n_factor, k_factor = mnk_factors
|
||||
|
||||
size_k = k_chunk * k_factor
|
||||
size_n = n_chunk * n_factor
|
||||
|
||||
# Filter act_order
|
||||
if act_order:
|
||||
if group_size == -1:
|
||||
return
|
||||
if group_size == size_k:
|
||||
return
|
||||
|
||||
# Normalize group_size
|
||||
if group_size == -1:
|
||||
group_size = size_k
|
||||
assert group_size <= size_k
|
||||
|
||||
if size_k % group_size != 0:
|
||||
pytest.skip("size_k must be divisible by group_size")
|
||||
|
||||
# Create input
|
||||
b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda")
|
||||
|
||||
# Quantize (and apply act_order if provided)
|
||||
w_ref, q_w, s, g_idx, rand_perm = gptq_quantize_weights(
|
||||
b_weight, quant_type, group_size, act_order
|
||||
)
|
||||
|
||||
q_w_gptq = pack_rows(q_w, quant_type.size_bits, size_k, size_n)
|
||||
|
||||
# For act_order, sort the "weights" and "g_idx" so that group ids are
|
||||
# increasing
|
||||
sort_indices = torch.empty(0, dtype=torch.int, device=b_weight.device)
|
||||
if act_order:
|
||||
q_w, g_idx, sort_indices = sort_weights(q_w, g_idx)
|
||||
|
||||
marlin_layout_perm = get_weight_perm(quant_type.size_bits)
|
||||
q_w_marlin_ref = marlin_weights(
|
||||
q_w, size_k, size_n, quant_type.size_bits, marlin_layout_perm
|
||||
)
|
||||
|
||||
# Run JIT repack kernel
|
||||
jit_output = gptq_marlin_repack(
|
||||
q_w_gptq, sort_indices, size_k, size_n, quant_type.size_bits
|
||||
)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# JIT should match the reference (computed from CPU marlin_weights)
|
||||
torch.testing.assert_close(jit_output, q_w_marlin_ref)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,210 @@
|
||||
import itertools
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.grouped_topk import grouped_topk as jit_grouped_topk
|
||||
from sglang.jit_kernel.utils import get_ci_test_range
|
||||
from sglang.srt.layers.moe.topk import biased_grouped_topk_impl
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
CORRECTNESS_CASES = get_ci_test_range(
|
||||
full_range=list(
|
||||
itertools.product(
|
||||
[1, 17, 128],
|
||||
[16, 32, 64, 128, 192, 256, 384, 512],
|
||||
[1, 2, 3, 4, 5, 6, 7, 8],
|
||||
)
|
||||
),
|
||||
ci_range=[
|
||||
(1, 16, 3), # smallest non-power-of-two topk
|
||||
(17, 128, 6), # Nemotron-3-Nano shape that exposed the bug
|
||||
(128, 192, 8), # Hunyuan-3 shape, power-of-two topk sanity case
|
||||
(33, 512, 7), # largest expert-count tier with non-power-of-two topk
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _make_inputs(num_tokens: int, num_experts: int, seed: int):
|
||||
torch.manual_seed(seed)
|
||||
hidden_states = torch.empty((num_tokens, 1), dtype=torch.float32, device="cuda")
|
||||
gating_output = torch.randn(
|
||||
(num_tokens, num_experts), dtype=torch.float32, device="cuda"
|
||||
)
|
||||
correction_bias = torch.randn(num_experts, dtype=torch.float32, device="cuda") * 0.1
|
||||
return hidden_states, gating_output, correction_bias
|
||||
|
||||
|
||||
def _scatter_by_expert(
|
||||
weights: torch.Tensor, ids: torch.Tensor, num_experts: int
|
||||
) -> torch.Tensor:
|
||||
dense = torch.zeros(
|
||||
(weights.shape[0], num_experts), dtype=torch.float32, device=weights.device
|
||||
)
|
||||
dense.scatter_(1, ids.long(), weights)
|
||||
return dense
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens,num_experts,topk", CORRECTNESS_CASES)
|
||||
def test_grouped_topk_renormalize_matches_reference(
|
||||
num_tokens: int, num_experts: int, topk: int
|
||||
) -> None:
|
||||
hidden_states, gating_output, correction_bias = _make_inputs(
|
||||
num_tokens, num_experts, seed=1000 + num_experts * 10 + topk
|
||||
)
|
||||
scaling_factor = 2.826 if (num_experts, topk) == (192, 8) else 1.0
|
||||
|
||||
topk_weights, topk_ids = jit_grouped_topk(
|
||||
gating_output,
|
||||
correction_bias,
|
||||
1,
|
||||
1,
|
||||
topk,
|
||||
True,
|
||||
scaling_factor,
|
||||
)
|
||||
ref_weights, ref_ids = biased_grouped_topk_impl(
|
||||
hidden_states,
|
||||
gating_output,
|
||||
correction_bias,
|
||||
topk,
|
||||
True,
|
||||
1,
|
||||
1,
|
||||
routed_scaling_factor=scaling_factor,
|
||||
apply_routed_scaling_factor_on_output=True,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(
|
||||
_scatter_by_expert(topk_weights, topk_ids, num_experts),
|
||||
_scatter_by_expert(ref_weights, ref_ids, num_experts),
|
||||
rtol=1e-5,
|
||||
atol=1e-6,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
topk_weights.sum(dim=-1),
|
||||
torch.full((num_tokens,), scaling_factor, dtype=torch.float32, device="cuda"),
|
||||
rtol=1e-5,
|
||||
atol=1e-6,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("topk", [3, 5, 6, 7])
|
||||
def test_grouped_topk_non_power_of_two_renormalize(topk: int) -> None:
|
||||
hidden_states, gating_output, correction_bias = _make_inputs(
|
||||
num_tokens=64, num_experts=128, seed=2000 + topk
|
||||
)
|
||||
|
||||
topk_weights, topk_ids = jit_grouped_topk(
|
||||
gating_output,
|
||||
correction_bias,
|
||||
1,
|
||||
1,
|
||||
topk,
|
||||
True,
|
||||
1.0,
|
||||
)
|
||||
ref_weights, ref_ids = biased_grouped_topk_impl(
|
||||
hidden_states,
|
||||
gating_output,
|
||||
correction_bias,
|
||||
topk,
|
||||
True,
|
||||
1,
|
||||
1,
|
||||
routed_scaling_factor=1.0,
|
||||
apply_routed_scaling_factor_on_output=True,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(
|
||||
_scatter_by_expert(topk_weights, topk_ids, 128),
|
||||
_scatter_by_expert(ref_weights, ref_ids, 128),
|
||||
rtol=1e-5,
|
||||
atol=1e-6,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
topk_weights.sum(dim=-1),
|
||||
torch.ones((64,), dtype=torch.float32, device="cuda"),
|
||||
rtol=1e-5,
|
||||
atol=1e-6,
|
||||
)
|
||||
|
||||
|
||||
def test_grouped_topk_negative_choice_scores_match_reference() -> None:
|
||||
hidden_states, gating_output, correction_bias = _make_inputs(
|
||||
num_tokens=64, num_experts=128, seed=23758
|
||||
)
|
||||
correction_bias.fill_(-2.0)
|
||||
|
||||
topk_weights, topk_ids = jit_grouped_topk(
|
||||
gating_output,
|
||||
correction_bias,
|
||||
1,
|
||||
1,
|
||||
6,
|
||||
True,
|
||||
1.0,
|
||||
)
|
||||
ref_weights, ref_ids = biased_grouped_topk_impl(
|
||||
hidden_states,
|
||||
gating_output,
|
||||
correction_bias,
|
||||
6,
|
||||
True,
|
||||
1,
|
||||
1,
|
||||
routed_scaling_factor=1.0,
|
||||
apply_routed_scaling_factor_on_output=True,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(
|
||||
_scatter_by_expert(topk_weights, topk_ids, 128),
|
||||
_scatter_by_expert(ref_weights, ref_ids, 128),
|
||||
rtol=1e-5,
|
||||
atol=1e-6,
|
||||
)
|
||||
|
||||
|
||||
def test_grouped_topk_without_renormalize_matches_reference() -> None:
|
||||
hidden_states, gating_output, correction_bias = _make_inputs(
|
||||
num_tokens=64, num_experts=128, seed=3006
|
||||
)
|
||||
|
||||
topk_weights, topk_ids = jit_grouped_topk(
|
||||
gating_output,
|
||||
correction_bias,
|
||||
1,
|
||||
1,
|
||||
6,
|
||||
False,
|
||||
1.0,
|
||||
)
|
||||
ref_weights, ref_ids = biased_grouped_topk_impl(
|
||||
hidden_states,
|
||||
gating_output,
|
||||
correction_bias,
|
||||
6,
|
||||
False,
|
||||
1,
|
||||
1,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(
|
||||
_scatter_by_expert(topk_weights, topk_ids, 128),
|
||||
_scatter_by_expert(ref_weights, ref_ids, 128),
|
||||
rtol=1e-5,
|
||||
atol=1e-6,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,428 @@
|
||||
import math
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from scipy.linalg import hadamard
|
||||
|
||||
from sglang.jit_kernel.hadamard import (
|
||||
hadamard_transform,
|
||||
hadamard_transform_12n,
|
||||
hadamard_transform_20n,
|
||||
hadamard_transform_28n,
|
||||
hadamard_transform_40n,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=128, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=512, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
# Exact M×N Hadamard matrices (±1 entries) copied from
|
||||
# python/sglang/jit_kernel/csrc/fast-hadamard-transform/code_gen.py.
|
||||
# These are non-power-of-2 Hadamard matrices constructed via Paley/Williamson methods.
|
||||
# "+" = +1, "-" = -1. Used by the _12n/_20n/_28n/_40n kernel variants.
|
||||
|
||||
_HAD_12_STR = """
|
||||
+-++++++++++
|
||||
--+-+-+-+-+-
|
||||
+++-++----++
|
||||
+---+--+-++-
|
||||
+++++-++----
|
||||
+-+---+--+-+
|
||||
++--+++-++--
|
||||
+--++---+--+
|
||||
++----+++-++
|
||||
+--+-++---+-
|
||||
++++----+++-
|
||||
+-+--+-++---
|
||||
"""
|
||||
|
||||
_HAD_20_STR = """
|
||||
+----+----++--++-++-
|
||||
-+----+---+++---+-++
|
||||
--+----+---+++-+-+-+
|
||||
---+----+---+++++-+-
|
||||
----+----++--++-++-+
|
||||
-+++++-----+--+++--+
|
||||
+-+++-+---+-+--+++--
|
||||
++-++--+---+-+--+++-
|
||||
+++-+---+---+-+--+++
|
||||
++++-----++--+-+--++
|
||||
--++-+-++-+-----++++
|
||||
---++-+-++-+---+-+++
|
||||
+---++-+-+--+--++-++
|
||||
++---++-+----+-+++-+
|
||||
-++---++-+----+++++-
|
||||
-+--+--++-+----+----
|
||||
+-+-----++-+----+---
|
||||
-+-+-+---+--+----+--
|
||||
--+-+++------+----+-
|
||||
+--+--++------+----+
|
||||
"""
|
||||
|
||||
_HAD_28_STR = """
|
||||
+------++----++-+--+-+--++--
|
||||
-+-----+++-----+-+--+-+--++-
|
||||
--+-----+++---+-+-+----+--++
|
||||
---+-----+++---+-+-+-+--+--+
|
||||
----+-----+++---+-+-+++--+--
|
||||
-----+-----++++--+-+--++--+-
|
||||
------++----++-+--+-+--++--+
|
||||
--++++-+-------++--+++-+--+-
|
||||
---++++-+-----+-++--+-+-+--+
|
||||
+---+++--+----++-++--+-+-+--
|
||||
++---++---+----++-++--+-+-+-
|
||||
+++---+----+----++-++--+-+-+
|
||||
++++--------+-+--++-++--+-+-
|
||||
-++++--------+++--++--+--+-+
|
||||
-+-++-++--++--+--------++++-
|
||||
+-+-++--+--++--+--------++++
|
||||
-+-+-++--+--++--+----+---+++
|
||||
+-+-+-++--+--+---+---++---++
|
||||
++-+-+-++--+------+--+++---+
|
||||
-++-+-+-++--+------+-++++---
|
||||
+-++-+---++--+------+-++++--
|
||||
-++--++-+-++-+++----++------
|
||||
+-++--++-+-++-+++-----+-----
|
||||
++-++---+-+-++-+++-----+----
|
||||
-++-++-+-+-+-+--+++-----+---
|
||||
--++-++++-+-+----+++-----+--
|
||||
+--++-+-++-+-+----+++-----+-
|
||||
++--++-+-++-+-+----++------+
|
||||
"""
|
||||
|
||||
_HAD_40_STR = """
|
||||
+-------------------+-------------------
|
||||
++-++----+-+-++++--+++-++----+-+-++++--+
|
||||
+++-++----+-+-++++--+++-++----+-+-++++--
|
||||
+-++-++----+-+-++++-+-++-++----+-+-++++-
|
||||
+--++-++----+-+-+++++--++-++----+-+-++++
|
||||
++--++-++----+-+-+++++--++-++----+-+-+++
|
||||
+++--++-++----+-+-+++++--++-++----+-+-++
|
||||
++++--++-++----+-+-+++++--++-++----+-+-+
|
||||
+++++--++-++----+-+-+++++--++-++----+-+-
|
||||
+-++++--++-++----+-++-++++--++-++----+-+
|
||||
++-++++--++-++----+-++-++++--++-++----+-
|
||||
+-+-++++--++-++----++-+-++++--++-++----+
|
||||
++-+-++++--++-++----++-+-++++--++-++----
|
||||
+-+-+-++++--++-++---+-+-+-++++--++-++---
|
||||
+--+-+-++++--++-++--+--+-+-++++--++-++--
|
||||
+---+-+-++++--++-++-+---+-+-++++--++-++-
|
||||
+----+-+-++++--++-+++----+-+-++++--++-++
|
||||
++----+-+-++++--++-+++----+-+-++++--++-+
|
||||
+++----+-+-++++--++-+++----+-+-++++--++-
|
||||
+-++----+-+-++++--+++-++----+-+-++++--++
|
||||
+--------------------+++++++++++++++++++
|
||||
++-++----+-+-++++--+--+--++++-+-+----++-
|
||||
+++-++----+-+-++++-----+--++++-+-+----++
|
||||
+-++-++----+-+-++++--+--+--++++-+-+----+
|
||||
+--++-++----+-+-++++-++--+--++++-+-+----
|
||||
++--++-++----+-+-+++--++--+--++++-+-+---
|
||||
+++--++-++----+-+-++---++--+--++++-+-+--
|
||||
++++--++-++----+-+-+----++--+--++++-+-+-
|
||||
+++++--++-++----+-+------++--+--++++-+-+
|
||||
+-++++--++-++----+-+-+----++--+--++++-+-
|
||||
++-++++--++-++----+---+----++--+--++++-+
|
||||
+-+-++++--++-++----+-+-+----++--+--++++-
|
||||
++-+-++++--++-++------+-+----++--+--++++
|
||||
+-+-+-++++--++-++----+-+-+----++--+--+++
|
||||
+--+-+-++++--++-++---++-+-+----++--+--++
|
||||
+---+-+-++++--++-++--+++-+-+----++--+--+
|
||||
+----+-+-++++--++-++-++++-+-+----++--+--
|
||||
++----+-+-++++--++-+--++++-+-+----++--+-
|
||||
+++----+-+-++++--++----++++-+-+----++--+
|
||||
+-++----+-+-++++--++-+--++++-+-+----++--
|
||||
"""
|
||||
|
||||
|
||||
def _parse_hadamard_str(s):
|
||||
"""Parse a ±1 string matrix definition into a numpy array."""
|
||||
s = s.strip().replace("+", "1").replace("-", "-1").split()
|
||||
return np.stack(
|
||||
[np.fromstring(" ".join(s[i]), dtype=np.int32, sep=" ") for i in range(len(s))]
|
||||
)
|
||||
|
||||
|
||||
# Parsed M×M special Hadamard matrices, keyed by M (the "multiple").
|
||||
# Copied from python/sglang/jit_kernel/csrc/fast-hadamard-transform/code_gen.py
|
||||
# (had_12_paley, had_20_will, had_28_will, had_40_tpal)
|
||||
_SPECIAL_MATRICES = {
|
||||
12: _parse_hadamard_str(_HAD_12_STR),
|
||||
20: _parse_hadamard_str(_HAD_20_STR),
|
||||
28: _parse_hadamard_str(_HAD_28_STR),
|
||||
40: _parse_hadamard_str(_HAD_40_STR),
|
||||
}
|
||||
|
||||
|
||||
def hadamard_transform_ref(x, scale=1.0):
|
||||
"""Reference impl for the general (power-of-2) hadamard_transform.
|
||||
|
||||
Pads dim to the next power of 2, multiplies by the full H matrix
|
||||
via F.linear, then truncates back to the original dim.
|
||||
"""
|
||||
x_shape = x.shape
|
||||
dim = x.shape[-1]
|
||||
x = x.reshape(-1, dim)
|
||||
log_dim = math.ceil(math.log2(dim)) if dim > 0 else 0
|
||||
dim_padded = 2**log_dim if dim > 0 else 1
|
||||
if dim != dim_padded:
|
||||
x = F.pad(x, (0, dim_padded - dim))
|
||||
H = torch.tensor(hadamard(dim_padded, dtype=float), dtype=x.dtype, device=x.device)
|
||||
out = F.linear(x, H)
|
||||
out = out * scale
|
||||
return out[..., :dim].reshape(*x_shape)
|
||||
|
||||
|
||||
def hadamard_transform_mn_ref(x, multiple, scale=1.0):
|
||||
"""Reference impl for the M×N hadamard variants (_12n, _20n, _28n, _40n).
|
||||
|
||||
The kernel computes (H_M ⊗ H_N) · x via two steps:
|
||||
1) H_N (power-of-2 Hadamard) along the N dimension
|
||||
2) H_M (special ±1 matrix) along the M dimension
|
||||
where dim = M * N, M = `multiple`, N = power of 2.
|
||||
"""
|
||||
x_shape = x.shape
|
||||
dim = x.shape[-1]
|
||||
x = x.reshape(-1, dim)
|
||||
|
||||
# The kernel requires dim % (4*M) == 0 (for vectorized memory access).
|
||||
# See python/sglang/jit_kernel/hadamard.py: pad_multiple = 4 * 12 / 4 * 20 / etc.
|
||||
pad_multiple = 4 * multiple
|
||||
if dim % pad_multiple != 0:
|
||||
pad_size = pad_multiple - dim % pad_multiple
|
||||
x = F.pad(x, (0, pad_size))
|
||||
dim_padded = dim + pad_size
|
||||
else:
|
||||
dim_padded = dim
|
||||
|
||||
# N = dim_padded / M, must be a power of 2
|
||||
n = dim_padded // multiple
|
||||
log_n = int(math.log2(n))
|
||||
assert 2**log_n == n, f"n={n} is not a power of 2"
|
||||
|
||||
batch = x.shape[0]
|
||||
x = x.reshape(batch, multiple, n) # (batch, M, N)
|
||||
|
||||
# Step 1: apply H_N (standard power-of-2 Hadamard) along the N dimension
|
||||
H_n = torch.tensor(hadamard(n, dtype=float), dtype=x.dtype, device=x.device)
|
||||
x = torch.einsum("bmn,kn->bmk", x, H_n)
|
||||
|
||||
# Step 2: apply H_M (special ±1 matrix) along the M dimension
|
||||
H_m = torch.tensor(
|
||||
_SPECIAL_MATRICES[multiple].astype(float), dtype=x.dtype, device=x.device
|
||||
)
|
||||
x = torch.einsum("bmn,km->bkn", x, H_m)
|
||||
|
||||
x = x.reshape(batch, -1) * scale
|
||||
return x[..., : x_shape[-1]].reshape(*x_shape)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
|
||||
@pytest.mark.parametrize(
|
||||
"dim",
|
||||
# Power-of-2 dims from sgl-kernel/tests/test_hadamard.py (old AOT test)
|
||||
[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768],
|
||||
)
|
||||
def test_hadamard_transform(dim, dtype):
|
||||
device = "cuda"
|
||||
|
||||
# Tolerances from sgl-kernel/tests/test_hadamard.py (old AOT test)
|
||||
if dtype == torch.float32:
|
||||
rtol, atol = 3e-4, 3e-3
|
||||
elif dtype == torch.bfloat16:
|
||||
rtol, atol = 1e-2, 5e-2
|
||||
else: # float16
|
||||
rtol, atol = 3e-3, 5e-3
|
||||
|
||||
torch.random.manual_seed(0)
|
||||
batch_size = 15
|
||||
|
||||
x = torch.randn(batch_size, dim, device=device, dtype=dtype)
|
||||
scale = 1.0 / math.sqrt(dim)
|
||||
|
||||
out = hadamard_transform(x, scale=scale)
|
||||
# Compute reference in float32 from a detached copy to avoid precision loss
|
||||
out_ref = hadamard_transform_ref(x.detach().clone().float(), scale=scale)
|
||||
|
||||
torch.testing.assert_close(out.float(), out_ref, rtol=rtol, atol=atol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
|
||||
@pytest.mark.parametrize(
|
||||
"dim",
|
||||
# Non-power-of-2 dims to test the padding path
|
||||
# (137 from sgl-kernel/tests/test_hadamard.py, 500/1000 added for coverage)
|
||||
[137, 500, 1000],
|
||||
)
|
||||
def test_hadamard_transform_non_power_of_two(dim, dtype):
|
||||
device = "cuda"
|
||||
|
||||
if dtype == torch.float32:
|
||||
rtol, atol = 3e-4, 3e-3
|
||||
elif dtype == torch.bfloat16:
|
||||
rtol, atol = 1e-2, 5e-2
|
||||
else:
|
||||
rtol, atol = 3e-3, 5e-3
|
||||
|
||||
torch.random.manual_seed(42)
|
||||
batch_size = 15
|
||||
|
||||
x = torch.randn(batch_size, dim, device=device, dtype=dtype)
|
||||
scale = 1.0 / math.sqrt(dim)
|
||||
|
||||
out = hadamard_transform(x, scale=scale)
|
||||
out_ref = hadamard_transform_ref(x.detach().clone().float(), scale=scale)
|
||||
|
||||
torch.testing.assert_close(out.float(), out_ref, rtol=rtol, atol=atol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
def test_hadamard_transform_3d_input(dtype):
|
||||
device = "cuda"
|
||||
|
||||
if dtype == torch.bfloat16:
|
||||
rtol, atol = 1e-2, 5e-2
|
||||
else:
|
||||
rtol, atol = 3e-3, 5e-3
|
||||
|
||||
torch.random.manual_seed(0)
|
||||
|
||||
x = torch.randn(4, 8, 256, device=device, dtype=dtype)
|
||||
scale = 1.0 / math.sqrt(256)
|
||||
|
||||
out = hadamard_transform(x, scale=scale)
|
||||
assert out.shape == x.shape
|
||||
|
||||
out_ref = hadamard_transform_ref(x.detach().clone().float(), scale=scale)
|
||||
torch.testing.assert_close(out.float(), out_ref, rtol=rtol, atol=atol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
def test_hadamard_transform_scale_one(dtype):
|
||||
device = "cuda"
|
||||
|
||||
if dtype == torch.bfloat16:
|
||||
rtol, atol = 1e-2, 5e-2
|
||||
else:
|
||||
rtol, atol = 3e-3, 5e-3
|
||||
|
||||
torch.random.manual_seed(0)
|
||||
|
||||
x = torch.randn(8, 64, device=device, dtype=dtype)
|
||||
|
||||
out = hadamard_transform(x, scale=1.0)
|
||||
out_ref = hadamard_transform_ref(x.detach().clone().float(), scale=1.0)
|
||||
|
||||
torch.testing.assert_close(out.float(), out_ref, rtol=rtol, atol=atol)
|
||||
|
||||
|
||||
# Test dimensions for M×N variants: dim = M * N where N = 2^k.
|
||||
# M = 12/20/28/40 are the non-power-of-2 Hadamard sizes registered in
|
||||
# python/sglang/jit_kernel/hadamard.py (Hadamard12NKernel, ..., Hadamard40NKernel).
|
||||
# range(2,9) gives N = 4,8,...,256 so dims cover a practical range.
|
||||
_12N_DIMS = [12 * (2**k) for k in range(2, 9)] # 48, 96, ... , 3072
|
||||
_20N_DIMS = [20 * (2**k) for k in range(2, 9)] # 80, 160, ... , 5120
|
||||
_28N_DIMS = [28 * (2**k) for k in range(2, 9)] # 112, 224, ... , 7168
|
||||
_40N_DIMS = [40 * (2**k) for k in range(2, 9)] # 160, 320, ... , 10240
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
|
||||
@pytest.mark.parametrize("dim", _12N_DIMS)
|
||||
def test_hadamard_transform_12n(dim, dtype):
|
||||
device = "cuda"
|
||||
|
||||
if dtype == torch.float32:
|
||||
rtol, atol = 3e-4, 3e-3
|
||||
elif dtype == torch.bfloat16:
|
||||
rtol, atol = 1e-2, 5e-2
|
||||
else:
|
||||
rtol, atol = 3e-3, 5e-3
|
||||
|
||||
torch.random.manual_seed(0)
|
||||
batch_size = 15
|
||||
|
||||
x = torch.randn(batch_size, dim, device=device, dtype=dtype)
|
||||
scale = 1.0 / math.sqrt(dim)
|
||||
|
||||
out = hadamard_transform_12n(x, scale=scale)
|
||||
out_ref = hadamard_transform_mn_ref(x.detach().clone().float(), 12, scale=scale)
|
||||
|
||||
torch.testing.assert_close(out.float(), out_ref, rtol=rtol, atol=atol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
|
||||
@pytest.mark.parametrize("dim", _20N_DIMS)
|
||||
def test_hadamard_transform_20n(dim, dtype):
|
||||
device = "cuda"
|
||||
|
||||
if dtype == torch.float32:
|
||||
rtol, atol = 3e-4, 3e-3
|
||||
elif dtype == torch.bfloat16:
|
||||
rtol, atol = 1e-2, 5e-2
|
||||
else:
|
||||
rtol, atol = 3e-3, 5e-3
|
||||
|
||||
torch.random.manual_seed(0)
|
||||
batch_size = 15
|
||||
|
||||
x = torch.randn(batch_size, dim, device=device, dtype=dtype)
|
||||
scale = 1.0 / math.sqrt(dim)
|
||||
|
||||
out = hadamard_transform_20n(x, scale=scale)
|
||||
out_ref = hadamard_transform_mn_ref(x.detach().clone().float(), 20, scale=scale)
|
||||
|
||||
torch.testing.assert_close(out.float(), out_ref, rtol=rtol, atol=atol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
|
||||
@pytest.mark.parametrize("dim", _28N_DIMS)
|
||||
def test_hadamard_transform_28n(dim, dtype):
|
||||
device = "cuda"
|
||||
|
||||
if dtype == torch.float32:
|
||||
rtol, atol = 3e-4, 3e-3
|
||||
elif dtype == torch.bfloat16:
|
||||
rtol, atol = 1e-2, 5e-2
|
||||
else:
|
||||
rtol, atol = 3e-3, 5e-3
|
||||
|
||||
torch.random.manual_seed(0)
|
||||
batch_size = 15
|
||||
|
||||
x = torch.randn(batch_size, dim, device=device, dtype=dtype)
|
||||
scale = 1.0 / math.sqrt(dim)
|
||||
|
||||
out = hadamard_transform_28n(x, scale=scale)
|
||||
out_ref = hadamard_transform_mn_ref(x.detach().clone().float(), 28, scale=scale)
|
||||
|
||||
torch.testing.assert_close(out.float(), out_ref, rtol=rtol, atol=atol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
|
||||
@pytest.mark.parametrize("dim", _40N_DIMS)
|
||||
def test_hadamard_transform_40n(dim, dtype):
|
||||
device = "cuda"
|
||||
|
||||
if dtype == torch.float32:
|
||||
rtol, atol = 3e-4, 3e-3
|
||||
elif dtype == torch.bfloat16:
|
||||
rtol, atol = 1e-2, 5e-2
|
||||
else:
|
||||
rtol, atol = 3e-3, 5e-3
|
||||
|
||||
torch.random.manual_seed(0)
|
||||
batch_size = 15
|
||||
|
||||
x = torch.randn(batch_size, dim, device=device, dtype=dtype)
|
||||
scale = 1.0 / math.sqrt(dim)
|
||||
|
||||
out = hadamard_transform_40n(x, scale=scale)
|
||||
out_ref = hadamard_transform_mn_ref(x.detach().clone().float(), 40, scale=scale)
|
||||
|
||||
torch.testing.assert_close(out.float(), out_ref, rtol=rtol, atol=atol)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -0,0 +1,247 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, MLATokenToKVPool
|
||||
from sglang.srt.mem_cache.memory_pool_host import (
|
||||
ALLOC_MEMORY_FUNCS,
|
||||
MHATokenToKVPoolHost,
|
||||
MLATokenToKVPoolHost,
|
||||
alloc_with_pin_memory,
|
||||
)
|
||||
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not torch.cuda.is_available()
|
||||
or is_npu()
|
||||
or is_xpu()
|
||||
or not (is_cuda() or is_hip()),
|
||||
reason="HiCache JIT tests require CUDA/ROCm.",
|
||||
)
|
||||
|
||||
DEVICE = "cuda"
|
||||
PAGE_SIZE = 1 if is_hip() else 16
|
||||
NUM_LAYERS = 2
|
||||
POOL_SIZE = PAGE_SIZE * 8
|
||||
MHA_ELEMENT_DIMS = [128, 256, 512, 1024]
|
||||
MLA_ELEMENT_DIMS = [576]
|
||||
LAYOUTS = ["layer_first", "page_first"]
|
||||
|
||||
|
||||
def _token_indices_for_pages(
|
||||
pages: torch.Tensor, page_size: int = PAGE_SIZE, device: str = DEVICE
|
||||
) -> torch.Tensor:
|
||||
parts = [
|
||||
torch.arange(
|
||||
int(page) * page_size,
|
||||
(int(page) + 1) * page_size,
|
||||
device=device,
|
||||
dtype=torch.int64,
|
||||
)
|
||||
for page in pages.tolist()
|
||||
]
|
||||
return torch.cat(parts, dim=0)
|
||||
|
||||
|
||||
def _pinned_host_pool(host_pool_cls, **kwargs):
|
||||
original_alloc = ALLOC_MEMORY_FUNCS[DEVICE]
|
||||
ALLOC_MEMORY_FUNCS[DEVICE] = alloc_with_pin_memory
|
||||
try:
|
||||
return host_pool_cls(
|
||||
host_to_device_ratio=2.0,
|
||||
host_size=0,
|
||||
page_size=PAGE_SIZE,
|
||||
pin_memory=True,
|
||||
device="cpu",
|
||||
**kwargs,
|
||||
)
|
||||
finally:
|
||||
ALLOC_MEMORY_FUNCS[DEVICE] = original_alloc
|
||||
|
||||
|
||||
def _copy_tensor_with_offset(tensor: torch.Tensor, offset: int) -> None:
|
||||
data = torch.arange(
|
||||
tensor.numel(), device=tensor.device, dtype=tensor.dtype
|
||||
).view_as(tensor)
|
||||
tensor.copy_(data + offset)
|
||||
|
||||
|
||||
def _run_transfer_roundtrip_mha(layout: str, element_dim: int) -> None:
|
||||
device_pool = MHATokenToKVPool(
|
||||
size=POOL_SIZE,
|
||||
page_size=PAGE_SIZE,
|
||||
head_num=element_dim // 128,
|
||||
head_dim=128,
|
||||
dtype=torch.bfloat16,
|
||||
layer_num=NUM_LAYERS,
|
||||
device=DEVICE,
|
||||
enable_memory_saver=False,
|
||||
)
|
||||
host_pool = _pinned_host_pool(
|
||||
MHATokenToKVPoolHost,
|
||||
device_pool=device_pool,
|
||||
layout=layout,
|
||||
)
|
||||
assert (
|
||||
host_pool.can_use_jit
|
||||
), f"Expected JIT HiCache kernel for MHA dim={element_dim}"
|
||||
|
||||
for layer_id in range(NUM_LAYERS):
|
||||
_copy_tensor_with_offset(device_pool.k_buffer[layer_id], layer_id)
|
||||
_copy_tensor_with_offset(device_pool.v_buffer[layer_id], layer_id + 100)
|
||||
|
||||
device_pages = torch.tensor([1, 2, 3], device=DEVICE, dtype=torch.int64)
|
||||
host_pages = torch.tensor([0, 1, 2], device=DEVICE, dtype=torch.int64)
|
||||
device_indices = _token_indices_for_pages(device_pages)
|
||||
host_indices = _token_indices_for_pages(host_pages)
|
||||
|
||||
host_pool.backup_from_device_all_layer(
|
||||
device_pool, host_indices, device_indices, "kernel"
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
for layer_id in range(NUM_LAYERS):
|
||||
for host_page, device_page in zip(host_pages.tolist(), device_pages.tolist()):
|
||||
host_start = host_page * PAGE_SIZE
|
||||
device_start = device_page * PAGE_SIZE
|
||||
assert torch.equal(
|
||||
host_pool.k_data_refs[layer_id][
|
||||
host_start : host_start + PAGE_SIZE
|
||||
].cpu(),
|
||||
device_pool.k_buffer[layer_id][
|
||||
device_start : device_start + PAGE_SIZE
|
||||
].cpu(),
|
||||
)
|
||||
assert torch.equal(
|
||||
host_pool.v_data_refs[layer_id][
|
||||
host_start : host_start + PAGE_SIZE
|
||||
].cpu(),
|
||||
device_pool.v_buffer[layer_id][
|
||||
device_start : device_start + PAGE_SIZE
|
||||
].cpu(),
|
||||
)
|
||||
|
||||
for layer_id in range(NUM_LAYERS):
|
||||
device_pool.k_buffer[layer_id].zero_()
|
||||
device_pool.v_buffer[layer_id].zero_()
|
||||
|
||||
load_pages = torch.tensor([4, 5, 6], device=DEVICE, dtype=torch.int64)
|
||||
load_indices = _token_indices_for_pages(load_pages)
|
||||
for layer_id in range(NUM_LAYERS):
|
||||
host_pool.load_to_device_per_layer(
|
||||
device_pool, host_indices, load_indices, layer_id, "kernel"
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
for layer_id in range(NUM_LAYERS):
|
||||
for host_page, device_page in zip(host_pages.tolist(), load_pages.tolist()):
|
||||
host_start = host_page * PAGE_SIZE
|
||||
device_start = device_page * PAGE_SIZE
|
||||
assert torch.equal(
|
||||
device_pool.k_buffer[layer_id][
|
||||
device_start : device_start + PAGE_SIZE
|
||||
].cpu(),
|
||||
host_pool.k_data_refs[layer_id][
|
||||
host_start : host_start + PAGE_SIZE
|
||||
].cpu(),
|
||||
)
|
||||
assert torch.equal(
|
||||
device_pool.v_buffer[layer_id][
|
||||
device_start : device_start + PAGE_SIZE
|
||||
].cpu(),
|
||||
host_pool.v_data_refs[layer_id][
|
||||
host_start : host_start + PAGE_SIZE
|
||||
].cpu(),
|
||||
)
|
||||
|
||||
|
||||
def _run_transfer_roundtrip_mla(layout: str, element_dim: int) -> None:
|
||||
device_pool = MLATokenToKVPool(
|
||||
size=POOL_SIZE,
|
||||
page_size=PAGE_SIZE,
|
||||
kv_lora_rank=element_dim - 64,
|
||||
qk_rope_head_dim=64,
|
||||
dtype=torch.bfloat16,
|
||||
layer_num=NUM_LAYERS,
|
||||
device=DEVICE,
|
||||
enable_memory_saver=False,
|
||||
)
|
||||
host_pool = _pinned_host_pool(
|
||||
MLATokenToKVPoolHost,
|
||||
device_pool=device_pool,
|
||||
layout=layout,
|
||||
)
|
||||
assert (
|
||||
host_pool.can_use_jit
|
||||
), f"Expected JIT HiCache kernel for MLA dim={element_dim}"
|
||||
|
||||
for layer_id in range(NUM_LAYERS):
|
||||
_copy_tensor_with_offset(device_pool.kv_buffer[layer_id], layer_id)
|
||||
|
||||
device_pages = torch.tensor([1, 2, 3], device=DEVICE, dtype=torch.int64)
|
||||
host_pages = torch.tensor([0, 1, 2], device=DEVICE, dtype=torch.int64)
|
||||
device_indices = _token_indices_for_pages(device_pages)
|
||||
host_indices = _token_indices_for_pages(host_pages)
|
||||
|
||||
host_pool.backup_from_device_all_layer(
|
||||
device_pool, host_indices, device_indices, "kernel"
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
for layer_id in range(NUM_LAYERS):
|
||||
for host_page, device_page in zip(host_pages.tolist(), device_pages.tolist()):
|
||||
host_start = host_page * PAGE_SIZE
|
||||
device_start = device_page * PAGE_SIZE
|
||||
assert torch.equal(
|
||||
host_pool.data_refs[layer_id][
|
||||
host_start : host_start + PAGE_SIZE
|
||||
].cpu(),
|
||||
device_pool.kv_buffer[layer_id][
|
||||
device_start : device_start + PAGE_SIZE
|
||||
].cpu(),
|
||||
)
|
||||
|
||||
for layer_id in range(NUM_LAYERS):
|
||||
device_pool.kv_buffer[layer_id].zero_()
|
||||
|
||||
load_pages = torch.tensor([4, 5, 6], device=DEVICE, dtype=torch.int64)
|
||||
load_indices = _token_indices_for_pages(load_pages)
|
||||
for layer_id in range(NUM_LAYERS):
|
||||
host_pool.load_to_device_per_layer(
|
||||
device_pool, host_indices, load_indices, layer_id, "kernel"
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
for layer_id in range(NUM_LAYERS):
|
||||
for host_page, device_page in zip(host_pages.tolist(), load_pages.tolist()):
|
||||
host_start = host_page * PAGE_SIZE
|
||||
device_start = device_page * PAGE_SIZE
|
||||
assert torch.equal(
|
||||
device_pool.kv_buffer[layer_id][
|
||||
device_start : device_start + PAGE_SIZE
|
||||
].cpu(),
|
||||
host_pool.data_refs[layer_id][
|
||||
host_start : host_start + PAGE_SIZE
|
||||
].cpu(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("layout", LAYOUTS)
|
||||
@pytest.mark.parametrize("element_dim", MHA_ELEMENT_DIMS)
|
||||
def test_hicache_transfer_mha(layout: str, element_dim: int) -> None:
|
||||
_run_transfer_roundtrip_mha(layout, element_dim)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("layout", LAYOUTS)
|
||||
@pytest.mark.parametrize("element_dim", MLA_ELEMENT_DIMS)
|
||||
def test_hicache_transfer_mla(layout: str, element_dim: int) -> None:
|
||||
_run_transfer_roundtrip_mla(layout, element_dim)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,426 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.hisparse import (
|
||||
load_cache_to_device_buffer_dsv4_mla,
|
||||
load_cache_to_device_buffer_mla,
|
||||
transfer_cache_dsv4_mla,
|
||||
)
|
||||
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not torch.cuda.is_available()
|
||||
or is_npu()
|
||||
or is_xpu()
|
||||
or not (is_cuda() or is_hip()),
|
||||
reason="HiSparse JIT tests require CUDA/ROCm.",
|
||||
)
|
||||
|
||||
DEVICE = "cuda"
|
||||
DTYPE = torch.float32
|
||||
KV_DIM = 8
|
||||
HOT_BUFFER_SIZE = 4
|
||||
PADDED_BUFFER_SIZE = HOT_BUFFER_SIZE + 1
|
||||
HOST_CACHE_SIZE = 16
|
||||
DEVICE_CACHE_SIZE = 16
|
||||
ITEM_SIZE_BYTES = KV_DIM * torch.empty((), dtype=DTYPE).element_size()
|
||||
DSV4_PAGE_SIZE = 64
|
||||
DSV4_VALUE_BYTES = 576
|
||||
DSV4_SCALE_BYTES = 8
|
||||
DSV4_ITEM_BYTES = DSV4_VALUE_BYTES + DSV4_SCALE_BYTES
|
||||
DSV4_PAGE_BYTES = ((DSV4_ITEM_BYTES * DSV4_PAGE_SIZE + 575) // 576) * 576
|
||||
DSV4_SCALE_OFFSET = DSV4_VALUE_BYTES * DSV4_PAGE_SIZE
|
||||
|
||||
|
||||
def _host_cache() -> torch.Tensor:
|
||||
host_cache = torch.empty(
|
||||
(HOST_CACHE_SIZE, 1, KV_DIM), dtype=DTYPE, device="cpu", pin_memory=True
|
||||
)
|
||||
host_cache.copy_(torch.arange(host_cache.numel(), dtype=DTYPE).view_as(host_cache))
|
||||
return host_cache
|
||||
|
||||
|
||||
def _dsv4_token_pattern(seed: int) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
value = (
|
||||
(torch.arange(DSV4_VALUE_BYTES, dtype=torch.int16) + seed)
|
||||
.remainder(256)
|
||||
.to(torch.uint8)
|
||||
)
|
||||
scale = (
|
||||
(torch.arange(DSV4_SCALE_BYTES, dtype=torch.int16) + seed + 17)
|
||||
.remainder(256)
|
||||
.to(torch.uint8)
|
||||
)
|
||||
return value, scale
|
||||
|
||||
|
||||
def _write_dsv4_token(cache: torch.Tensor, loc: int, seed: int) -> None:
|
||||
page = loc // DSV4_PAGE_SIZE
|
||||
offset = loc % DSV4_PAGE_SIZE
|
||||
value, scale = _dsv4_token_pattern(seed)
|
||||
cache[page, offset * DSV4_VALUE_BYTES : (offset + 1) * DSV4_VALUE_BYTES].copy_(
|
||||
value.to(cache.device)
|
||||
)
|
||||
scale_start = DSV4_SCALE_OFFSET + offset * DSV4_SCALE_BYTES
|
||||
cache[page, scale_start : scale_start + DSV4_SCALE_BYTES].copy_(
|
||||
scale.to(cache.device)
|
||||
)
|
||||
|
||||
|
||||
def _read_dsv4_token(cache: torch.Tensor, loc: int) -> torch.Tensor:
|
||||
page = loc // DSV4_PAGE_SIZE
|
||||
offset = loc % DSV4_PAGE_SIZE
|
||||
value = cache[page, offset * DSV4_VALUE_BYTES : (offset + 1) * DSV4_VALUE_BYTES]
|
||||
scale_start = DSV4_SCALE_OFFSET + offset * DSV4_SCALE_BYTES
|
||||
scale = cache[page, scale_start : scale_start + DSV4_SCALE_BYTES]
|
||||
return torch.cat([value, scale])
|
||||
|
||||
|
||||
def _dsv4_ptrs(cache: torch.Tensor) -> torch.Tensor:
|
||||
return torch.tensor([cache.data_ptr()], dtype=torch.uint64, device=DEVICE)
|
||||
|
||||
|
||||
def _run_kernel(
|
||||
*,
|
||||
top_k_tokens: torch.Tensor,
|
||||
device_buffer_tokens: torch.Tensor,
|
||||
host_cache_locs: torch.Tensor,
|
||||
device_buffer_locs: torch.Tensor,
|
||||
host_cache: torch.Tensor,
|
||||
device_buffer: torch.Tensor,
|
||||
lru_slots: torch.Tensor,
|
||||
seq_len: int | None = None,
|
||||
seq_lens: torch.Tensor | None = None,
|
||||
seq_lens_dtype: torch.dtype = torch.int32,
|
||||
req_pool_indices: torch.Tensor | None = None,
|
||||
num_real_reqs: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
batch_size = top_k_tokens.shape[0]
|
||||
if req_pool_indices is None:
|
||||
req_pool_indices = torch.arange(batch_size, dtype=torch.int64, device=DEVICE)
|
||||
if seq_lens is None:
|
||||
seq_lens = torch.full(
|
||||
(batch_size,), seq_len, dtype=seq_lens_dtype, device=DEVICE
|
||||
)
|
||||
if num_real_reqs is None:
|
||||
num_real_reqs = batch_size
|
||||
|
||||
out = torch.full_like(top_k_tokens, -1)
|
||||
load_cache_to_device_buffer_mla(
|
||||
top_k_tokens=top_k_tokens,
|
||||
device_buffer_tokens=device_buffer_tokens,
|
||||
host_cache_locs=host_cache_locs,
|
||||
device_buffer_locs=device_buffer_locs,
|
||||
host_cache=host_cache,
|
||||
device_buffer=device_buffer,
|
||||
top_k_device_locs=out,
|
||||
req_pool_indices=req_pool_indices,
|
||||
seq_lens=seq_lens,
|
||||
lru_slots=lru_slots,
|
||||
item_size_bytes=ITEM_SIZE_BYTES,
|
||||
num_top_k=top_k_tokens.shape[1],
|
||||
hot_buffer_size=HOT_BUFFER_SIZE,
|
||||
page_size=1,
|
||||
block_size=256,
|
||||
num_real_reqs=torch.tensor([num_real_reqs], dtype=torch.int32, device=DEVICE),
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
return out
|
||||
|
||||
|
||||
def _make_state(
|
||||
device_buffer_locs_rows: list[list[int]],
|
||||
device_buffer_tokens_rows: list[list[int]],
|
||||
newest_tokens: list[int],
|
||||
):
|
||||
host_cache = _host_cache()
|
||||
device_buffer = torch.full(
|
||||
(DEVICE_CACHE_SIZE, 1, KV_DIM), -1, dtype=DTYPE, device=DEVICE
|
||||
)
|
||||
device_buffer_locs = torch.tensor(
|
||||
device_buffer_locs_rows, dtype=torch.int32, device=DEVICE
|
||||
)
|
||||
device_buffer_tokens = torch.tensor(
|
||||
device_buffer_tokens_rows, dtype=torch.int32, device=DEVICE
|
||||
)
|
||||
lru_slots = (
|
||||
torch.arange(HOT_BUFFER_SIZE, dtype=torch.int16, device=DEVICE)
|
||||
.view(1, -1)
|
||||
.repeat(device_buffer_locs.shape[0], 1)
|
||||
)
|
||||
host_cache_locs = (
|
||||
torch.arange(HOST_CACHE_SIZE, dtype=torch.int64, device=DEVICE)
|
||||
.view(1, -1)
|
||||
.repeat(device_buffer_locs.shape[0], 1)
|
||||
)
|
||||
|
||||
# Slots 0..3 participate in LRU; slot 4 is the reserved newest slot.
|
||||
for rid, newest_token in enumerate(newest_tokens):
|
||||
for slot, token in enumerate(device_buffer_tokens_rows[rid][:HOT_BUFFER_SIZE]):
|
||||
if token >= 0:
|
||||
device_buffer[device_buffer_locs[rid, slot]].copy_(
|
||||
host_cache[token].to(DEVICE, non_blocking=True)
|
||||
)
|
||||
device_buffer[device_buffer_locs[rid, HOT_BUFFER_SIZE]].copy_(
|
||||
host_cache[newest_token].to(DEVICE, non_blocking=True)
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
return {
|
||||
"host_cache": host_cache,
|
||||
"device_buffer": device_buffer,
|
||||
"device_buffer_locs": device_buffer_locs,
|
||||
"device_buffer_tokens": device_buffer_tokens,
|
||||
"lru_slots": lru_slots,
|
||||
"host_cache_locs": host_cache_locs,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_hip(), reason="DSV4 paged-layout HiSparse test is CUDA-only.")
|
||||
def test_transfer_cache_dsv4_mla_copies_paged_token() -> None:
|
||||
src_cache = torch.zeros((2, DSV4_PAGE_BYTES), dtype=torch.uint8, device=DEVICE)
|
||||
dst_cache = torch.zeros(
|
||||
(2, DSV4_PAGE_BYTES), dtype=torch.uint8, device="cpu", pin_memory=True
|
||||
)
|
||||
src_loc = DSV4_PAGE_SIZE + 6
|
||||
dst_loc = DSV4_PAGE_SIZE + 1
|
||||
_write_dsv4_token(src_cache, src_loc, seed=41)
|
||||
|
||||
transfer_cache_dsv4_mla(
|
||||
src_ptrs=_dsv4_ptrs(src_cache),
|
||||
dst_ptrs=_dsv4_ptrs(dst_cache),
|
||||
src_indices=torch.tensor([src_loc], dtype=torch.int64, device=DEVICE),
|
||||
dst_indices=torch.tensor([dst_loc], dtype=torch.int64, device=DEVICE),
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
assert torch.equal(
|
||||
_read_dsv4_token(dst_cache, dst_loc).to(DEVICE),
|
||||
_read_dsv4_token(src_cache, src_loc),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_hip(), reason="DSV4 paged-layout HiSparse test is CUDA-only.")
|
||||
def test_dsv4_swap_in_reads_paged_host_layout() -> None:
|
||||
host_cache = torch.zeros(
|
||||
(2, DSV4_PAGE_BYTES), dtype=torch.uint8, device="cpu", pin_memory=True
|
||||
)
|
||||
device_buffer = torch.zeros((2, DSV4_PAGE_BYTES), dtype=torch.uint8, device=DEVICE)
|
||||
host_loc = DSV4_PAGE_SIZE + 1
|
||||
swap_loc = DSV4_PAGE_SIZE + 12
|
||||
_write_dsv4_token(host_cache, host_loc, seed=41)
|
||||
|
||||
top_k_tokens = torch.tensor([[3]], dtype=torch.int32, device=DEVICE)
|
||||
device_buffer_tokens = torch.full(
|
||||
(1, PADDED_BUFFER_SIZE), -1, dtype=torch.int32, device=DEVICE
|
||||
)
|
||||
host_cache_locs = torch.zeros((1, 8), dtype=torch.int64, device=DEVICE)
|
||||
host_cache_locs[0, 3] = host_loc
|
||||
device_buffer_locs = torch.tensor(
|
||||
[[swap_loc, swap_loc + 1, swap_loc + 2, swap_loc + 3, swap_loc + 4]],
|
||||
dtype=torch.int32,
|
||||
device=DEVICE,
|
||||
)
|
||||
lru_slots = torch.arange(HOT_BUFFER_SIZE, dtype=torch.int16, device=DEVICE).view(
|
||||
1, -1
|
||||
)
|
||||
out = torch.full_like(top_k_tokens, -1)
|
||||
|
||||
load_cache_to_device_buffer_dsv4_mla(
|
||||
top_k_tokens=top_k_tokens,
|
||||
device_buffer_tokens=device_buffer_tokens,
|
||||
host_cache_locs=host_cache_locs,
|
||||
device_buffer_locs=device_buffer_locs,
|
||||
host_cache=host_cache,
|
||||
device_buffer=device_buffer,
|
||||
top_k_device_locs=out,
|
||||
req_pool_indices=torch.tensor([0], dtype=torch.int64, device=DEVICE),
|
||||
seq_lens=torch.tensor([8], dtype=torch.int32, device=DEVICE),
|
||||
lru_slots=lru_slots,
|
||||
item_size_bytes=DSV4_ITEM_BYTES,
|
||||
num_top_k=1,
|
||||
hot_buffer_size=HOT_BUFFER_SIZE,
|
||||
page_size=1,
|
||||
block_size=256,
|
||||
num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE),
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
assert out.item() == swap_loc
|
||||
assert torch.equal(
|
||||
_read_dsv4_token(device_buffer, swap_loc),
|
||||
_read_dsv4_token(host_cache, host_loc).to(DEVICE),
|
||||
)
|
||||
|
||||
|
||||
def _long_case():
|
||||
# One-request baseline used by the stateful cases below:
|
||||
# req 0 LRU slots : [0, 1, 2, 3]
|
||||
# req 0 cached tokens : slot0->1, slot1->4, slot2->2, slot3->5
|
||||
# req 0 physical locs : slot0->9, slot1->7, slot2->3, slot3->5
|
||||
# req 0 newest slot : slot4/newest -> token 7 at physical loc 11
|
||||
return _make_state([[9, 7, 3, 5, 11]], [[1, 4, 2, 5, -1]], [7])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seq_lens_dtype", [torch.int32, torch.int64])
|
||||
def test_load_cache_to_device_buffer_fast_path(seq_lens_dtype: torch.dtype) -> None:
|
||||
host_cache = _host_cache()
|
||||
device_buffer = torch.arange(
|
||||
DEVICE_CACHE_SIZE * KV_DIM, dtype=DTYPE, device=DEVICE
|
||||
).view(DEVICE_CACHE_SIZE, 1, KV_DIM)
|
||||
device_buffer_before = device_buffer.clone()
|
||||
device_buffer_locs = torch.tensor(
|
||||
[[13, 9, 5, 1, 15]], dtype=torch.int32, device=DEVICE
|
||||
)
|
||||
device_buffer_tokens = torch.tensor(
|
||||
[[10, 11, 12, 13, -1]], dtype=torch.int32, device=DEVICE
|
||||
)
|
||||
device_buffer_tokens_before = device_buffer_tokens.clone()
|
||||
lru_slots = torch.tensor([[0, 1, 2, 3]], dtype=torch.int16, device=DEVICE)
|
||||
lru_slots_before = lru_slots.clone()
|
||||
|
||||
# Short-sequence layout:
|
||||
# token position 0 -> physical loc 13
|
||||
# token position 1 -> physical loc 9
|
||||
# token position 2 -> physical loc 5
|
||||
#
|
||||
# seq_len <= HOT_BUFFER_SIZE should skip host loads and LRU mutations,
|
||||
# so top_k_tokens acts like direct indexing into device_buffer_locs.
|
||||
out = _run_kernel(
|
||||
top_k_tokens=torch.tensor([[2, 0, 1]], dtype=torch.int32, device=DEVICE),
|
||||
device_buffer_tokens=device_buffer_tokens,
|
||||
host_cache_locs=torch.arange(
|
||||
HOST_CACHE_SIZE, dtype=torch.int64, device=DEVICE
|
||||
).view(1, -1),
|
||||
device_buffer_locs=device_buffer_locs,
|
||||
host_cache=host_cache,
|
||||
device_buffer=device_buffer,
|
||||
lru_slots=lru_slots,
|
||||
seq_len=3,
|
||||
seq_lens_dtype=seq_lens_dtype,
|
||||
)
|
||||
|
||||
assert torch.equal(out.cpu(), torch.tensor([[5, 13, 9]], dtype=torch.int32))
|
||||
assert torch.equal(device_buffer_tokens.cpu(), device_buffer_tokens_before.cpu())
|
||||
assert torch.equal(lru_slots.cpu(), lru_slots_before.cpu())
|
||||
assert torch.equal(device_buffer.cpu(), device_buffer_before.cpu())
|
||||
|
||||
|
||||
def test_load_cache_to_device_buffer_hits_newest_and_updates_lru() -> None:
|
||||
state = _long_case()
|
||||
|
||||
# Query [4, 2, 7]:
|
||||
# 4 hits slot1 -> loc 7
|
||||
# 2 hits slot2 -> loc 3
|
||||
# 7 is the newest token -> reserved newest loc 11
|
||||
#
|
||||
# Hits move to the MRU tail, so [0, 1, 2, 3] becomes [0, 3, 1, 2].
|
||||
out = _run_kernel(
|
||||
top_k_tokens=torch.tensor([[4, 2, 7]], dtype=torch.int32, device=DEVICE),
|
||||
seq_len=8,
|
||||
**state,
|
||||
)
|
||||
|
||||
assert torch.equal(out.cpu(), torch.tensor([[7, 3, 11]], dtype=torch.int32))
|
||||
assert torch.equal(
|
||||
state["device_buffer_tokens"].cpu(),
|
||||
torch.tensor([[1, 4, 2, 5, -1]], dtype=torch.int32),
|
||||
)
|
||||
assert torch.equal(
|
||||
state["lru_slots"].cpu(), torch.tensor([[0, 3, 1, 2]], dtype=torch.int16)
|
||||
)
|
||||
|
||||
|
||||
def test_load_cache_to_device_buffer_miss_uses_updated_lru_slot() -> None:
|
||||
state = _long_case()
|
||||
|
||||
# Step 1: touch tokens [4, 2], so LRU becomes [0, 3, 1, 2].
|
||||
# Step 2: query token 6, which is a miss.
|
||||
# The kernel should reuse the new LRU head slot0, whose physical loc is 9.
|
||||
# This round has no regular hits, so the freshly loaded miss slot ends up at the tail.
|
||||
_run_kernel(
|
||||
top_k_tokens=torch.tensor([[4, 2]], dtype=torch.int32, device=DEVICE),
|
||||
seq_len=8,
|
||||
**state,
|
||||
)
|
||||
out = _run_kernel(
|
||||
top_k_tokens=torch.tensor([[6]], dtype=torch.int32, device=DEVICE),
|
||||
seq_len=8,
|
||||
**state,
|
||||
)
|
||||
|
||||
assert torch.equal(out.cpu(), torch.tensor([[9]], dtype=torch.int32))
|
||||
assert torch.equal(
|
||||
state["device_buffer_tokens"].cpu(),
|
||||
torch.tensor([[6, 4, 2, 5, -1]], dtype=torch.int32),
|
||||
)
|
||||
assert torch.equal(
|
||||
state["lru_slots"].cpu(), torch.tensor([[3, 1, 2, 0]], dtype=torch.int16)
|
||||
)
|
||||
assert torch.equal(state["device_buffer"][9].cpu(), state["host_cache"][6])
|
||||
|
||||
|
||||
def test_load_cache_to_device_buffer_batched_with_padding() -> None:
|
||||
state = _make_state(
|
||||
[
|
||||
[9, 7, 3, 5, 11],
|
||||
[12, 10, 8, 6, 14],
|
||||
[15, 4, 2, 1, 13],
|
||||
],
|
||||
[
|
||||
[1, 4, 2, 5, -1],
|
||||
[0, 1, 2, 3, -1],
|
||||
[9, 8, 7, 6, -1],
|
||||
],
|
||||
[7, 4, 5],
|
||||
)
|
||||
padded_tokens_before = state["device_buffer_tokens"][2].clone()
|
||||
padded_lru_before = state["lru_slots"][2].clone()
|
||||
|
||||
# req 0: long path
|
||||
# cached tokens/locs : 1@9, 4@7, 2@3, 5@5, newest 7@11
|
||||
# query [4, 6, 7] : hit loc 7, miss into slot0/loc 9, newest loc 11
|
||||
# LRU update : remaining evictables [2, 3], then miss [0], then hit [1]
|
||||
# : [0, 1, 2, 3] -> [2, 3, 0, 1]
|
||||
#
|
||||
# req 1: fast path
|
||||
# seq_len = 3 <= HOT_BUFFER_SIZE, so [2, 1, 0] maps directly to locs [8, 10, 12]
|
||||
#
|
||||
# req 2: padded block
|
||||
# num_real_reqs = 2 means this row must be ignored entirely.
|
||||
out = _run_kernel(
|
||||
top_k_tokens=torch.tensor(
|
||||
[[4, 6, 7], [2, 1, 0], [9, 8, 7]], dtype=torch.int32, device=DEVICE
|
||||
),
|
||||
seq_lens=torch.tensor([8, 3, 8], dtype=torch.int32, device=DEVICE),
|
||||
num_real_reqs=2,
|
||||
**state,
|
||||
)
|
||||
|
||||
assert torch.equal(
|
||||
out.cpu(),
|
||||
torch.tensor([[7, 9, 11], [8, 10, 12], [-1, -1, -1]], dtype=torch.int32),
|
||||
)
|
||||
assert torch.equal(
|
||||
state["device_buffer_tokens"][:2].cpu(),
|
||||
torch.tensor([[6, 4, 2, 5, -1], [0, 1, 2, 3, -1]], dtype=torch.int32),
|
||||
)
|
||||
assert torch.equal(
|
||||
state["lru_slots"][:2].cpu(),
|
||||
torch.tensor([[2, 3, 0, 1], [0, 1, 2, 3]], dtype=torch.int16),
|
||||
)
|
||||
assert torch.equal(
|
||||
state["device_buffer_tokens"][2].cpu(), padded_tokens_before.cpu()
|
||||
)
|
||||
assert torch.equal(state["lru_slots"][2].cpu(), padded_lru_before.cpu())
|
||||
assert torch.equal(state["device_buffer"][9].cpu(), state["host_cache"][6])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,104 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.mla_kv_pack_quantize_fp8 import mla_kv_pack_quantize_fp8
|
||||
from sglang.jit_kernel.utils import get_ci_test_range
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=60, suite="base-b-kernel-unit-1-gpu-large")
|
||||
|
||||
DEVICE = "cuda"
|
||||
|
||||
SHAPES = get_ci_test_range(
|
||||
[(128, 64, 128), (64, 32, 64)],
|
||||
[(128, 64, 128)],
|
||||
)
|
||||
NUM_HEADS = get_ci_test_range([8, 16, 32, 64], [16, 32])
|
||||
BATCH_SIZES = get_ci_test_range(
|
||||
[1, 4, 17, 64, 257, 1024, 4096, 16384],
|
||||
[1, 64, 1024, 16384],
|
||||
)
|
||||
|
||||
|
||||
def _ref(k_nope, k_pe, v, k_scale_inv, v_scale_inv, fp8_dtype):
|
||||
s, h, qk_nope = k_nope.shape
|
||||
qk_rope = k_pe.shape[-1]
|
||||
v_head = v.shape[-1]
|
||||
if k_pe.dim() == 3:
|
||||
k_pe = k_pe.squeeze(1)
|
||||
|
||||
k_bf16 = torch.empty(
|
||||
(s, h, qk_nope + qk_rope), dtype=k_nope.dtype, device=k_nope.device
|
||||
)
|
||||
k_bf16[..., :qk_nope] = k_nope
|
||||
k_bf16[..., qk_nope:] = k_pe.unsqueeze(1).expand(-1, h, -1)
|
||||
|
||||
k_fp8 = (k_bf16.float() * k_scale_inv).to(fp8_dtype)
|
||||
v_fp8 = (v.float() * v_scale_inv).to(fp8_dtype)
|
||||
return k_fp8, v_fp8
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.parametrize("shape", SHAPES)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("batch_size", BATCH_SIZES)
|
||||
def test_correctness(dtype, shape, num_heads, batch_size):
|
||||
qk_nope, qk_rope, v_head = shape
|
||||
|
||||
torch.manual_seed(0)
|
||||
k_nope = torch.randn((batch_size, num_heads, qk_nope), dtype=dtype, device=DEVICE)
|
||||
k_pe = torch.randn((batch_size, 1, qk_rope), dtype=dtype, device=DEVICE)
|
||||
v = torch.randn((batch_size, num_heads, v_head), dtype=dtype, device=DEVICE)
|
||||
|
||||
k_scale_inv = 0.7
|
||||
v_scale_inv = 1.3
|
||||
|
||||
k_fp8, v_fp8 = mla_kv_pack_quantize_fp8(
|
||||
k_nope, k_pe, v, k_scale_inv=k_scale_inv, v_scale_inv=v_scale_inv
|
||||
)
|
||||
|
||||
k_ref, v_ref = _ref(k_nope, k_pe, v, k_scale_inv, v_scale_inv, torch.float8_e4m3fn)
|
||||
|
||||
torch.testing.assert_close(k_fp8.float(), k_ref.float(), rtol=1e-2, atol=0.5)
|
||||
torch.testing.assert_close(v_fp8.float(), v_ref.float(), rtol=1e-2, atol=0.5)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
def test_strided_inputs(dtype):
|
||||
s, h = 16, 32
|
||||
qk_nope, qk_rope, v_head = 128, 64, 128
|
||||
|
||||
full = torch.randn(
|
||||
(s, h, qk_nope * 2), dtype=dtype, device=DEVICE, requires_grad=False
|
||||
)
|
||||
k_nope = full[..., qk_nope:]
|
||||
assert k_nope.stride(-1) == 1
|
||||
|
||||
k_pe = torch.randn((s, 1, qk_rope), dtype=dtype, device=DEVICE)
|
||||
v = torch.randn((s, h, v_head), dtype=dtype, device=DEVICE)
|
||||
|
||||
k_fp8, v_fp8 = mla_kv_pack_quantize_fp8(k_nope, k_pe, v)
|
||||
k_ref, v_ref = _ref(k_nope, k_pe, v, 1.0, 1.0, torch.float8_e4m3fn)
|
||||
torch.testing.assert_close(k_fp8.float(), k_ref.float(), rtol=1e-2, atol=0.5)
|
||||
torch.testing.assert_close(v_fp8.float(), v_ref.float(), rtol=1e-2, atol=0.5)
|
||||
|
||||
|
||||
def test_kpe_2d_accepted():
|
||||
s, h = 8, 16
|
||||
qk_nope, qk_rope, v_head = 128, 64, 128
|
||||
dtype = torch.bfloat16
|
||||
|
||||
k_nope = torch.randn((s, h, qk_nope), dtype=dtype, device=DEVICE)
|
||||
k_pe = torch.randn((s, qk_rope), dtype=dtype, device=DEVICE)
|
||||
v = torch.randn((s, h, v_head), dtype=dtype, device=DEVICE)
|
||||
|
||||
k_fp8, v_fp8 = mla_kv_pack_quantize_fp8(k_nope, k_pe, v)
|
||||
k_ref, v_ref = _ref(k_nope, k_pe.unsqueeze(1), v, 1.0, 1.0, torch.float8_e4m3fn)
|
||||
torch.testing.assert_close(k_fp8.float(), k_ref.float(), rtol=1e-2, atol=0.5)
|
||||
torch.testing.assert_close(v_fp8.float(), v_ref.float(), rtol=1e-2, atol=0.5)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,349 @@
|
||||
import itertools
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.jit_kernel.moe_align import moe_align_block_size
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=28, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
def ceil_div(a, b):
|
||||
return (a + b - 1) // b
|
||||
|
||||
|
||||
@triton.jit
|
||||
def moe_align_block_size_stage1(
|
||||
topk_ids_ptr,
|
||||
tokens_cnts_ptr,
|
||||
num_experts: tl.constexpr,
|
||||
numel: tl.constexpr,
|
||||
tokens_per_thread: tl.constexpr,
|
||||
):
|
||||
pid = tl.program_id(0)
|
||||
start_idx = pid * tokens_per_thread
|
||||
off_c = (pid + 1) * num_experts
|
||||
|
||||
for i in range(tokens_per_thread):
|
||||
if start_idx + i < numel:
|
||||
idx = tl.load(topk_ids_ptr + start_idx + i)
|
||||
token_cnt = tl.load(tokens_cnts_ptr + off_c + idx)
|
||||
tl.store(tokens_cnts_ptr + off_c + idx, token_cnt + 1)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def moe_align_block_size_stage2(
|
||||
tokens_cnts_ptr,
|
||||
num_experts: tl.constexpr,
|
||||
):
|
||||
pid = tl.program_id(0)
|
||||
last_cnt = 0
|
||||
for i in range(1, num_experts + 1):
|
||||
token_cnt = tl.load(tokens_cnts_ptr + i * num_experts + pid)
|
||||
last_cnt = last_cnt + token_cnt
|
||||
tl.store(tokens_cnts_ptr + i * num_experts + pid, last_cnt)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def moe_align_block_size_stage3(
|
||||
total_tokens_post_pad_ptr,
|
||||
tokens_cnts_ptr,
|
||||
cumsum_ptr,
|
||||
num_experts: tl.constexpr,
|
||||
block_size: tl.constexpr,
|
||||
):
|
||||
last_cumsum = 0
|
||||
off_cnt = num_experts * num_experts
|
||||
for i in range(1, num_experts + 1):
|
||||
token_cnt = tl.load(tokens_cnts_ptr + off_cnt + i - 1)
|
||||
last_cumsum = last_cumsum + tl.cdiv(token_cnt, block_size) * block_size
|
||||
tl.store(cumsum_ptr + i, last_cumsum)
|
||||
tl.store(total_tokens_post_pad_ptr, last_cumsum)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def moe_align_block_size_stage4(
|
||||
topk_ids_ptr,
|
||||
sorted_token_ids_ptr,
|
||||
expert_ids_ptr,
|
||||
tokens_cnts_ptr,
|
||||
cumsum_ptr,
|
||||
num_experts: tl.constexpr,
|
||||
block_size: tl.constexpr,
|
||||
numel: tl.constexpr,
|
||||
tokens_per_thread: tl.constexpr,
|
||||
):
|
||||
pid = tl.program_id(0)
|
||||
start_idx = tl.load(cumsum_ptr + pid)
|
||||
end_idx = tl.load(cumsum_ptr + pid + 1)
|
||||
|
||||
for i in range(start_idx, end_idx, block_size):
|
||||
tl.store(expert_ids_ptr + i // block_size, pid)
|
||||
|
||||
start_idx = pid * tokens_per_thread
|
||||
off_t = pid * num_experts
|
||||
|
||||
for i in range(start_idx, tl.minimum(start_idx + tokens_per_thread, numel)):
|
||||
expert_id = tl.load(topk_ids_ptr + i)
|
||||
token_cnt = tl.load(tokens_cnts_ptr + off_t + expert_id)
|
||||
rank_post_pad = token_cnt + tl.load(cumsum_ptr + expert_id)
|
||||
tl.store(sorted_token_ids_ptr + rank_post_pad, i)
|
||||
tl.store(tokens_cnts_ptr + off_t + expert_id, token_cnt + 1)
|
||||
|
||||
|
||||
def moe_align_block_size_triton(
|
||||
topk_ids: torch.Tensor,
|
||||
num_experts: int,
|
||||
block_size: int,
|
||||
sorted_token_ids: torch.Tensor,
|
||||
expert_ids: torch.Tensor,
|
||||
num_tokens_post_pad: torch.Tensor,
|
||||
) -> None:
|
||||
numel = topk_ids.numel()
|
||||
grid = (num_experts,)
|
||||
tokens_cnts = torch.zeros(
|
||||
(num_experts + 1, num_experts), dtype=torch.int32, device=topk_ids.device
|
||||
)
|
||||
cumsum = torch.zeros((num_experts + 1,), dtype=torch.int32, device=topk_ids.device)
|
||||
tokens_per_thread = ceil_div(numel, num_experts)
|
||||
|
||||
moe_align_block_size_stage1[grid](
|
||||
topk_ids,
|
||||
tokens_cnts,
|
||||
num_experts,
|
||||
numel,
|
||||
tokens_per_thread,
|
||||
)
|
||||
moe_align_block_size_stage2[grid](
|
||||
tokens_cnts,
|
||||
num_experts,
|
||||
)
|
||||
moe_align_block_size_stage3[(1,)](
|
||||
num_tokens_post_pad,
|
||||
tokens_cnts,
|
||||
cumsum,
|
||||
num_experts,
|
||||
block_size,
|
||||
)
|
||||
moe_align_block_size_stage4[grid](
|
||||
topk_ids,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
tokens_cnts,
|
||||
cumsum,
|
||||
num_experts,
|
||||
block_size,
|
||||
numel,
|
||||
tokens_per_thread,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"block_size,num_tokens,topk,num_experts,pad_sorted_token_ids",
|
||||
list(
|
||||
itertools.product(
|
||||
[32, 64, 128, 256], # block_size
|
||||
[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096], # num_tokens
|
||||
[1, 2, 4, 8, 16, 32, 64], # topk
|
||||
[64, 160, 256, 257, 260, 264], # num_experts
|
||||
[True, False], # pad_sorted_token_ids
|
||||
)
|
||||
),
|
||||
)
|
||||
def test_moe_align_block_size_compare_implementations(
|
||||
block_size, num_tokens, topk, num_experts, pad_sorted_token_ids
|
||||
):
|
||||
topk_ids = torch.argsort(torch.rand(num_tokens, num_experts, device="cuda"), dim=1)[
|
||||
:, :topk
|
||||
]
|
||||
|
||||
max_num_tokens_padded = topk_ids.numel() + (num_experts + 1) * (block_size - 1)
|
||||
if topk_ids.numel() < num_experts + 1:
|
||||
max_num_tokens_padded = topk_ids.numel() * block_size
|
||||
|
||||
sorted_ids_cuda = torch.empty(
|
||||
(max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device
|
||||
)
|
||||
if not pad_sorted_token_ids:
|
||||
sorted_ids_cuda.fill_(topk_ids.numel())
|
||||
max_num_m_blocks = max_num_tokens_padded // block_size
|
||||
expert_ids_cuda = torch.zeros(
|
||||
(max_num_m_blocks,), dtype=torch.int32, device=topk_ids.device
|
||||
)
|
||||
num_tokens_post_pad_cuda = torch.empty(
|
||||
(1), dtype=torch.int32, device=topk_ids.device
|
||||
)
|
||||
cumsum_buffer = torch.empty(
|
||||
num_experts + 2, dtype=torch.int32, device=topk_ids.device
|
||||
)
|
||||
|
||||
sorted_ids_triton = torch.empty_like(sorted_ids_cuda)
|
||||
sorted_ids_triton.fill_(topk_ids.numel())
|
||||
expert_ids_triton = torch.zeros_like(expert_ids_cuda)
|
||||
num_tokens_post_pad_triton = torch.empty_like(num_tokens_post_pad_cuda)
|
||||
|
||||
moe_align_block_size(
|
||||
topk_ids,
|
||||
num_experts + 1,
|
||||
block_size,
|
||||
sorted_ids_cuda,
|
||||
expert_ids_cuda,
|
||||
num_tokens_post_pad_cuda,
|
||||
cumsum_buffer,
|
||||
pad_sorted_token_ids,
|
||||
)
|
||||
|
||||
moe_align_block_size_triton(
|
||||
topk_ids,
|
||||
num_experts + 1,
|
||||
block_size,
|
||||
sorted_ids_triton,
|
||||
expert_ids_triton,
|
||||
num_tokens_post_pad_triton,
|
||||
)
|
||||
|
||||
assert torch.allclose(expert_ids_cuda, expert_ids_triton, atol=0, rtol=0), (
|
||||
f"Expert IDs mismatch for block_size={block_size}, "
|
||||
f"num_tokens={num_tokens}, topk={topk}\n"
|
||||
f"CUDA expert_ids: {expert_ids_cuda}\n"
|
||||
f"Triton expert_ids: {expert_ids_triton}"
|
||||
)
|
||||
|
||||
assert torch.allclose(
|
||||
num_tokens_post_pad_cuda, num_tokens_post_pad_triton, atol=0, rtol=0
|
||||
), (
|
||||
f"Num tokens post pad mismatch for block_size={block_size}, "
|
||||
f"num_tokens={num_tokens}, topk={topk}\n"
|
||||
f"CUDA num_tokens_post_pad: {num_tokens_post_pad_cuda}\n"
|
||||
f"Triton num_tokens_post_pad: {num_tokens_post_pad_triton}"
|
||||
)
|
||||
|
||||
# Select an expert to check
|
||||
expert_idx = expert_ids_cuda.max().item()
|
||||
|
||||
# Get the first and last block id where expert_ids_cuda == expert_idx
|
||||
matching_indices = torch.where(expert_ids_cuda == expert_idx)[0]
|
||||
block_sorted_start = matching_indices[0].item() * block_size
|
||||
block_sorted_end = min(
|
||||
(matching_indices[-1].item() + 1) * block_size,
|
||||
num_tokens_post_pad_cuda.item(),
|
||||
)
|
||||
|
||||
selected_sorted_ids_cuda = sorted_ids_cuda[
|
||||
block_sorted_start:block_sorted_end
|
||||
].sort()[0]
|
||||
selected_sorted_ids_triton = sorted_ids_triton[
|
||||
block_sorted_start:block_sorted_end
|
||||
].sort()[0]
|
||||
|
||||
assert torch.allclose(
|
||||
selected_sorted_ids_cuda,
|
||||
selected_sorted_ids_triton,
|
||||
atol=0,
|
||||
rtol=0,
|
||||
), (
|
||||
f"Sorted IDs mismatch for block_size={block_size}, "
|
||||
f"num_tokens={num_tokens}, topk={topk}\n"
|
||||
f"CUDA sorted_ids: {selected_sorted_ids_cuda}\n"
|
||||
f"Triton sorted_ids: {selected_sorted_ids_triton}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"block_size,num_tokens,topk,num_experts",
|
||||
list(
|
||||
itertools.product(
|
||||
[64, 128], # block_size
|
||||
[1, 8, 32, 256], # num_tokens
|
||||
[8], # topk
|
||||
[
|
||||
1025,
|
||||
2048,
|
||||
4095,
|
||||
], # num_experts (>1024 to exercise v2 kernel, max 4095 real experts)
|
||||
)
|
||||
),
|
||||
)
|
||||
def test_moe_align_block_size_v2_large_num_experts(
|
||||
block_size, num_tokens, topk, num_experts
|
||||
):
|
||||
"""Test moe_align_block_size v2 kernel for >1024 experts against Triton reference."""
|
||||
topk_ids = torch.randint(
|
||||
0, num_experts, (num_tokens, topk), dtype=torch.int32, device="cuda"
|
||||
)
|
||||
|
||||
max_num_tokens_padded = topk_ids.numel() + (num_experts + 1) * (block_size - 1)
|
||||
if topk_ids.numel() < num_experts + 1:
|
||||
max_num_tokens_padded = topk_ids.numel() * block_size
|
||||
|
||||
sorted_ids_cuda = torch.empty(
|
||||
(max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device
|
||||
)
|
||||
sorted_ids_cuda.fill_(topk_ids.numel())
|
||||
max_num_m_blocks = max_num_tokens_padded // block_size
|
||||
expert_ids_cuda = torch.zeros(
|
||||
(max_num_m_blocks,), dtype=torch.int32, device=topk_ids.device
|
||||
)
|
||||
num_tokens_post_pad_cuda = torch.empty(
|
||||
(1), dtype=torch.int32, device=topk_ids.device
|
||||
)
|
||||
cumsum_buffer = torch.empty(
|
||||
num_experts + 2, dtype=torch.int32, device=topk_ids.device
|
||||
)
|
||||
|
||||
sorted_ids_triton = torch.empty_like(sorted_ids_cuda)
|
||||
sorted_ids_triton.fill_(topk_ids.numel())
|
||||
expert_ids_triton = torch.zeros_like(expert_ids_cuda)
|
||||
num_tokens_post_pad_triton = torch.empty_like(num_tokens_post_pad_cuda)
|
||||
|
||||
moe_align_block_size(
|
||||
topk_ids,
|
||||
num_experts + 1,
|
||||
block_size,
|
||||
sorted_ids_cuda,
|
||||
expert_ids_cuda,
|
||||
num_tokens_post_pad_cuda,
|
||||
cumsum_buffer,
|
||||
True,
|
||||
)
|
||||
|
||||
moe_align_block_size_triton(
|
||||
topk_ids,
|
||||
num_experts + 1,
|
||||
block_size,
|
||||
sorted_ids_triton,
|
||||
expert_ids_triton,
|
||||
num_tokens_post_pad_triton,
|
||||
)
|
||||
|
||||
assert torch.equal(num_tokens_post_pad_cuda, num_tokens_post_pad_triton), (
|
||||
f"Num tokens post pad mismatch: CUDA={num_tokens_post_pad_cuda.item()}, "
|
||||
f"Triton={num_tokens_post_pad_triton.item()}"
|
||||
)
|
||||
|
||||
ntp = num_tokens_post_pad_cuda.item()
|
||||
num_blocks = ntp // block_size
|
||||
|
||||
assert torch.equal(expert_ids_cuda[:num_blocks], expert_ids_triton[:num_blocks]), (
|
||||
f"Expert IDs mismatch for block_size={block_size}, "
|
||||
f"num_tokens={num_tokens}, topk={topk}, num_experts={num_experts}"
|
||||
)
|
||||
|
||||
# Compare sorted_token_ids per expert block (order within block may differ)
|
||||
for b in range(num_blocks):
|
||||
s, e = b * block_size, (b + 1) * block_size
|
||||
block_cuda = sorted_ids_cuda[s:e].sort().values
|
||||
block_triton = sorted_ids_triton[s:e].sort().values
|
||||
assert torch.equal(block_cuda, block_triton), (
|
||||
f"Block {b} sorted_ids mismatch for num_experts={num_experts}, "
|
||||
f"num_tokens={num_tokens}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -0,0 +1,168 @@
|
||||
# Temporarily adapted from https://github.com/vllm-project/vllm/blob/main/tests/lora/test_moe_lora_align_sum.py, will optimize in future refactor
|
||||
import random
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
# ---------------------------------------------------------
|
||||
# IMPORT PREBUILT KERNEL
|
||||
# ---------------------------------------------------------
|
||||
from sglang.jit_kernel.moe_lora_align import moe_lora_align_block_size
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=28, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
def round_up(x, base):
|
||||
return ((x + base - 1) // base) * base
|
||||
|
||||
|
||||
def CEILDIV(x, y):
|
||||
return (x + y - 1) // y
|
||||
|
||||
|
||||
def sample_data(num_experts, max_loras, num_tokens, topk_num):
|
||||
# 1. Generate TopK IDs (Flattened tokens)
|
||||
topk_ids = torch.zeros((num_tokens, topk_num), dtype=torch.int32)
|
||||
for i in range(num_tokens):
|
||||
pool = list(range(num_experts))
|
||||
random.shuffle(pool)
|
||||
for j in range(topk_num):
|
||||
topk_ids[i, j] = pool[j]
|
||||
|
||||
# 2. Generate Random Requests (Segments)
|
||||
# We split num_tokens into random chunks to simulate a batch of requests
|
||||
remaining_tokens = num_tokens
|
||||
seg_lens = []
|
||||
while remaining_tokens > 0:
|
||||
# Random length between 1 and remaining
|
||||
length = random.randint(1, min(32, remaining_tokens))
|
||||
if remaining_tokens - length < 0:
|
||||
length = remaining_tokens
|
||||
seg_lens.append(length)
|
||||
remaining_tokens -= length
|
||||
|
||||
# Ensure we cover the full range exactly (cleanup last segment)
|
||||
if sum(seg_lens) < num_tokens:
|
||||
seg_lens.append(num_tokens - sum(seg_lens))
|
||||
|
||||
# 3. Build seg_indptr [0, len1, len1+len2, ...]
|
||||
seg_indptr = torch.cumsum(
|
||||
torch.tensor([0] + seg_lens, dtype=torch.int32), dim=0
|
||||
).to(dtype=torch.int32)
|
||||
|
||||
# 4. Assign a LoRA ID to each Request
|
||||
num_reqs = len(seg_lens)
|
||||
req_to_lora = torch.randint(0, max_loras, (num_reqs,), dtype=torch.int32)
|
||||
|
||||
return (topk_ids.to("cuda"), seg_indptr.to("cuda"), req_to_lora.to("cuda"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [100, 200, 1024, 4096])
|
||||
@pytest.mark.parametrize("topk_num", [6])
|
||||
@pytest.mark.parametrize("num_experts", [64, 128, 256, 512])
|
||||
@pytest.mark.parametrize("max_loras", [2, 32])
|
||||
@pytest.mark.parametrize("block_size", [16])
|
||||
def test_moe_lora_align_block_size(
|
||||
num_tokens, topk_num, num_experts, max_loras, block_size
|
||||
):
|
||||
# sample data
|
||||
random.seed(1)
|
||||
torch.manual_seed(1)
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA is not available, skipping moe_lora_align_block_size test.")
|
||||
# UPDATED: Get the new 3-step mapping tensors
|
||||
topk_ids, seg_indptr, req_to_lora = sample_data(
|
||||
num_experts, max_loras, num_tokens, topk_num
|
||||
)
|
||||
|
||||
# compute paddings
|
||||
max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
|
||||
max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
|
||||
max_num_m_blocks = CEILDIV(max_num_tokens_padded, block_size)
|
||||
|
||||
# init output tensors
|
||||
sorted_token_ids = torch.full(
|
||||
(max_loras * max_num_tokens_padded,),
|
||||
topk_ids.numel(),
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
expert_ids = torch.full(
|
||||
(max_loras * max_num_m_blocks,), num_experts, dtype=torch.int32, device="cuda"
|
||||
)
|
||||
num_tokens_post_pad = torch.zeros((max_loras,), dtype=torch.int32, device="cuda")
|
||||
adapter_enabled = torch.ones((max_loras + 1,), dtype=torch.int32, device="cuda")
|
||||
lora_ids = torch.arange(max_loras, dtype=torch.int32, device="cuda")
|
||||
|
||||
# UPDATED: Call kernel with new signature
|
||||
moe_lora_align_block_size(
|
||||
topk_ids,
|
||||
seg_indptr, # Arg 2: Pointers
|
||||
req_to_lora, # Arg 3: Request Map
|
||||
num_experts,
|
||||
block_size,
|
||||
max_loras,
|
||||
max_num_tokens_padded,
|
||||
max_num_m_blocks,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_pad,
|
||||
adapter_enabled,
|
||||
lora_ids,
|
||||
None,
|
||||
)
|
||||
|
||||
# verify values
|
||||
expert_ids = expert_ids.view(max_loras, -1)
|
||||
sorted_token_ids = sorted_token_ids.view(max_loras, -1, block_size)
|
||||
|
||||
# Reconstruct token-level ownership for verification logic
|
||||
# We expand req_to_lora back to [num_tokens] on CPU just to check correctness
|
||||
# This proves the kernel (which used the compressed format) produced the right result
|
||||
cpu_seg_indptr = seg_indptr.cpu()
|
||||
cpu_req_to_lora = req_to_lora.cpu()
|
||||
token_ownership = torch.zeros(num_tokens, dtype=torch.int32)
|
||||
|
||||
for r in range(len(cpu_req_to_lora)):
|
||||
start = cpu_seg_indptr[r]
|
||||
end = cpu_seg_indptr[r + 1]
|
||||
token_ownership[start:end] = cpu_req_to_lora[r]
|
||||
|
||||
token_ownership = token_ownership.to("cuda")
|
||||
|
||||
for lora_idx in range(max_loras):
|
||||
# Count how many tokens actually belong to this LoRA
|
||||
expected_count = (token_ownership == lora_idx).sum().item()
|
||||
|
||||
# Verify the kernel processed a reasonable number of tokens (sanity check)
|
||||
# Note: num_tokens_post_pad includes padding, so it might be larger than expected_count
|
||||
assert num_tokens_post_pad[lora_idx].item() >= expected_count * topk_num
|
||||
|
||||
for token_idx in range(sorted_token_ids.size(1)):
|
||||
block = sorted_token_ids[lora_idx][token_idx]
|
||||
# Valid indices are those less than total numel
|
||||
indices = block[block != topk_ids.numel()]
|
||||
|
||||
if indices.numel() > 0:
|
||||
# 1. Verify routing: Does the token actually route to this expert?
|
||||
expert_id = expert_ids[lora_idx][token_idx]
|
||||
assert torch.all(topk_ids.view(-1)[indices] == expert_id)
|
||||
|
||||
# 2. Verify ownership: Did the kernel grab the correct tokens for this LoRA?
|
||||
# The indices in 'sorted_token_ids' point to the flattened [token, topk] array.
|
||||
# We divide by topk_num to get the original token index.
|
||||
original_token_indices = indices // topk_num
|
||||
|
||||
# Check that all tokens in this block truly belong to 'lora_idx'
|
||||
actual_owners = token_ownership[original_token_indices]
|
||||
assert torch.all(
|
||||
actual_owners == lora_idx
|
||||
), f"Kernel put tokens from LoRA {actual_owners} into block for LoRA {lora_idx}"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -0,0 +1,615 @@
|
||||
import itertools
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from sgl_kernel.scalar_type import scalar_types
|
||||
|
||||
from sglang.jit_kernel.moe_wna16_marlin import moe_wna16_marlin_gemm
|
||||
from sglang.srt.layers.moe.fused_moe_triton import moe_align_block_size
|
||||
from sglang.srt.layers.moe.fused_moe_triton.fused_marlin_moe import fused_marlin_moe
|
||||
from sglang.srt.layers.quantization.marlin_utils_fp4 import (
|
||||
prepare_moe_nvfp4_layer_for_marlin,
|
||||
)
|
||||
from sglang.srt.utils.common import is_sm80_supported, is_sm90_supported
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_marlin_utils import (
|
||||
awq_marlin_quantize,
|
||||
make_nvfp4_weight_and_ref,
|
||||
marlin_quantize,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
def _has_aot_moe_wna16_marlin_gemm() -> bool:
|
||||
return hasattr(torch.ops.sgl_kernel, "moe_wna16_marlin_gemm") and hasattr(
|
||||
torch.ops.sgl_kernel.moe_wna16_marlin_gemm, "default"
|
||||
)
|
||||
|
||||
|
||||
AOT_AVAILABLE = _has_aot_moe_wna16_marlin_gemm()
|
||||
|
||||
|
||||
def stack_and_dev(tensors: list[torch.Tensor]):
|
||||
dev = tensors[0].device
|
||||
return torch.stack(tensors, dim=0).to(dev)
|
||||
|
||||
|
||||
def _get_scalar_type(num_bits: int, has_zp: bool):
|
||||
if has_zp:
|
||||
assert num_bits == 4
|
||||
return scalar_types.uint4
|
||||
else:
|
||||
return scalar_types.uint4b8 if num_bits == 4 else scalar_types.uint8b128
|
||||
|
||||
|
||||
def _setup_moe_weights(e, n, k, quant_type, group_size, act_order, dtype):
|
||||
"""Set up quantized MoE weights for a single gate (e experts, output n, input k)."""
|
||||
has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8]
|
||||
|
||||
w = torch.randn((e, n, k), device="cuda", dtype=dtype) / 20
|
||||
|
||||
w_ref_l = []
|
||||
qweight_l = []
|
||||
scales_l = []
|
||||
zeros_l = []
|
||||
g_idx_l = []
|
||||
sort_indices_l = []
|
||||
|
||||
for i in range(e):
|
||||
if has_zp:
|
||||
w_ref, qweight, scales, zeros = awq_marlin_quantize(
|
||||
w[i].transpose(1, 0), quant_type, group_size
|
||||
)
|
||||
w_ref_l.append(w_ref.T)
|
||||
qweight_l.append(qweight)
|
||||
scales_l.append(scales)
|
||||
zeros_l.append(zeros)
|
||||
else:
|
||||
test_perm = torch.randperm(k)
|
||||
w_ref, qweight, scales, g_idx, sort_indices, _ = marlin_quantize(
|
||||
w[i].transpose(1, 0), quant_type, group_size, act_order, test_perm
|
||||
)
|
||||
w_ref_l.append(w_ref.T)
|
||||
qweight_l.append(qweight)
|
||||
scales_l.append(scales)
|
||||
g_idx_l.append(g_idx)
|
||||
sort_indices_l.append(sort_indices)
|
||||
|
||||
w_ref = stack_and_dev(w_ref_l)
|
||||
qweight = stack_and_dev(qweight_l).contiguous()
|
||||
scales = stack_and_dev(scales_l)
|
||||
g_idx = stack_and_dev(g_idx_l) if g_idx_l else None
|
||||
sort_indices = stack_and_dev(sort_indices_l) if sort_indices_l else None
|
||||
zeros = stack_and_dev(zeros_l) if zeros_l else None
|
||||
|
||||
return w_ref, qweight, scales, zeros, g_idx, sort_indices
|
||||
|
||||
|
||||
def _run_single_gemm(
|
||||
fn,
|
||||
a,
|
||||
c,
|
||||
qweight,
|
||||
scales,
|
||||
zeros,
|
||||
g_idx,
|
||||
sort_indices,
|
||||
workspace,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_padded,
|
||||
topk_weights,
|
||||
quant_type,
|
||||
block_size_m,
|
||||
topk,
|
||||
size_m,
|
||||
size_n,
|
||||
size_k,
|
||||
mul_topk_weights,
|
||||
is_k_full,
|
||||
use_atomic_add,
|
||||
):
|
||||
return fn(
|
||||
a,
|
||||
c,
|
||||
qweight,
|
||||
None, # b_bias
|
||||
scales,
|
||||
None, # global_scale
|
||||
zeros,
|
||||
g_idx,
|
||||
sort_indices,
|
||||
workspace,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_padded,
|
||||
topk_weights,
|
||||
moe_block_size=block_size_m,
|
||||
top_k=topk,
|
||||
mul_topk_weights=mul_topk_weights,
|
||||
is_ep=False,
|
||||
b_q_type=quant_type,
|
||||
size_m=size_m,
|
||||
size_n=size_n,
|
||||
size_k=size_k,
|
||||
is_k_full=is_k_full,
|
||||
use_atomic_add=use_atomic_add,
|
||||
use_fp32_reduce=True,
|
||||
is_zp_float=False,
|
||||
)
|
||||
|
||||
|
||||
def _run_single_gemm_aot(
|
||||
a,
|
||||
c,
|
||||
qweight,
|
||||
scales,
|
||||
zeros,
|
||||
g_idx,
|
||||
sort_indices,
|
||||
workspace,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_padded,
|
||||
topk_weights,
|
||||
quant_type,
|
||||
block_size_m,
|
||||
topk,
|
||||
size_m,
|
||||
size_n,
|
||||
size_k,
|
||||
mul_topk_weights,
|
||||
is_k_full,
|
||||
use_atomic_add,
|
||||
):
|
||||
return torch.ops.sgl_kernel.moe_wna16_marlin_gemm.default(
|
||||
a,
|
||||
c,
|
||||
qweight,
|
||||
None, # b_bias
|
||||
scales,
|
||||
None, # global_scale
|
||||
zeros,
|
||||
g_idx,
|
||||
sort_indices,
|
||||
workspace,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_padded,
|
||||
topk_weights,
|
||||
moe_block_size=block_size_m,
|
||||
top_k=topk,
|
||||
mul_topk_weights=mul_topk_weights,
|
||||
is_ep=False,
|
||||
b_q_type_id=quant_type.id,
|
||||
size_m=size_m,
|
||||
size_n=size_n,
|
||||
size_k=size_k,
|
||||
is_k_full=is_k_full,
|
||||
use_atomic_add=use_atomic_add,
|
||||
use_fp32_reduce=True,
|
||||
is_zp_float=False,
|
||||
)
|
||||
|
||||
|
||||
def generate_test_cases():
|
||||
m_list = [1, 123]
|
||||
n_list = [128, 1024]
|
||||
k_list = [256]
|
||||
e_list = [4]
|
||||
topk_list = [2]
|
||||
dtype_list = [torch.float16, torch.bfloat16]
|
||||
group_size_list = [128]
|
||||
act_order_list = [False, True]
|
||||
quant_type_list = [scalar_types.uint4, scalar_types.uint4b8]
|
||||
|
||||
all_combinations = itertools.product(
|
||||
m_list,
|
||||
n_list,
|
||||
k_list,
|
||||
e_list,
|
||||
topk_list,
|
||||
dtype_list,
|
||||
group_size_list,
|
||||
act_order_list,
|
||||
quant_type_list,
|
||||
)
|
||||
|
||||
def is_valid(m, n, k, e, topk, dtype, group_size, act_order, quant_type):
|
||||
has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8]
|
||||
if act_order:
|
||||
if group_size == -1 or group_size == k:
|
||||
return False
|
||||
if has_zp:
|
||||
return False
|
||||
if group_size > 0 and k % group_size != 0:
|
||||
return False
|
||||
return True
|
||||
|
||||
return [case for case in all_combinations if is_valid(*case)]
|
||||
|
||||
|
||||
TEST_CASES = generate_test_cases()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"m,n,k,e,topk,dtype,group_size,act_order,quant_type",
|
||||
TEST_CASES,
|
||||
ids=[
|
||||
f"m{c[0]}_n{c[1]}_k{c[2]}_e{c[3]}_t{c[4]}_{c[5].__name__ if hasattr(c[5], '__name__') else str(c[5]).split('.')[-1]}_g{c[6]}_act{c[7]}_{c[8]}"
|
||||
for c in TEST_CASES
|
||||
],
|
||||
)
|
||||
def test_moe_wna16_marlin_gemm(
|
||||
m, n, k, e, topk, dtype, group_size, act_order, quant_type
|
||||
):
|
||||
if not AOT_AVAILABLE:
|
||||
pytest.skip("sgl_kernel moe_wna16_marlin_gemm AOT op not available")
|
||||
|
||||
torch.manual_seed(0)
|
||||
|
||||
has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8]
|
||||
|
||||
a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
|
||||
|
||||
# Set up quantized weights for first gemm (gate_up: output 2*n, input k)
|
||||
w_ref1, qweight1, scales1, zeros1, g_idx1, sort_indices1 = _setup_moe_weights(
|
||||
e, 2 * n, k, quant_type, group_size, act_order, dtype
|
||||
)
|
||||
|
||||
# Compute block_size_m
|
||||
for block_size_m in [8, 16, 32, 48, 64]:
|
||||
if m * topk / e / block_size_m < 0.9:
|
||||
break
|
||||
|
||||
# Align tokens
|
||||
score = torch.randn((m, e), device="cuda", dtype=dtype)
|
||||
score_softmax = torch.softmax(score, dim=-1, dtype=torch.float32)
|
||||
topk_weights, topk_ids = torch.topk(score_softmax, topk)
|
||||
|
||||
sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
|
||||
topk_ids, block_size_m, e
|
||||
)
|
||||
|
||||
# Workspace
|
||||
sms = torch.cuda.get_device_properties("cuda").multi_processor_count
|
||||
max_workspace_size = (max(2 * n, k) // 64) * (
|
||||
sorted_token_ids.size(0) // block_size_m
|
||||
)
|
||||
max_workspace_size = min(max_workspace_size, sms * 4)
|
||||
workspace = torch.zeros(
|
||||
max_workspace_size, dtype=torch.int, device="cuda", requires_grad=False
|
||||
)
|
||||
|
||||
use_atomic_add = (
|
||||
dtype == torch.half or torch.cuda.get_device_capability("cuda")[0] >= 9
|
||||
)
|
||||
|
||||
scalar_type = _get_scalar_type(4, has_zp)
|
||||
|
||||
# --- Run JIT kernel ---
|
||||
c_jit = torch.empty((m * topk, 2 * n), dtype=dtype, device="cuda")
|
||||
c_jit = _run_single_gemm(
|
||||
moe_wna16_marlin_gemm,
|
||||
a,
|
||||
c_jit,
|
||||
qweight1,
|
||||
scales1,
|
||||
zeros1,
|
||||
g_idx1,
|
||||
sort_indices1,
|
||||
workspace,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_padded,
|
||||
topk_weights,
|
||||
scalar_type,
|
||||
block_size_m,
|
||||
topk,
|
||||
m,
|
||||
2 * n,
|
||||
k,
|
||||
False,
|
||||
True,
|
||||
use_atomic_add,
|
||||
)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# --- Check bitwise equality with AOT kernel ---
|
||||
c_aot = torch.empty((m * topk, 2 * n), dtype=dtype, device="cuda")
|
||||
c_aot = _run_single_gemm_aot(
|
||||
a,
|
||||
c_aot,
|
||||
qweight1,
|
||||
scales1,
|
||||
zeros1,
|
||||
g_idx1,
|
||||
sort_indices1,
|
||||
workspace,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_padded,
|
||||
topk_weights,
|
||||
scalar_type,
|
||||
block_size_m,
|
||||
topk,
|
||||
m,
|
||||
2 * n,
|
||||
k,
|
||||
False,
|
||||
True,
|
||||
use_atomic_add,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
torch.testing.assert_close(c_jit, c_aot, rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Skip, test pass locally but compiling takes too long in CI")
|
||||
@pytest.mark.skipif(
|
||||
not (is_sm80_supported() or is_sm90_supported()),
|
||||
reason="Non-gated NVFP4 Marlin fallback test requires CUDA SM8X/SM9X",
|
||||
)
|
||||
def test_fused_marlin_moe_non_gated_relu2():
|
||||
torch.manual_seed(0)
|
||||
|
||||
m = 17
|
||||
n = 128
|
||||
k = 256
|
||||
e = 4
|
||||
topk = 2
|
||||
dtype = torch.float16
|
||||
group_size = 128
|
||||
quant_type = scalar_types.uint4b8
|
||||
|
||||
hidden_states = torch.randn((m, k), device="cuda", dtype=dtype) / 10
|
||||
w_ref1, qweight1, scales1, zeros1, g_idx1, sort_indices1 = _setup_moe_weights(
|
||||
e, n, k, quant_type, group_size, False, dtype
|
||||
)
|
||||
w_ref2, qweight2, scales2, zeros2, g_idx2, sort_indices2 = _setup_moe_weights(
|
||||
e, k, n, quant_type, group_size, False, dtype
|
||||
)
|
||||
|
||||
router_logits = torch.randn((m, e), device="cuda", dtype=dtype)
|
||||
score_softmax = torch.softmax(router_logits, dim=-1, dtype=torch.float32)
|
||||
topk_weights, topk_ids = torch.topk(score_softmax, topk)
|
||||
|
||||
output = fused_marlin_moe(
|
||||
hidden_states=hidden_states,
|
||||
w1=qweight1,
|
||||
w2=qweight2,
|
||||
w1_scale=scales1,
|
||||
w2_scale=scales2,
|
||||
gating_output=router_logits,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
g_idx1=g_idx1,
|
||||
g_idx2=g_idx2,
|
||||
sort_indices1=sort_indices1,
|
||||
sort_indices2=sort_indices2,
|
||||
w1_zeros=zeros1,
|
||||
w2_zeros=zeros2,
|
||||
num_bits=4,
|
||||
is_k_full=True,
|
||||
routed_scaling_factor=1.0,
|
||||
activation="relu2",
|
||||
is_gated=False,
|
||||
)
|
||||
|
||||
output_ref = torch.zeros_like(hidden_states)
|
||||
for token_idx in range(m):
|
||||
for route_idx in range(topk):
|
||||
expert_id = topk_ids[token_idx, route_idx]
|
||||
intermediate = hidden_states[token_idx] @ w_ref1[expert_id].T
|
||||
intermediate = torch.square(torch.relu(intermediate))
|
||||
routed = intermediate @ w_ref2[expert_id].T
|
||||
output_ref[token_idx] += routed * topk_weights[token_idx, route_idx]
|
||||
|
||||
torch.cuda.synchronize()
|
||||
torch.testing.assert_close(output, output_ref, rtol=0.04, atol=0.04)
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Skip, test pass locally but compiling takes too long in CI")
|
||||
@pytest.mark.skipif(
|
||||
not (is_sm80_supported() or is_sm90_supported()),
|
||||
reason="NVFP4 Marlin MoE padding test requires CUDA SM8X/SM9X",
|
||||
)
|
||||
def test_fused_marlin_moe_nvfp4_non_gated_padded_intermediate_launches():
|
||||
torch.manual_seed(0)
|
||||
|
||||
m = 17
|
||||
intermediate_size = 192
|
||||
hidden_size = 256
|
||||
e = 4
|
||||
topk = 2
|
||||
dtype = torch.bfloat16
|
||||
nvfp4_group_size = 16
|
||||
|
||||
layer = torch.nn.Module()
|
||||
layer.quant_config = SimpleNamespace(group_size=nvfp4_group_size)
|
||||
layer.moe_runner_config = SimpleNamespace(is_gated=False)
|
||||
layer.params_dtype = dtype
|
||||
layer.intermediate_size_per_partition = intermediate_size
|
||||
layer.w13_weight = torch.nn.Parameter(
|
||||
torch.randint(
|
||||
0,
|
||||
256,
|
||||
(e, intermediate_size, hidden_size // 2),
|
||||
device="cuda",
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w2_weight = torch.nn.Parameter(
|
||||
torch.randint(
|
||||
0,
|
||||
256,
|
||||
(e, hidden_size, intermediate_size // 2),
|
||||
device="cuda",
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w13_weight_scale = torch.nn.Parameter(
|
||||
torch.rand(
|
||||
(e, intermediate_size, hidden_size // nvfp4_group_size),
|
||||
device="cuda",
|
||||
dtype=dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w2_weight_scale = torch.nn.Parameter(
|
||||
torch.rand(
|
||||
(e, hidden_size, intermediate_size // nvfp4_group_size),
|
||||
device="cuda",
|
||||
dtype=dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w13_weight_scale_2 = torch.nn.Parameter(
|
||||
torch.ones((e,), device="cuda", dtype=dtype), requires_grad=False
|
||||
)
|
||||
layer.w2_weight_scale_2 = torch.nn.Parameter(
|
||||
torch.ones((e,), device="cuda", dtype=dtype), requires_grad=False
|
||||
)
|
||||
prepare_moe_nvfp4_layer_for_marlin(layer)
|
||||
|
||||
assert layer.w13_weight.shape[1] * 16 == 256
|
||||
assert layer.w2_weight.shape[1] * 16 == 256
|
||||
|
||||
hidden_states = torch.randn((m, hidden_size), device="cuda", dtype=dtype) / 10
|
||||
|
||||
score = torch.randn((m, e), device="cuda", dtype=dtype)
|
||||
score_softmax = torch.softmax(score, dim=-1, dtype=torch.float32)
|
||||
topk_weights, topk_ids = torch.topk(score_softmax, topk)
|
||||
|
||||
out = fused_marlin_moe(
|
||||
hidden_states=hidden_states,
|
||||
w1=layer.w13_weight,
|
||||
w2=layer.w2_weight,
|
||||
w1_scale=layer.w13_weight_scale,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
gating_output=score,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
w1_global_scale=layer.w13_weight_scale_2,
|
||||
w2_global_scale=layer.w2_weight_scale_2,
|
||||
workspace=layer.workspace,
|
||||
num_bits=4,
|
||||
is_k_full=True,
|
||||
routed_scaling_factor=1.0,
|
||||
activation="relu2",
|
||||
is_gated=False,
|
||||
)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
assert out.shape == (m, hidden_size)
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Skip, test pass locally but compiling takes too long in CI")
|
||||
@pytest.mark.skipif(
|
||||
not (is_sm80_supported() or is_sm90_supported()),
|
||||
reason="NVFP4 Marlin MoE numeric test requires CUDA SM80, SM86, or SM90",
|
||||
)
|
||||
def test_fused_marlin_moe_nvfp4_non_gated_matches_dequant_reference():
|
||||
torch.manual_seed(0)
|
||||
|
||||
m = 17
|
||||
intermediate_size = 192
|
||||
hidden_size = 256
|
||||
e = 4
|
||||
topk = 2
|
||||
dtype = torch.bfloat16
|
||||
group_size = 16
|
||||
routed_scaling_factor = 1.0
|
||||
|
||||
w13_packed_l, w13_scales_l, w13_gscale_l, w13_ref_l = [], [], [], []
|
||||
w2_packed_l, w2_scales_l, w2_gscale_l, w2_ref_l = [], [], [], []
|
||||
for _ in range(e):
|
||||
packed, scales, gscale, ref = make_nvfp4_weight_and_ref(
|
||||
intermediate_size, hidden_size, dtype, group_size=group_size
|
||||
)
|
||||
w13_packed_l.append(packed)
|
||||
w13_scales_l.append(scales)
|
||||
w13_gscale_l.append(gscale)
|
||||
w13_ref_l.append(ref)
|
||||
|
||||
packed, scales, gscale, ref = make_nvfp4_weight_and_ref(
|
||||
hidden_size, intermediate_size, dtype, group_size=group_size
|
||||
)
|
||||
w2_packed_l.append(packed)
|
||||
w2_scales_l.append(scales)
|
||||
w2_gscale_l.append(gscale)
|
||||
w2_ref_l.append(ref)
|
||||
|
||||
layer = torch.nn.Module()
|
||||
layer.quant_config = SimpleNamespace(group_size=group_size)
|
||||
layer.moe_runner_config = SimpleNamespace(is_gated=False)
|
||||
layer.params_dtype = dtype
|
||||
layer.intermediate_size_per_partition = intermediate_size
|
||||
layer.w13_weight = torch.nn.Parameter(
|
||||
torch.stack(w13_packed_l), requires_grad=False
|
||||
)
|
||||
layer.w2_weight = torch.nn.Parameter(torch.stack(w2_packed_l), requires_grad=False)
|
||||
layer.w13_weight_scale = torch.nn.Parameter(
|
||||
torch.stack(w13_scales_l), requires_grad=False
|
||||
)
|
||||
layer.w2_weight_scale = torch.nn.Parameter(
|
||||
torch.stack(w2_scales_l), requires_grad=False
|
||||
)
|
||||
layer.w13_weight_scale_2 = torch.nn.Parameter(
|
||||
torch.stack(w13_gscale_l), requires_grad=False
|
||||
)
|
||||
layer.w2_weight_scale_2 = torch.nn.Parameter(
|
||||
torch.stack(w2_gscale_l), requires_grad=False
|
||||
)
|
||||
prepare_moe_nvfp4_layer_for_marlin(layer)
|
||||
|
||||
# Scale activations down so relu² doesn't blow up intermediate magnitudes;
|
||||
# this keeps output values small so tighter element-wise tolerance is realistic.
|
||||
hidden_states = torch.randn((m, hidden_size), device="cuda", dtype=dtype) / 20
|
||||
router_logits = torch.randn((m, e), device="cuda", dtype=dtype)
|
||||
score_softmax = torch.softmax(router_logits, dim=-1, dtype=torch.float32)
|
||||
topk_weights, topk_ids = torch.topk(score_softmax, topk)
|
||||
|
||||
output = fused_marlin_moe(
|
||||
hidden_states=hidden_states,
|
||||
w1=layer.w13_weight,
|
||||
w2=layer.w2_weight,
|
||||
w1_scale=layer.w13_weight_scale,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
gating_output=router_logits,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
w1_global_scale=layer.w13_weight_scale_2,
|
||||
w2_global_scale=layer.w2_weight_scale_2,
|
||||
workspace=layer.workspace,
|
||||
num_bits=4,
|
||||
is_k_full=True,
|
||||
routed_scaling_factor=routed_scaling_factor,
|
||||
activation="relu2",
|
||||
is_gated=False,
|
||||
)
|
||||
|
||||
w13_ref = torch.stack(w13_ref_l)
|
||||
w2_ref = torch.stack(w2_ref_l)
|
||||
output_ref = torch.zeros_like(hidden_states)
|
||||
for token_idx in range(m):
|
||||
for route_idx in range(topk):
|
||||
expert_id = topk_ids[token_idx, route_idx]
|
||||
intermediate = hidden_states[token_idx] @ w13_ref[expert_id].T
|
||||
intermediate = torch.square(torch.relu(intermediate))
|
||||
routed = intermediate @ w2_ref[expert_id].T
|
||||
output_ref[token_idx] += routed * topk_weights[token_idx, route_idx]
|
||||
output_ref *= routed_scaling_factor
|
||||
|
||||
torch.cuda.synchronize()
|
||||
torch.testing.assert_close(output, output_ref, rtol=0.05, atol=0.25)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,153 @@
|
||||
import random
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.mxfp8 import (
|
||||
es_sm100_mxfp8_blockscaled_grouped_quant,
|
||||
es_sm100_mxfp8_blockscaled_moe_grouped_gemm,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=5, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
def align(val: int, alignment: int = 128) -> int:
|
||||
return int((val + alignment - 1) // alignment * alignment)
|
||||
|
||||
|
||||
# Copy from: https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/utils.py
|
||||
def calc_diff(x, y):
|
||||
x, y = x.double(), y.double()
|
||||
denominator = (x * x + y * y).sum()
|
||||
sim = 2 * (x * y).sum() / denominator
|
||||
return 1 - sim
|
||||
|
||||
|
||||
def is_sm100_supported(device=None) -> bool:
|
||||
return (torch.cuda.get_device_capability(device)[0] == 10) and (
|
||||
torch.version.cuda >= "12.8"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not is_sm100_supported(),
|
||||
reason="test_mxfp8_moe at jit kernen is only supported on sm100",
|
||||
)
|
||||
@pytest.mark.parametrize("num_experts", [8, 16, 32, 64])
|
||||
@pytest.mark.parametrize("out_dtype", [torch.half, torch.bfloat16])
|
||||
def test_es_sm100_mxfp8_blockscaled_grouped_mm(num_experts, out_dtype):
|
||||
device = "cuda"
|
||||
alignment = 128
|
||||
n_g = random.randint(1, 64) * alignment
|
||||
k_g = random.randint(1, 64) * alignment
|
||||
|
||||
expert_offset = 0
|
||||
expert_offsets = []
|
||||
aux_expert_offset = 0
|
||||
aux_expert_offsets = []
|
||||
a_blockscale_offset = 0
|
||||
a_blockscale_offsets = []
|
||||
b_blockscale_offset = 0
|
||||
b_blockscale_offsets = []
|
||||
a_list = []
|
||||
b_list = []
|
||||
ref_d_list = []
|
||||
tokens_per_expert = []
|
||||
|
||||
for g in range(num_experts):
|
||||
m_g = random.randint(1, 512)
|
||||
tokens_per_expert.append(m_g)
|
||||
expert_offsets.append(expert_offset)
|
||||
expert_offset += m_g
|
||||
aux_expert_offsets.append(aux_expert_offset)
|
||||
aux_expert_offset += n_g
|
||||
a_blockscale_offsets.append(a_blockscale_offset)
|
||||
a_blockscale_offset += align(m_g, 128)
|
||||
b_blockscale_offsets.append(b_blockscale_offset)
|
||||
b_blockscale_offset += n_g # n_g already align to 128
|
||||
|
||||
a = torch.normal(
|
||||
0.0, std=1.0, size=(m_g, k_g), device=device, dtype=out_dtype
|
||||
) # (M, K):(K, 1)
|
||||
b = torch.normal(
|
||||
0.0, std=1.0, size=(n_g, k_g), device=device, dtype=out_dtype
|
||||
) # (N, K):(K, 1)
|
||||
|
||||
a_list.append(a)
|
||||
b_list.append(b)
|
||||
ref_d = a @ b.T
|
||||
ref_d_list.append(ref_d)
|
||||
a = torch.concat(a_list, dim=0)
|
||||
b = torch.concat(b_list, dim=0)
|
||||
|
||||
_expert_offsets = torch.tensor(expert_offsets).to(device=device, dtype=torch.int32)
|
||||
_aux_expert_offsets = torch.tensor(aux_expert_offsets).to(
|
||||
device=device, dtype=torch.int32
|
||||
)
|
||||
_a_blockscale_offsets = torch.tensor(a_blockscale_offsets).to(
|
||||
device=device, dtype=torch.int32
|
||||
)
|
||||
_b_blockscale_offsets = torch.tensor(b_blockscale_offsets).to(
|
||||
device=device, dtype=torch.int32
|
||||
)
|
||||
|
||||
a_quant = torch.zeros_like(a, dtype=torch.float8_e4m3fn, device=device)
|
||||
a_scale_factor = torch.zeros(
|
||||
(a_blockscale_offset, k_g // 32), dtype=torch.uint8, device=device
|
||||
)
|
||||
|
||||
b_quant = torch.zeros_like(b, dtype=torch.float8_e4m3fn, device=device)
|
||||
b_scale_factor = torch.zeros(
|
||||
(num_experts * n_g, k_g // 32), dtype=torch.uint8, device=device
|
||||
)
|
||||
tokens_per_expert = torch.tensor(tokens_per_expert).to(
|
||||
device=device, dtype=torch.int32
|
||||
)
|
||||
workspace = torch.empty((1024, 1024, 1024), dtype=torch.uint8, device=device)
|
||||
|
||||
es_sm100_mxfp8_blockscaled_grouped_quant(
|
||||
a,
|
||||
tokens_per_expert,
|
||||
_expert_offsets,
|
||||
_a_blockscale_offsets,
|
||||
a_quant,
|
||||
a_scale_factor,
|
||||
)
|
||||
es_sm100_mxfp8_blockscaled_grouped_quant(
|
||||
b,
|
||||
torch.ones_like(tokens_per_expert) * n_g,
|
||||
_aux_expert_offsets,
|
||||
_b_blockscale_offsets,
|
||||
b_quant,
|
||||
b_scale_factor,
|
||||
)
|
||||
|
||||
b_quant = b_quant.view(num_experts, n_g, k_g)
|
||||
b_scale_factor = b_scale_factor.view(num_experts, n_g, k_g // 32)
|
||||
d = es_sm100_mxfp8_blockscaled_moe_grouped_gemm(
|
||||
b_quant,
|
||||
a_quant,
|
||||
b_scale_factor,
|
||||
a_scale_factor,
|
||||
_expert_offsets,
|
||||
_a_blockscale_offsets,
|
||||
tokens_per_expert,
|
||||
workspace,
|
||||
a.dtype,
|
||||
)
|
||||
|
||||
for g in range(num_experts):
|
||||
baseline = ref_d_list[g]
|
||||
actual = d[expert_offsets[g] : (expert_offsets[g] + tokens_per_expert[g])]
|
||||
diff = calc_diff(actual, baseline)
|
||||
assert diff < 0.001
|
||||
print(
|
||||
f"m_g={baseline.shape[0]} n_g={n_g} k_g={k_g} num_experts={num_experts}, out_dtype={out_dtype}, diff={diff:.5f}: OK"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -0,0 +1,135 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.ngram_embedding import (
|
||||
compute_n_gram_ids,
|
||||
compute_n_gram_ids_decode,
|
||||
update_token_table,
|
||||
update_token_table_decode,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
|
||||
|
||||
|
||||
def _make_ngram_params(ne_n: int, ne_k: int, vocab_size: int):
|
||||
ne_weights = torch.zeros([ne_n - 1, ne_k, ne_n], dtype=torch.int32)
|
||||
ne_mods = torch.zeros([ne_n - 1, ne_k], dtype=torch.int32)
|
||||
exclusive_sums = torch.zeros([(ne_n - 1) * ne_k + 1], dtype=torch.int32)
|
||||
|
||||
for n in range(2, ne_n + 1):
|
||||
for k in range(ne_k):
|
||||
config_id = (n - 2) * ne_k + k
|
||||
mod = 65537 + 2 * config_id
|
||||
ne_mods[n - 2][k] = mod
|
||||
exclusive_sums[config_id + 1] = exclusive_sums[config_id] + mod
|
||||
for delta in range(ne_n):
|
||||
ne_weights[n - 2][k][delta] = pow(vocab_size, delta, mod)
|
||||
|
||||
return (
|
||||
ne_weights.cuda(),
|
||||
ne_mods.cuda(),
|
||||
exclusive_sums.cuda(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", [1, 2, 17, 128, 1024])
|
||||
def test_compute_n_gram_ids_decode_matches_general(batch_size: int) -> None:
|
||||
ne_n = 8
|
||||
ne_k = 2
|
||||
vocab_size = 32000
|
||||
max_context_len = 1024
|
||||
max_running_reqs = batch_size + 8
|
||||
num_configs = (ne_n - 1) * ne_k
|
||||
|
||||
ne_weights, ne_mods, exclusive_sums = _make_ngram_params(ne_n, ne_k, vocab_size)
|
||||
ne_token_table = torch.randint(
|
||||
0,
|
||||
vocab_size,
|
||||
(max_running_reqs, max_context_len),
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
row_indices = torch.randperm(max_running_reqs, device="cuda")[:batch_size].to(
|
||||
torch.int64
|
||||
)
|
||||
column_starts = torch.randint(
|
||||
0, max_context_len, (batch_size,), dtype=torch.int32, device="cuda"
|
||||
)
|
||||
tokens = torch.empty(batch_size, dtype=torch.int32, device="cuda")
|
||||
exclusive_req_len_sums = torch.arange(
|
||||
batch_size + 1, dtype=torch.int32, device="cuda"
|
||||
)
|
||||
n_gram_ids_general = torch.empty(
|
||||
(batch_size, num_configs), dtype=torch.int32, device="cuda"
|
||||
)
|
||||
n_gram_ids_decode = torch.empty_like(n_gram_ids_general)
|
||||
|
||||
compute_n_gram_ids(
|
||||
ne_n=ne_n,
|
||||
ne_k=ne_k,
|
||||
ne_weights=ne_weights,
|
||||
ne_mods=ne_mods,
|
||||
exclusive_ne_embedder_size_sums=exclusive_sums,
|
||||
tokens=tokens,
|
||||
exclusive_req_len_sums=exclusive_req_len_sums,
|
||||
ne_token_table=ne_token_table,
|
||||
row_indices=row_indices,
|
||||
column_starts=column_starts,
|
||||
n_gram_ids=n_gram_ids_general,
|
||||
)
|
||||
compute_n_gram_ids_decode(
|
||||
ne_n=ne_n,
|
||||
ne_k=ne_k,
|
||||
ne_weights=ne_weights,
|
||||
ne_mods=ne_mods,
|
||||
exclusive_ne_embedder_size_sums=exclusive_sums,
|
||||
ne_token_table=ne_token_table,
|
||||
row_indices=row_indices,
|
||||
column_starts=column_starts,
|
||||
n_gram_ids=n_gram_ids_decode,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(n_gram_ids_decode, n_gram_ids_general, atol=0, rtol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", [1, 2, 17, 128, 1024])
|
||||
def test_update_token_table_decode_matches_general(batch_size: int) -> None:
|
||||
max_context_len = 4096
|
||||
max_running_reqs = batch_size + 8
|
||||
tokens = torch.arange(batch_size, dtype=torch.int32, device="cuda") + 100
|
||||
row_indices = torch.randperm(max_running_reqs, device="cuda")[:batch_size].to(
|
||||
torch.int64
|
||||
)
|
||||
column_starts = torch.randint(
|
||||
0, max_context_len, (batch_size,), dtype=torch.int32, device="cuda"
|
||||
)
|
||||
req_lens = torch.ones(batch_size, dtype=torch.int32, device="cuda")
|
||||
|
||||
token_table_general = torch.full(
|
||||
(max_running_reqs, max_context_len), -1, dtype=torch.int32, device="cuda"
|
||||
)
|
||||
token_table_decode = token_table_general.clone()
|
||||
|
||||
update_token_table(
|
||||
tokens=tokens,
|
||||
ne_token_table=token_table_general,
|
||||
row_indices=row_indices,
|
||||
column_starts=column_starts,
|
||||
req_lens=req_lens,
|
||||
ignore_tokens=None,
|
||||
)
|
||||
update_token_table_decode(
|
||||
tokens=tokens,
|
||||
ne_token_table=token_table_decode,
|
||||
row_indices=row_indices,
|
||||
column_starts=column_starts,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(token_table_decode, token_table_general, atol=0, rtol=0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,137 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.nvfp4 import (
|
||||
cutlass_fp4_group_mm,
|
||||
scaled_fp4_experts_quant,
|
||||
scaled_fp4_quant,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=5, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
FLOAT4_E2M1_MAX = 6.0
|
||||
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
|
||||
|
||||
|
||||
def _nvfp4_supported() -> bool:
|
||||
return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0)
|
||||
|
||||
|
||||
def _round_up(x: int, y: int) -> int:
|
||||
return ((x + y - 1) // y) * y
|
||||
|
||||
|
||||
def _build_expert_offsets(
|
||||
m_per_expert: list[int], device: torch.device
|
||||
) -> torch.Tensor:
|
||||
offsets = [0]
|
||||
for m in m_per_expert:
|
||||
offsets.append(offsets[-1] + m)
|
||||
return torch.tensor(offsets, dtype=torch.int32, device=device)
|
||||
|
||||
|
||||
def _build_blockscale_offsets(
|
||||
m_per_expert: list[int], device: torch.device
|
||||
) -> torch.Tensor:
|
||||
offsets = [0]
|
||||
for m in m_per_expert:
|
||||
offsets.append(offsets[-1] + _round_up(m, 128))
|
||||
return torch.tensor(offsets, dtype=torch.int32, device=device)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0"
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
def test_nvfp4_blockwise_moe_grouped_mm(dtype: torch.dtype) -> None:
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda")
|
||||
|
||||
num_experts = 4
|
||||
m_per_expert = [33, 17, 48, 29]
|
||||
n = 256
|
||||
k = 128
|
||||
|
||||
expert_offsets_full = _build_expert_offsets(m_per_expert, device)
|
||||
blockscale_offsets_full = _build_blockscale_offsets(m_per_expert, device)
|
||||
|
||||
total_m = int(expert_offsets_full[-1].item())
|
||||
a = torch.randn((total_m, k), device=device, dtype=dtype) * 0.1
|
||||
b = torch.randn((num_experts, n, k), device=device, dtype=dtype) * 0.1
|
||||
|
||||
a_global_scale = torch.empty((num_experts,), device=device, dtype=torch.float32)
|
||||
for i in range(num_experts):
|
||||
start = int(expert_offsets_full[i].item())
|
||||
end = int(expert_offsets_full[i + 1].item())
|
||||
amax = a[start:end].abs().max().to(torch.float32)
|
||||
a_global_scale[i] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / amax
|
||||
|
||||
b_global_scale = torch.empty((num_experts,), device=device, dtype=torch.float32)
|
||||
for i in range(num_experts):
|
||||
bmax = b[i].abs().max().to(torch.float32)
|
||||
b_global_scale[i] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / bmax
|
||||
|
||||
a_fp4, a_blockscale = scaled_fp4_experts_quant(
|
||||
a,
|
||||
a_global_scale,
|
||||
expert_offsets_full,
|
||||
blockscale_offsets_full,
|
||||
topk=1,
|
||||
)
|
||||
|
||||
b_fp4 = torch.empty((num_experts, n, k // 2), device=device, dtype=torch.uint8)
|
||||
b_blockscale = torch.empty(
|
||||
(num_experts, _round_up(n, 128), _round_up(k // 16, 4)),
|
||||
device=device,
|
||||
dtype=torch.float8_e4m3fn,
|
||||
)
|
||||
for i in range(num_experts):
|
||||
b_fp4_i, b_scale_i = scaled_fp4_quant(b[i], b_global_scale[i])
|
||||
b_fp4[i].copy_(b_fp4_i)
|
||||
b_blockscale[i].copy_(b_scale_i)
|
||||
|
||||
alphas = (1.0 / (a_global_scale * b_global_scale)).to(torch.float32)
|
||||
|
||||
params = {
|
||||
"ab_strides": torch.full((num_experts,), k, dtype=torch.int64, device=device),
|
||||
"c_strides": torch.full((num_experts,), n, dtype=torch.int64, device=device),
|
||||
"problem_sizes": torch.tensor(
|
||||
[[m, n, k] for m in m_per_expert], dtype=torch.int32, device=device
|
||||
),
|
||||
"expert_offsets": expert_offsets_full[:-1].contiguous(),
|
||||
"blockscale_offsets": blockscale_offsets_full[:-1].contiguous(),
|
||||
"a_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
|
||||
"b_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
|
||||
"out_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
|
||||
"a_scales_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
|
||||
"b_scales_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
|
||||
"alpha_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
|
||||
"layout_sfa": torch.empty((num_experts, 5), dtype=torch.int64, device=device),
|
||||
"layout_sfb": torch.empty((num_experts, 5), dtype=torch.int64, device=device),
|
||||
}
|
||||
|
||||
out = cutlass_fp4_group_mm(
|
||||
a_fp4,
|
||||
b_fp4,
|
||||
a_blockscale,
|
||||
b_blockscale,
|
||||
alphas,
|
||||
dtype,
|
||||
params,
|
||||
)
|
||||
|
||||
ref = torch.empty((total_m, n), device=device, dtype=dtype)
|
||||
for i in range(num_experts):
|
||||
start = int(expert_offsets_full[i].item())
|
||||
end = int(expert_offsets_full[i + 1].item())
|
||||
ref[start:end] = torch.matmul(a[start:end], b[i].t())
|
||||
|
||||
torch.testing.assert_close(out, ref, atol=1e-1, rtol=1e-1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,152 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=5, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
def _nvfp4_supported() -> bool:
|
||||
return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0)
|
||||
|
||||
|
||||
DTYPES = [torch.float16, torch.bfloat16]
|
||||
SHAPES = [
|
||||
(128, 128, 64),
|
||||
(128, 128, 128),
|
||||
(256, 128, 64),
|
||||
(128, 256, 128),
|
||||
(150, 128, 64),
|
||||
]
|
||||
|
||||
FLOAT4_E2M1_MAX = 6.0
|
||||
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
|
||||
|
||||
K_E2M1_TO_FLOAT = [
|
||||
0.0,
|
||||
0.5,
|
||||
1.0,
|
||||
1.5,
|
||||
2.0,
|
||||
3.0,
|
||||
4.0,
|
||||
6.0,
|
||||
]
|
||||
|
||||
|
||||
def e2m1_to_fp32(int4_value: int) -> float:
|
||||
sign_bit = int4_value & 0x8
|
||||
int4_abs_value = int4_value & 0x7
|
||||
float_result = K_E2M1_TO_FLOAT[int4_abs_value]
|
||||
return -float_result if sign_bit else float_result
|
||||
|
||||
|
||||
def break_fp4_bytes(a: torch.Tensor) -> torch.Tensor:
|
||||
assert a.dtype == torch.uint8
|
||||
m, n = a.shape
|
||||
a = a.flatten()
|
||||
high_half_byte = (a & 0xF0) >> 4
|
||||
low_half_byte = a & 0x0F
|
||||
f_h = torch.tensor([e2m1_to_fp32(x) for x in high_half_byte], device=a.device)
|
||||
f_l = torch.tensor([e2m1_to_fp32(x) for x in low_half_byte], device=a.device)
|
||||
return torch.stack((f_l, f_h), dim=-1).reshape(m, n * 2)
|
||||
|
||||
|
||||
def convert_swizzled_to_linear(
|
||||
a_sf_swizzled: torch.Tensor, m: int, k: int, block_size: int
|
||||
) -> torch.Tensor:
|
||||
sf_m, sf_k = a_sf_swizzled.shape
|
||||
del sf_m, sf_k
|
||||
m_tiles = (m + 128 - 1) // 128
|
||||
f = block_size * 4
|
||||
k_tiles = (k + f - 1) // f
|
||||
tmp = torch.reshape(a_sf_swizzled, (1, m_tiles, k_tiles, 32, 4, 4))
|
||||
tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5))
|
||||
out = tmp.reshape(m_tiles * 128, k_tiles * f // block_size)
|
||||
return out[0:m, 0 : k // block_size]
|
||||
|
||||
|
||||
def dequantize_to_dtype(
|
||||
tensor_fp4: torch.Tensor,
|
||||
tensor_sf: torch.Tensor,
|
||||
global_scale: torch.Tensor,
|
||||
block_size: int = 16,
|
||||
) -> torch.Tensor:
|
||||
assert tensor_fp4.dtype == torch.uint8
|
||||
m, packed_k = tensor_fp4.shape
|
||||
k = packed_k * 2
|
||||
tensor_f32 = break_fp4_bytes(tensor_fp4)
|
||||
tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size)
|
||||
tensor_sf = tensor_sf.view(torch.float8_e4m3fn)
|
||||
tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size)
|
||||
tensor_sf_dtype = tensor_sf.to(torch.float32) / global_scale
|
||||
return (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k)
|
||||
|
||||
|
||||
def get_ref_results(
|
||||
a_fp4: torch.Tensor,
|
||||
b_fp4: torch.Tensor,
|
||||
a_sf: torch.Tensor,
|
||||
b_sf: torch.Tensor,
|
||||
a_global_scale: torch.Tensor,
|
||||
b_global_scale: torch.Tensor,
|
||||
block_size: int,
|
||||
) -> torch.Tensor:
|
||||
a_in_dtype = dequantize_to_dtype(a_fp4, a_sf, a_global_scale, block_size=block_size)
|
||||
b_in_dtype = dequantize_to_dtype(b_fp4, b_sf, b_global_scale, block_size=block_size)
|
||||
return torch.matmul(a_in_dtype, b_in_dtype.t())
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0"
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("shape", SHAPES)
|
||||
def test_nvfp4_gemm(dtype: torch.dtype, shape: tuple[int, int, int]) -> None:
|
||||
m, n, packed_k = shape
|
||||
k = packed_k * 2
|
||||
block_size = 16
|
||||
|
||||
a_dtype = torch.randn((m, k), dtype=dtype, device="cuda")
|
||||
b_dtype = torch.randn((n, k), dtype=dtype, device="cuda")
|
||||
|
||||
a_global_scale = (
|
||||
(FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a_dtype.flatten(), dim=-1)
|
||||
).to(torch.float32)
|
||||
b_global_scale = (
|
||||
(FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(b_dtype.flatten(), dim=-1)
|
||||
).to(torch.float32)
|
||||
|
||||
alpha = 1.0 / (a_global_scale * b_global_scale)
|
||||
|
||||
a_fp4, a_scale_interleaved = scaled_fp4_quant(a_dtype, a_global_scale)
|
||||
b_fp4, b_scale_interleaved = scaled_fp4_quant(b_dtype, b_global_scale)
|
||||
|
||||
expected_out = get_ref_results(
|
||||
a_fp4,
|
||||
b_fp4,
|
||||
a_scale_interleaved,
|
||||
b_scale_interleaved,
|
||||
a_global_scale,
|
||||
b_global_scale,
|
||||
block_size,
|
||||
)
|
||||
|
||||
out = cutlass_scaled_fp4_mm(
|
||||
a_fp4,
|
||||
b_fp4,
|
||||
a_scale_interleaved,
|
||||
b_scale_interleaved,
|
||||
alpha,
|
||||
dtype,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(out, expected_out.to(dtype=dtype), atol=1e-1, rtol=1e-1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,225 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.nvfp4 import (
|
||||
scaled_fp4_grouped_quant,
|
||||
scaled_fp4_quant,
|
||||
silu_and_mul_scaled_fp4_grouped_quant,
|
||||
)
|
||||
|
||||
try:
|
||||
from sgl_kernel import silu_and_mul as _sgl_silu_and_mul
|
||||
except Exception:
|
||||
_sgl_silu_and_mul = None
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=5, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
def _nvfp4_supported() -> bool:
|
||||
return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0)
|
||||
|
||||
|
||||
def _silu_and_mul_reference(x: torch.Tensor) -> torch.Tensor:
|
||||
if _sgl_silu_and_mul is not None:
|
||||
return _sgl_silu_and_mul(x)
|
||||
k = x.shape[-1] // 2
|
||||
return torch.nn.functional.silu(x[:, :, :k]) * x[:, :, k:]
|
||||
|
||||
|
||||
DTYPES = [torch.float16, torch.bfloat16]
|
||||
SHAPES = [(128, 64), (128, 128), (256, 64), (256, 128)]
|
||||
PAD_SHAPES = [
|
||||
(90, 64),
|
||||
(150, 64),
|
||||
(128, 48),
|
||||
(128, 80),
|
||||
]
|
||||
|
||||
FLOAT4_E2M1_MAX = 6.0
|
||||
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
|
||||
BLOCK_SIZE = 16
|
||||
|
||||
E2M1_TO_FLOAT32 = [
|
||||
0.0,
|
||||
0.5,
|
||||
1.0,
|
||||
1.5,
|
||||
2.0,
|
||||
3.0,
|
||||
4.0,
|
||||
6.0,
|
||||
0.0,
|
||||
-0.5,
|
||||
-1.0,
|
||||
-1.5,
|
||||
-2.0,
|
||||
-3.0,
|
||||
-4.0,
|
||||
-6.0,
|
||||
]
|
||||
|
||||
|
||||
def cast_from_fp4(x: torch.Tensor, m: int, n: int) -> torch.Tensor:
|
||||
v_2nd = (x & 0xF).to(torch.long)
|
||||
v_1st = ((x >> 4) & 0xF).to(torch.long)
|
||||
c = torch.stack((v_2nd, v_1st), dim=-1).flatten()
|
||||
lut = torch.tensor(E2M1_TO_FLOAT32, device=x.device, dtype=torch.float32)
|
||||
return lut[c].reshape(m, n)
|
||||
|
||||
|
||||
def cast_to_fp4(x: torch.Tensor) -> torch.Tensor:
|
||||
sign = torch.sign(x)
|
||||
x = torch.abs(x)
|
||||
x[(x >= 0.0) & (x <= 0.25)] = 0.0
|
||||
x[(x > 0.25) & (x < 0.75)] = 0.5
|
||||
x[(x >= 0.75) & (x <= 1.25)] = 1.0
|
||||
x[(x > 1.25) & (x < 1.75)] = 1.5
|
||||
x[(x >= 1.75) & (x <= 2.5)] = 2.0
|
||||
x[(x > 2.5) & (x < 3.5)] = 3.0
|
||||
x[(x >= 3.5) & (x <= 5.0)] = 4.0
|
||||
x[x > 5.0] = 6.0
|
||||
return x * sign
|
||||
|
||||
|
||||
def get_reciprocal(x):
|
||||
if isinstance(x, torch.Tensor):
|
||||
return torch.where(x == 0, torch.tensor(0.0, dtype=x.dtype), 1.0 / x)
|
||||
return 0.0 if x == 0 else 1.0 / x
|
||||
|
||||
|
||||
def ref_nvfp4_quant(x: torch.Tensor, global_scale: torch.Tensor):
|
||||
assert global_scale.dtype == torch.float32
|
||||
assert x.ndim == 2
|
||||
m, n = x.shape
|
||||
x = torch.reshape(x, (m, n // BLOCK_SIZE, BLOCK_SIZE))
|
||||
vec_max = torch.max(torch.abs(x), dim=-1, keepdim=True)[0].to(torch.float32)
|
||||
scale = global_scale * (vec_max * get_reciprocal(FLOAT4_E2M1_MAX))
|
||||
scale = scale.to(torch.float8_e4m3fn).to(torch.float32)
|
||||
output_scale = get_reciprocal(scale * get_reciprocal(global_scale))
|
||||
|
||||
scaled_x = x.to(torch.float32) * output_scale
|
||||
clipped_x = torch.clamp(scaled_x, -6.0, 6.0).reshape(m, n)
|
||||
return cast_to_fp4(clipped_x), scale.squeeze(-1)
|
||||
|
||||
|
||||
def recover_swizzled_scales(scale: torch.Tensor, m: int, n: int) -> torch.Tensor:
|
||||
rounded_m = ((m + 128 - 1) // 128) * 128
|
||||
scale_n = n // BLOCK_SIZE
|
||||
rounded_n = ((scale_n + 4 - 1) // 4) * 4
|
||||
tmp = torch.reshape(scale, (1, rounded_m // 128, rounded_n // 4, 32, 4, 4))
|
||||
tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5))
|
||||
result = torch.reshape(tmp, (rounded_m, rounded_n)).to(torch.float32)
|
||||
return result[:m, :scale_n]
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0"
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("shape", SHAPES)
|
||||
def test_quantize_to_fp4(dtype: torch.dtype, shape: tuple[int, int]) -> None:
|
||||
torch.manual_seed(42)
|
||||
m, n = shape
|
||||
|
||||
x = torch.randn((m, n), dtype=dtype, device="cuda")
|
||||
tensor_amax = torch.abs(x).max().to(torch.float32)
|
||||
global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
|
||||
out_ref, scale_ref = ref_nvfp4_quant(x, global_scale)
|
||||
|
||||
out, out_scale = scaled_fp4_quant(x, global_scale)
|
||||
scale_ans = recover_swizzled_scales(out_scale, m, n)
|
||||
out_ans = cast_from_fp4(out, m, n)
|
||||
|
||||
torch.testing.assert_close(out_ans, out_ref)
|
||||
torch.testing.assert_close(scale_ans, scale_ref)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0"
|
||||
)
|
||||
@pytest.mark.parametrize("shape", PAD_SHAPES)
|
||||
def test_quantize_to_fp4_padded(shape: tuple[int, int]) -> None:
|
||||
torch.manual_seed(42)
|
||||
m, n = shape
|
||||
x = torch.randn((m, n), dtype=torch.float16, device="cuda")
|
||||
|
||||
tensor_amax = torch.abs(x).max().to(torch.float32)
|
||||
global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
|
||||
out_ref, scale_ref = ref_nvfp4_quant(x, global_scale)
|
||||
|
||||
out, out_scale = scaled_fp4_quant(x, global_scale)
|
||||
scale_ans = recover_swizzled_scales(out_scale, m, n)
|
||||
out_ans = cast_from_fp4(out, m, n)
|
||||
|
||||
torch.testing.assert_close(out_ans, out_ref)
|
||||
torch.testing.assert_close(scale_ans, scale_ref)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0"
|
||||
)
|
||||
@pytest.mark.parametrize("shape", [(2, 128, 512), (2, 100, 128)])
|
||||
def test_quantize_to_fp4_grouped(shape: tuple[int, int, int]) -> None:
|
||||
torch.manual_seed(42)
|
||||
l, m, k = shape
|
||||
|
||||
x = torch.randn((l, m, k), dtype=torch.bfloat16, device="cuda")
|
||||
mask = torch.randint(1, max(2, m // 2), (l,), dtype=torch.int32, device="cuda")
|
||||
tensor_amax = x.abs().amax(dim=(1, 2)).to(torch.float32)
|
||||
x_sf_global = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
|
||||
|
||||
output, output_scales = scaled_fp4_grouped_quant(x, x_sf_global, mask)
|
||||
output = output.permute(2, 0, 1)
|
||||
padded_m = ((m + 128 - 1) // 128) * 128
|
||||
output_scales = output_scales.permute(5, 2, 4, 0, 1, 3).view(l, padded_m, -1)
|
||||
|
||||
for i in range(l):
|
||||
a_fp4, a_scale_interleaved = scaled_fp4_quant(x[i], x_sf_global[i])
|
||||
torch.testing.assert_close(a_fp4[: mask[i]], output[i][: mask[i]])
|
||||
scale_ref = recover_swizzled_scales(a_scale_interleaved, m, k)
|
||||
scale_ans = recover_swizzled_scales(output_scales[i], m, k)
|
||||
torch.testing.assert_close(scale_ref[: mask[i]], scale_ans[: mask[i]])
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0"
|
||||
)
|
||||
@pytest.mark.parametrize("shape", [(4, 96, 256), (8, 128, 512)])
|
||||
def test_silu_and_mul_quantize_to_fp4_grouped(shape: tuple[int, int, int]) -> None:
|
||||
torch.manual_seed(42)
|
||||
l, m, k = shape
|
||||
|
||||
x = torch.randn((l, m, k * 2), dtype=torch.bfloat16, device="cuda")
|
||||
mask = torch.randint(1, max(2, m // 2), (l,), dtype=torch.int32, device="cuda")
|
||||
|
||||
ref_y = _silu_and_mul_reference(x)
|
||||
|
||||
tensor_amax = ref_y.abs().amax(dim=(1, 2)).to(torch.float32)
|
||||
y_sf_global = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
|
||||
|
||||
ref_output, ref_output_scales = scaled_fp4_grouped_quant(ref_y, y_sf_global, mask)
|
||||
output, output_scales = silu_and_mul_scaled_fp4_grouped_quant(x, y_sf_global, mask)
|
||||
|
||||
output = output.permute(2, 0, 1)
|
||||
ref_output = ref_output.permute(2, 0, 1)
|
||||
|
||||
padded_m = ((m + 128 - 1) // 128) * 128
|
||||
output_scales = output_scales.permute(5, 2, 4, 0, 1, 3).view(l, padded_m, -1)
|
||||
ref_output_scales = ref_output_scales.permute(5, 2, 4, 0, 1, 3).view(
|
||||
l, padded_m, -1
|
||||
)
|
||||
|
||||
for i in range(l):
|
||||
torch.testing.assert_close(ref_output[i, : mask[i]], output[i, : mask[i]])
|
||||
scale_ref = recover_swizzled_scales(ref_output_scales[i], m, k)
|
||||
scale_ans = recover_swizzled_scales(output_scales[i], m, k)
|
||||
torch.testing.assert_close(scale_ref[: mask[i]], scale_ans[: mask[i]])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,91 @@
|
||||
import itertools
|
||||
import sys
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=16, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
try:
|
||||
from sglang.srt.utils import is_hip
|
||||
|
||||
_is_hip = is_hip()
|
||||
except ImportError:
|
||||
_is_hip = False
|
||||
|
||||
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
|
||||
|
||||
|
||||
def sglang_scaled_fp8_quant(
|
||||
input: torch.Tensor,
|
||||
scale: Optional[torch.Tensor] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
fp8_type_: torch.dtype = torch.float8_e4m3fn
|
||||
output = torch.empty_like(input, device=input.device, dtype=fp8_type_)
|
||||
is_static = True
|
||||
if scale is None:
|
||||
scale = torch.zeros(1, device=input.device, dtype=torch.float32)
|
||||
is_static = False
|
||||
per_tensor_quant_fp8(input, output, scale, is_static)
|
||||
|
||||
return output, scale
|
||||
|
||||
|
||||
def torch_scaled_fp8_quant(tensor, inv_scale):
|
||||
finfo = torch.finfo(torch.float8_e4m3fn)
|
||||
scale = inv_scale.reciprocal()
|
||||
qweight = (tensor.to(torch.float32) * scale).clamp(min=finfo.min, max=finfo.max)
|
||||
qweight = qweight.to(torch.float8_e4m3fn)
|
||||
return qweight
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"num_tokens,hidden_dim",
|
||||
list(itertools.product([128, 256, 512], [512, 2048, 4096])),
|
||||
)
|
||||
def test_jit_per_tensor_quant_compare_implementations(
|
||||
num_tokens: int,
|
||||
hidden_dim: int,
|
||||
):
|
||||
device = torch.device("cuda")
|
||||
x = torch.rand((num_tokens, hidden_dim), dtype=torch.float16, device=device)
|
||||
|
||||
sglang_out, sglang_scale = sglang_scaled_fp8_quant(x)
|
||||
torch_out = torch_scaled_fp8_quant(x, sglang_scale)
|
||||
|
||||
torch.testing.assert_close(
|
||||
sglang_out.float(), torch_out.float(), rtol=1e-3, atol=1e-3
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", [(4, 8, 64), (2, 16, 128), (19260817, 1, 1)])
|
||||
def test_jit_per_tensor_quant_supports_3d(shape):
|
||||
device = torch.device("cuda")
|
||||
x = torch.rand(shape, dtype=torch.bfloat16, device=device)
|
||||
out = torch.empty_like(x, device=x.device, dtype=fp8_type_)
|
||||
scale = torch.zeros(1, device=x.device, dtype=torch.float32)
|
||||
|
||||
per_tensor_quant_fp8(x, out, scale, is_static=False)
|
||||
|
||||
x_2d = x.flatten(0, -2)
|
||||
out_ref_2d = torch_scaled_fp8_quant(x_2d, scale)
|
||||
out_ref = out_ref_2d.reshape(shape)
|
||||
|
||||
torch.testing.assert_close(out.float(), out_ref.float(), rtol=1e-3, atol=1e-3)
|
||||
|
||||
scale = torch.rand(1, dtype=torch.float32, device=device)
|
||||
sglang_out, _ = sglang_scaled_fp8_quant(x, scale)
|
||||
torch_out = torch_scaled_fp8_quant(x, scale)
|
||||
|
||||
torch.testing.assert_close(
|
||||
sglang_out.float(), torch_out.float(), rtol=1e-3, atol=1e-3
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,210 @@
|
||||
import itertools
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.utils import is_hip
|
||||
|
||||
_is_hip = is_hip()
|
||||
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
|
||||
|
||||
from sgl_kernel.test_utils import (
|
||||
assert_all_close_or_tiny_diff,
|
||||
create_per_token_group_quant_test_data,
|
||||
)
|
||||
|
||||
from sglang.jit_kernel.per_token_group_quant_8bit import (
|
||||
per_token_group_quant_8bit as sglang_per_token_group_quant_8bit,
|
||||
)
|
||||
from sglang.srt.layers.quantization.fp8_kernel import (
|
||||
create_per_token_group_quant_fp8_output_scale,
|
||||
)
|
||||
from sglang.srt.layers.quantization.fp8_kernel import (
|
||||
per_token_group_quant_8bit as triton_per_token_group_quant_8bit,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=16, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
configs = list(
|
||||
itertools.product(
|
||||
[1, 4, 16, 64, 127, 128, 512, 1024, 4096, 8192], # num_tokens
|
||||
[128, 256, 384, 512, 1024, 1536, 1664, 2048, 4096, 7168, 16384], # hidden_dim
|
||||
[16, 32, 64, 128], # group_size
|
||||
[None], # num_ranks
|
||||
[fp8_type_], # dtype
|
||||
[
|
||||
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,
|
||||
),
|
||||
],
|
||||
)
|
||||
) + list(
|
||||
itertools.product(
|
||||
[1, 4, 1 * 8, 4 * 8, 64 * 8, 256 * 8, 768 * 8],
|
||||
[2048],
|
||||
[128],
|
||||
[8, 16, 32, 48],
|
||||
[fp8_type_],
|
||||
[
|
||||
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",
|
||||
),
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@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
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,494 @@
|
||||
import sys
|
||||
import time
|
||||
from typing import Optional, Tuple, Union
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.jit_kernel.rope import rotary_embedding
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=18, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def burn_kernel(out_ptr, iters: tl.constexpr):
|
||||
pid = tl.program_id(0)
|
||||
x = tl.full((), pid + 1, dtype=tl.uint32)
|
||||
|
||||
a = tl.full((), 1664525, dtype=tl.uint32)
|
||||
c = tl.full((), 1013904223, dtype=tl.uint32)
|
||||
sh = tl.full((), 13, dtype=tl.uint32)
|
||||
|
||||
for _ in range(iters):
|
||||
x = x * a + c
|
||||
x = x ^ (x >> sh)
|
||||
|
||||
if pid == 0:
|
||||
tl.store(out_ptr, x)
|
||||
|
||||
|
||||
def triton_burn(ms: float, grid=(256,)):
|
||||
iters = int(ms * 20000)
|
||||
out = torch.empty((), device="cuda", dtype=torch.uint32)
|
||||
burn_kernel[grid](out, iters=iters)
|
||||
return out
|
||||
|
||||
|
||||
def create_test_inputs(
|
||||
head_size, batch_size, seq_len, device, dtype, num_q_heads, num_kv_heads
|
||||
):
|
||||
"""Create test inputs."""
|
||||
total_tokens = batch_size * seq_len
|
||||
|
||||
query = torch.randn(
|
||||
batch_size, seq_len, num_q_heads, head_size, dtype=dtype, device=device
|
||||
)
|
||||
key = torch.randn(
|
||||
batch_size, seq_len, num_kv_heads, head_size, dtype=dtype, device=device
|
||||
)
|
||||
|
||||
pos_ids = torch.randint(
|
||||
0, min(seq_len * 2, 100), (total_tokens,), dtype=torch.long, device=device
|
||||
)
|
||||
|
||||
query = query.view(total_tokens, num_q_heads, head_size)
|
||||
key = key.view(total_tokens, num_kv_heads, head_size)
|
||||
|
||||
return query, key, pos_ids
|
||||
|
||||
|
||||
def create_cos_sin_cache(rotary_dim, max_position_embeddings, base, dtype, device):
|
||||
"""Create cos/sin cache for rotary embedding."""
|
||||
max_pos = max_position_embeddings
|
||||
extended_max_pos = max(max_pos, 100)
|
||||
cos_sin_cache = torch.zeros(
|
||||
extended_max_pos, rotary_dim, dtype=dtype, device=device
|
||||
)
|
||||
|
||||
inv_freq = 1.0 / (
|
||||
base
|
||||
** (
|
||||
torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=device)
|
||||
/ rotary_dim
|
||||
)
|
||||
)
|
||||
t = torch.arange(extended_max_pos, dtype=torch.float32, device=device)
|
||||
freqs = torch.outer(t, inv_freq)
|
||||
cos_cache = torch.cos(freqs).to(dtype)
|
||||
sin_cache = torch.sin(freqs).to(dtype)
|
||||
|
||||
cos_sin_cache[:, : rotary_dim // 2] = cos_cache
|
||||
cos_sin_cache[:, rotary_dim // 2 :] = sin_cache
|
||||
|
||||
return cos_sin_cache
|
||||
|
||||
|
||||
# vLLM torch native
|
||||
def _apply_rotary_emb(
|
||||
x: torch.Tensor,
|
||||
cos: torch.Tensor,
|
||||
sin: torch.Tensor,
|
||||
is_neox_style: bool,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
x: [num_tokens, num_heads, head_size]
|
||||
cos: [num_tokens, head_size // 2]
|
||||
sin: [num_tokens, head_size // 2]
|
||||
is_neox_style: Whether to use the Neox-style or GPT-J-style rotary
|
||||
positional embeddings.
|
||||
"""
|
||||
cos = cos.unsqueeze(-2).to(x.dtype)
|
||||
sin = sin.unsqueeze(-2).to(x.dtype)
|
||||
if is_neox_style:
|
||||
x1, x2 = torch.chunk(x, 2, dim=-1)
|
||||
else:
|
||||
x1 = x[..., ::2]
|
||||
x2 = x[..., 1::2]
|
||||
o1 = x1 * cos - x2 * sin
|
||||
o2 = x2 * cos + x1 * sin
|
||||
if is_neox_style:
|
||||
return torch.cat((o1, o2), dim=-1)
|
||||
else:
|
||||
return torch.stack((o1, o2), dim=-1).flatten(-2)
|
||||
|
||||
|
||||
class RotaryEmbedding(torch.nn.Module):
|
||||
# Reference: https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/rotary_embedding.py
|
||||
def __init__(
|
||||
self,
|
||||
head_size: int,
|
||||
rotary_dim: int,
|
||||
max_position_embeddings: int,
|
||||
base: int,
|
||||
is_neox_style: bool,
|
||||
dtype: torch.dtype,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.head_size = head_size
|
||||
self.rotary_dim = rotary_dim
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.base = base
|
||||
self.is_neox_style = is_neox_style
|
||||
self.dtype = dtype
|
||||
|
||||
cache = self._compute_cos_sin_cache()
|
||||
self.cos_sin_cache: torch.Tensor
|
||||
self.register_buffer("cos_sin_cache", cache, persistent=False)
|
||||
|
||||
def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor:
|
||||
inv_freq = 1.0 / (
|
||||
base
|
||||
** (
|
||||
torch.arange(0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim
|
||||
)
|
||||
)
|
||||
return inv_freq
|
||||
|
||||
def _compute_cos_sin_cache(self) -> torch.Tensor:
|
||||
"""Compute the cos and sin cache."""
|
||||
inv_freq = self._compute_inv_freq(self.base)
|
||||
t = torch.arange(self.max_position_embeddings, dtype=torch.float)
|
||||
|
||||
freqs = torch.einsum("i,j -> ij", t, inv_freq)
|
||||
cos = freqs.cos()
|
||||
sin = freqs.sin()
|
||||
cache = torch.cat((cos, sin), dim=-1)
|
||||
return cache
|
||||
|
||||
def forward_native(
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: Optional[torch.Tensor] = None,
|
||||
offsets: Optional[torch.Tensor] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""A PyTorch-native implementation of forward()."""
|
||||
|
||||
if offsets is not None:
|
||||
positions = positions + offsets
|
||||
|
||||
positions = positions.flatten()
|
||||
num_tokens = positions.shape[0]
|
||||
cos_sin = self.cos_sin_cache.index_select(0, positions)
|
||||
|
||||
cos, sin = cos_sin.chunk(2, dim=-1)
|
||||
|
||||
query_shape = query.shape
|
||||
query = query.view(num_tokens, -1, self.head_size)
|
||||
query_rot = query[..., : self.rotary_dim]
|
||||
query_pass = query[..., self.rotary_dim :]
|
||||
query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)
|
||||
query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
|
||||
|
||||
# Modification: convert to the correct dtype
|
||||
query = query.to(self.dtype)
|
||||
|
||||
if key is not None:
|
||||
key_shape = key.shape
|
||||
key = key.view(num_tokens, -1, self.head_size)
|
||||
key_rot = key[..., : self.rotary_dim]
|
||||
key_pass = key[..., self.rotary_dim :]
|
||||
key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style)
|
||||
key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
|
||||
|
||||
key = key.to(self.dtype)
|
||||
|
||||
return query, key
|
||||
|
||||
|
||||
def get_torch_rotary_embedding(
|
||||
head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype, device
|
||||
):
|
||||
"""Initialize Torch Native RotaryEmbedding based on vLLM implementation."""
|
||||
return RotaryEmbedding(
|
||||
head_size=head_size,
|
||||
rotary_dim=rotary_dim,
|
||||
max_position_embeddings=max_position_embeddings,
|
||||
base=base,
|
||||
is_neox_style=is_neox_style,
|
||||
dtype=dtype,
|
||||
).to(device)
|
||||
|
||||
|
||||
def get_sgl_rotary_embedding(
|
||||
head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype, device
|
||||
):
|
||||
"""Initialize SglKernelRotaryEmbedding."""
|
||||
try:
|
||||
from sgl_kernel.testing.rotary_embedding import SglKernelRotaryEmbedding
|
||||
except ImportError:
|
||||
pytest.skip(
|
||||
"SglKernelRotaryEmbedding is not available. Test case can be removed."
|
||||
)
|
||||
|
||||
return SglKernelRotaryEmbedding(
|
||||
head_size=head_size,
|
||||
rotary_dim=rotary_dim,
|
||||
max_position_embeddings=max_position_embeddings,
|
||||
base=base,
|
||||
is_neox_style=is_neox_style,
|
||||
dtype=dtype,
|
||||
).to(device)
|
||||
|
||||
|
||||
def compare_results(jit_out, sgl_out, dtype):
|
||||
"""Compare results between JIT and SGL implementations."""
|
||||
if jit_out is None:
|
||||
assert sgl_out is None
|
||||
return
|
||||
|
||||
assert sgl_out is not None
|
||||
|
||||
# Check for NaN values
|
||||
assert not torch.isnan(jit_out).any(), "NaN in JIT results"
|
||||
assert not torch.isnan(sgl_out).any(), "NaN in SGL results"
|
||||
|
||||
# Compare results
|
||||
atol = 4e-2 if dtype != torch.float32 else 1e-5
|
||||
rtol = 4e-2 if dtype != torch.float32 else 1e-5
|
||||
|
||||
torch.testing.assert_close(jit_out, sgl_out, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype, device, batch_size, seq_len, num_q_heads, num_kv_heads",
|
||||
[
|
||||
# GPT-OSS cases
|
||||
*[
|
||||
(64, 64, 4096, 8000, True, torch.bfloat16, "cuda", bs, sl, 8, 8)
|
||||
for bs, sl in [(1, 1), (32, 1), (128, 1), (512, 1), (2, 512), (4, 4096)]
|
||||
],
|
||||
# Other cases
|
||||
(64, 64, 32, 8000, True, torch.bfloat16, "cuda", 32, 32, 1, 1),
|
||||
(256, 128, 4096, 10000, True, torch.bfloat16, "cuda", 2, 512, 4, 2),
|
||||
(512, 128, 311, 10000, True, torch.bfloat16, "cuda", 3, 39, 4, 2),
|
||||
(128, 128, 2048, 10000, False, torch.bfloat16, "cuda", 2, 512, 32, 8),
|
||||
(128, 128, 2048, 10000, False, torch.bfloat16, "cuda", 2, 512, 16, 4),
|
||||
(512, 128, 311, 10000, False, torch.bfloat16, "cuda", 3, 39, 4, 2),
|
||||
(64, 64, 32, 8000, True, torch.float32, "cuda", 32, 32, 1, 1),
|
||||
(256, 128, 4096, 10000, True, torch.float32, "cuda", 2, 512, 4, 2),
|
||||
(512, 128, 311, 10000, True, torch.float32, "cuda", 3, 39, 4, 2),
|
||||
(128, 128, 2048, 10000, False, torch.float32, "cuda", 2, 512, 32, 8),
|
||||
(128, 128, 2048, 10000, False, torch.float32, "cuda", 2, 512, 16, 4),
|
||||
(512, 128, 311, 10000, False, torch.float32, "cuda", 3, 39, 4, 2),
|
||||
# Additional test cases for different head sizes and dtypes
|
||||
(64, 32, 1024, 10000, True, torch.float16, "cuda", 16, 64, 8, 4),
|
||||
(128, 64, 2048, 10000, True, torch.float16, "cuda", 8, 128, 16, 8),
|
||||
(256, 128, 4096, 10000, True, torch.float16, "cuda", 4, 256, 8, 4),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"key_is_none",
|
||||
[True, False],
|
||||
)
|
||||
def test_correctness(
|
||||
head_size,
|
||||
rotary_dim,
|
||||
max_position_embeddings,
|
||||
base,
|
||||
is_neox_style,
|
||||
dtype,
|
||||
device,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_q_heads,
|
||||
num_kv_heads,
|
||||
key_is_none,
|
||||
):
|
||||
"""Test correctness of JIT rotary embedding implementation."""
|
||||
# Create inputs and caches
|
||||
query, key, pos_ids = create_test_inputs(
|
||||
head_size, batch_size, seq_len, device, dtype, num_q_heads, num_kv_heads
|
||||
)
|
||||
cos_sin_cache = create_cos_sin_cache(
|
||||
rotary_dim, max_position_embeddings, base, dtype, device
|
||||
)
|
||||
|
||||
# Initialize torch kernel
|
||||
torch_rotary_emb = get_torch_rotary_embedding(
|
||||
head_size,
|
||||
rotary_dim,
|
||||
max_position_embeddings,
|
||||
base,
|
||||
is_neox_style,
|
||||
dtype,
|
||||
device,
|
||||
)
|
||||
torch_rotary_emb.cos_sin_cache = cos_sin_cache
|
||||
r = torch.randn_like(query)
|
||||
|
||||
# Apply rotary embeddings
|
||||
query_jit, key_jit = query.clone(), key.clone()
|
||||
query_torch, key_torch = query.clone(), key.clone()
|
||||
stream_jit = torch.get_device_module("cuda").Stream()
|
||||
stream_kernel = torch.get_device_module("cuda").Stream()
|
||||
|
||||
if key_is_none:
|
||||
key_jit = None
|
||||
key_torch = None
|
||||
triton_burn(100.0, grid=(1024,))
|
||||
|
||||
r_jit, r_torch = r.clone(), r.clone()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
with torch.cuda.stream(stream_jit):
|
||||
# Test if rotary_embedding runs on stream_jit
|
||||
triton_burn(100.0, grid=(1024,))
|
||||
query_jit = query_jit + r_jit
|
||||
query_jit_out, key_jit_out = rotary_embedding(
|
||||
positions=pos_ids,
|
||||
query=query_jit,
|
||||
key=key_jit,
|
||||
head_size=head_size,
|
||||
cos_sin_cache=cos_sin_cache,
|
||||
is_neox=is_neox_style,
|
||||
)
|
||||
|
||||
with torch.cuda.stream(stream_kernel):
|
||||
triton_burn(100.0, grid=(1024,))
|
||||
query_torch = query_torch + r_torch
|
||||
query_torch_out, key_torch_out = torch_rotary_emb.forward_native(
|
||||
positions=pos_ids, query=query_torch, key=key_torch
|
||||
)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
compare_results(query_jit_out, query_torch_out, dtype)
|
||||
compare_results(key_jit_out, key_torch_out, dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype, device, batch_size, seq_len, num_q_heads, num_kv_heads",
|
||||
[
|
||||
# Small scale
|
||||
(64, 64, 4096, 8000, True, torch.bfloat16, "cuda", 1, 1, 8, 8),
|
||||
(64, 64, 4096, 8000, True, torch.bfloat16, "cuda", 4, 16, 8, 8),
|
||||
# Medium scale
|
||||
(64, 64, 4096, 8000, True, torch.bfloat16, "cuda", 8, 64, 8, 8),
|
||||
(64, 64, 4096, 8000, True, torch.bfloat16, "cuda", 16, 128, 8, 8),
|
||||
# Large scale
|
||||
(64, 64, 4096, 8000, True, torch.bfloat16, "cuda", 32, 512, 8, 8),
|
||||
(64, 64, 4096, 8000, True, torch.bfloat16, "cuda", 64, 1024, 8, 8),
|
||||
],
|
||||
)
|
||||
def test_performance(
|
||||
head_size: int,
|
||||
rotary_dim: int,
|
||||
max_position_embeddings: int,
|
||||
base: int,
|
||||
is_neox_style,
|
||||
dtype,
|
||||
device,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_q_heads,
|
||||
num_kv_heads,
|
||||
):
|
||||
"""Performance test comparing JIT and SGL implementations with accuracy validation."""
|
||||
# Create inputs and caches
|
||||
query, key, pos_ids = create_test_inputs(
|
||||
head_size, batch_size, seq_len, device, dtype, num_q_heads, num_kv_heads
|
||||
)
|
||||
cos_sin_cache = create_cos_sin_cache(
|
||||
rotary_dim, max_position_embeddings, base, dtype, device
|
||||
)
|
||||
|
||||
# Initialize SGL kernel
|
||||
sgl_rotary_emb = get_sgl_rotary_embedding(
|
||||
head_size,
|
||||
rotary_dim,
|
||||
max_position_embeddings,
|
||||
base,
|
||||
is_neox_style,
|
||||
dtype,
|
||||
device,
|
||||
)
|
||||
sgl_rotary_emb.cos_sin_cache = cos_sin_cache
|
||||
|
||||
warmup = 3
|
||||
|
||||
# Warmup runs
|
||||
for _ in range(warmup):
|
||||
query_warm, key_warm = query.clone(), key.clone()
|
||||
rotary_embedding(
|
||||
positions=pos_ids,
|
||||
query=query_warm,
|
||||
key=key_warm,
|
||||
head_size=head_size,
|
||||
cos_sin_cache=cos_sin_cache,
|
||||
is_neox=is_neox_style,
|
||||
)
|
||||
|
||||
query_sgl_warm, key_sgl_warm = query.clone(), key.clone()
|
||||
sgl_rotary_emb.forward_cuda(
|
||||
positions=pos_ids, query=query_sgl_warm, key=key_sgl_warm
|
||||
)
|
||||
|
||||
iteration = 100
|
||||
|
||||
# Time JIT implementation
|
||||
torch.cuda.synchronize()
|
||||
start_time = time.time()
|
||||
for _ in range(iteration):
|
||||
query_jit, key_jit = query.clone(), key.clone()
|
||||
rotary_embedding(
|
||||
positions=pos_ids,
|
||||
query=query_jit,
|
||||
key=key_jit,
|
||||
head_size=head_size,
|
||||
cos_sin_cache=cos_sin_cache,
|
||||
is_neox=is_neox_style,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
jit_time = (time.time() - start_time) / iteration
|
||||
|
||||
# Time SGL implementation
|
||||
torch.cuda.synchronize()
|
||||
start_time = time.time()
|
||||
for _ in range(iteration):
|
||||
query_sgl, key_sgl = query.clone(), key.clone()
|
||||
sgl_rotary_emb.forward_cuda(positions=pos_ids, query=query_sgl, key=key_sgl)
|
||||
torch.cuda.synchronize()
|
||||
sgl_time = (time.time() - start_time) / iteration
|
||||
|
||||
# Accuracy validation during performance test
|
||||
# Run one more time to get outputs for comparison
|
||||
query_jit_final, key_jit_final = query.clone(), key.clone()
|
||||
query_sgl_final, key_sgl_final = query.clone(), key.clone()
|
||||
|
||||
query_jit_out, key_jit_out = rotary_embedding(
|
||||
positions=pos_ids,
|
||||
query=query_jit_final,
|
||||
key=key_jit_final,
|
||||
head_size=head_size,
|
||||
cos_sin_cache=cos_sin_cache,
|
||||
is_neox=is_neox_style,
|
||||
)
|
||||
|
||||
query_sgl_out, key_sgl_out = sgl_rotary_emb.forward_cuda(
|
||||
positions=pos_ids, query=query_sgl_final, key=key_sgl_final
|
||||
)
|
||||
|
||||
# Validate accuracy
|
||||
compare_results(query_jit_out, query_sgl_out, dtype)
|
||||
compare_results(key_jit_out, key_sgl_out, dtype)
|
||||
|
||||
# Print results
|
||||
total_tokens = batch_size * seq_len
|
||||
print(
|
||||
f"\nPerformance Test - Batch={batch_size}, SeqLen={seq_len}, Tokens={total_tokens}"
|
||||
)
|
||||
print(f"JIT: {jit_time*1000:.9f}ms, SGL: {sgl_time*1000:.9f}ms")
|
||||
if sgl_time > 0:
|
||||
speedup = sgl_time / jit_time if jit_time > 0 else float("inf")
|
||||
print(f"Speedup (SGL/JIT): {speedup:.2f}x")
|
||||
|
||||
assert jit_time >= 0 and sgl_time >= 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,101 @@
|
||||
import itertools
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.jit_kernel.utils import get_ci_test_range
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=37, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=148, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
def sglang_aot_qknorm(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
) -> None:
|
||||
from sgl_kernel import rmsnorm
|
||||
|
||||
head_dim = q.shape[-1]
|
||||
q = q.view(-1, head_dim)
|
||||
k = k.view(-1, head_dim)
|
||||
rmsnorm(q, q_weight, out=q)
|
||||
rmsnorm(k, k_weight, out=k)
|
||||
|
||||
|
||||
def sglang_jit_qknorm(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
) -> None:
|
||||
from sglang.jit_kernel.norm import fused_inplace_qknorm
|
||||
|
||||
fused_inplace_qknorm(q, k, q_weight, k_weight)
|
||||
|
||||
|
||||
def flashinfer_qknorm(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
) -> None:
|
||||
from flashinfer.norm import rmsnorm
|
||||
|
||||
rmsnorm(q, q_weight, out=q)
|
||||
rmsnorm(k, k_weight, out=k)
|
||||
|
||||
|
||||
@torch.compile()
|
||||
def torch_impl_qknorm(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
eps: float = 1e-6,
|
||||
) -> None:
|
||||
q_mean = q.float().pow(2).mean(dim=-1, keepdim=True)
|
||||
k_mean = k.float().pow(2).mean(dim=-1, keepdim=True)
|
||||
q_norm = (q_mean + eps).rsqrt()
|
||||
k_norm = (k_mean + eps).rsqrt()
|
||||
q.copy_(q.float() * q_norm * q_weight.float())
|
||||
k.copy_(k.float() * k_norm * k_weight.float())
|
||||
|
||||
|
||||
BS_LIST = [2**n for n in range(0, 14)]
|
||||
BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)]
|
||||
BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 256, 4109])
|
||||
N_K_LIST = get_ci_test_range([2, 4], [2, 4])
|
||||
N_Q_LIST = get_ci_test_range([8, 16], [8, 16])
|
||||
HEAD_DIM_LIST = get_ci_test_range([64, 128, 256, 512, 1024], [64, 256, 1024])
|
||||
DEVICE = "cuda"
|
||||
DTYPE = torch.bfloat16
|
||||
|
||||
# NOTE(dark): sgl_kernel use flashinfer template, which is bitwise identical to flashinfer impl.
|
||||
# However, sgl-jit-kernel, flashinfer, torch_impl, may have small numerical differences.
|
||||
# so we allow a small rel/abs tolerance in correctness test.
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_size,n_k,n_q,head_dim",
|
||||
list(itertools.product(BS_LIST, N_K_LIST, N_Q_LIST, HEAD_DIM_LIST)),
|
||||
)
|
||||
def test_qknorm(batch_size: int, n_k: int, n_q: int, head_dim: int) -> None:
|
||||
q = torch.randn(batch_size, n_q, head_dim, device=DEVICE, dtype=DTYPE)
|
||||
k = torch.randn(batch_size, n_k, head_dim, device=DEVICE, dtype=DTYPE)
|
||||
q_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
|
||||
k_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
|
||||
q_k_aot = (q.clone(), k.clone())
|
||||
q_k_jit = (q.clone(), k.clone())
|
||||
sglang_aot_qknorm(q_k_aot[0], q_k_aot[1], q_weight, k_weight)
|
||||
sglang_jit_qknorm(q_k_jit[0], q_k_jit[1], q_weight, k_weight)
|
||||
triton.testing.assert_close(q_k_aot[0], q_k_jit[0], atol=1e-2, rtol=1e-2)
|
||||
triton.testing.assert_close(q_k_aot[1], q_k_jit[1], atol=1e-2, rtol=1e-2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,83 @@
|
||||
import itertools
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.jit_kernel.utils import get_ci_test_range
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=15, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
def sglang_jit_qknorm_across_heads(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
) -> None:
|
||||
from sglang.jit_kernel.norm import fused_inplace_qknorm_across_heads
|
||||
|
||||
fused_inplace_qknorm_across_heads(q, k, q_weight, k_weight)
|
||||
|
||||
|
||||
def sglang_aot_qknorm_across_heads(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
) -> None:
|
||||
from sgl_kernel import rmsnorm
|
||||
|
||||
rmsnorm(q, q_weight, out=q)
|
||||
rmsnorm(k, k_weight, out=k)
|
||||
|
||||
|
||||
@torch.compile()
|
||||
def torch_impl_qknorm_across_heads(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
eps: float = 1e-6,
|
||||
) -> None:
|
||||
q_mean = q.float().pow(2).mean(dim=-1, keepdim=True)
|
||||
k_mean = k.float().pow(2).mean(dim=-1, keepdim=True)
|
||||
q_norm = (q_mean + eps).rsqrt()
|
||||
k_norm = (k_mean + eps).rsqrt()
|
||||
q.copy_(q.float() * q_norm * q_weight.float())
|
||||
k.copy_(k.float() * k_norm * k_weight.float())
|
||||
|
||||
|
||||
BS_LIST = [2**n for n in range(0, 14)]
|
||||
BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)]
|
||||
BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 256, 4109])
|
||||
HIDDEN_DIM_LIST = get_ci_test_range([512, 1024, 2048, 4096], [512, 2048, 4096])
|
||||
DEVICE = "cuda"
|
||||
DTYPE = torch.bfloat16
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_size,hidden_dim",
|
||||
list(itertools.product(BS_LIST, HIDDEN_DIM_LIST)),
|
||||
)
|
||||
def test_qknorm_across_heads(batch_size: int, hidden_dim: int) -> None:
|
||||
q = torch.randn(batch_size, hidden_dim, device=DEVICE, dtype=DTYPE)
|
||||
k = torch.randn(batch_size, hidden_dim, device=DEVICE, dtype=DTYPE)
|
||||
q_weight = torch.randn(hidden_dim, device=DEVICE, dtype=DTYPE)
|
||||
k_weight = torch.randn(hidden_dim, device=DEVICE, dtype=DTYPE)
|
||||
|
||||
q_k_jit = (q.clone(), k.clone())
|
||||
q_k_aot = (q.clone(), k.clone())
|
||||
|
||||
sglang_jit_qknorm_across_heads(q_k_jit[0], q_k_jit[1], q_weight, k_weight)
|
||||
sglang_aot_qknorm_across_heads(q_k_aot[0], q_k_aot[1], q_weight, k_weight)
|
||||
|
||||
triton.testing.assert_close(q_k_jit[0], q_k_aot[0], atol=1e-2, rtol=1e-2)
|
||||
triton.testing.assert_close(q_k_jit[1], q_k_aot[1], atol=1e-2, rtol=1e-2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,86 @@
|
||||
# Adapted from https://github.com/flashinfer-ai/flashinfer/blob/main/tests/test_sampling.py
|
||||
# and /sgl-workspace/sglang/sgl-kernel/tests/test_sampling.py
|
||||
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import sgl_kernel
|
||||
import torch
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=6, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", [1, 99, 989])
|
||||
@pytest.mark.parametrize("vocab_size", [111, 32000, 128256])
|
||||
@pytest.mark.parametrize("k", [10, 100, 500])
|
||||
def test_top_k_renorm_probs(batch_size, vocab_size, k):
|
||||
"""Test top_k_renorm_probs kernel for correctness.
|
||||
|
||||
This test validates that the kernel correctly:
|
||||
1. Identifies the top-k probabilities
|
||||
2. Masks out non-top-k values
|
||||
3. Renormalizes the remaining probabilities to sum to 1
|
||||
"""
|
||||
if k > vocab_size:
|
||||
pytest.skip("k should be less than vocab_size")
|
||||
torch.manual_seed(42)
|
||||
pre_norm_prob = torch.rand(batch_size, vocab_size, device="cuda:0")
|
||||
normalized_prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
|
||||
sorted_prob, _ = torch.sort(normalized_prob, descending=True)
|
||||
pivot = sorted_prob[:, k - 1]
|
||||
mask = (normalized_prob >= pivot.unsqueeze(-1)).int()
|
||||
renorm_prob_ground_truth = normalized_prob.clone()
|
||||
renorm_prob_ground_truth[mask == 0] = 0
|
||||
renorm_prob_ground_truth = renorm_prob_ground_truth / renorm_prob_ground_truth.sum(
|
||||
dim=-1, keepdim=True
|
||||
)
|
||||
|
||||
renorm_prob = sgl_kernel.top_k_renorm_prob(normalized_prob, k)
|
||||
for i in range(batch_size):
|
||||
torch.testing.assert_close(
|
||||
renorm_prob_ground_truth[i],
|
||||
renorm_prob[i],
|
||||
rtol=1e-3,
|
||||
atol=1e-3,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", [1, 99, 989])
|
||||
@pytest.mark.parametrize("vocab_size", [111, 32000, 128256])
|
||||
@pytest.mark.parametrize("p", [0.1, 0.5, 0.9])
|
||||
def test_top_p_renorm_probs(batch_size, vocab_size, p):
|
||||
"""Test top_p_renorm_probs kernel for correctness.
|
||||
|
||||
This test validates that the kernel correctly:
|
||||
1. Computes the cumulative probability distribution
|
||||
2. Identifies tokens in the top-p threshold
|
||||
3. Masks out tokens outside the threshold
|
||||
4. Renormalizes the remaining probabilities to sum to 1
|
||||
"""
|
||||
torch.manual_seed(42)
|
||||
pre_norm_prob = torch.rand(batch_size, vocab_size, device="cuda:0")
|
||||
normalized_prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
|
||||
sorted_prob, indices = torch.sort(normalized_prob, descending=False)
|
||||
cdf = torch.cumsum(sorted_prob, dim=-1)
|
||||
mask = torch.zeros(batch_size, vocab_size, dtype=torch.int32, device="cuda:0")
|
||||
mask.scatter_add_(1, indices, (cdf >= (1 - p)).int())
|
||||
renorm_prob_ground_truth = normalized_prob.clone()
|
||||
renorm_prob_ground_truth[mask == 0] = 0
|
||||
renorm_prob_ground_truth = renorm_prob_ground_truth / renorm_prob_ground_truth.sum(
|
||||
dim=-1, keepdim=True
|
||||
)
|
||||
|
||||
renorm_prob = sgl_kernel.top_p_renorm_prob(normalized_prob, p)
|
||||
torch.testing.assert_close(
|
||||
renorm_prob_ground_truth,
|
||||
renorm_prob,
|
||||
rtol=1e-3,
|
||||
atol=1e-3,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -0,0 +1,69 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.resolve_future_token_ids import resolve_future_token_ids_cuda
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=9, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
def _reference_resolve(input_ids, future_map):
|
||||
"""Reference implementation using plain torch."""
|
||||
result = input_ids.clone()
|
||||
result[:] = torch.where(
|
||||
result < 0,
|
||||
future_map[torch.clamp(-result, min=0)],
|
||||
result,
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
@pytest.mark.parametrize("size", [1, 2, 127, 128, 255, 256, 1024, 4097])
|
||||
@pytest.mark.parametrize("dtype", [torch.int32, torch.int64])
|
||||
class TestResolveFutureTokenIds:
|
||||
def test_all_negative(self, size: int, dtype: torch.dtype) -> None:
|
||||
map_size = 8192
|
||||
future_map = torch.randint(0, 50000, (map_size,), dtype=dtype, device="cuda")
|
||||
# Negative indices in range [-map_size+1, -1]
|
||||
input_ids = -torch.randint(1, map_size, (size,), dtype=dtype, device="cuda")
|
||||
|
||||
expected = _reference_resolve(input_ids, future_map)
|
||||
resolve_future_token_ids_cuda(input_ids, future_map)
|
||||
assert torch.equal(input_ids, expected)
|
||||
|
||||
def test_all_non_negative(self, size: int, dtype: torch.dtype) -> None:
|
||||
map_size = 16
|
||||
future_map = torch.randint(0, 50000, (map_size,), dtype=dtype, device="cuda")
|
||||
input_ids = torch.randint(0, 50000, (size,), dtype=dtype, device="cuda")
|
||||
|
||||
expected = input_ids.clone()
|
||||
resolve_future_token_ids_cuda(input_ids, future_map)
|
||||
assert torch.equal(input_ids, expected)
|
||||
|
||||
def test_mixed(self, size: int, dtype: torch.dtype) -> None:
|
||||
map_size = 8192
|
||||
future_map = torch.randint(0, 50000, (map_size,), dtype=dtype, device="cuda")
|
||||
# Mix of negative and non-negative
|
||||
input_ids = torch.randint(
|
||||
-map_size + 1, 50000, (size,), dtype=dtype, device="cuda"
|
||||
)
|
||||
|
||||
expected = _reference_resolve(input_ids, future_map)
|
||||
resolve_future_token_ids_cuda(input_ids, future_map)
|
||||
assert torch.equal(input_ids, expected)
|
||||
|
||||
def test_zeros(self, size: int, dtype: torch.dtype) -> None:
|
||||
map_size = 16
|
||||
future_map = torch.randint(0, 50000, (map_size,), dtype=dtype, device="cuda")
|
||||
input_ids = torch.zeros(size, dtype=dtype, device="cuda")
|
||||
|
||||
expected = input_ids.clone()
|
||||
resolve_future_token_ids_cuda(input_ids, future_map)
|
||||
assert torch.equal(input_ids, expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,107 @@
|
||||
import itertools
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.utils import get_ci_test_range
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=45, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=240, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
EPS = 1e-6
|
||||
DEVICE = "cuda"
|
||||
DTYPES = [torch.float16, torch.bfloat16]
|
||||
|
||||
|
||||
def sglang_jit_rmsnorm(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
*,
|
||||
output: torch.Tensor | None = None,
|
||||
eps: float = EPS,
|
||||
) -> None:
|
||||
from sglang.jit_kernel.norm import rmsnorm
|
||||
|
||||
rmsnorm(input, weight, out=output, eps=eps)
|
||||
|
||||
|
||||
def flashinfer_rmsnorm(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
*,
|
||||
output: torch.Tensor,
|
||||
eps: float = EPS,
|
||||
) -> None:
|
||||
from flashinfer.norm import rmsnorm
|
||||
|
||||
rmsnorm(input, weight, out=output, eps=eps)
|
||||
|
||||
|
||||
BS_LIST = [2**n for n in range(0, 14)]
|
||||
BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)]
|
||||
BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 256, 4109])
|
||||
SUPPORTED_HIDDEN_SIZE_LIST = get_ci_test_range(
|
||||
[64, 128, 256, 512, *range(1024, 8192 + 1, 1024), 2304, 2560, 12288, 16384],
|
||||
[256, 1024, 16384],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_size,hidden_size",
|
||||
list(itertools.product(BS_LIST, SUPPORTED_HIDDEN_SIZE_LIST)),
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("specify_out", [True, False])
|
||||
def test_rmsnorm(
|
||||
batch_size: int, hidden_size: int, dtype: torch.dtype, specify_out: bool
|
||||
) -> None:
|
||||
input = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
|
||||
weight = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
|
||||
|
||||
input_flashinfer = input.clone()
|
||||
output_flashinfer = torch.empty_like(input)
|
||||
flashinfer_rmsnorm(input_flashinfer, weight, output=output_flashinfer)
|
||||
|
||||
if specify_out:
|
||||
output_sglang = torch.empty_like(input)
|
||||
sglang_jit_rmsnorm(input, weight, output=output_sglang)
|
||||
else:
|
||||
output_sglang = input.clone()
|
||||
sglang_jit_rmsnorm(output_sglang, weight, output=output_sglang)
|
||||
|
||||
torch.testing.assert_close(output_sglang, output_flashinfer, atol=1e-2, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("hidden_size", [64, 128, 256, 512, 8192, 8704, 16384])
|
||||
def test_rmsnorm_hidden_size_support(hidden_size: int) -> None:
|
||||
from sglang.jit_kernel.norm import _is_supported_rmsnorm_hidden_size
|
||||
|
||||
assert _is_supported_rmsnorm_hidden_size(hidden_size)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("hidden_size", "expected"),
|
||||
[
|
||||
(64, "RMSNormWarpKernel"),
|
||||
(128, "RMSNormWarpKernel"),
|
||||
(256, "RMSNormWarpKernel"),
|
||||
(512, "RMSNormHalfKernel"),
|
||||
(1536, "RMSNormKernel"),
|
||||
(2048, "RMSNormHalfKernel"),
|
||||
(2304, "RMSNormKernel"), # NOTE: not 512 aligned
|
||||
(8192, "RMSNormHalfKernel"),
|
||||
(8704, "RMSNormHalfKernel"),
|
||||
(16384, "RMSNormHalfKernel"),
|
||||
],
|
||||
)
|
||||
def test_rmsnorm_kernel_dispatch(hidden_size: int, expected: str) -> None:
|
||||
from sglang.jit_kernel.norm import _rmsnorm_kernel_class
|
||||
|
||||
assert _rmsnorm_kernel_class(hidden_size) == expected
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user