1068 lines
38 KiB
Python
1068 lines
38 KiB
Python
"""
|
|
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
|
|
"""
|
|
|
|
import time
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
# =============================================================================
|
|
# Helper Functions
|
|
# =============================================================================
|
|
|
|
|
|
def create_test_metadata(
|
|
bs: int,
|
|
max_len: int,
|
|
max_seqlen_k: int,
|
|
seqlens_expanded_size: int,
|
|
has_real_page_table: bool = False,
|
|
has_flashmla: bool = False,
|
|
device: str = "cuda",
|
|
):
|
|
"""Create test metadata tensors matching NSA backend structure."""
|
|
# Basic tensors (always present)
|
|
cache_seqlens_src = torch.randint(
|
|
1, max_len, (bs,), dtype=torch.int32, device=device
|
|
)
|
|
cu_seqlens_k_src = torch.zeros(bs + 1, dtype=torch.int32, device=device)
|
|
cu_seqlens_k_src[1:] = torch.cumsum(cache_seqlens_src, dim=0)
|
|
|
|
page_indices_src = torch.randint(
|
|
0, 1000, (bs, max_len), dtype=torch.int32, device=device
|
|
)
|
|
nsa_cache_seqlens_src = torch.randint(
|
|
1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device
|
|
)
|
|
seqlens_expanded_src = torch.randint(
|
|
1, max_seqlen_k, (seqlens_expanded_size,), dtype=torch.int32, device=device
|
|
)
|
|
nsa_cu_seqlens_k_src = torch.zeros(
|
|
seqlens_expanded_size + 1, dtype=torch.int32, device=device
|
|
)
|
|
nsa_cu_seqlens_k_src[1:] = torch.cumsum(nsa_cache_seqlens_src, dim=0)
|
|
|
|
# Destination tensors
|
|
cache_seqlens_dst = torch.zeros(bs, dtype=torch.int32, device=device)
|
|
cu_seqlens_k_dst = torch.zeros(bs + 1, dtype=torch.int32, device=device)
|
|
page_table_1_dst = torch.zeros((bs, max_len + 16), dtype=torch.int32, device=device)
|
|
nsa_cache_seqlens_dst = torch.zeros(
|
|
seqlens_expanded_size, dtype=torch.int32, device=device
|
|
)
|
|
nsa_seqlens_expanded_dst = torch.zeros(
|
|
seqlens_expanded_size, dtype=torch.int32, device=device
|
|
)
|
|
nsa_cu_seqlens_k_dst = torch.zeros(
|
|
seqlens_expanded_size + 1, dtype=torch.int32, device=device
|
|
)
|
|
|
|
# Optional tensors
|
|
real_page_table_src = None
|
|
real_page_table_dst = None
|
|
if has_real_page_table:
|
|
real_page_table_cols = max_len // 2
|
|
real_page_table_src = torch.randint(
|
|
0, 1000, (bs, real_page_table_cols), dtype=torch.int32, device=device
|
|
)
|
|
real_page_table_dst = torch.zeros(
|
|
(bs, real_page_table_cols + 8), dtype=torch.int32, device=device
|
|
)
|
|
|
|
flashmla_num_splits_src = None
|
|
flashmla_num_splits_dst = None
|
|
flashmla_metadata_src = None
|
|
flashmla_metadata_dst = None
|
|
if has_flashmla:
|
|
flashmla_num_splits_src = torch.randint(
|
|
1, 10, (seqlens_expanded_size + 1,), dtype=torch.int32, device=device
|
|
)
|
|
flashmla_num_splits_dst = torch.zeros(
|
|
seqlens_expanded_size + 1, dtype=torch.int32, device=device
|
|
)
|
|
# FlashMLA metadata is typically (num_sm_parts, TileSchedulerMetaDataSize)
|
|
# For testing, we use a simplified size
|
|
flashmla_metadata_size = 128
|
|
flashmla_metadata_src = torch.randint(
|
|
0, 100, (flashmla_metadata_size,), dtype=torch.int32, device=device
|
|
)
|
|
flashmla_metadata_dst = torch.zeros(
|
|
flashmla_metadata_size, dtype=torch.int32, device=device
|
|
)
|
|
|
|
return {
|
|
"src": {
|
|
"cache_seqlens": cache_seqlens_src,
|
|
"cu_seqlens_k": cu_seqlens_k_src,
|
|
"page_indices": page_indices_src,
|
|
"nsa_cache_seqlens": nsa_cache_seqlens_src,
|
|
"seqlens_expanded": seqlens_expanded_src,
|
|
"nsa_cu_seqlens_k": nsa_cu_seqlens_k_src,
|
|
"real_page_table": real_page_table_src,
|
|
"flashmla_num_splits": flashmla_num_splits_src,
|
|
"flashmla_metadata": flashmla_metadata_src,
|
|
},
|
|
"dst": {
|
|
"cache_seqlens": cache_seqlens_dst,
|
|
"cu_seqlens_k": cu_seqlens_k_dst,
|
|
"page_table_1": page_table_1_dst,
|
|
"nsa_cache_seqlens": nsa_cache_seqlens_dst,
|
|
"nsa_seqlens_expanded": nsa_seqlens_expanded_dst,
|
|
"nsa_cu_seqlens_k": nsa_cu_seqlens_k_dst,
|
|
"real_page_table": real_page_table_dst,
|
|
"flashmla_num_splits": flashmla_num_splits_dst,
|
|
"flashmla_metadata": flashmla_metadata_dst,
|
|
},
|
|
}
|
|
|
|
|
|
def reference_copy_decode(src, dst, max_len):
|
|
"""Reference implementation: individual .copy_() for DECODE 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:])
|
|
dst["page_table_1"][:, :max_len].copy_(src["page_indices"])
|
|
dst["nsa_cache_seqlens"].copy_(src["nsa_cache_seqlens"])
|
|
dst["nsa_cu_seqlens_k"][1 : bs + 1].copy_(src["nsa_cu_seqlens_k"][1 : bs + 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 = bs + 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_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["nsa_seqlens_expanded"][:seqlens_expanded_size].copy_(src["seqlens_expanded"])
|
|
dst["nsa_cache_seqlens"][:seqlens_expanded_size].copy_(src["nsa_cache_seqlens"])
|
|
dst["nsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1].copy_(
|
|
src["nsa_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["nsa_seqlens_expanded"][:seqlens_expanded_size].copy_(src["seqlens_expanded"])
|
|
dst["nsa_cache_seqlens"][:seqlens_expanded_size].copy_(src["nsa_cache_seqlens"])
|
|
dst["nsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1].copy_(
|
|
src["nsa_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
|
|
# =============================================================================
|
|
|
|
|
|
def test_fused_metadata_copy_dtype_validation():
|
|
"""Test that dtype validation rejects non-int32 tensors."""
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA not available")
|
|
|
|
from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_cuda
|
|
|
|
bs = 2
|
|
max_len = 128
|
|
max_seqlen_k = 256
|
|
seqlens_expanded_size = bs
|
|
device = "cuda"
|
|
|
|
# Create tensors with WRONG dtype (int64 instead of int32)
|
|
cache_seqlens_src_wrong = torch.randint(
|
|
1, max_len, (bs,), dtype=torch.int64, device=device
|
|
)
|
|
cu_seqlens_k_src = torch.zeros(bs + 1, dtype=torch.int32, device=device)
|
|
page_indices_src = torch.randint(
|
|
0, 1000, (bs, max_len), dtype=torch.int32, device=device
|
|
)
|
|
nsa_cache_seqlens_src = torch.randint(
|
|
1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device
|
|
)
|
|
seqlens_expanded_src = torch.randint(
|
|
1, max_seqlen_k, (seqlens_expanded_size,), dtype=torch.int32, device=device
|
|
)
|
|
nsa_cu_seqlens_k_src = torch.zeros(
|
|
seqlens_expanded_size + 1, dtype=torch.int32, device=device
|
|
)
|
|
|
|
# Destination tensors (correct dtype)
|
|
cache_seqlens_dst = torch.zeros(bs, dtype=torch.int32, device=device)
|
|
cu_seqlens_k_dst = torch.zeros(bs + 1, dtype=torch.int32, device=device)
|
|
page_table_1_dst = torch.zeros((bs, max_len + 16), dtype=torch.int32, device=device)
|
|
nsa_cache_seqlens_dst = torch.zeros(
|
|
seqlens_expanded_size, dtype=torch.int32, device=device
|
|
)
|
|
nsa_seqlens_expanded_dst = torch.zeros(
|
|
seqlens_expanded_size, dtype=torch.int32, device=device
|
|
)
|
|
nsa_cu_seqlens_k_dst = torch.zeros(
|
|
seqlens_expanded_size + 1, dtype=torch.int32, device=device
|
|
)
|
|
|
|
# Test 1: Wrong dtype for source tensor should raise RuntimeError
|
|
with pytest.raises(RuntimeError, match="must have dtype int32"):
|
|
fused_metadata_copy_cuda(
|
|
cache_seqlens_src_wrong, # Wrong dtype: int64
|
|
cu_seqlens_k_src,
|
|
page_indices_src,
|
|
nsa_cache_seqlens_src,
|
|
seqlens_expanded_src,
|
|
nsa_cu_seqlens_k_src,
|
|
None, # real_page_table_src
|
|
None, # flashmla_num_splits_src
|
|
None, # flashmla_metadata_src
|
|
cache_seqlens_dst,
|
|
cu_seqlens_k_dst,
|
|
page_table_1_dst,
|
|
nsa_cache_seqlens_dst,
|
|
nsa_seqlens_expanded_dst,
|
|
nsa_cu_seqlens_k_dst,
|
|
None, # real_page_table_dst
|
|
None, # flashmla_num_splits_dst
|
|
None, # flashmla_metadata_dst
|
|
0, # forward_mode
|
|
bs,
|
|
max_len,
|
|
max_seqlen_k,
|
|
seqlens_expanded_size,
|
|
)
|
|
|
|
# Test 2: Wrong dtype for destination tensor should also raise RuntimeError
|
|
cache_seqlens_src = torch.randint(
|
|
1, max_len, (bs,), dtype=torch.int32, device=device
|
|
)
|
|
cache_seqlens_dst_wrong = torch.zeros(bs, dtype=torch.int64, device=device)
|
|
|
|
with pytest.raises(RuntimeError, match="must have dtype int32"):
|
|
fused_metadata_copy_cuda(
|
|
cache_seqlens_src,
|
|
cu_seqlens_k_src,
|
|
page_indices_src,
|
|
nsa_cache_seqlens_src,
|
|
seqlens_expanded_src,
|
|
nsa_cu_seqlens_k_src,
|
|
None,
|
|
None,
|
|
None,
|
|
cache_seqlens_dst_wrong, # Wrong dtype: int64
|
|
cu_seqlens_k_dst,
|
|
page_table_1_dst,
|
|
nsa_cache_seqlens_dst,
|
|
nsa_seqlens_expanded_dst,
|
|
nsa_cu_seqlens_k_dst,
|
|
None,
|
|
None,
|
|
None,
|
|
0,
|
|
bs,
|
|
max_len,
|
|
max_seqlen_k,
|
|
seqlens_expanded_size,
|
|
)
|
|
|
|
|
|
@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):
|
|
"""Test fused metadata copy kernel against reference implementation."""
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA not available")
|
|
|
|
from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_cuda
|
|
|
|
max_len = 128
|
|
max_seqlen_k = 256
|
|
seqlens_expanded_size = bs if forward_mode == 0 else bs * 2
|
|
|
|
# Create test data
|
|
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=has_real_page_table,
|
|
has_flashmla=has_flashmla,
|
|
)
|
|
|
|
# Create separate destination tensors for reference and fused kernel
|
|
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()
|
|
}
|
|
|
|
# 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
|
|
)
|
|
|
|
# Run fused kernel
|
|
fused_metadata_copy_cuda(
|
|
data["src"]["cache_seqlens"],
|
|
data["src"]["cu_seqlens_k"],
|
|
data["src"]["page_indices"],
|
|
data["src"]["nsa_cache_seqlens"],
|
|
data["src"]["seqlens_expanded"],
|
|
data["src"]["nsa_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["nsa_cache_seqlens"],
|
|
dst_fused["nsa_seqlens_expanded"],
|
|
dst_fused["nsa_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,
|
|
)
|
|
|
|
# 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["nsa_cache_seqlens"], dst_fused["nsa_cache_seqlens"]
|
|
), "nsa_cache_seqlens mismatch"
|
|
assert torch.equal(
|
|
dst_ref["nsa_seqlens_expanded"], dst_fused["nsa_seqlens_expanded"]
|
|
), "nsa_seqlens_expanded mismatch"
|
|
assert torch.equal(
|
|
dst_ref["nsa_cu_seqlens_k"], dst_fused["nsa_cu_seqlens_k"]
|
|
), "nsa_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.jit_kernel.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"]["nsa_cache_seqlens"],
|
|
data["src"]["seqlens_expanded"],
|
|
data["src"]["nsa_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["nsa_cache_seqlens"],
|
|
dst_fused["nsa_seqlens_expanded"],
|
|
dst_fused["nsa_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"
|
|
|
|
|
|
# =============================================================================
|
|
# Multi-Backend Kernel Tests
|
|
# =============================================================================
|
|
|
|
|
|
def create_test_metadata_multi(
|
|
bs: int,
|
|
max_len: int,
|
|
seqlens_expanded_size: int,
|
|
has_real_page_table: bool = False,
|
|
has_flashmla: bool = False,
|
|
device: str = "cuda",
|
|
):
|
|
"""Create test metadata tensors for multi-backend testing."""
|
|
# Source tensors (precomputed metadata)
|
|
cache_seqlens_src = torch.randint(
|
|
1, max_len, (bs,), dtype=torch.int32, device=device
|
|
)
|
|
cu_seqlens_k_src = torch.zeros(bs + 1, dtype=torch.int32, device=device)
|
|
cu_seqlens_k_src[1:] = torch.cumsum(cache_seqlens_src, dim=0)
|
|
|
|
page_indices_src = torch.randint(
|
|
0, 1000, (bs, max_len), dtype=torch.int32, device=device
|
|
)
|
|
nsa_cache_seqlens_src = torch.randint(
|
|
1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device
|
|
)
|
|
nsa_cu_seqlens_k_src = torch.zeros(
|
|
seqlens_expanded_size + 1, dtype=torch.int32, device=device
|
|
)
|
|
nsa_cu_seqlens_k_src[1:] = torch.cumsum(nsa_cache_seqlens_src, dim=0)
|
|
|
|
# Optional tensors
|
|
real_page_table_src = None
|
|
if has_real_page_table:
|
|
real_page_table_cols = max_len // 2
|
|
real_page_table_src = torch.randint(
|
|
0, 1000, (bs, real_page_table_cols), dtype=torch.int32, device=device
|
|
)
|
|
|
|
flashmla_num_splits_src = None
|
|
flashmla_metadata_src = None
|
|
if has_flashmla:
|
|
flashmla_num_splits_src = torch.randint(
|
|
1, 10, (seqlens_expanded_size + 1,), dtype=torch.int32, device=device
|
|
)
|
|
flashmla_metadata_size = 128
|
|
flashmla_metadata_src = torch.randint(
|
|
0, 100, (flashmla_metadata_size,), dtype=torch.int32, device=device
|
|
)
|
|
|
|
# Create destination tensors for 3 backends
|
|
def create_dst_tensors():
|
|
cache_seqlens_dst = torch.zeros(bs, dtype=torch.int32, device=device)
|
|
cu_seqlens_k_dst = torch.zeros(bs + 1, dtype=torch.int32, device=device)
|
|
page_table_1_dst = torch.zeros(
|
|
(bs, max_len + 16), dtype=torch.int32, device=device
|
|
)
|
|
nsa_cache_seqlens_dst = torch.zeros(
|
|
seqlens_expanded_size, dtype=torch.int32, device=device
|
|
)
|
|
nsa_cu_seqlens_k_dst = torch.zeros(
|
|
seqlens_expanded_size + 1, dtype=torch.int32, device=device
|
|
)
|
|
|
|
real_page_table_dst = None
|
|
if has_real_page_table:
|
|
real_page_table_cols = max_len // 2
|
|
real_page_table_dst = torch.zeros(
|
|
(bs, real_page_table_cols + 8), dtype=torch.int32, device=device
|
|
)
|
|
|
|
flashmla_num_splits_dst = None
|
|
flashmla_metadata_dst = None
|
|
if has_flashmla:
|
|
flashmla_num_splits_dst = torch.zeros(
|
|
seqlens_expanded_size + 1, dtype=torch.int32, device=device
|
|
)
|
|
flashmla_metadata_size = 128
|
|
flashmla_metadata_dst = torch.zeros(
|
|
flashmla_metadata_size, dtype=torch.int32, device=device
|
|
)
|
|
|
|
return {
|
|
"cache_seqlens_int32": cache_seqlens_dst,
|
|
"cu_seqlens_k": cu_seqlens_k_dst,
|
|
"page_table_1": page_table_1_dst,
|
|
"nsa_cache_seqlens_int32": nsa_cache_seqlens_dst,
|
|
"nsa_cu_seqlens_k": nsa_cu_seqlens_k_dst,
|
|
"real_page_table": real_page_table_dst,
|
|
"flashmla_num_splits": flashmla_num_splits_dst,
|
|
"flashmla_metadata": flashmla_metadata_dst,
|
|
}
|
|
|
|
return {
|
|
"src": {
|
|
"cache_seqlens": cache_seqlens_src,
|
|
"cu_seqlens_k": cu_seqlens_k_src,
|
|
"page_indices": page_indices_src,
|
|
"nsa_cache_seqlens": nsa_cache_seqlens_src,
|
|
"nsa_cu_seqlens_k": nsa_cu_seqlens_k_src,
|
|
"real_page_table": real_page_table_src,
|
|
"flashmla_num_splits": flashmla_num_splits_src,
|
|
"flashmla_metadata": flashmla_metadata_src,
|
|
},
|
|
"dst0": create_dst_tensors(),
|
|
"dst1": create_dst_tensors(),
|
|
"dst2": create_dst_tensors(),
|
|
}
|
|
|
|
|
|
def reference_copy_for_loop(src, dst_list, bs, max_len):
|
|
"""Reference implementation: for-loop calling copy for each backend."""
|
|
for dst in dst_list:
|
|
# Simulate what init_forward_metadata_replay_cuda_graph_from_precomputed does
|
|
dst["cache_seqlens_int32"].copy_(src["cache_seqlens"])
|
|
dst["cu_seqlens_k"][1:].copy_(src["cu_seqlens_k"][1:])
|
|
dst["page_table_1"][:, :max_len].copy_(src["page_indices"])
|
|
dst["nsa_cache_seqlens_int32"].copy_(src["nsa_cache_seqlens"])
|
|
dst["nsa_cu_seqlens_k"][1 : bs + 1].copy_(src["nsa_cu_seqlens_k"][1 : bs + 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 = bs + 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 test_fused_metadata_copy_multi_dtype_validation():
|
|
"""Test that dtype validation rejects non-int32 tensors for multi-backend kernel."""
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA not available")
|
|
|
|
from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_multi_cuda
|
|
|
|
bs = 2
|
|
max_len = 128
|
|
seqlens_expanded_size = bs
|
|
device = "cuda"
|
|
|
|
# Create source tensors - one with WRONG dtype
|
|
cache_seqlens_src_wrong = torch.randint(
|
|
1, max_len, (bs,), dtype=torch.int64, device=device # Wrong dtype!
|
|
)
|
|
cu_seqlens_k_src = torch.zeros(bs + 1, dtype=torch.int32, device=device)
|
|
page_indices_src = torch.randint(
|
|
0, 1000, (bs, max_len), dtype=torch.int32, device=device
|
|
)
|
|
nsa_cache_seqlens_src = torch.randint(
|
|
1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device
|
|
)
|
|
nsa_cu_seqlens_k_src = torch.zeros(
|
|
seqlens_expanded_size + 1, dtype=torch.int32, device=device
|
|
)
|
|
|
|
# Create destination tensors for 3 backends (all correct dtype)
|
|
def create_dst():
|
|
return {
|
|
"cache_seqlens": torch.zeros(bs, dtype=torch.int32, device=device),
|
|
"cu_seqlens_k": torch.zeros(bs + 1, dtype=torch.int32, device=device),
|
|
"page_table_1": torch.zeros(
|
|
(bs, max_len + 16), dtype=torch.int32, device=device
|
|
),
|
|
"nsa_cache_seqlens": torch.zeros(
|
|
seqlens_expanded_size, dtype=torch.int32, device=device
|
|
),
|
|
"nsa_cu_seqlens_k": torch.zeros(
|
|
seqlens_expanded_size + 1, dtype=torch.int32, device=device
|
|
),
|
|
}
|
|
|
|
dst0 = create_dst()
|
|
dst1 = create_dst()
|
|
dst2 = create_dst()
|
|
|
|
# Test: Wrong dtype for source tensor should raise RuntimeError
|
|
with pytest.raises(RuntimeError, match="must have dtype int32"):
|
|
fused_metadata_copy_multi_cuda(
|
|
cache_seqlens_src_wrong, # Wrong dtype: int64
|
|
cu_seqlens_k_src,
|
|
page_indices_src,
|
|
nsa_cache_seqlens_src,
|
|
nsa_cu_seqlens_k_src,
|
|
None, # real_page_table_src
|
|
None, # flashmla_num_splits_src
|
|
None, # flashmla_metadata_src
|
|
# Backend 0
|
|
dst0["cache_seqlens"],
|
|
dst0["cu_seqlens_k"],
|
|
dst0["page_table_1"],
|
|
dst0["nsa_cache_seqlens"],
|
|
dst0["nsa_cu_seqlens_k"],
|
|
None,
|
|
None,
|
|
None,
|
|
# Backend 1
|
|
dst1["cache_seqlens"],
|
|
dst1["cu_seqlens_k"],
|
|
dst1["page_table_1"],
|
|
dst1["nsa_cache_seqlens"],
|
|
dst1["nsa_cu_seqlens_k"],
|
|
None,
|
|
None,
|
|
None,
|
|
# Backend 2
|
|
dst2["cache_seqlens"],
|
|
dst2["cu_seqlens_k"],
|
|
dst2["page_table_1"],
|
|
dst2["nsa_cache_seqlens"],
|
|
dst2["nsa_cu_seqlens_k"],
|
|
None,
|
|
None,
|
|
None,
|
|
# Parameters
|
|
bs,
|
|
max_len,
|
|
seqlens_expanded_size,
|
|
)
|
|
|
|
|
|
@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])
|
|
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():
|
|
pytest.skip("CUDA not available")
|
|
|
|
from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_multi_cuda
|
|
|
|
max_len = 128
|
|
seqlens_expanded_size = bs
|
|
|
|
# Create test data
|
|
data = create_test_metadata_multi(
|
|
bs=bs,
|
|
max_len=max_len,
|
|
seqlens_expanded_size=seqlens_expanded_size,
|
|
has_real_page_table=has_real_page_table,
|
|
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_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"],
|
|
data["src"]["cu_seqlens_k"],
|
|
data["src"]["page_indices"],
|
|
data["src"]["nsa_cache_seqlens"],
|
|
data["src"]["nsa_cu_seqlens_k"],
|
|
data["src"]["real_page_table"],
|
|
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["nsa_cache_seqlens_int32"],
|
|
dst_fused_0["nsa_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["nsa_cache_seqlens_int32"],
|
|
dst_fused_1["nsa_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["nsa_cache_seqlens_int32"],
|
|
dst_fused_2["nsa_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",
|
|
"nsa_cache_seqlens_int32",
|
|
"nsa_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.jit_kernel.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"]["nsa_cache_seqlens"],
|
|
data["src"]["nsa_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["nsa_cache_seqlens_int32"],
|
|
dst_fused_0["nsa_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["nsa_cache_seqlens_int32"],
|
|
dst_fused_1["nsa_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["nsa_cache_seqlens_int32"],
|
|
dst_fused_2["nsa_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"]["nsa_cache_seqlens"],
|
|
data["src"]["nsa_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["nsa_cache_seqlens_int32"],
|
|
dst_fused_0["nsa_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["nsa_cache_seqlens_int32"],
|
|
dst_fused_1["nsa_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["nsa_cache_seqlens_int32"],
|
|
dst_fused_2["nsa_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"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v", "-s"])
|