[CI] Move JIT kernel tests + benchmarks to test/registered/jit; add in-package guard (#27644)
This commit is contained in:
Liangsheng Yin
@@ -1,260 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from typing import Tuple, Union
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.jit_kernel.benchmark.bench_activation import register_cuda_ci
|
||||
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,
|
||||
)
|
||||
|
||||
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"]))
|
||||
@@ -1,337 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from typing import Tuple, Union
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.jit_kernel.benchmark.bench_activation import register_cuda_ci
|
||||
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,
|
||||
)
|
||||
|
||||
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"]))
|
||||
@@ -1,228 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,142 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,460 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,133 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,252 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,104 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,153 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,226 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,193 +0,0 @@
|
||||
"""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"]))
|
||||
@@ -1,156 +0,0 @@
|
||||
"""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"]))
|
||||
@@ -1,70 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
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")
|
||||
|
||||
|
||||
_CONSTS_CUH: Path = (
|
||||
Path(__file__).resolve().parents[2] / "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"]))
|
||||
@@ -1,440 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,842 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,222 +0,0 @@
|
||||
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
@@ -1,278 +0,0 @@
|
||||
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()
|
||||
@@ -1,64 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,214 +0,0 @@
|
||||
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
@@ -1,256 +0,0 @@
|
||||
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
@@ -1,205 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,41 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,169 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,125 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,111 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,46 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,175 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,239 +0,0 @@
|
||||
"""
|
||||
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)
|
||||
@@ -1,312 +0,0 @@
|
||||
"""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
@@ -1,96 +0,0 @@
|
||||
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
@@ -1,459 +0,0 @@
|
||||
"""
|
||||
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"]))
|
||||
@@ -1,238 +0,0 @@
|
||||
"""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"]))
|
||||
@@ -1,194 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,96 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,210 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,428 +0,0 @@
|
||||
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__]))
|
||||
@@ -1,247 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,426 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,104 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,349 +0,0 @@
|
||||
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__]))
|
||||
@@ -1,168 +0,0 @@
|
||||
# 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__]))
|
||||
@@ -1,615 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,153 +0,0 @@
|
||||
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__]))
|
||||
@@ -1,135 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,137 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,152 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,225 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,91 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,210 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,494 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,101 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,83 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,86 +0,0 @@
|
||||
# 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__]))
|
||||
@@ -1,69 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,107 +0,0 @@
|
||||
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"]))
|
||||
@@ -1,136 +0,0 @@
|
||||
"""Tests for the JIT rmsnorm_hf kernel (HF LlamaRMSNorm semantics)."""
|
||||
|
||||
import itertools
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.rmsnorm_hf import (
|
||||
is_supported_rmsnorm_hf_hidden_size,
|
||||
rmsnorm_hf,
|
||||
)
|
||||
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=30, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
EPS = 1e-5
|
||||
DEVICE = "cuda"
|
||||
DTYPES = [torch.float16, torch.bfloat16]
|
||||
|
||||
|
||||
def hf_rmsnorm_reference(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor:
|
||||
"""HF LlamaRMSNorm: normalize fp32, cast normalized x to dtype, then multiply weight."""
|
||||
x_fp32 = x.to(torch.float32)
|
||||
variance = x_fp32.pow(2).mean(-1, keepdim=True)
|
||||
x_normed = x_fp32 * torch.rsqrt(variance + eps)
|
||||
return w * x_normed.to(x.dtype)
|
||||
|
||||
|
||||
def sgl_rmsnorm_reference(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor:
|
||||
"""Old sgl_kernel.rmsnorm semantics — weight multiply in fp32, cast at the end."""
|
||||
x_fp32 = x.to(torch.float32)
|
||||
variance = x_fp32.pow(2).mean(-1, keepdim=True)
|
||||
x_normed = x_fp32 * torch.rsqrt(variance + eps)
|
||||
return (x_normed * w.to(torch.float32)).to(x.dtype)
|
||||
|
||||
|
||||
BS_LIST = get_ci_test_range(
|
||||
[1, 2, 4, 7, 16, 64, 128, 512, 1024, 4096],
|
||||
[1, 16, 1024],
|
||||
)
|
||||
HIDDEN_SIZE_LIST = get_ci_test_range(
|
||||
# Warp-kernel shapes (q/k RMSNorm head dims) + CTA-kernel shapes.
|
||||
[32, 64, 96, 128, 256, 512, 1024, 2048, 3072, 4096, 8192, 16384],
|
||||
[128, 512, 4096, 16384],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_size,hidden_size",
|
||||
list(itertools.product(BS_LIST, HIDDEN_SIZE_LIST)),
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
def test_rmsnorm_hf_correctness(
|
||||
batch_size: int, hidden_size: int, dtype: torch.dtype
|
||||
) -> None:
|
||||
torch.manual_seed(0)
|
||||
x = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
|
||||
w = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
|
||||
out = rmsnorm_hf(x, w, EPS)
|
||||
ref = hf_rmsnorm_reference(x, w, EPS)
|
||||
# Loose atol — the kernel's block-reduce order differs from PyTorch's
|
||||
# `mean`, producing ~1 fp16 ULP of drift on some shapes.
|
||||
# The SGL-semantics regression guard below is what catches the cast-order
|
||||
# bug this PR fixes; it's reduction-order-invariant.
|
||||
torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
def test_rmsnorm_hf_out_param(dtype: torch.dtype) -> None:
|
||||
torch.manual_seed(0)
|
||||
x = torch.randn(8, 4096, device=DEVICE, dtype=dtype)
|
||||
w = torch.randn(4096, device=DEVICE, dtype=dtype)
|
||||
out = torch.empty_like(x)
|
||||
result = rmsnorm_hf(x, w, EPS, out=out)
|
||||
assert result.data_ptr() == out.data_ptr()
|
||||
torch.testing.assert_close(
|
||||
out, hf_rmsnorm_reference(x, w, EPS), atol=1e-2, rtol=1e-2
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
def test_rmsnorm_hf_matches_hf_not_sgl(dtype: torch.dtype) -> None:
|
||||
"""Regression guard: kernel must follow HF (cast-before-mul), not the old
|
||||
sgl_kernel.rmsnorm semantics (fp32-mul-then-cast). Reduction-order drift
|
||||
prevents a bit-exact assert against HF, so instead assert the kernel is
|
||||
strictly closer to HF than to the SGL reference."""
|
||||
torch.manual_seed(0)
|
||||
x = torch.randn(64, 4096, device=DEVICE, dtype=dtype)
|
||||
w = torch.randn(4096, device=DEVICE, dtype=dtype)
|
||||
out = rmsnorm_hf(x, w, EPS).float()
|
||||
hf_ref = hf_rmsnorm_reference(x, w, EPS).float()
|
||||
sgl_ref = sgl_rmsnorm_reference(x, w, EPS).float()
|
||||
assert (sgl_ref - hf_ref).abs().max() > 0, "inputs don't exercise the difference"
|
||||
diff_hf = (out - hf_ref).abs().max().item()
|
||||
diff_sgl = (out - sgl_ref).abs().max().item()
|
||||
assert (
|
||||
diff_hf < diff_sgl
|
||||
), f"kernel closer to SGL than HF (hf={diff_hf}, sgl={diff_sgl})"
|
||||
|
||||
|
||||
def test_rmsnorm_hf_empty_input() -> None:
|
||||
"""Empty input must short-circuit: the C++ launcher rejects num_tokens=0."""
|
||||
x = torch.empty(0, 4096, device=DEVICE, dtype=torch.float16)
|
||||
w = torch.randn(4096, device=DEVICE, dtype=torch.float16)
|
||||
out = rmsnorm_hf(x, w, EPS)
|
||||
assert out.shape == x.shape and out.numel() == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("hidden_size", "expected"),
|
||||
[
|
||||
(16, False),
|
||||
(32, True),
|
||||
(64, True),
|
||||
(96, True),
|
||||
(128, True),
|
||||
(256, True),
|
||||
(288, True),
|
||||
(384, True),
|
||||
(500, False),
|
||||
(512, True),
|
||||
(3072, True),
|
||||
(4096, True),
|
||||
(8192, True),
|
||||
(4097, False),
|
||||
],
|
||||
)
|
||||
def test_is_supported_hidden_size(hidden_size: int, expected: bool) -> None:
|
||||
assert is_supported_rmsnorm_hf_hidden_size(hidden_size) is expected
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -1,257 +0,0 @@
|
||||
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=64, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=256, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
DEVICE = "cuda"
|
||||
DTYPE = torch.bfloat16
|
||||
MAX_SEQ_LEN = 131072 # common seq length
|
||||
ROPE_BASE = 10000.0
|
||||
CACHE_SIZE = 1024 * 128
|
||||
|
||||
|
||||
def create_cos_sin_cache(
|
||||
rotary_dim: int,
|
||||
max_position: int = MAX_SEQ_LEN,
|
||||
base: float = ROPE_BASE,
|
||||
) -> torch.Tensor:
|
||||
"""Create cos/sin cache compatible with SGLang layout: [max_pos, rotary_dim]."""
|
||||
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)
|
||||
cos = freqs.cos()
|
||||
sin = freqs.sin()
|
||||
cache = torch.cat((cos, sin), dim=-1) # [max_pos, rotary_dim]
|
||||
return cache
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Implementation wrappers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def sglang_jit_rope(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
cos_sin_cache: 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)
|
||||
|
||||
|
||||
def flashinfer_rope(
|
||||
q: torch.Tensor,
|
||||
k: 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
|
||||
|
||||
head_size = q.shape[-1]
|
||||
# flashinfer expects [nnz, num_heads * head_size]
|
||||
q_2d = q.view(q.shape[0], -1)
|
||||
k_2d = k.view(k.shape[0], -1)
|
||||
apply_rope_with_cos_sin_cache_inplace(
|
||||
positions=positions,
|
||||
query=q_2d,
|
||||
key=k_2d,
|
||||
head_size=head_size,
|
||||
cos_sin_cache=cos_sin_cache,
|
||||
is_neox=is_neox,
|
||||
)
|
||||
|
||||
|
||||
def torch_impl_rope(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
cos_sin_cache: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
is_neox: bool,
|
||||
) -> None:
|
||||
# TODO: implement a pure-PyTorch reference for extra coverage
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test parameters
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
BS_LIST = [2**x for x in range(12)]
|
||||
BS_LIST += [x + 1 for x in BS_LIST] # odd sizes to stress non-aligned paths
|
||||
BS_LIST = get_ci_test_range(BS_LIST, [1, 129, 2048, 2049])
|
||||
NUM_KV_HEADS_LIST = get_ci_test_range([1, 2, 8], [1, 8])
|
||||
GQA_RATIO = get_ci_test_range([1, 4, 8], [1, 8])
|
||||
ROPE_DIM_LIST = get_ci_test_range([64, 128, 256, 512], [64, 256])
|
||||
IS_NEOX_LIST = [False, True]
|
||||
DTYPE_LIST = get_ci_test_range(
|
||||
[torch.bfloat16, torch.float16], [torch.bfloat16, torch.float16]
|
||||
)
|
||||
PARTIAL_ROPE_DIM_LIST = get_ci_test_range([64, 80, 96, 128], [64, 96])
|
||||
HEAD_DIM_LIST = get_ci_test_range([64, 128, 256], [64, 256])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", BS_LIST)
|
||||
@pytest.mark.parametrize("gqa_ratio", GQA_RATIO)
|
||||
@pytest.mark.parametrize("num_kv_heads", NUM_KV_HEADS_LIST)
|
||||
@pytest.mark.parametrize("rope_dim", ROPE_DIM_LIST)
|
||||
@pytest.mark.parametrize("is_neox", IS_NEOX_LIST)
|
||||
@pytest.mark.parametrize("dtype", DTYPE_LIST)
|
||||
def test_rope(
|
||||
batch_size: int,
|
||||
gqa_ratio: int,
|
||||
num_kv_heads: int,
|
||||
rope_dim: int,
|
||||
is_neox: bool,
|
||||
dtype: torch.dtype,
|
||||
) -> None:
|
||||
num_qo_heads = num_kv_heads * gqa_ratio
|
||||
q = torch.randn(batch_size, num_qo_heads, rope_dim, device=DEVICE, dtype=dtype)
|
||||
k = torch.randn(batch_size, num_kv_heads, rope_dim, device=DEVICE, dtype=dtype)
|
||||
positions = torch.randint(
|
||||
0, MAX_SEQ_LEN, (batch_size,), device=DEVICE, dtype=torch.int64
|
||||
)
|
||||
cos_sin_cache = create_cos_sin_cache(rope_dim)
|
||||
|
||||
q_fi, k_fi = q.clone(), k.clone()
|
||||
q_jit, k_jit = q.clone(), k.clone()
|
||||
|
||||
flashinfer_rope(q_fi, k_fi, cos_sin_cache, positions, is_neox)
|
||||
sglang_jit_rope(q_jit, k_jit, cos_sin_cache, positions, is_neox)
|
||||
|
||||
atol = rtol = 1e-2
|
||||
triton.testing.assert_close(q_fi, q_jit, atol=atol, rtol=rtol)
|
||||
triton.testing.assert_close(k_fi, k_jit, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.int32, torch.int64])
|
||||
def test_rope_position_dtypes(dtype: torch.dtype) -> None:
|
||||
"""Ensure both int32 and int64 position tensors work correctly."""
|
||||
batch_size, num_qo_heads, num_kv_heads, rope_dim = 16384, 16, 2, 128
|
||||
is_neox = True
|
||||
|
||||
q = torch.randn(batch_size, num_qo_heads, rope_dim, device=DEVICE, dtype=DTYPE)
|
||||
k = torch.randn(batch_size, num_kv_heads, rope_dim, device=DEVICE, dtype=DTYPE)
|
||||
positions = torch.randint(0, MAX_SEQ_LEN, (batch_size,), device=DEVICE, dtype=dtype)
|
||||
cos_sin_cache = create_cos_sin_cache(rope_dim)
|
||||
|
||||
q_fi, k_fi = q.clone(), k.clone()
|
||||
q_jit, k_jit = q.clone(), k.clone()
|
||||
|
||||
flashinfer_rope(q_fi, k_fi, cos_sin_cache, positions.long(), is_neox)
|
||||
sglang_jit_rope(q_jit, k_jit, cos_sin_cache, positions, is_neox)
|
||||
|
||||
atol = rtol = 1e-2
|
||||
triton.testing.assert_close(q_fi, q_jit, atol=atol, rtol=rtol)
|
||||
triton.testing.assert_close(k_fi, k_jit, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", BS_LIST)
|
||||
@pytest.mark.parametrize("is_neox", IS_NEOX_LIST)
|
||||
@pytest.mark.parametrize("rope_dim", PARTIAL_ROPE_DIM_LIST)
|
||||
@pytest.mark.parametrize("head_dim", HEAD_DIM_LIST)
|
||||
def test_partial_rope(batch_size: int, is_neox: bool, rope_dim: int, head_dim: int):
|
||||
if head_dim < rope_dim:
|
||||
pytest.skip("Invalid config: head_dim must be >= rope_dim.")
|
||||
num_qo_heads, num_kv_heads = 8, 2
|
||||
|
||||
q = torch.randn(batch_size, num_qo_heads, head_dim, device=DEVICE, dtype=DTYPE)
|
||||
k = torch.randn(batch_size, num_kv_heads, head_dim, device=DEVICE, dtype=DTYPE)
|
||||
positions = torch.randint(0, MAX_SEQ_LEN, (batch_size,), device=DEVICE)
|
||||
cos_sin_cache = create_cos_sin_cache(rope_dim)
|
||||
|
||||
q_fi, k_fi = q.clone(), k.clone()
|
||||
q_jit, k_jit = q.clone(), k.clone()
|
||||
rope = ..., slice(rope_dim) # NOTE: flashinfer by default apply to first rope_dim
|
||||
|
||||
flashinfer_rope(q_fi, k_fi, cos_sin_cache, positions.long(), is_neox)
|
||||
sglang_jit_rope(q_jit[rope], k_jit[rope], cos_sin_cache, positions, is_neox)
|
||||
|
||||
atol = rtol = 1e-2
|
||||
triton.testing.assert_close(q_fi, q_jit, atol=atol, rtol=rtol)
|
||||
triton.testing.assert_close(k_fi, k_jit, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", BS_LIST)
|
||||
@pytest.mark.parametrize("gqa_ratio", GQA_RATIO)
|
||||
@pytest.mark.parametrize("num_kv_heads", NUM_KV_HEADS_LIST)
|
||||
@pytest.mark.parametrize("rope_dim", ROPE_DIM_LIST)
|
||||
@pytest.mark.parametrize("is_neox", IS_NEOX_LIST)
|
||||
def test_fused_rope_store(
|
||||
batch_size: int,
|
||||
gqa_ratio: int,
|
||||
num_kv_heads: int,
|
||||
rope_dim: int,
|
||||
is_neox: bool,
|
||||
) -> None:
|
||||
"""Test fused RoPE + KV cache store against separate RoPE + manual store."""
|
||||
from sglang.jit_kernel.rope import apply_rope_inplace_with_kvcache
|
||||
|
||||
num_qo_heads = num_kv_heads * gqa_ratio
|
||||
dtype = DTYPE
|
||||
|
||||
q = torch.randn(batch_size, num_qo_heads, rope_dim, device=DEVICE, dtype=dtype)
|
||||
k = torch.randn(batch_size, num_kv_heads, rope_dim, device=DEVICE, dtype=dtype)
|
||||
v = torch.randn(batch_size, num_kv_heads, rope_dim, device=DEVICE, dtype=dtype)
|
||||
positions = torch.randint(
|
||||
0, MAX_SEQ_LEN, (batch_size,), device=DEVICE, dtype=torch.int64
|
||||
)
|
||||
out_loc = torch.randperm(CACHE_SIZE, device=DEVICE, dtype=torch.int64)[:batch_size]
|
||||
cos_sin_cache = create_cos_sin_cache(rope_dim)
|
||||
|
||||
row_size = num_kv_heads * rope_dim
|
||||
k_cache_ref = torch.zeros(CACHE_SIZE, row_size, device=DEVICE, dtype=dtype)
|
||||
v_cache_ref = torch.zeros(CACHE_SIZE, row_size, device=DEVICE, dtype=dtype)
|
||||
k_cache_fused = torch.zeros(CACHE_SIZE, row_size, device=DEVICE, dtype=dtype)
|
||||
v_cache_fused = torch.zeros(CACHE_SIZE, row_size, device=DEVICE, dtype=dtype)
|
||||
|
||||
# --- reference: separate RoPE then manual scatter ---
|
||||
q_ref, k_ref = q.clone(), k.clone()
|
||||
flashinfer_rope(q_ref, k_ref, cos_sin_cache, positions, is_neox)
|
||||
k_cache_ref[out_loc] = k_ref.view(batch_size, -1)
|
||||
v_cache_ref[out_loc] = v.view(batch_size, -1)
|
||||
|
||||
# --- fused kernel ---
|
||||
q_fused, k_fused = q.clone(), k.clone()
|
||||
v_fused = v.clone()
|
||||
apply_rope_inplace_with_kvcache(
|
||||
q_fused,
|
||||
k_fused,
|
||||
v_fused,
|
||||
k_cache_fused,
|
||||
v_cache_fused,
|
||||
cos_sin_cache,
|
||||
positions,
|
||||
out_loc,
|
||||
is_neox=is_neox,
|
||||
)
|
||||
|
||||
atol = rtol = 1e-2
|
||||
# q should match RoPE-only result
|
||||
triton.testing.assert_close(q_ref, q_fused, atol=atol, rtol=rtol)
|
||||
# k_cache should contain the rotated k
|
||||
triton.testing.assert_close(
|
||||
k_cache_ref[out_loc], k_cache_fused[out_loc], atol=atol, rtol=rtol
|
||||
)
|
||||
# v_cache should be an exact copy
|
||||
assert torch.all(v_cache_ref[out_loc] == v_cache_fused[out_loc]), "v_cache mismatch"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -1,126 +0,0 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.set_mla_kv_buffer import (
|
||||
can_use_set_mla_kv_buffer,
|
||||
set_mla_kv_buffer,
|
||||
)
|
||||
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=30, suite="base-b-kernel-unit-1-gpu-large")
|
||||
|
||||
DEVICE = "cuda"
|
||||
CACHE_SIZE = 4096
|
||||
|
||||
# (nope_dim, rope_dim) pairs: standard MLA, MLA scale buffer, FP8 nope-extended layout.
|
||||
SHAPES = get_ci_test_range(
|
||||
[(512, 64), (512, 32), (256, 64), (128, 64), (528, 64)],
|
||||
[(512, 64), (528, 64)],
|
||||
)
|
||||
BATCH_SIZES = get_ci_test_range([1, 7, 64, 257, 1024], [1, 64, 1024])
|
||||
|
||||
|
||||
def _ref(kv_buffer, loc, cache_k_nope, cache_k_rope):
|
||||
nope_dim = cache_k_nope.shape[-1]
|
||||
n_loc = loc.shape[0]
|
||||
src_nope = cache_k_nope.reshape(n_loc, -1)
|
||||
src_rope = cache_k_rope.reshape(n_loc, -1)
|
||||
kv_view = kv_buffer.view(kv_buffer.shape[0], -1)
|
||||
kv_view[loc.long(), :nope_dim] = src_nope
|
||||
kv_view[loc.long(), nope_dim : nope_dim + src_rope.shape[-1]] = src_rope
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
@pytest.mark.parametrize("shape", SHAPES)
|
||||
@pytest.mark.parametrize("batch_size", BATCH_SIZES)
|
||||
def test_set_mla_kv_buffer_correctness(dtype, shape, batch_size):
|
||||
nope_dim, rope_dim = shape
|
||||
total_dim = nope_dim + rope_dim
|
||||
|
||||
cache_k_nope = torch.randn((batch_size, 1, nope_dim), dtype=dtype, device=DEVICE)
|
||||
cache_k_rope = torch.randn((batch_size, 1, rope_dim), dtype=dtype, device=DEVICE)
|
||||
kv_buffer = torch.randn((CACHE_SIZE, 1, total_dim), dtype=dtype, device=DEVICE)
|
||||
kv_ref = kv_buffer.clone()
|
||||
|
||||
loc = torch.randperm(CACHE_SIZE, device=DEVICE)[:batch_size]
|
||||
|
||||
set_mla_kv_buffer(kv_buffer, loc, cache_k_nope, cache_k_rope)
|
||||
_ref(kv_ref, loc, cache_k_nope, cache_k_rope)
|
||||
|
||||
assert torch.equal(kv_buffer, kv_ref)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("loc_dtype", [torch.int32, torch.int64])
|
||||
def test_set_mla_kv_buffer_loc_dtypes(loc_dtype):
|
||||
nope_dim, rope_dim = 512, 64
|
||||
batch_size = 128
|
||||
dtype = torch.bfloat16
|
||||
|
||||
cache_k_nope = torch.randn((batch_size, 1, nope_dim), dtype=dtype, device=DEVICE)
|
||||
cache_k_rope = torch.randn((batch_size, 1, rope_dim), dtype=dtype, device=DEVICE)
|
||||
kv_buffer = torch.randn(
|
||||
(CACHE_SIZE, 1, nope_dim + rope_dim), dtype=dtype, device=DEVICE
|
||||
)
|
||||
kv_ref = kv_buffer.clone()
|
||||
|
||||
loc = torch.randperm(CACHE_SIZE, device=DEVICE)[:batch_size].to(loc_dtype)
|
||||
|
||||
set_mla_kv_buffer(kv_buffer, loc, cache_k_nope, cache_k_rope)
|
||||
_ref(kv_ref, loc, cache_k_nope, cache_k_rope)
|
||||
|
||||
assert torch.equal(kv_buffer, kv_ref)
|
||||
|
||||
|
||||
def test_set_mla_kv_buffer_uint8_byte_layout():
|
||||
"""FP8 DSA byte-layout: cache_k_nope is uint8 with [fp8(512) | scales(16)] = 528,
|
||||
cache_k_rope is uint8 [128]; total payload = 656 bytes."""
|
||||
nope_bytes, rope_bytes = 528, 128
|
||||
batch_size = 64
|
||||
dtype = torch.uint8
|
||||
|
||||
cache_k_nope = torch.randint(
|
||||
0, 256, (batch_size, 1, nope_bytes), dtype=dtype, device=DEVICE
|
||||
)
|
||||
cache_k_rope = torch.randint(
|
||||
0, 256, (batch_size, 1, rope_bytes), dtype=dtype, device=DEVICE
|
||||
)
|
||||
kv_buffer = torch.randint(
|
||||
0, 256, (CACHE_SIZE, 1, nope_bytes + rope_bytes), dtype=dtype, device=DEVICE
|
||||
)
|
||||
kv_ref = kv_buffer.clone()
|
||||
|
||||
loc = torch.randperm(CACHE_SIZE, device=DEVICE)[:batch_size]
|
||||
|
||||
set_mla_kv_buffer(kv_buffer, loc, cache_k_nope, cache_k_rope)
|
||||
_ref(kv_ref, loc, cache_k_nope, cache_k_rope)
|
||||
|
||||
assert torch.equal(kv_buffer, kv_ref)
|
||||
|
||||
|
||||
def test_set_mla_kv_buffer_empty_loc():
|
||||
nope_dim, rope_dim = 512, 64
|
||||
dtype = torch.bfloat16
|
||||
cache_k_nope = torch.empty((0, 1, nope_dim), dtype=dtype, device=DEVICE)
|
||||
cache_k_rope = torch.empty((0, 1, rope_dim), dtype=dtype, device=DEVICE)
|
||||
kv_buffer = torch.randn(
|
||||
(CACHE_SIZE, 1, nope_dim + rope_dim), dtype=dtype, device=DEVICE
|
||||
)
|
||||
kv_before = kv_buffer.clone()
|
||||
|
||||
loc = torch.empty((0,), dtype=torch.int64, device=DEVICE)
|
||||
set_mla_kv_buffer(kv_buffer, loc, cache_k_nope, cache_k_rope)
|
||||
|
||||
assert torch.equal(kv_buffer, kv_before)
|
||||
|
||||
|
||||
def test_can_use_set_mla_kv_buffer():
|
||||
assert can_use_set_mla_kv_buffer(1024, 128) # bf16 (512,64)
|
||||
assert can_use_set_mla_kv_buffer(528, 128) # fp8 byte layout
|
||||
assert not can_use_set_mla_kv_buffer(13, 8) # not multiple of 4
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -1,133 +0,0 @@
|
||||
import itertools
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kvcache import can_use_store_cache, store_cache
|
||||
from sglang.jit_kernel.utils import get_ci_test_range
|
||||
from sglang.test.ci.ci_register import register_amd_ci, 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)
|
||||
register_amd_ci(est_time=55, suite="jit-kernel-unit-test-amd")
|
||||
|
||||
BS_LIST = [2**n for n in range(0, 15)]
|
||||
BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)]
|
||||
BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 256, 16399])
|
||||
HIDDEN_DIMS = get_ci_test_range(
|
||||
[64, 128, 256, 512, 1024, 96, 98, 100], [64, 512, 1024, 98]
|
||||
)
|
||||
CACHE_SIZE = 1024 * 1024
|
||||
DTYPE = torch.bfloat16
|
||||
DEVICE = "cuda"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_size,element_dim",
|
||||
list(itertools.product(BS_LIST, HIDDEN_DIMS)),
|
||||
)
|
||||
def test_store_cache(batch_size: int, element_dim: int) -> None:
|
||||
k = torch.randn((batch_size, element_dim), dtype=DTYPE, device=DEVICE)
|
||||
v = torch.randn((batch_size, element_dim), dtype=DTYPE, device=DEVICE)
|
||||
k_cache = torch.randn((CACHE_SIZE, element_dim), dtype=DTYPE, device=DEVICE)
|
||||
v_cache = torch.randn((CACHE_SIZE, element_dim), dtype=DTYPE, device=DEVICE)
|
||||
indices = torch.randperm(CACHE_SIZE, device=DEVICE)[:batch_size]
|
||||
|
||||
# AOT store cache
|
||||
store_cache(k, v, k_cache, v_cache, indices)
|
||||
|
||||
assert torch.all(k_cache[indices] == k)
|
||||
assert torch.all(v_cache[indices] == v)
|
||||
|
||||
|
||||
# Smaller subset for targeted tests below
|
||||
REPR_BS = get_ci_test_range([1, 7, 128], [1, 128])
|
||||
REPR_DIMS = get_ci_test_range([64, 128, 512, 1024, 96], [64, 1024, 96])
|
||||
SMALL_CACHE = 4096
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
|
||||
@pytest.mark.parametrize(
|
||||
"batch_size,element_dim",
|
||||
list(itertools.product(REPR_BS, REPR_DIMS)),
|
||||
)
|
||||
def test_store_cache_dtypes(
|
||||
batch_size: int, element_dim: int, dtype: torch.dtype
|
||||
) -> None:
|
||||
k = torch.randn((batch_size, element_dim), dtype=dtype, device=DEVICE)
|
||||
v = torch.randn((batch_size, element_dim), dtype=dtype, device=DEVICE)
|
||||
k_cache = torch.randn((SMALL_CACHE, element_dim), dtype=dtype, device=DEVICE)
|
||||
v_cache = torch.randn((SMALL_CACHE, element_dim), dtype=dtype, device=DEVICE)
|
||||
indices = torch.randperm(SMALL_CACHE, device=DEVICE)[:batch_size]
|
||||
|
||||
store_cache(k, v, k_cache, v_cache, indices)
|
||||
|
||||
assert torch.all(k_cache[indices] == k)
|
||||
assert torch.all(v_cache[indices] == v)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_size,element_dim",
|
||||
list(itertools.product(REPR_BS, REPR_DIMS)),
|
||||
)
|
||||
def test_store_cache_int32_indices(batch_size: int, element_dim: int) -> None:
|
||||
k = torch.randn((batch_size, element_dim), dtype=DTYPE, device=DEVICE)
|
||||
v = torch.randn((batch_size, element_dim), dtype=DTYPE, device=DEVICE)
|
||||
k_cache = torch.randn((SMALL_CACHE, element_dim), dtype=DTYPE, device=DEVICE)
|
||||
v_cache = torch.randn((SMALL_CACHE, element_dim), dtype=DTYPE, device=DEVICE)
|
||||
# int32 indices exercise a different CUDA template instantiation than default int64
|
||||
indices = torch.randperm(SMALL_CACHE, device=DEVICE)[:batch_size].to(torch.int32)
|
||||
|
||||
store_cache(k, v, k_cache, v_cache, indices)
|
||||
|
||||
assert torch.all(k_cache[indices.long()] == k)
|
||||
assert torch.all(v_cache[indices.long()] == v)
|
||||
|
||||
|
||||
def _valid_num_splits(element_dim: int, dtype: torch.dtype) -> list:
|
||||
"""Return the list of valid num_split values for a given element_dim/dtype."""
|
||||
row_bytes = element_dim * dtype.itemsize
|
||||
splits = [1]
|
||||
if row_bytes % (2 * 128) == 0:
|
||||
splits.append(2)
|
||||
if row_bytes % (4 * 128) == 0:
|
||||
splits.append(4)
|
||||
return splits
|
||||
|
||||
|
||||
_NUM_SPLIT_CASES = [
|
||||
(_dim, _ns, _dtype)
|
||||
for _dtype in [torch.float16, torch.bfloat16, torch.float32]
|
||||
for _dim in REPR_DIMS
|
||||
for _ns in _valid_num_splits(_dim, _dtype)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("element_dim,num_split,dtype", _NUM_SPLIT_CASES)
|
||||
def test_store_cache_num_split(
|
||||
element_dim: int, num_split: int, dtype: torch.dtype
|
||||
) -> None:
|
||||
batch_size = 128
|
||||
k = torch.randn((batch_size, element_dim), dtype=dtype, device=DEVICE)
|
||||
v = torch.randn((batch_size, element_dim), dtype=dtype, device=DEVICE)
|
||||
k_cache = torch.randn((SMALL_CACHE, element_dim), dtype=dtype, device=DEVICE)
|
||||
v_cache = torch.randn((SMALL_CACHE, element_dim), dtype=dtype, device=DEVICE)
|
||||
indices = torch.randperm(SMALL_CACHE, device=DEVICE)[:batch_size]
|
||||
|
||||
# Verify each num_split kernel path (1, 2, 4) produces correct results
|
||||
store_cache(k, v, k_cache, v_cache, indices, num_split=num_split)
|
||||
|
||||
assert torch.all(k_cache[indices] == k)
|
||||
assert torch.all(v_cache[indices] == v)
|
||||
|
||||
|
||||
def test_can_use_store_cache() -> None:
|
||||
assert can_use_store_cache(128)
|
||||
assert can_use_store_cache(256)
|
||||
assert can_use_store_cache(1024)
|
||||
assert can_use_store_cache(2048)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -1,183 +0,0 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
try:
|
||||
import tabulate
|
||||
except Exception:
|
||||
tabulate = None
|
||||
|
||||
from sglang.jit_kernel.timestep_embedding import (
|
||||
timestep_embedding as timestep_embedding_cuda,
|
||||
)
|
||||
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=16, suite="base-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
CORRECTNESS_BATCH_SIZES = get_ci_test_range(
|
||||
[1, 2, 8, 128, 256, 512, 1536, 2048, 4096, 11008, 16384],
|
||||
[1, 128, 2048, 16384],
|
||||
)
|
||||
CORRECTNESS_DIMS = get_ci_test_range(
|
||||
[32, 128, 256, 512, 1536, 2048, 4096, 8192],
|
||||
[32, 512, 8192],
|
||||
)
|
||||
DIFFUSERS_BATCH_SIZES = get_ci_test_range(
|
||||
[1, 2, 8, 128, 256, 512, 1536, 2048, 16384],
|
||||
[1, 512, 16384],
|
||||
)
|
||||
DIFFUSERS_DIMS = get_ci_test_range([32, 256, 512, 1536, 8192], [32, 512, 8192])
|
||||
DTYPES = get_ci_test_range(
|
||||
[torch.float16, torch.bfloat16, torch.float32],
|
||||
[torch.float16, torch.bfloat16],
|
||||
)
|
||||
SCALES = get_ci_test_range([1, 0.01], [1, 0.01])
|
||||
|
||||
|
||||
def get_timestep_embedding_reference(
|
||||
timesteps: torch.Tensor,
|
||||
dim: int,
|
||||
*,
|
||||
flip_sin_to_cos: bool = False,
|
||||
downscale_freq_shift: float = 1,
|
||||
scale: float = 1,
|
||||
max_period: int = 10000,
|
||||
):
|
||||
assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
|
||||
|
||||
timesteps = timesteps.to(torch.float32)
|
||||
half_dim = dim // 2
|
||||
exponent = -torch.log(
|
||||
torch.tensor(max_period, dtype=torch.float32, device=timesteps.device)
|
||||
) * torch.arange(
|
||||
start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
|
||||
)
|
||||
exponent = exponent / (half_dim - downscale_freq_shift)
|
||||
|
||||
emb = torch.exp(exponent)
|
||||
emb = timesteps[:, None].float() * emb[None, :]
|
||||
|
||||
emb = scale * emb
|
||||
|
||||
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
|
||||
if flip_sin_to_cos:
|
||||
emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
|
||||
if dim % 2 == 1:
|
||||
emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
|
||||
return emb
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", CORRECTNESS_BATCH_SIZES)
|
||||
@pytest.mark.parametrize("dim", CORRECTNESS_DIMS)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
def test_timestep_embedding_correctness_with_sgld(batch_size, dim, dtype):
|
||||
device = "cuda"
|
||||
t = torch.randint(low=0, high=1000, size=(batch_size,), device=device).to(dtype)
|
||||
torch_output = get_timestep_embedding_reference(
|
||||
t, dim, flip_sin_to_cos=True, downscale_freq_shift=0
|
||||
)
|
||||
cuda_output = timestep_embedding_cuda(
|
||||
t, dim, flip_sin_to_cos=True, downscale_freq_shift=0
|
||||
)
|
||||
torch.testing.assert_close(torch_output, cuda_output, atol=1e-3, rtol=1e-3)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", DIFFUSERS_BATCH_SIZES)
|
||||
@pytest.mark.parametrize("dim", DIFFUSERS_DIMS)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("flip_sin_to_cos", [False, True])
|
||||
@pytest.mark.parametrize("downscale_freq_shift", [0, 1])
|
||||
@pytest.mark.parametrize("scale", SCALES)
|
||||
def test_timestep_embedding_correctness_with_diffusers(
|
||||
batch_size, dim, flip_sin_to_cos, downscale_freq_shift, scale, dtype
|
||||
):
|
||||
device = "cuda"
|
||||
t = torch.randint(low=0, high=1000, size=(batch_size,), device=device).to(dtype)
|
||||
torch_output = get_timestep_embedding_reference(
|
||||
t,
|
||||
dim,
|
||||
flip_sin_to_cos=flip_sin_to_cos,
|
||||
downscale_freq_shift=downscale_freq_shift,
|
||||
scale=scale,
|
||||
max_period=10000,
|
||||
)
|
||||
cuda_output = timestep_embedding_cuda(
|
||||
t,
|
||||
dim,
|
||||
flip_sin_to_cos=flip_sin_to_cos,
|
||||
downscale_freq_shift=downscale_freq_shift,
|
||||
scale=scale,
|
||||
max_period=10000,
|
||||
)
|
||||
torch.testing.assert_close(torch_output, cuda_output, atol=1e-3, rtol=1e-3)
|
||||
|
||||
|
||||
def test_timestep_embedding_perf():
|
||||
if os.environ.get("SGLANG_RUN_JIT_KERNEL_PERF_TESTS") != "1":
|
||||
pytest.skip("Perf test disabled by default")
|
||||
if tabulate is None:
|
||||
pytest.skip("Optional dependency 'tabulate' is not installed")
|
||||
|
||||
NUM_BATCH = [1, 2, 8, 63, 256, 512, 613, 1024, 1536]
|
||||
NUM_DIM = [32, 64, 128, 256, 512, 1024, 2048, 4096]
|
||||
|
||||
def perf_kernel_fn(kernel_fn: callable, *args, **kwargs):
|
||||
warmup_times = 4
|
||||
repeat_times = 20
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
|
||||
for _ in range(warmup_times):
|
||||
output_fn = kernel_fn(*args, **kwargs)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
start.record()
|
||||
for _ in range(repeat_times):
|
||||
output_fn = kernel_fn(*args, **kwargs)
|
||||
end.record()
|
||||
end.synchronize()
|
||||
return start.elapsed_time(end) / repeat_times
|
||||
|
||||
device = "cuda"
|
||||
results = []
|
||||
|
||||
cuda_speedups = []
|
||||
for B in NUM_BATCH:
|
||||
for dim in NUM_DIM:
|
||||
t = torch.linspace(0, max(100000, B), steps=B, device=device).to(
|
||||
torch.float32
|
||||
)
|
||||
time_torch = perf_kernel_fn(get_timestep_embedding_reference, t, dim)
|
||||
time_cuda = perf_kernel_fn(timestep_embedding_cuda, t, dim)
|
||||
speedup_cuda = time_torch / time_cuda
|
||||
|
||||
results.append(
|
||||
{
|
||||
"Batch Size": B,
|
||||
"Dimension": dim,
|
||||
"Torch Time (ms)": time_torch,
|
||||
"CUDA Time (ms)": time_cuda,
|
||||
"Speedup (CUDA)": speedup_cuda,
|
||||
}
|
||||
)
|
||||
cuda_speedups.append(speedup_cuda)
|
||||
|
||||
print("=== Timestep Embedding Benchmark Results ===")
|
||||
print(
|
||||
tabulate.tabulate(
|
||||
results,
|
||||
headers="keys",
|
||||
tablefmt="fancy_grid",
|
||||
floatfmt=(".0f", ".0f", ".6f", ".6f", ".5f"),
|
||||
)
|
||||
)
|
||||
print(f"Average Speedup(cuda): {np.mean(cuda_speedups):.4f}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -1,168 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import triton
|
||||
|
||||
from sglang.jit_kernel.all_reduce import fused_parallel_qknorm
|
||||
from sglang.jit_kernel.tests.test_custom_all_reduce import multiprocess_test
|
||||
from sglang.jit_kernel.tests.utils import multiprocess_main
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
Q_K_DIMS = [(6144, 1024)]
|
||||
EPS = 1e-6
|
||||
BATCH_SIZES = [2**n for n in range(0, 14)]
|
||||
DTYPES = [torch.float16, torch.bfloat16, torch.float32]
|
||||
TEST_CONFIG = list(itertools.product(Q_K_DIMS, BATCH_SIZES, DTYPES))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("nproc", [2, 4, 8])
|
||||
def test_tp_qknorm(nproc: int) -> None:
|
||||
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)
|
||||
|
||||
|
||||
def init_distributed():
|
||||
import sglang.srt.distributed.parallel_state as ps
|
||||
from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
|
||||
CustomAllReduceV2,
|
||||
)
|
||||
|
||||
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_pull_size = 0
|
||||
max_push_size = 8 * max(BATCH_SIZES)
|
||||
comm = CustomAllReduceV2(cpu_group, device, max_pull_size, max_push_size)
|
||||
if comm.disabled:
|
||||
raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system")
|
||||
|
||||
return rank, world_size, device, cpu_group, nccl_group, comm
|
||||
|
||||
|
||||
def _all_gather_cat(x: torch.Tensor, group: dist.ProcessGroup) -> torch.Tensor:
|
||||
gathered = [torch.empty_like(x) for _ in range(dist.get_world_size(group=group))]
|
||||
dist.all_gather(gathered, x, group=group)
|
||||
return torch.cat(gathered, dim=-1)
|
||||
|
||||
|
||||
def _rmsnorm_ref(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
|
||||
x_fp32 = x.float()
|
||||
scale = (x_fp32.pow(2).mean(dim=-1, keepdim=True) + eps).rsqrt()
|
||||
return (x_fp32 * scale * weight.float()).to(x.dtype)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def worker_test(
|
||||
rank: int,
|
||||
world_size: int,
|
||||
device: torch.device,
|
||||
nccl_group: dist.ProcessGroup,
|
||||
comm,
|
||||
q_k_dim: tuple[int, int],
|
||||
batch_size: int,
|
||||
dtype: torch.dtype,
|
||||
) -> Optional[RuntimeError]:
|
||||
q_dim, k_dim = q_k_dim
|
||||
local_q_dim = q_dim // world_size
|
||||
local_k_dim = k_dim // world_size
|
||||
|
||||
q = torch.randn(batch_size, local_q_dim, device=device, dtype=dtype)
|
||||
k = torch.randn(batch_size, local_k_dim, device=device, dtype=dtype)
|
||||
q_weight = torch.randn(local_q_dim, device=device, dtype=dtype)
|
||||
k_weight = torch.randn(local_k_dim, device=device, dtype=dtype)
|
||||
|
||||
q_ref = _all_gather_cat(q, nccl_group)
|
||||
k_ref = _all_gather_cat(k, nccl_group)
|
||||
q_weight_ref = _all_gather_cat(q_weight.unsqueeze(0), nccl_group).squeeze(0)
|
||||
k_weight_ref = _all_gather_cat(k_weight.unsqueeze(0), nccl_group).squeeze(0)
|
||||
|
||||
q_expected = _rmsnorm_ref(q_ref, q_weight_ref, EPS)
|
||||
k_expected = _rmsnorm_ref(k_ref, k_weight_ref, EPS)
|
||||
q_expected = q_expected[:, rank * local_q_dim : (rank + 1) * local_q_dim]
|
||||
k_expected = k_expected[:, rank * local_k_dim : (rank + 1) * local_k_dim]
|
||||
|
||||
fused_parallel_qknorm(
|
||||
comm.obj,
|
||||
q,
|
||||
k,
|
||||
q_weight,
|
||||
k_weight,
|
||||
EPS,
|
||||
)
|
||||
|
||||
try:
|
||||
triton.testing.assert_close(q, q_expected, atol=1e-2, rtol=1e-2)
|
||||
triton.testing.assert_close(k, k_expected, atol=1e-2, rtol=1e-2)
|
||||
except AssertionError as err:
|
||||
return RuntimeError(
|
||||
f"TP QKNorm mismatch for {batch_size=}, {dtype=}, {world_size=}, {rank=}: {err}"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def worker_main() -> None:
|
||||
rank, world_size, device, cpu_group, nccl_group, comm = init_distributed()
|
||||
torch.cuda.set_stream(torch.cuda.Stream())
|
||||
|
||||
for q_k_dim, batch_size, dtype in TEST_CONFIG:
|
||||
error = worker_test(
|
||||
rank,
|
||||
world_size,
|
||||
device,
|
||||
nccl_group,
|
||||
comm,
|
||||
q_k_dim,
|
||||
batch_size,
|
||||
dtype,
|
||||
)
|
||||
result = torch.tensor([int(error is not None)])
|
||||
dist.all_reduce(result, group=cpu_group)
|
||||
if error is not None:
|
||||
print(str(error))
|
||||
if bool(result.item()):
|
||||
raise RuntimeError(
|
||||
f"TP QKNorm test failed for {q_k_dim=}, {batch_size=}, {dtype=}, {world_size=}"
|
||||
)
|
||||
|
||||
print(f"Rank {rank} passed all tests.")
|
||||
comm.close()
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
multiprocess_main(__file__, worker_main)
|
||||
Reference in New Issue
Block a user