419 lines
15 KiB
Python
419 lines
15 KiB
Python
"""
|
|
Unit tests for the fused Triton kernel in normal_decode_set_metadata.
|
|
|
|
This test suite verifies:
|
|
1. Correctness against reference PyTorch implementation
|
|
2. Different page sizes (1, 16, 64)
|
|
3. With and without Sliding Window Attention (SWA)
|
|
4. Various batch sizes and sequence lengths
|
|
5. Edge cases
|
|
"""
|
|
|
|
import unittest
|
|
|
|
import torch
|
|
|
|
from sglang.srt.layers.attention.flashattention_backend import (
|
|
normal_decode_set_metadata,
|
|
)
|
|
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
|
|
from sglang.test.ci.ci_register import register_cuda_ci
|
|
from sglang.test.test_utils import CustomTestCase
|
|
|
|
# Register this test for CUDA CI in stage-b (fast attention/kernel tests)
|
|
register_cuda_ci(est_time=25, suite="stage-b-test-1-gpu-large")
|
|
|
|
|
|
def reference_normal_decode_set_metadata(
|
|
cache_seqlens_int32: torch.Tensor,
|
|
cu_seqlens_k: torch.Tensor,
|
|
page_table: torch.Tensor,
|
|
req_to_token: torch.Tensor,
|
|
req_pool_indices: torch.Tensor,
|
|
strided_indices: torch.Tensor,
|
|
max_seq_pages: int,
|
|
seq_lens: torch.Tensor,
|
|
seq_len_delta: int,
|
|
page_size: int,
|
|
swa_page_table: torch.Tensor = None,
|
|
token_to_kv_pool=None,
|
|
):
|
|
"""
|
|
Reference implementation using original PyTorch operations.
|
|
This is the pre-Triton version for correctness comparison.
|
|
"""
|
|
cache_seqlens_int32.copy_(seq_lens + seq_len_delta)
|
|
cu_seqlens_k[1:].copy_(torch.cumsum(cache_seqlens_int32, dim=0, dtype=torch.int32))
|
|
page_indices = req_to_token[
|
|
req_pool_indices[:, None],
|
|
strided_indices[:max_seq_pages][None, :],
|
|
]
|
|
page_table[:, :max_seq_pages].copy_(page_indices // page_size)
|
|
|
|
if swa_page_table is not None and token_to_kv_pool is not None:
|
|
swa_page_indices = token_to_kv_pool.translate_loc_from_full_to_swa(page_indices)
|
|
swa_page_table[:, :max_seq_pages].copy_(swa_page_indices // page_size)
|
|
|
|
|
|
@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
|
|
class TestNormalDecodeSetMetadata(CustomTestCase):
|
|
"""Test fused Triton kernel in normal_decode_set_metadata."""
|
|
|
|
def setUp(self):
|
|
self.device = "cuda"
|
|
self.dtype = torch.int32
|
|
|
|
def _create_test_data(
|
|
self,
|
|
batch_size: int,
|
|
max_seq_len: int,
|
|
page_size: int,
|
|
has_swa: bool = False,
|
|
seq_len_delta: int = 0,
|
|
):
|
|
"""Create test data for normal_decode_set_metadata."""
|
|
# Random sequence lengths for each batch
|
|
seq_lens = torch.randint(
|
|
max_seq_len // 2,
|
|
max_seq_len + 1,
|
|
(batch_size,),
|
|
dtype=torch.int64,
|
|
device=self.device,
|
|
)
|
|
|
|
# Calculate max_seq_pages
|
|
max_len = seq_lens.max().item()
|
|
max_seq_pages = (max_len + seq_len_delta + page_size - 1) // page_size
|
|
|
|
# Create req_pool_indices (maps batch index to pool index)
|
|
req_pool_indices = torch.arange(
|
|
batch_size, dtype=torch.int32, device=self.device
|
|
)
|
|
|
|
# Create strided_indices for page table indexing
|
|
if page_size == 1:
|
|
strided_indices = torch.arange(
|
|
max_seq_len * 2, dtype=torch.int32, device=self.device
|
|
)
|
|
else:
|
|
strided_indices = torch.arange(
|
|
0, max_seq_len * 2, page_size, dtype=torch.int32, device=self.device
|
|
)
|
|
|
|
# Create req_to_token pool (simulates token locations in KV cache)
|
|
pool_size = batch_size
|
|
max_tokens = max_seq_len * 2
|
|
req_to_token = torch.randint(
|
|
0, 10000, (pool_size, max_tokens), dtype=torch.int32, device=self.device
|
|
)
|
|
|
|
# Output tensors (to be filled by the function)
|
|
cache_seqlens_int32 = torch.zeros(
|
|
batch_size, dtype=torch.int32, device=self.device
|
|
)
|
|
cu_seqlens_k = torch.zeros(
|
|
batch_size + 1, dtype=torch.int32, device=self.device
|
|
)
|
|
page_table = torch.zeros(
|
|
(batch_size, max_seq_pages + 10), dtype=torch.int32, device=self.device
|
|
)
|
|
|
|
# SWA setup if needed
|
|
swa_page_table = None
|
|
token_to_kv_pool = None
|
|
if has_swa:
|
|
swa_page_table = torch.zeros(
|
|
(batch_size, max_seq_pages + 10), dtype=torch.int32, device=self.device
|
|
)
|
|
# Create a simple SWA KV pool for testing
|
|
token_to_kv_pool = self._create_swa_kv_pool(10000, page_size)
|
|
|
|
return {
|
|
"cache_seqlens_int32": cache_seqlens_int32,
|
|
"cu_seqlens_k": cu_seqlens_k,
|
|
"page_table": page_table,
|
|
"req_to_token": req_to_token,
|
|
"req_pool_indices": req_pool_indices,
|
|
"strided_indices": strided_indices,
|
|
"max_seq_pages": max_seq_pages,
|
|
"seq_lens": seq_lens,
|
|
"seq_len_delta": seq_len_delta,
|
|
"page_size": page_size,
|
|
"swa_page_table": swa_page_table,
|
|
"token_to_kv_pool": token_to_kv_pool,
|
|
}
|
|
|
|
def _create_swa_kv_pool(self, size: int, page_size: int):
|
|
"""Create a mock SWA KV pool for testing that inherits from SWAKVPool."""
|
|
|
|
# Create a minimal mock that inherits from SWAKVPool to pass isinstance check
|
|
class MinimalSWAKVPool(SWAKVPool):
|
|
def __init__(self, size, device):
|
|
# Don't call super().__init__() to avoid complex initialization
|
|
# Just set the minimal attributes needed for the test
|
|
self.full_to_swa_index_mapping = torch.arange(
|
|
size, dtype=torch.int32, device=device
|
|
)
|
|
# Add some randomness to simulate real SWA mapping
|
|
self.full_to_swa_index_mapping = (
|
|
self.full_to_swa_index_mapping
|
|
+ torch.randint(0, 100, (size,), device=device)
|
|
) % size
|
|
self.device = device
|
|
|
|
def translate_loc_from_full_to_swa(self, page_indices):
|
|
"""Mock translation method."""
|
|
return self.full_to_swa_index_mapping[page_indices]
|
|
|
|
return MinimalSWAKVPool(size, self.device)
|
|
|
|
def _run_test(
|
|
self,
|
|
batch_size: int,
|
|
max_seq_len: int,
|
|
page_size: int,
|
|
has_swa: bool = False,
|
|
seq_len_delta: int = 0,
|
|
):
|
|
"""Run a single test configuration."""
|
|
# Create test data
|
|
test_data = self._create_test_data(
|
|
batch_size, max_seq_len, page_size, has_swa, seq_len_delta
|
|
)
|
|
|
|
# Clone data for reference implementation
|
|
ref_data = {
|
|
"cache_seqlens_int32": test_data["cache_seqlens_int32"].clone(),
|
|
"cu_seqlens_k": test_data["cu_seqlens_k"].clone(),
|
|
"page_table": test_data["page_table"].clone(),
|
|
"swa_page_table": test_data["swa_page_table"].clone() if has_swa else None,
|
|
}
|
|
|
|
# Run reference implementation
|
|
reference_normal_decode_set_metadata(
|
|
ref_data["cache_seqlens_int32"],
|
|
ref_data["cu_seqlens_k"],
|
|
ref_data["page_table"],
|
|
test_data["req_to_token"],
|
|
test_data["req_pool_indices"],
|
|
test_data["strided_indices"],
|
|
test_data["max_seq_pages"],
|
|
test_data["seq_lens"],
|
|
test_data["seq_len_delta"],
|
|
test_data["page_size"],
|
|
ref_data["swa_page_table"],
|
|
test_data["token_to_kv_pool"],
|
|
)
|
|
|
|
# Run fused Triton implementation
|
|
normal_decode_set_metadata(
|
|
test_data["cache_seqlens_int32"],
|
|
test_data["cu_seqlens_k"],
|
|
test_data["page_table"],
|
|
test_data["req_to_token"],
|
|
test_data["req_pool_indices"],
|
|
test_data["strided_indices"],
|
|
test_data["max_seq_pages"],
|
|
test_data["seq_lens"],
|
|
test_data["seq_len_delta"],
|
|
test_data["page_size"],
|
|
test_data["swa_page_table"],
|
|
test_data["token_to_kv_pool"],
|
|
)
|
|
|
|
# Compare results
|
|
self.assertTrue(
|
|
torch.equal(
|
|
test_data["cache_seqlens_int32"], ref_data["cache_seqlens_int32"]
|
|
),
|
|
f"cache_seqlens_int32 mismatch. Expected:\n{ref_data['cache_seqlens_int32']}\nGot:\n{test_data['cache_seqlens_int32']}",
|
|
)
|
|
|
|
self.assertTrue(
|
|
torch.equal(test_data["cu_seqlens_k"], ref_data["cu_seqlens_k"]),
|
|
f"cu_seqlens_k mismatch. Expected:\n{ref_data['cu_seqlens_k']}\nGot:\n{test_data['cu_seqlens_k']}",
|
|
)
|
|
|
|
self.assertTrue(
|
|
torch.equal(test_data["page_table"], ref_data["page_table"]),
|
|
f"page_table mismatch at bs={batch_size}, page_size={page_size}",
|
|
)
|
|
|
|
if has_swa:
|
|
self.assertTrue(
|
|
torch.equal(test_data["swa_page_table"], ref_data["swa_page_table"]),
|
|
f"swa_page_table mismatch at bs={batch_size}, page_size={page_size}",
|
|
)
|
|
|
|
# Test cases for page_size=1 (uses specialized kernel _fused_metadata_kernel_ps1_no_swa)
|
|
def test_page_size_1_small_batch(self):
|
|
"""Test with page_size=1, small batch."""
|
|
self._run_test(batch_size=2, max_seq_len=128, page_size=1, has_swa=False)
|
|
|
|
def test_page_size_1_medium_batch(self):
|
|
"""Test with page_size=1, medium batch."""
|
|
self._run_test(batch_size=16, max_seq_len=256, page_size=1, has_swa=False)
|
|
|
|
def test_page_size_1_large_batch(self):
|
|
"""Test with page_size=1, large batch."""
|
|
self._run_test(batch_size=64, max_seq_len=512, page_size=1, has_swa=False)
|
|
|
|
def test_page_size_1_with_seq_len_delta(self):
|
|
"""Test with page_size=1 and seq_len_delta > 0."""
|
|
self._run_test(
|
|
batch_size=8, max_seq_len=200, page_size=1, has_swa=False, seq_len_delta=5
|
|
)
|
|
|
|
# Test cases for page_size > 1 (uses general kernel _fused_metadata_kernel_general)
|
|
def test_page_size_16_small_batch(self):
|
|
"""Test with page_size=16, small batch."""
|
|
self._run_test(batch_size=4, max_seq_len=256, page_size=16, has_swa=False)
|
|
|
|
def test_page_size_16_medium_batch(self):
|
|
"""Test with page_size=16, medium batch."""
|
|
self._run_test(batch_size=16, max_seq_len=512, page_size=16, has_swa=False)
|
|
|
|
def test_page_size_64_small_batch(self):
|
|
"""Test with page_size=64, small batch."""
|
|
self._run_test(batch_size=4, max_seq_len=512, page_size=64, has_swa=False)
|
|
|
|
def test_page_size_64_medium_batch(self):
|
|
"""Test with page_size=64, medium batch."""
|
|
self._run_test(batch_size=32, max_seq_len=1024, page_size=64, has_swa=False)
|
|
|
|
def test_page_size_64_with_seq_len_delta(self):
|
|
"""Test with page_size=64 and seq_len_delta > 0."""
|
|
self._run_test(
|
|
batch_size=8, max_seq_len=512, page_size=64, has_swa=False, seq_len_delta=3
|
|
)
|
|
|
|
# Test cases with Sliding Window Attention (SWA)
|
|
def test_page_size_16_with_swa(self):
|
|
"""Test with page_size=16 and SWA enabled."""
|
|
self._run_test(batch_size=8, max_seq_len=256, page_size=16, has_swa=True)
|
|
|
|
def test_page_size_64_with_swa(self):
|
|
"""Test with page_size=64 and SWA enabled."""
|
|
self._run_test(batch_size=16, max_seq_len=512, page_size=64, has_swa=True)
|
|
|
|
def test_page_size_64_with_swa_and_delta(self):
|
|
"""Test with page_size=64, SWA, and seq_len_delta."""
|
|
self._run_test(
|
|
batch_size=8, max_seq_len=400, page_size=64, has_swa=True, seq_len_delta=2
|
|
)
|
|
|
|
# Edge cases
|
|
def test_batch_size_1(self):
|
|
"""Test with single batch."""
|
|
self._run_test(batch_size=1, max_seq_len=128, page_size=1, has_swa=False)
|
|
self._run_test(batch_size=1, max_seq_len=256, page_size=64, has_swa=False)
|
|
|
|
def test_max_seq_pages_small(self):
|
|
"""Test edge case where max_seq_pages could be very small."""
|
|
# This tests when sequences are very short
|
|
test_data = self._create_test_data(
|
|
batch_size=2, max_seq_len=10, page_size=64, has_swa=False
|
|
)
|
|
|
|
# Run fused implementation (should handle gracefully)
|
|
normal_decode_set_metadata(
|
|
test_data["cache_seqlens_int32"],
|
|
test_data["cu_seqlens_k"],
|
|
test_data["page_table"],
|
|
test_data["req_to_token"],
|
|
test_data["req_pool_indices"],
|
|
test_data["strided_indices"],
|
|
test_data["max_seq_pages"],
|
|
test_data["seq_lens"],
|
|
test_data["seq_len_delta"],
|
|
test_data["page_size"],
|
|
test_data["swa_page_table"],
|
|
test_data["token_to_kv_pool"],
|
|
)
|
|
|
|
# Verify no crashes and basic properties
|
|
self.assertEqual(
|
|
test_data["cache_seqlens_int32"].sum().item(),
|
|
test_data["seq_lens"].sum().item(),
|
|
)
|
|
|
|
def test_power_of_two_page_sizes(self):
|
|
"""Test various power-of-2 page sizes."""
|
|
page_sizes = [1, 2, 4, 8, 16, 32, 64, 128]
|
|
for page_size in page_sizes:
|
|
with self.subTest(page_size=page_size):
|
|
self._run_test(
|
|
batch_size=4, max_seq_len=256, page_size=page_size, has_swa=False
|
|
)
|
|
|
|
def test_varied_sequence_lengths(self):
|
|
"""Test with highly varied sequence lengths in the same batch."""
|
|
batch_size = 8
|
|
max_seq_len = 512
|
|
page_size = 64
|
|
|
|
test_data = self._create_test_data(
|
|
batch_size, max_seq_len, page_size, has_swa=False
|
|
)
|
|
|
|
# Manually set varied sequence lengths
|
|
test_data["seq_lens"] = torch.tensor(
|
|
[10, 50, 100, 200, 300, 450, 500, 512],
|
|
dtype=torch.int64,
|
|
device=self.device,
|
|
)
|
|
test_data["max_seq_pages"] = (
|
|
test_data["seq_lens"].max().item() + page_size - 1
|
|
) // page_size
|
|
|
|
# Run both implementations
|
|
ref_data = {
|
|
"cache_seqlens_int32": test_data["cache_seqlens_int32"].clone(),
|
|
"cu_seqlens_k": test_data["cu_seqlens_k"].clone(),
|
|
"page_table": test_data["page_table"].clone(),
|
|
}
|
|
|
|
reference_normal_decode_set_metadata(
|
|
ref_data["cache_seqlens_int32"],
|
|
ref_data["cu_seqlens_k"],
|
|
ref_data["page_table"],
|
|
test_data["req_to_token"],
|
|
test_data["req_pool_indices"],
|
|
test_data["strided_indices"],
|
|
test_data["max_seq_pages"],
|
|
test_data["seq_lens"],
|
|
0,
|
|
page_size,
|
|
None,
|
|
None,
|
|
)
|
|
|
|
normal_decode_set_metadata(
|
|
test_data["cache_seqlens_int32"],
|
|
test_data["cu_seqlens_k"],
|
|
test_data["page_table"],
|
|
test_data["req_to_token"],
|
|
test_data["req_pool_indices"],
|
|
test_data["strided_indices"],
|
|
test_data["max_seq_pages"],
|
|
test_data["seq_lens"],
|
|
0,
|
|
page_size,
|
|
None,
|
|
None,
|
|
)
|
|
|
|
self.assertTrue(
|
|
torch.equal(
|
|
test_data["cache_seqlens_int32"], ref_data["cache_seqlens_int32"]
|
|
)
|
|
)
|
|
self.assertTrue(
|
|
torch.equal(test_data["cu_seqlens_k"], ref_data["cu_seqlens_k"])
|
|
)
|
|
self.assertTrue(torch.equal(test_data["page_table"], ref_data["page_table"]))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|