minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Cleanup] Deduplicate kernel tests, diffusion fixtures and benchmark helpers (#40265)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-09-19 19:45:38 +08:00
committed by GitHub
parent 0b0d2c257a
commit cb22f2451e
25 changed files with 512 additions and 1936 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kernels/ops/attention/test_cute_dsl_fp8_paged_mqa_logits.py
+12 -133
View File
@@ -7,141 +7,18 @@ import pytest
import torch
from sglang.kernels.ops.attention.dsa import cutedsl_paged_mqa_logits, pick_dsl_expand
from sglang.srt.layers.attention.dsa.utils import (
fp8_mqa_logits_ceil_to_ue8m0,
fp8_mqa_logits_make_fused_kv,
)
from sglang.srt.utils import is_sm100_supported
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kernels.paged_mqa import (
BLOCK_KV,
HEAD_DIM,
assert_paged_mqa_matches_ref,
generate_paged_mqa_test_data,
ref_fp8_paged_mqa_logits,
)
register_cuda_ci(est_time=180, stage="nightly", runner_config="4-gpu-b200")
BLOCK_KV = 64
HEAD_DIM = 128
def _ref_fp8_paged_mqa_logits(
q_fp8,
kv_fp8,
kv_scales,
weights,
context_lens,
block_table,
max_model_len,
block_kv,
):
B, next_n, H, D = q_fp8.shape
device = q_fp8.device
logits = torch.full(
(B * next_n, max_model_len), float("-inf"), device=device, dtype=torch.float32
)
q_f32 = q_fp8.float()
for b in range(B):
ctx_len = context_lens[b].item()
q_positions = torch.arange(ctx_len - next_n, ctx_len, device=device)
w = weights[b * next_n : (b + 1) * next_n, :]
for blk_idx in range((ctx_len + block_kv - 1) // block_kv):
phys_blk = block_table[b, blk_idx].item()
k_f32 = kv_fp8[phys_blk].float()
scales = kv_scales[phys_blk]
k_positions = torch.arange(
blk_idx * block_kv, (blk_idx + 1) * block_kv, device=device
)
mask = (k_positions[None, :] < ctx_len) & (
k_positions[None, :] <= q_positions[:, None]
)
qk = torch.matmul(q_f32[b].permute(1, 0, 2), k_f32.T)
qk = torch.where(mask[None, :, :], qk, torch.zeros(1, device=device))
qk = torch.relu(qk)
weighted = (w.T[:, :, None] * qk).sum(dim=0)
weighted = weighted * scales[None, :]
start_pos = blk_idx * block_kv
end_pos = start_pos + block_kv
logits[b * next_n : (b + 1) * next_n, start_pos:end_pos] = torch.where(
mask,
weighted,
torch.tensor(float("-inf"), device=device, dtype=torch.float32),
)
return logits
def _generate_test_data(
batch_size,
next_n,
num_heads,
avg_context_len,
max_model_len,
device="cuda",
):
torch.manual_seed(42)
torch.cuda.manual_seed(42)
context_lens = torch.randint(
max(BLOCK_KV, int(0.7 * avg_context_len)),
int(1.3 * avg_context_len) + 1,
(batch_size,),
dtype=torch.int32,
device="cpu",
).clamp(max=max_model_len)
max_blocks_per_seq = (max_model_len + BLOCK_KV - 1) // BLOCK_KV
total_blocks = ((context_lens + BLOCK_KV - 1) // BLOCK_KV).sum().item()
num_phys_blocks = total_blocks + batch_size * 2
block_table = torch.full(
(batch_size, max_blocks_per_seq), 0, dtype=torch.int32, device=device
)
blk_offset = 0
for i in range(batch_size):
n_blks = (context_lens[i].item() + BLOCK_KV - 1) // BLOCK_KV
block_table[i, :n_blks] = torch.arange(
blk_offset, blk_offset + n_blks, dtype=torch.int32, device=device
)
blk_offset += n_blks
q_bf16 = torch.randn(batch_size, next_n, num_heads, HEAD_DIM, device=device)
q_fp8 = q_bf16.to(torch.float8_e4m3fn)
kv_bf16 = torch.randn(num_phys_blocks, BLOCK_KV, HEAD_DIM, device=device)
kv_amax = kv_bf16.abs().float().amax(dim=-1, keepdim=True).clamp(1e-4)
kv_scale = fp8_mqa_logits_ceil_to_ue8m0(kv_amax / 448.0).squeeze(-1)
kv_fp8 = (kv_bf16 / kv_scale.unsqueeze(-1)).to(torch.float8_e4m3fn)
weights = torch.randn(
batch_size * next_n, num_heads, device=device, dtype=torch.float32
)
kv_fused = fp8_mqa_logits_make_fused_kv(kv_fp8, kv_scale, BLOCK_KV, HEAD_DIM)
return {
"q_fp8": q_fp8,
"kv_fp8": kv_fp8,
"kv_scales": kv_scale,
"kv_fused": kv_fused,
"weights": weights,
"context_lens": context_lens.to(device),
"block_table": block_table,
}
def _assert_matches_ref(logits, ref_logits, context_lens, B, next_n, max_model_len):
device = logits.device
positions = torch.arange(max_model_len, device=device).unsqueeze(0)
row_indices = torch.arange(B * next_n, device=device) // next_n
next_n_offset = torch.arange(B * next_n, device=device) % next_n
end_pos = context_lens[row_indices] - next_n + next_n_offset
mask = positions <= end_pos.unsqueeze(1)
logits_masked = logits.float().masked_fill(~mask, 0)
ref_masked = ref_logits.float().masked_fill(~mask, 0)
torch.testing.assert_close(logits_masked, ref_masked, atol=5e-5, rtol=1e-5)
def _run_cutedsl_paged_mqa_logits(
data, batch_size, next_n, num_heads, max_model_len, is_target_verify
@@ -204,7 +81,9 @@ def _run_cutedsl_paged_mqa_logits(
@pytest.mark.parametrize("avg_ctx", [128, 1024, 4096, 16384])
def test_cutedsl_paged_mqa_logits(batch_size, next_n, num_heads, avg_ctx):
max_model_len = max(avg_ctx * 2, 2048)
data = _generate_test_data(batch_size, next_n, num_heads, avg_ctx, max_model_len)
data = generate_paged_mqa_test_data(
batch_size, next_n, num_heads, avg_ctx, max_model_len
)
logits = _run_cutedsl_paged_mqa_logits(
data,
@@ -215,7 +94,7 @@ def test_cutedsl_paged_mqa_logits(batch_size, next_n, num_heads, avg_ctx):
is_target_verify=next_n >= 2,
)
ref_logits = _ref_fp8_paged_mqa_logits(
ref_logits = ref_fp8_paged_mqa_logits(
data["q_fp8"],
data["kv_fp8"],
data["kv_scales"],
@@ -225,7 +104,7 @@ def test_cutedsl_paged_mqa_logits(batch_size, next_n, num_heads, avg_ctx):
max_model_len,
BLOCK_KV,
)
_assert_matches_ref(
assert_paged_mqa_matches_ref(
logits, ref_logits, data["context_lens"], batch_size, next_n, max_model_len
)
test/registered/kernels/ops/attention/test_deepgemm_paged_mqa_logits.py
+12 -133
View File
@@ -10,141 +10,18 @@ from sglang.kernels.ops.attention.dsa import (
deepgemm_paged_mqa_logits_native,
deepgemm_paged_mqa_logits_split,
)
from sglang.srt.layers.attention.dsa.utils import (
fp8_mqa_logits_ceil_to_ue8m0,
fp8_mqa_logits_make_fused_kv,
)
from sglang.srt.utils import is_sm90_supported, is_sm100_supported
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kernels.paged_mqa import (
BLOCK_KV,
HEAD_DIM,
assert_paged_mqa_matches_ref,
generate_paged_mqa_test_data,
ref_fp8_paged_mqa_logits,
)
register_cuda_ci(est_time=40, stage="nightly", runner_config="4-gpu-b200")
BLOCK_KV = 64
HEAD_DIM = 128
def _ref_fp8_paged_mqa_logits(
q_fp8,
kv_fp8,
kv_scales,
weights,
context_lens,
block_table,
max_model_len,
block_kv,
):
B, next_n, H, D = q_fp8.shape
device = q_fp8.device
logits = torch.full(
(B * next_n, max_model_len), float("-inf"), device=device, dtype=torch.float32
)
q_f32 = q_fp8.float()
for b in range(B):
ctx_len = context_lens[b].item()
q_positions = torch.arange(ctx_len - next_n, ctx_len, device=device)
w = weights[b * next_n : (b + 1) * next_n, :]
for blk_idx in range((ctx_len + block_kv - 1) // block_kv):
phys_blk = block_table[b, blk_idx].item()
k_f32 = kv_fp8[phys_blk].float()
scales = kv_scales[phys_blk]
k_positions = torch.arange(
blk_idx * block_kv, (blk_idx + 1) * block_kv, device=device
)
mask = (k_positions[None, :] < ctx_len) & (
k_positions[None, :] <= q_positions[:, None]
)
qk = torch.matmul(q_f32[b].permute(1, 0, 2), k_f32.T)
qk = torch.where(mask[None, :, :], qk, torch.zeros(1, device=device))
qk = torch.relu(qk)
weighted = (w.T[:, :, None] * qk).sum(dim=0)
weighted = weighted * scales[None, :]
start_pos = blk_idx * block_kv
end_pos = start_pos + block_kv
logits[b * next_n : (b + 1) * next_n, start_pos:end_pos] = torch.where(
mask,
weighted,
torch.tensor(float("-inf"), device=device, dtype=torch.float32),
)
return logits
def _generate_test_data(
batch_size,
next_n,
num_heads,
avg_context_len,
max_model_len,
device="cuda",
):
torch.manual_seed(42)
torch.cuda.manual_seed(42)
context_lens = torch.randint(
max(BLOCK_KV, int(0.7 * avg_context_len)),
int(1.3 * avg_context_len) + 1,
(batch_size,),
dtype=torch.int32,
device="cpu",
).clamp(max=max_model_len)
max_blocks_per_seq = (max_model_len + BLOCK_KV - 1) // BLOCK_KV
total_blocks = ((context_lens + BLOCK_KV - 1) // BLOCK_KV).sum().item()
num_phys_blocks = total_blocks + batch_size * 2
block_table = torch.full(
(batch_size, max_blocks_per_seq), 0, dtype=torch.int32, device=device
)
blk_offset = 0
for i in range(batch_size):
n_blks = (context_lens[i].item() + BLOCK_KV - 1) // BLOCK_KV
block_table[i, :n_blks] = torch.arange(
blk_offset, blk_offset + n_blks, dtype=torch.int32, device=device
)
blk_offset += n_blks
q_bf16 = torch.randn(batch_size, next_n, num_heads, HEAD_DIM, device=device)
q_fp8 = q_bf16.to(torch.float8_e4m3fn)
kv_bf16 = torch.randn(num_phys_blocks, BLOCK_KV, HEAD_DIM, device=device)
kv_amax = kv_bf16.abs().float().amax(dim=-1, keepdim=True).clamp(1e-4)
kv_scale = fp8_mqa_logits_ceil_to_ue8m0(kv_amax / 448.0).squeeze(-1)
kv_fp8 = (kv_bf16 / kv_scale.unsqueeze(-1)).to(torch.float8_e4m3fn)
weights = torch.randn(
batch_size * next_n, num_heads, device=device, dtype=torch.float32
)
kv_fused = fp8_mqa_logits_make_fused_kv(kv_fp8, kv_scale, BLOCK_KV, HEAD_DIM)
return {
"q_fp8": q_fp8,
"kv_fp8": kv_fp8,
"kv_scales": kv_scale,
"kv_fused": kv_fused,
"weights": weights,
"context_lens": context_lens.to(device),
"block_table": block_table,
}
def _assert_matches_ref(logits, ref_logits, context_lens, B, next_n, max_model_len):
device = logits.device
positions = torch.arange(max_model_len, device=device).unsqueeze(0)
row_indices = torch.arange(B * next_n, device=device) // next_n
next_n_offset = torch.arange(B * next_n, device=device) % next_n
end_pos = context_lens[row_indices] - next_n + next_n_offset
mask = positions <= end_pos.unsqueeze(1)
logits_masked = logits.float().masked_fill(~mask, 0)
ref_masked = ref_logits.float().masked_fill(~mask, 0)
torch.testing.assert_close(logits_masked, ref_masked, atol=5e-5, rtol=1e-5)
def _run_deepgemm_paged_mqa_logits(data, batch_size, next_n, num_heads, max_model_len):
"""Mirrors the DEEPGEMM dispatch in
@@ -208,13 +85,15 @@ def _run_deepgemm_paged_mqa_logits(data, batch_size, next_n, num_heads, max_mode
@pytest.mark.parametrize("avg_ctx", [128, 1024, 4096, 16384])
def test_deepgemm_paged_mqa_logits(batch_size, next_n, num_heads, avg_ctx):
max_model_len = max(avg_ctx * 2, 2048)
data = _generate_test_data(batch_size, next_n, num_heads, avg_ctx, max_model_len)
data = generate_paged_mqa_test_data(
batch_size, next_n, num_heads, avg_ctx, max_model_len
)
logits = _run_deepgemm_paged_mqa_logits(
data, batch_size, next_n, num_heads, max_model_len
)
ref_logits = _ref_fp8_paged_mqa_logits(
ref_logits = ref_fp8_paged_mqa_logits(
data["q_fp8"],
data["kv_fp8"],
data["kv_scales"],
@@ -224,7 +103,7 @@ def test_deepgemm_paged_mqa_logits(batch_size, next_n, num_heads, avg_ctx):
max_model_len,
BLOCK_KV,
)
_assert_matches_ref(
assert_paged_mqa_matches_ref(
logits, ref_logits, data["context_lens"], batch_size, next_n, max_model_len
)
test/registered/kernels/ops/attention/test_fused_metadata_copy.py
+71 -496
View File
@@ -1,15 +1,6 @@
"""
Comprehensive tests for JIT-compiled fused metadata copy kernels.
This test suite verifies:
1. Single-backend fused kernel (fused_metadata_copy_cuda) - all forward modes
2. Multi-backend fused kernel (fused_metadata_copy_multi_cuda) - 3 backends at once
3. Correctness against reference implementations
4. Performance benchmarks and speedup measurements
"""
"""Compare single- and multi-backend metadata copies with PyTorch references."""
import sys
import time
import pytest
import torch
@@ -151,62 +142,6 @@ def reference_copy_decode(src, dst, max_len):
dst["flashmla_metadata"].copy_(src["flashmla_metadata"])
def reference_copy_target_verify(src, dst, max_seqlen_k, seqlens_expanded_size):
"""Reference implementation: individual .copy_() for TARGET_VERIFY mode."""
bs = src["cache_seqlens"].shape[0]
dst["cache_seqlens"].copy_(src["cache_seqlens"])
dst["cu_seqlens_k"][1:].copy_(src["cu_seqlens_k"][1:])
rows, cols = src["page_indices"].shape
dst["page_table_1"][:rows, :cols].copy_(src["page_indices"])
dst["dsa_seqlens_expanded"][:seqlens_expanded_size].copy_(src["seqlens_expanded"])
dst["dsa_cache_seqlens"][:seqlens_expanded_size].copy_(src["dsa_cache_seqlens"])
dst["dsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1].copy_(
src["dsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1]
)
if src["real_page_table"] is not None:
rows, cols = src["real_page_table"].shape
dst["real_page_table"][:rows, :cols].copy_(src["real_page_table"])
if src["flashmla_num_splits"] is not None:
flashmla_size = seqlens_expanded_size + 1
dst["flashmla_num_splits"][:flashmla_size].copy_(
src["flashmla_num_splits"][:flashmla_size]
)
if src["flashmla_metadata"] is not None:
dst["flashmla_metadata"].copy_(src["flashmla_metadata"])
def reference_copy_draft_extend(src, dst, max_seqlen_k, seqlens_expanded_size):
"""Reference implementation: individual .copy_() for DRAFT_EXTEND mode."""
bs = src["cache_seqlens"].shape[0]
dst["cache_seqlens"].copy_(src["cache_seqlens"])
dst["cu_seqlens_k"][1:].copy_(src["cu_seqlens_k"][1:])
rows, cols = src["page_indices"].shape
dst["page_table_1"][:rows, :cols].copy_(src["page_indices"])
dst["dsa_seqlens_expanded"][:seqlens_expanded_size].copy_(src["seqlens_expanded"])
dst["dsa_cache_seqlens"][:seqlens_expanded_size].copy_(src["dsa_cache_seqlens"])
dst["dsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1].copy_(
src["dsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1]
)
if src["real_page_table"] is not None:
rows, cols = src["real_page_table"].shape
dst["real_page_table"][:rows, :cols].copy_(src["real_page_table"])
if src["flashmla_num_splits"] is not None:
flashmla_size = seqlens_expanded_size + 1
dst["flashmla_num_splits"][:flashmla_size].copy_(
src["flashmla_num_splits"][:flashmla_size]
)
if src["flashmla_metadata"] is not None:
dst["flashmla_metadata"].copy_(src["flashmla_metadata"])
# =============================================================================
# Single-Backend Kernel Tests
# =============================================================================
@@ -321,13 +256,17 @@ def test_fused_metadata_copy_dtype_validation():
)
@pytest.mark.parametrize("bs", [1, 2, 4, 8])
@pytest.mark.parametrize(
"forward_mode", [0]
) # DECODE mode only (other modes not fully tested yet)
@pytest.mark.parametrize("has_real_page_table", [False, True])
@pytest.mark.parametrize("has_flashmla", [False, True])
def test_fused_metadata_copy(bs, forward_mode, has_real_page_table, has_flashmla):
"bs,has_real_page_table,has_flashmla",
[
(bs, page, mla)
for bs in (1, 2, 4, 8)
for page in (False, True)
for mla in (False, True)
]
+ [(16, True, True), (32, True, True)],
)
def test_fused_metadata_copy(bs, has_real_page_table, has_flashmla):
"""Test fused metadata copy kernel against reference implementation."""
if not torch.cuda.is_available():
pytest.skip("CUDA not available")
@@ -336,9 +275,10 @@ def test_fused_metadata_copy(bs, forward_mode, has_real_page_table, has_flashmla
fused_metadata_copy_cuda,
)
forward_mode = 0 # DECODE
max_len = 128
max_seqlen_k = 256
seqlens_expanded_size = bs if forward_mode == 0 else bs * 2
seqlens_expanded_size = bs
# Create test data
data = create_test_metadata(
@@ -356,17 +296,7 @@ def test_fused_metadata_copy(bs, forward_mode, has_real_page_table, has_flashmla
k: v.clone() if v is not None else None for k, v in data["dst"].items()
}
# Run reference implementation
if forward_mode == 0: # DECODE
reference_copy_decode(data["src"], dst_ref, max_len)
elif forward_mode == 1: # TARGET_VERIFY
reference_copy_target_verify(
data["src"], dst_ref, max_seqlen_k, seqlens_expanded_size
)
else: # DRAFT_EXTEND
reference_copy_draft_extend(
data["src"], dst_ref, max_seqlen_k, seqlens_expanded_size
)
reference_copy_decode(data["src"], dst_ref, max_len)
# Run fused kernel
fused_metadata_copy_cuda(
@@ -395,101 +325,11 @@ def test_fused_metadata_copy(bs, forward_mode, has_real_page_table, has_flashmla
seqlens_expanded_size,
)
# Compare results
assert torch.equal(dst_ref["cache_seqlens"], dst_fused["cache_seqlens"]), (
"cache_seqlens mismatch"
)
assert torch.equal(dst_ref["cu_seqlens_k"], dst_fused["cu_seqlens_k"]), (
"cu_seqlens_k mismatch"
)
assert torch.equal(dst_ref["page_table_1"], dst_fused["page_table_1"]), (
"page_table_1 mismatch"
)
assert torch.equal(dst_ref["dsa_cache_seqlens"], dst_fused["dsa_cache_seqlens"]), (
"dsa_cache_seqlens mismatch"
)
assert torch.equal(
dst_ref["dsa_seqlens_expanded"], dst_fused["dsa_seqlens_expanded"]
), "dsa_seqlens_expanded mismatch"
assert torch.equal(dst_ref["dsa_cu_seqlens_k"], dst_fused["dsa_cu_seqlens_k"]), (
"dsa_cu_seqlens_k mismatch"
)
if has_real_page_table:
assert torch.equal(dst_ref["real_page_table"], dst_fused["real_page_table"]), (
"real_page_table mismatch"
)
if has_flashmla:
assert torch.equal(
dst_ref["flashmla_num_splits"], dst_fused["flashmla_num_splits"]
), "flashmla_num_splits mismatch"
assert torch.equal(
dst_ref["flashmla_metadata"], dst_fused["flashmla_metadata"]
), "flashmla_metadata mismatch"
@pytest.mark.parametrize("bs", [16, 32])
def test_fused_metadata_copy_large_batch(bs):
"""Test with larger batch sizes."""
if not torch.cuda.is_available():
pytest.skip("CUDA not available")
from sglang.kernels.ops.attention.fused_metadata_copy import (
fused_metadata_copy_cuda,
)
forward_mode = 0 # DECODE
max_len = 128
max_seqlen_k = 256
seqlens_expanded_size = bs
data = create_test_metadata(
bs=bs,
max_len=max_len,
max_seqlen_k=max_seqlen_k,
seqlens_expanded_size=seqlens_expanded_size,
has_real_page_table=True,
has_flashmla=True,
)
dst_ref = {k: v.clone() if v is not None else None for k, v in data["dst"].items()}
dst_fused = {
k: v.clone() if v is not None else None for k, v in data["dst"].items()
}
reference_copy_decode(data["src"], dst_ref, max_len)
fused_metadata_copy_cuda(
data["src"]["cache_seqlens"],
data["src"]["cu_seqlens_k"],
data["src"]["page_indices"],
data["src"]["dsa_cache_seqlens"],
data["src"]["seqlens_expanded"],
data["src"]["dsa_cu_seqlens_k"],
data["src"]["real_page_table"],
data["src"]["flashmla_num_splits"],
data["src"]["flashmla_metadata"],
dst_fused["cache_seqlens"],
dst_fused["cu_seqlens_k"],
dst_fused["page_table_1"],
dst_fused["dsa_cache_seqlens"],
dst_fused["dsa_seqlens_expanded"],
dst_fused["dsa_cu_seqlens_k"],
dst_fused["real_page_table"],
dst_fused["flashmla_num_splits"],
dst_fused["flashmla_metadata"],
forward_mode,
bs,
max_len,
max_seqlen_k,
seqlens_expanded_size,
)
# Verify all tensors match
for key in dst_ref:
if dst_ref[key] is not None:
assert torch.equal(dst_ref[key], dst_fused[key]), f"{key} mismatch"
for key, expected in dst_ref.items():
if expected is not None:
torch.testing.assert_close(
dst_fused[key], expected, rtol=0, atol=0, msg=key
)
# =============================================================================
@@ -725,9 +565,16 @@ def test_fused_metadata_copy_multi_dtype_validation():
)
@pytest.mark.parametrize("bs", [1, 2, 4, 8, 16])
@pytest.mark.parametrize("has_real_page_table", [False, True])
@pytest.mark.parametrize("has_flashmla", [False, True])
@pytest.mark.parametrize(
"bs,has_real_page_table,has_flashmla",
[
(bs, page, mla)
for bs in (1, 2, 4, 8, 16)
for page in (False, True)
for mla in (False, True)
]
+ [(32, True, True), (64, True, True)],
)
def test_fused_metadata_copy_multi(bs, has_real_page_table, has_flashmla):
"""Test fused multi-backend metadata copy kernel against for-loop version."""
if not torch.cuda.is_available():
@@ -749,38 +596,16 @@ def test_fused_metadata_copy_multi(bs, has_real_page_table, has_flashmla):
has_flashmla=has_flashmla,
)
# Create separate destination tensors for reference (for-loop) and fused kernel
dst_ref_0 = {
k: v.clone() if v is not None else None for k, v in data["dst0"].items()
}
dst_ref_1 = {
k: v.clone() if v is not None else None for k, v in data["dst1"].items()
}
dst_ref_2 = {
k: v.clone() if v is not None else None for k, v in data["dst2"].items()
}
dst_ref = [
{k: v.clone() if v is not None else None for k, v in data[f"dst{i}"].items()}
for i in range(3)
]
dst_fused = [
{k: v.clone() if v is not None else None for k, v in data[f"dst{i}"].items()}
for i in range(3)
]
reference_copy_for_loop(data["src"], dst_ref, bs, max_len)
dst_fused_0 = {
k: v.clone() if v is not None else None for k, v in data["dst0"].items()
}
dst_fused_1 = {
k: v.clone() if v is not None else None for k, v in data["dst1"].items()
}
dst_fused_2 = {
k: v.clone() if v is not None else None for k, v in data["dst2"].items()
}
# Run reference implementation (for-loop)
torch.cuda.synchronize()
loop_start = time.perf_counter()
reference_copy_for_loop(data["src"], [dst_ref_0, dst_ref_1, dst_ref_2], bs, max_len)
torch.cuda.synchronize()
loop_end = time.perf_counter()
loop_time = loop_end - loop_start
# Run fused kernel
torch.cuda.synchronize()
fused_start = time.perf_counter()
fused_metadata_copy_multi_cuda(
# Source tensors
data["src"]["cache_seqlens"],
@@ -792,296 +617,46 @@ def test_fused_metadata_copy_multi(bs, has_real_page_table, has_flashmla):
data["src"]["flashmla_num_splits"],
data["src"]["flashmla_metadata"],
# Destination tensors for backend 0
dst_fused_0["cache_seqlens_int32"],
dst_fused_0["cu_seqlens_k"],
dst_fused_0["page_table_1"],
dst_fused_0["dsa_cache_seqlens_int32"],
dst_fused_0["dsa_cu_seqlens_k"],
dst_fused_0["real_page_table"],
dst_fused_0["flashmla_num_splits"],
dst_fused_0["flashmla_metadata"],
dst_fused[0]["cache_seqlens_int32"],
dst_fused[0]["cu_seqlens_k"],
dst_fused[0]["page_table_1"],
dst_fused[0]["dsa_cache_seqlens_int32"],
dst_fused[0]["dsa_cu_seqlens_k"],
dst_fused[0]["real_page_table"],
dst_fused[0]["flashmla_num_splits"],
dst_fused[0]["flashmla_metadata"],
# Destination tensors for backend 1
dst_fused_1["cache_seqlens_int32"],
dst_fused_1["cu_seqlens_k"],
dst_fused_1["page_table_1"],
dst_fused_1["dsa_cache_seqlens_int32"],
dst_fused_1["dsa_cu_seqlens_k"],
dst_fused_1["real_page_table"],
dst_fused_1["flashmla_num_splits"],
dst_fused_1["flashmla_metadata"],
dst_fused[1]["cache_seqlens_int32"],
dst_fused[1]["cu_seqlens_k"],
dst_fused[1]["page_table_1"],
dst_fused[1]["dsa_cache_seqlens_int32"],
dst_fused[1]["dsa_cu_seqlens_k"],
dst_fused[1]["real_page_table"],
dst_fused[1]["flashmla_num_splits"],
dst_fused[1]["flashmla_metadata"],
# Destination tensors for backend 2
dst_fused_2["cache_seqlens_int32"],
dst_fused_2["cu_seqlens_k"],
dst_fused_2["page_table_1"],
dst_fused_2["dsa_cache_seqlens_int32"],
dst_fused_2["dsa_cu_seqlens_k"],
dst_fused_2["real_page_table"],
dst_fused_2["flashmla_num_splits"],
dst_fused_2["flashmla_metadata"],
dst_fused[2]["cache_seqlens_int32"],
dst_fused[2]["cu_seqlens_k"],
dst_fused[2]["page_table_1"],
dst_fused[2]["dsa_cache_seqlens_int32"],
dst_fused[2]["dsa_cu_seqlens_k"],
dst_fused[2]["real_page_table"],
dst_fused[2]["flashmla_num_splits"],
dst_fused[2]["flashmla_metadata"],
# Parameters
bs,
max_len,
seqlens_expanded_size,
)
torch.cuda.synchronize()
fused_end = time.perf_counter()
fused_time = fused_end - fused_start
# Compare results for all 3 backends
speedup = loop_time / fused_time if fused_time > 0 else 0
print(
f"\n[VERIFY] bs={bs}, real_page_table={has_real_page_table}, flashmla={has_flashmla}"
)
print(
f"[VERIFY] Fused time: {fused_time * 1000:.3f}ms, Loop time: {loop_time * 1000:.3f}ms, Speedup: {speedup:.2f}x"
)
max_diff = 0.0
all_match = True
for backend_idx, (dst_ref, dst_fused) in enumerate(
[
(dst_ref_0, dst_fused_0),
(dst_ref_1, dst_fused_1),
(dst_ref_2, dst_fused_2),
]
):
for key in [
"cache_seqlens_int32",
"cu_seqlens_k",
"page_table_1",
"dsa_cache_seqlens_int32",
"dsa_cu_seqlens_k",
]:
if not torch.equal(dst_ref[key], dst_fused[key]):
diff = (
(dst_ref[key].float() - dst_fused[key].float()).abs().max().item()
)
max_diff = max(max_diff, diff)
all_match = False
print(
f"[ERROR] Backend {backend_idx} {key}: MISMATCH! Max diff: {diff}"
)
if has_real_page_table and dst_ref["real_page_table"] is not None:
if not torch.equal(
dst_ref["real_page_table"], dst_fused["real_page_table"]
):
diff = (
(
dst_ref["real_page_table"].float()
- dst_fused["real_page_table"].float()
)
.abs()
.max()
.item()
)
max_diff = max(max_diff, diff)
all_match = False
print(
f"[ERROR] Backend {backend_idx} real_page_table: MISMATCH! Max diff: {diff}"
)
if has_flashmla:
if dst_ref["flashmla_num_splits"] is not None and not torch.equal(
dst_ref["flashmla_num_splits"], dst_fused["flashmla_num_splits"]
):
diff = (
(
dst_ref["flashmla_num_splits"].float()
- dst_fused["flashmla_num_splits"].float()
)
.abs()
.max()
.item()
)
max_diff = max(max_diff, diff)
all_match = False
print(
f"[ERROR] Backend {backend_idx} flashmla_num_splits: MISMATCH! Max diff: {diff}"
)
if dst_ref["flashmla_metadata"] is not None and not torch.equal(
dst_ref["flashmla_metadata"], dst_fused["flashmla_metadata"]
):
diff = (
(
dst_ref["flashmla_metadata"].float()
- dst_fused["flashmla_metadata"].float()
)
.abs()
.max()
.item()
)
max_diff = max(max_diff, diff)
all_match = False
print(
f"[ERROR] Backend {backend_idx} flashmla_metadata: MISMATCH! Max diff: {diff}"
)
if not all_match:
error_msg = (
f"Fused metadata copy verification FAILED! "
f"Maximum difference: {max_diff}. "
f"The fused kernel produces different results than the for-loop version."
)
print(f"[ERROR] {error_msg}")
raise AssertionError(error_msg)
print(f"[VERIFY] Verification PASSED - all tensors match!")
@pytest.mark.parametrize("bs", [32, 64])
def test_fused_metadata_copy_multi_large_batch(bs):
"""Test with larger batch sizes and timing comparison."""
if not torch.cuda.is_available():
pytest.skip("CUDA not available")
from sglang.kernels.ops.attention.fused_metadata_copy import (
fused_metadata_copy_multi_cuda,
)
max_len = 128
seqlens_expanded_size = bs
data = create_test_metadata_multi(
bs=bs,
max_len=max_len,
seqlens_expanded_size=seqlens_expanded_size,
has_real_page_table=True,
has_flashmla=True,
)
dst_ref_0 = {
k: v.clone() if v is not None else None for k, v in data["dst0"].items()
}
dst_ref_1 = {
k: v.clone() if v is not None else None for k, v in data["dst1"].items()
}
dst_ref_2 = {
k: v.clone() if v is not None else None for k, v in data["dst2"].items()
}
dst_fused_0 = {
k: v.clone() if v is not None else None for k, v in data["dst0"].items()
}
dst_fused_1 = {
k: v.clone() if v is not None else None for k, v in data["dst1"].items()
}
dst_fused_2 = {
k: v.clone() if v is not None else None for k, v in data["dst2"].items()
}
# Warmup
for _ in range(5):
reference_copy_for_loop(
data["src"], [dst_ref_0, dst_ref_1, dst_ref_2], bs, max_len
)
fused_metadata_copy_multi_cuda(
data["src"]["cache_seqlens"],
data["src"]["cu_seqlens_k"],
data["src"]["page_indices"],
data["src"]["dsa_cache_seqlens"],
data["src"]["dsa_cu_seqlens_k"],
data["src"]["real_page_table"],
data["src"]["flashmla_num_splits"],
data["src"]["flashmla_metadata"],
dst_fused_0["cache_seqlens_int32"],
dst_fused_0["cu_seqlens_k"],
dst_fused_0["page_table_1"],
dst_fused_0["dsa_cache_seqlens_int32"],
dst_fused_0["dsa_cu_seqlens_k"],
dst_fused_0["real_page_table"],
dst_fused_0["flashmla_num_splits"],
dst_fused_0["flashmla_metadata"],
dst_fused_1["cache_seqlens_int32"],
dst_fused_1["cu_seqlens_k"],
dst_fused_1["page_table_1"],
dst_fused_1["dsa_cache_seqlens_int32"],
dst_fused_1["dsa_cu_seqlens_k"],
dst_fused_1["real_page_table"],
dst_fused_1["flashmla_num_splits"],
dst_fused_1["flashmla_metadata"],
dst_fused_2["cache_seqlens_int32"],
dst_fused_2["cu_seqlens_k"],
dst_fused_2["page_table_1"],
dst_fused_2["dsa_cache_seqlens_int32"],
dst_fused_2["dsa_cu_seqlens_k"],
dst_fused_2["real_page_table"],
dst_fused_2["flashmla_num_splits"],
dst_fused_2["flashmla_metadata"],
bs,
max_len,
seqlens_expanded_size,
)
torch.cuda.synchronize()
# Actual timing
torch.cuda.synchronize()
loop_start = time.perf_counter()
reference_copy_for_loop(data["src"], [dst_ref_0, dst_ref_1, dst_ref_2], bs, max_len)
torch.cuda.synchronize()
loop_time = time.perf_counter() - loop_start
torch.cuda.synchronize()
fused_start = time.perf_counter()
fused_metadata_copy_multi_cuda(
data["src"]["cache_seqlens"],
data["src"]["cu_seqlens_k"],
data["src"]["page_indices"],
data["src"]["dsa_cache_seqlens"],
data["src"]["dsa_cu_seqlens_k"],
data["src"]["real_page_table"],
data["src"]["flashmla_num_splits"],
data["src"]["flashmla_metadata"],
dst_fused_0["cache_seqlens_int32"],
dst_fused_0["cu_seqlens_k"],
dst_fused_0["page_table_1"],
dst_fused_0["dsa_cache_seqlens_int32"],
dst_fused_0["dsa_cu_seqlens_k"],
dst_fused_0["real_page_table"],
dst_fused_0["flashmla_num_splits"],
dst_fused_0["flashmla_metadata"],
dst_fused_1["cache_seqlens_int32"],
dst_fused_1["cu_seqlens_k"],
dst_fused_1["page_table_1"],
dst_fused_1["dsa_cache_seqlens_int32"],
dst_fused_1["dsa_cu_seqlens_k"],
dst_fused_1["real_page_table"],
dst_fused_1["flashmla_num_splits"],
dst_fused_1["flashmla_metadata"],
dst_fused_2["cache_seqlens_int32"],
dst_fused_2["cu_seqlens_k"],
dst_fused_2["page_table_1"],
dst_fused_2["dsa_cache_seqlens_int32"],
dst_fused_2["dsa_cu_seqlens_k"],
dst_fused_2["real_page_table"],
dst_fused_2["flashmla_num_splits"],
dst_fused_2["flashmla_metadata"],
bs,
max_len,
seqlens_expanded_size,
)
torch.cuda.synchronize()
fused_time = time.perf_counter() - fused_start
speedup = loop_time / fused_time if fused_time > 0 else 0
print(
f"\n[PERF] Large batch (bs={bs}): Fused={fused_time * 1000:.3f}ms, Loop={loop_time * 1000:.3f}ms, Speedup={speedup:.2f}x"
)
# Verify correctness
for backend_idx, (dst_ref, dst_fused) in enumerate(
[
(dst_ref_0, dst_fused_0),
(dst_ref_1, dst_fused_1),
(dst_ref_2, dst_fused_2),
]
):
for key in dst_ref:
if dst_ref[key] is not None and dst_fused[key] is not None:
assert torch.equal(dst_ref[key], dst_fused[key]), (
f"Backend {backend_idx} {key} mismatch"
for backend_idx, (expected, actual) in enumerate(zip(dst_ref, dst_fused)):
for key, tensor in expected.items():
if tensor is not None:
torch.testing.assert_close(
actual[key],
tensor,
rtol=0,
atol=0,
msg=f"Backend {backend_idx} {key}",
)
test/registered/kernels/ops/diffusion/test_flux2_token_cat_fp8.py
+5 -10
View File
@@ -36,19 +36,14 @@ def test_flux2_token_cat_fp8_is_bit_exact(tokens: int) -> None:
assert torch.equal(actual, expected)
def test_flux2_token_cat_fp8_rejects_compile() -> None:
@pytest.mark.parametrize(
"guard", ["torch.compiler.is_compiling", "torch.cuda.is_current_stream_capturing"]
)
def test_flux2_token_cat_fp8_rejects_capture(guard) -> None:
attention = torch.empty((1, 1, 16), device="cuda", dtype=torch.bfloat16)
mlp = torch.empty((1, 1, 48), device="cuda", dtype=torch.bfloat16)
scale = torch.ones((1,), device="cuda", dtype=torch.float32)
with patch("torch.compiler.is_compiling", return_value=True):
assert try_flux2_token_cat_fp8(attention, mlp, scale) is None
def test_flux2_token_cat_fp8_rejects_cuda_graph_capture() -> None:
attention = torch.empty((1, 1, 16), device="cuda", dtype=torch.bfloat16)
mlp = torch.empty((1, 1, 48), device="cuda", dtype=torch.bfloat16)
scale = torch.ones((1,), device="cuda", dtype=torch.float32)
with patch("torch.cuda.is_current_stream_capturing", return_value=True):
with patch(guard, return_value=True):
assert try_flux2_token_cat_fp8(attention, mlp, scale) is None
test/registered/kernels/ops/diffusion/test_modulate.py
+23 -47
View File
@@ -146,26 +146,6 @@ GATE_CASES = [
]
def _assert_gate_add(out, ref):
if ref.dtype == torch.float32:
# fp32 has no rounding boundary to reproduce; the kernel keeps the
# accumulation in fp32 and only order may differ.
torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
else:
assert torch.equal(out, ref)
@pytest.mark.parametrize("residual_shape,gate_shape", GATE_CASES)
def test_residual_gate_add_matches_torch(residual_shape, gate_shape):
residual = torch.randn(residual_shape, device=DEVICE, dtype=torch.bfloat16)
update = torch.randn_like(residual)
gate = torch.randn(gate_shape, device=DEVICE, dtype=torch.bfloat16)
ref = residual + update * gate
_assert_gate_add(residual_gate_add_cuda(residual, update, gate), ref)
assert torch.equal(residual_gate_add(residual, update, gate), ref)
# LingBot per-token gates are [B, S, 1]: one scalar per token, broadcast
# along the hidden dimension.
PER_TOKEN_GATE_CASES = [
@@ -175,8 +155,17 @@ PER_TOKEN_GATE_CASES = [
]
@pytest.mark.parametrize("residual_shape,gate_shape", PER_TOKEN_GATE_CASES)
def test_residual_gate_add_per_token_matches_torch(residual_shape, gate_shape):
def _assert_gate_add(out, ref):
if ref.dtype == torch.float32:
# fp32 has no rounding boundary to reproduce; the kernel keeps the
# accumulation in fp32 and only order may differ.
torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
else:
assert torch.equal(out, ref)
@pytest.mark.parametrize("residual_shape,gate_shape", GATE_CASES + PER_TOKEN_GATE_CASES)
def test_residual_gate_add_matches_torch(residual_shape, gate_shape):
residual = torch.randn(residual_shape, device=DEVICE, dtype=torch.bfloat16)
update = torch.randn_like(residual)
gate = torch.randn(gate_shape, device=DEVICE, dtype=torch.bfloat16)
@@ -188,19 +177,12 @@ def test_residual_gate_add_per_token_matches_torch(residual_shape, gate_shape):
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
def test_residual_gate_add_per_token_dtypes(dtype):
residual = torch.randn((1, 2560, 512), device=DEVICE, dtype=dtype)
update = torch.randn_like(residual)
gate = torch.randn((1, 2560, 1), device=DEVICE, dtype=dtype)
_assert_gate_add(
residual_gate_add_cuda(residual, update, gate), residual + update * gate
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
@pytest.mark.parametrize("gate_shape", [(1, 1, 64), (1, 9, 64)])
def test_residual_gate_add_dtypes(dtype, gate_shape):
residual = torch.randn((1, 9, 64), device=DEVICE, dtype=dtype)
@pytest.mark.parametrize(
"shape,gate_shape",
[((1, 9, 64), (1, 1, 64)), ((1, 9, 64), (1, 9, 64)), PER_TOKEN_GATE_CASES[0]],
)
def test_residual_gate_add_dtypes(dtype, shape, gate_shape):
residual = torch.randn(shape, device=DEVICE, dtype=dtype)
update = torch.randn_like(residual)
gate = torch.randn(gate_shape, device=DEVICE, dtype=dtype)
_assert_gate_add(
@@ -248,10 +230,12 @@ def test_residual_gate_add_transposed_storage_offsets():
assert torch.equal(out, residual + update * gate)
def test_residual_gate_add_transposed_torch_compile_fullgraph():
residual = torch.randn((1, 128, 32), device=DEVICE, dtype=torch.bfloat16).transpose(
1, 2
)
@pytest.mark.parametrize("transposed", [False, True])
def test_residual_gate_add_torch_compile_fullgraph(transposed):
shape = (1, 128, 32) if transposed else (1, 32, 128)
residual = torch.randn(shape, device=DEVICE, dtype=torch.bfloat16)
if transposed:
residual = residual.transpose(1, 2)
update = torch.randn_like(residual, memory_format=torch.contiguous_format)
gate = torch.randn((1, 1, 128), device=DEVICE, dtype=torch.bfloat16)
compiled = torch.compile(residual_gate_add, fullgraph=True)
@@ -309,14 +293,6 @@ def test_residual_gate_add_guards_and_eager_fallback():
)
def test_residual_gate_add_torch_compile_fullgraph():
residual = torch.randn((1, 32, 128), device=DEVICE, dtype=torch.bfloat16)
update = torch.randn_like(residual)
gate = torch.randn((1, 1, 128), device=DEVICE, dtype=torch.bfloat16)
compiled = torch.compile(residual_gate_add, fullgraph=True)
assert torch.equal(compiled(residual, update, gate), residual + update * gate)
@torch.no_grad()
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_scaled_residual_add_is_bit_exact(dtype):
test/registered/kernels/ops/quantization/test_awq_marlin_moe_repack.py
+2 -23
View File
@@ -1,6 +1,5 @@
import sys
import numpy as np
import pytest
import torch
from sgl_kernel.scalar_type import scalar_types
@@ -8,8 +7,9 @@ from sgl_kernel.scalar_type import scalar_types
from sglang.kernels.ops.quantization.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.srt.layers.quantization.utils import quantize_weights
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_marlin_utils import awq_pack
register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
@@ -23,27 +23,6 @@ def _has_aot_awq_marlin_moe_repack() -> bool:
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)])
test/registered/kernels/ops/quantization/test_awq_marlin_repack.py
+2 -24
View File
@@ -1,6 +1,5 @@
import sys
import numpy as np
import pytest
import torch
from sgl_kernel.scalar_type import scalar_types
@@ -8,9 +7,9 @@ from sgl_kernel.scalar_type import scalar_types
from sglang.kernels.ops.quantization.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.srt.layers.quantization.utils import quantize_weights
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_marlin_utils import get_weight_perm, marlin_weights
from sglang.test.test_marlin_utils import awq_pack, get_weight_perm, marlin_weights
register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
@@ -24,27 +23,6 @@ def _has_aot_awq_marlin_repack() -> bool:
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])