diff --git a/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py b/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py index eea98c401..d6c499df0 100644 --- a/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py +++ b/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py @@ -16,7 +16,7 @@ s: scale, 1 item per token, fp32 class GetK: @classmethod def execute(cls, *args, **kwargs): - return cls.torch_fast(*args, **kwargs) + return cls.triton(*args, **kwargs) @classmethod def slow( @@ -67,11 +67,28 @@ class GetK: out = flat_buf[flat_indices] return out.view(-1, 128) + @classmethod + def triton( + cls, pool: "NSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor + ): + """ + Triton implementation for gathering K data from paged buffer. + :param page_indices: (num_pages,), int32/int64 + :return: (seq_len, index_head_dim), uint8 + """ + return _get_k_triton( + buf=buf, + page_indices=page_indices, + seq_len=seq_len, + page_size=pool.page_size, + index_head_dim=pool.index_head_dim, + ) + class GetS: @classmethod def execute(cls, *args, **kwargs): - return cls.torch_fast(*args, **kwargs) + return cls.triton(*args, **kwargs) @classmethod def slow( @@ -119,6 +136,48 @@ class GetS: out = flat_buf[flat_indices] return out.view(-1, 4) + @classmethod + def triton( + cls, pool: "NSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor + ): + """ + Triton implementation for gathering S (scale) data from paged buffer. + :param page_indices: (num_pages,), int32/int64 + :return: (seq_len, 4), uint8 + """ + return _get_s_triton( + buf=buf, + page_indices=page_indices, + seq_len=seq_len, + page_size=pool.page_size, + index_head_dim=pool.index_head_dim, + ) + + +class GetKAndS: + @classmethod + def execute(cls, *args, **kwargs): + return cls.triton(*args, **kwargs) + + @classmethod + def triton( + cls, pool: "NSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor + ): + """ + Triton implementation for gathering both K and S data from paged buffer in a single call. + :param page_indices: (num_pages,), int32/int64 + :return: tuple of (k_fp8, k_scale) where + k_fp8: (seq_len, index_head_dim), uint8 + k_scale: (seq_len, 4), uint8 + """ + return _get_k_and_s_triton( + buf=buf, + page_indices=page_indices, + seq_len=seq_len, + page_size=pool.page_size, + index_head_dim=pool.index_head_dim, + ) + class SetK: @classmethod @@ -363,3 +422,260 @@ def _set_k_and_s_triton_kernel( tl.store(buf_fp8_ptr + out_k_offsets, k) tl.store(buf_fp32_ptr + out_s_offset, k_scale) + + +def _get_k_triton( + buf: torch.Tensor, + page_indices: torch.Tensor, + seq_len: int, + page_size: int, + index_head_dim: int, +): + """ + Gather K (key) data from paged buffer using Triton. + + :param buf: (num_pages, page_size * 128 + page_size * 4), uint8 + :param page_indices: (num_pages,), int32/int64 + :param seq_len: int, number of tokens to gather + :param page_size: int, typically 64 + :param index_head_dim: int, typically 128 + :return: (seq_len, index_head_dim), uint8 + """ + num_pages, buf_numel_per_page = buf.shape + + # Allocate output + out = torch.empty((seq_len, index_head_dim), dtype=torch.uint8, device=buf.device) + + # Launch kernel with one thread per token + grid = (seq_len,) + _get_k_triton_kernel[grid]( + buf, + page_indices, + out, + seq_len, + page_size, + buf_numel_per_page, + index_head_dim, + BLOCK_SIZE=128, + ) + + return out + + +@triton.jit +def _get_k_triton_kernel( + buf_ptr, + page_indices_ptr, + out_ptr, + seq_len: tl.constexpr, + page_size: tl.constexpr, + buf_numel_per_page: tl.constexpr, + index_head_dim: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """ + Each program handles one token (seq_len tokens total). + Loads 128 bytes from the appropriate page. + """ + token_id = tl.program_id(0) + + # Calculate which page and offset within page + page_idx = token_id // page_size + token_offset_in_page = token_id % page_size + + # Load the page index from page_indices + page_index = tl.load(page_indices_ptr + page_idx) + + # Calculate source offset in buf + # buf[page_index, token_offset_in_page * index_head_dim : ...] + src_base_offset = ( + page_index * buf_numel_per_page + token_offset_in_page * index_head_dim + ) + + # Load 128 bytes (index_head_dim elements) + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < index_head_dim + data = tl.load(buf_ptr + src_base_offset + offsets, mask=mask) + + # Store to output + dst_offset = token_id * index_head_dim + tl.store(out_ptr + dst_offset + offsets, data, mask=mask) + + +def _get_s_triton( + buf: torch.Tensor, + page_indices: torch.Tensor, + seq_len: int, + page_size: int, + index_head_dim: int, +): + """ + Gather S (scale) data from paged buffer using Triton. + + :param buf: (num_pages, page_size * 128 + page_size * 4), uint8 + :param page_indices: (num_pages,), int32/int64 + :param seq_len: int, number of tokens to gather + :param page_size: int, typically 64 + :param index_head_dim: int, typically 128 + :return: (seq_len, 4), uint8 (representing fp32 scale) + """ + num_pages, buf_numel_per_page = buf.shape + s_offset_in_page = page_size * index_head_dim # Scales start after K data + + # Allocate output + out = torch.empty((seq_len, 4), dtype=torch.uint8, device=buf.device) + + # Launch kernel with one thread per token + grid = (seq_len,) + _get_s_triton_kernel[grid]( + buf, + page_indices, + out, + seq_len, + page_size, + buf_numel_per_page, + s_offset_in_page, + ) + + return out + + +@triton.jit +def _get_s_triton_kernel( + buf_ptr, + page_indices_ptr, + out_ptr, + seq_len: tl.constexpr, + page_size: tl.constexpr, + buf_numel_per_page: tl.constexpr, + s_offset_in_page: tl.constexpr, +): + """ + Each program handles one token (seq_len tokens total). + Loads 4 bytes (fp32 scale) from the appropriate page. + """ + token_id = tl.program_id(0) + + # Calculate which page and offset within page + page_idx = token_id // page_size + token_offset_in_page = token_id % page_size + + # Load the page index from page_indices + page_index = tl.load(page_indices_ptr + page_idx) + + # Calculate source offset in buf + # Scales are stored after K data: page_size * index_head_dim offset + # buf[page_index, s_offset_in_page + token_offset_in_page * 4 : ...] + src_base_offset = ( + page_index * buf_numel_per_page + s_offset_in_page + token_offset_in_page * 4 + ) + + # Load 4 bytes (fp32 scale) + offsets = tl.arange(0, 4) + data = tl.load(buf_ptr + src_base_offset + offsets) + + # Store to output + dst_offset = token_id * 4 + tl.store(out_ptr + dst_offset + offsets, data) + + +def _get_k_and_s_triton( + buf: torch.Tensor, + page_indices: torch.Tensor, + seq_len: int, + page_size: int, + index_head_dim: int, +): + """ + Fused gather of both K (key) and S (scale) data from paged buffer using Triton. + This is more efficient than calling GetK and GetS separately. + + :param buf: (num_pages, page_size * 128 + page_size * 4), uint8 + :param page_indices: (num_pages,), int32/int64 + :param seq_len: int, number of tokens to gather + :param page_size: int, typically 64 + :param index_head_dim: int, typically 128 + :return: tuple of (k_out, s_out) where + k_out: (seq_len, index_head_dim), uint8 + s_out: (seq_len, 4), uint8 + """ + num_pages, buf_numel_per_page = buf.shape + s_offset_in_page = page_size * index_head_dim # Scales start after K data + + # Allocate outputs + k_out = torch.empty((seq_len, index_head_dim), dtype=torch.uint8, device=buf.device) + s_out = torch.empty((seq_len, 4), dtype=torch.uint8, device=buf.device) + + # Launch kernel with one thread per token + grid = (seq_len,) + _get_k_and_s_triton_kernel[grid]( + buf, + page_indices, + k_out, + s_out, + seq_len, + page_size, + buf_numel_per_page, + index_head_dim, + s_offset_in_page, + BLOCK_SIZE_K=128, + ) + + return k_out, s_out + + +@triton.jit +def _get_k_and_s_triton_kernel( + buf_ptr, + page_indices_ptr, + k_out_ptr, + s_out_ptr, + seq_len: tl.constexpr, + page_size: tl.constexpr, + buf_numel_per_page: tl.constexpr, + index_head_dim: tl.constexpr, + s_offset_in_page: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, +): + """ + Fused kernel that gathers both K and S data in a single pass. + Each program handles one token (seq_len tokens total). + Loads 128 bytes (K) + 4 bytes (S) from the appropriate page. + """ + token_id = tl.program_id(0) + + # Calculate which page and offset within page + page_idx = token_id // page_size + token_offset_in_page = token_id % page_size + + # Load the page index from page_indices + page_index = tl.load(page_indices_ptr + page_idx) + + # ===== Load K data (128 bytes) ===== + # Calculate source offset for K in buf + k_src_base_offset = ( + page_index * buf_numel_per_page + token_offset_in_page * index_head_dim + ) + + # Load 128 bytes (index_head_dim elements) + k_offsets = tl.arange(0, BLOCK_SIZE_K) + k_mask = k_offsets < index_head_dim + k_data = tl.load(buf_ptr + k_src_base_offset + k_offsets, mask=k_mask) + + # Store K to output + k_dst_offset = token_id * index_head_dim + tl.store(k_out_ptr + k_dst_offset + k_offsets, k_data, mask=k_mask) + + # ===== Load S data (4 bytes) ===== + # Calculate source offset for S in buf + s_src_base_offset = ( + page_index * buf_numel_per_page + s_offset_in_page + token_offset_in_page * 4 + ) + + # Load 4 bytes (fp32 scale) + s_offsets = tl.arange(0, 4) + s_data = tl.load(buf_ptr + s_src_base_offset + s_offsets) + + # Store S to output + s_dst_offset = token_id * 4 + tl.store(s_out_ptr + s_dst_offset + s_offsets, s_data) diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 55aaadca2..7518c909a 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -385,12 +385,8 @@ class Indexer(CustomOp): for i in range(forward_batch.batch_size): seq_len = forward_batch.seq_lens_cpu[i].item() assert isinstance(seq_len, int) - k_fp8 = forward_batch.token_to_kv_pool.get_index_k_continuous( - layer_id, - seq_len, - block_tables[i], - ) - k_scale = forward_batch.token_to_kv_pool.get_index_k_scale_continuous( + # Use fused Triton kernel to get both K and scale in a single call + k_fp8, k_scale = forward_batch.token_to_kv_pool.get_index_k_scale_buffer( layer_id, seq_len, block_tables[i], diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index e61a5540e..64a1ba366 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1826,6 +1826,28 @@ class NSATokenToKVPool(MLATokenToKVPool): self, buf, seq_len=seq_len, page_indices=page_indices ) + def get_index_k_scale_buffer( + self, + layer_id: int, + seq_len: int, + page_indices: torch.Tensor, + ): + """ + Fused method to get both index K and scale data in a single call using Triton. + More efficient than calling get_index_k_continuous and get_index_k_scale_continuous separately. + + :param layer_id: Layer index + :param seq_len: Sequence length + :param page_indices: Page indices tensor + :return: tuple of (k_fp8, k_scale) where + k_fp8: (seq_len, index_head_dim), uint8 + k_scale: (seq_len, 4), uint8 + """ + buf = self.index_k_with_scale_buffer[layer_id - self.start_layer] + return index_buf_accessor.GetKAndS.execute( + self, buf, seq_len=seq_len, page_indices=page_indices + ) + def set_index_k_scale_buffer( self, layer_id: int, diff --git a/test/manual/layers/attention/nsa/test_index_buf_accessor.py b/test/manual/layers/attention/nsa/test_index_buf_accessor.py new file mode 100644 index 000000000..5e17e1859 --- /dev/null +++ b/test/manual/layers/attention/nsa/test_index_buf_accessor.py @@ -0,0 +1,554 @@ +""" +Correctness tests for NSA Indexer K/S Buffer Access with Fused Triton Kernels. + +This test verifies that the optimized Triton implementations (GetK, GetS, GetKAndS) +produce identical results to the torch_fast baseline implementations. + +Test coverage: +- GetK.triton() vs GetK.torch_fast() +- GetS.triton() vs GetS.torch_fast() +- GetKAndS.triton() vs separate GetK.torch_fast() + GetS.torch_fast() +""" + +import pytest +import torch + +from sglang.srt.layers.attention.nsa.index_buf_accessor import GetK, GetKAndS, GetS + + +class MockNSATokenToKVPool: + """Mock pool object that mimics NSATokenToKVPool for testing.""" + + def __init__( + self, + page_size: int = 64, + index_head_dim: int = 128, + quant_block_size: int = 128, + device: str = "cuda", + ): + self.page_size = page_size + self.index_head_dim = index_head_dim + self.quant_block_size = quant_block_size + self.device = device + + +def create_test_buffer( + num_pages: int, + page_size: int = 64, + index_head_dim: int = 128, + device: str = "cuda", +) -> torch.Tensor: + """ + Create a test buffer mimicking the K/S buffer structure. + + Buffer layout per page: + - First page_size * index_head_dim bytes: K data (fp8, stored as uint8) + - Next page_size * 4 bytes: S data (fp32 scales, stored as uint8) + + Args: + num_pages: Number of pages to allocate + page_size: Tokens per page (typically 64) + index_head_dim: Dimension of K vectors (typically 128) + device: Device to allocate on + + Returns: + Buffer of shape (num_pages, page_size * index_head_dim + page_size * 4) + """ + buf_numel_per_page = page_size * index_head_dim + page_size * 4 + buf = torch.randint( + 0, 256, (num_pages, buf_numel_per_page), dtype=torch.uint8, device=device + ) + return buf + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +class TestGetK: + """Test cases for GetK.triton() correctness.""" + + @pytest.mark.parametrize("num_pages", [1, 2, 4, 8, 16]) + @pytest.mark.parametrize("seq_len", [64, 128, 256, 512, 1024]) + @pytest.mark.parametrize("page_size", [64]) + @pytest.mark.parametrize("index_head_dim", [128]) + def test_getk_correctness(self, num_pages, seq_len, page_size, index_head_dim): + """Test GetK.triton() produces same output as GetK.torch_fast().""" + device = torch.device("cuda") + + # Ensure seq_len doesn't exceed available pages + max_seq_len = num_pages * page_size + seq_len = min(seq_len, max_seq_len) + + # Create mock pool + pool = MockNSATokenToKVPool( + page_size=page_size, index_head_dim=index_head_dim, device=device + ) + + # Create test buffer + buf = create_test_buffer( + num_pages=num_pages, + page_size=page_size, + index_head_dim=index_head_dim, + device=device, + ) + + # Create page indices + num_pages_needed = (seq_len + page_size - 1) // page_size + page_indices = torch.randint( + 0, num_pages, (num_pages_needed,), dtype=torch.int32, device=device + ) + + # Run both implementations + output_torch = GetK.torch_fast(pool, buf, seq_len, page_indices) + output_triton = GetK.triton(pool, buf, seq_len, page_indices) + + # Verify shapes + assert output_torch.shape == (seq_len, index_head_dim) + assert output_triton.shape == (seq_len, index_head_dim) + assert output_torch.dtype == torch.uint8 + assert output_triton.dtype == torch.uint8 + + # Compare results (should be exact match) + torch.testing.assert_close( + output_triton, output_torch, rtol=0, atol=0, msg="GetK outputs differ" + ) + + def test_getk_sequential_pages(self): + """Test GetK with sequential page indices.""" + device = torch.device("cuda") + page_size = 64 + index_head_dim = 128 + num_pages = 10 + seq_len = 320 # 5 pages + + pool = MockNSATokenToKVPool( + page_size=page_size, index_head_dim=index_head_dim, device=device + ) + buf = create_test_buffer(num_pages, page_size, index_head_dim, device) + + # Sequential page indices [0, 1, 2, 3, 4] + page_indices = torch.arange(5, dtype=torch.int32, device=device) + + output_torch = GetK.torch_fast(pool, buf, seq_len, page_indices) + output_triton = GetK.triton(pool, buf, seq_len, page_indices) + + torch.testing.assert_close(output_triton, output_torch, rtol=0, atol=0) + + def test_getk_repeated_pages(self): + """Test GetK with repeated page indices.""" + device = torch.device("cuda") + page_size = 64 + index_head_dim = 128 + num_pages = 5 + seq_len = 192 # 3 pages + + pool = MockNSATokenToKVPool( + page_size=page_size, index_head_dim=index_head_dim, device=device + ) + buf = create_test_buffer(num_pages, page_size, index_head_dim, device) + + # Repeated page indices [2, 2, 2] + page_indices = torch.full((3,), 2, dtype=torch.int32, device=device) + + output_torch = GetK.torch_fast(pool, buf, seq_len, page_indices) + output_triton = GetK.triton(pool, buf, seq_len, page_indices) + + torch.testing.assert_close(output_triton, output_torch, rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +class TestGetS: + """Test cases for GetS.triton() correctness.""" + + @pytest.mark.parametrize("num_pages", [1, 2, 4, 8, 16]) + @pytest.mark.parametrize("seq_len", [64, 128, 256, 512, 1024]) + @pytest.mark.parametrize("page_size", [64]) + @pytest.mark.parametrize("index_head_dim", [128]) + def test_gets_correctness(self, num_pages, seq_len, page_size, index_head_dim): + """Test GetS.triton() produces same output as GetS.torch_fast().""" + device = torch.device("cuda") + + # Ensure seq_len doesn't exceed available pages + max_seq_len = num_pages * page_size + seq_len = min(seq_len, max_seq_len) + + # Create mock pool + pool = MockNSATokenToKVPool( + page_size=page_size, index_head_dim=index_head_dim, device=device + ) + + # Create test buffer + buf = create_test_buffer( + num_pages=num_pages, + page_size=page_size, + index_head_dim=index_head_dim, + device=device, + ) + + # Create page indices + num_pages_needed = (seq_len + page_size - 1) // page_size + page_indices = torch.randint( + 0, num_pages, (num_pages_needed,), dtype=torch.int32, device=device + ) + + # Run both implementations + output_torch = GetS.torch_fast(pool, buf, seq_len, page_indices) + output_triton = GetS.triton(pool, buf, seq_len, page_indices) + + # Verify shapes + assert output_torch.shape == (seq_len, 4) + assert output_triton.shape == (seq_len, 4) + assert output_torch.dtype == torch.uint8 + assert output_triton.dtype == torch.uint8 + + # Compare results (should be exact match) + torch.testing.assert_close( + output_triton, output_torch, rtol=0, atol=0, msg="GetS outputs differ" + ) + + def test_gets_sequential_pages(self): + """Test GetS with sequential page indices.""" + device = torch.device("cuda") + page_size = 64 + index_head_dim = 128 + num_pages = 10 + seq_len = 320 # 5 pages + + pool = MockNSATokenToKVPool( + page_size=page_size, index_head_dim=index_head_dim, device=device + ) + buf = create_test_buffer(num_pages, page_size, index_head_dim, device) + + # Sequential page indices [0, 1, 2, 3, 4] + page_indices = torch.arange(5, dtype=torch.int32, device=device) + + output_torch = GetS.torch_fast(pool, buf, seq_len, page_indices) + output_triton = GetS.triton(pool, buf, seq_len, page_indices) + + torch.testing.assert_close(output_triton, output_torch, rtol=0, atol=0) + + def test_gets_repeated_pages(self): + """Test GetS with repeated page indices.""" + device = torch.device("cuda") + page_size = 64 + index_head_dim = 128 + num_pages = 5 + seq_len = 192 # 3 pages + + pool = MockNSATokenToKVPool( + page_size=page_size, index_head_dim=index_head_dim, device=device + ) + buf = create_test_buffer(num_pages, page_size, index_head_dim, device) + + # Repeated page indices [2, 2, 2] + page_indices = torch.full((3,), 2, dtype=torch.int32, device=device) + + output_torch = GetS.torch_fast(pool, buf, seq_len, page_indices) + output_triton = GetS.triton(pool, buf, seq_len, page_indices) + + torch.testing.assert_close(output_triton, output_torch, rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +class TestGetKAndS: + """Test cases for GetKAndS.triton() correctness.""" + + @pytest.mark.parametrize("num_pages", [1, 2, 4, 8, 16]) + @pytest.mark.parametrize("seq_len", [64, 128, 256, 512, 1024]) + @pytest.mark.parametrize("page_size", [64]) + @pytest.mark.parametrize("index_head_dim", [128]) + def test_get_k_and_s_correctness( + self, num_pages, seq_len, page_size, index_head_dim + ): + """Test GetKAndS.triton() produces same output as separate torch_fast calls.""" + device = torch.device("cuda") + + # Ensure seq_len doesn't exceed available pages + max_seq_len = num_pages * page_size + seq_len = min(seq_len, max_seq_len) + + # Create mock pool + pool = MockNSATokenToKVPool( + page_size=page_size, index_head_dim=index_head_dim, device=device + ) + + # Create test buffer + buf = create_test_buffer( + num_pages=num_pages, + page_size=page_size, + index_head_dim=index_head_dim, + device=device, + ) + + # Create page indices + num_pages_needed = (seq_len + page_size - 1) // page_size + page_indices = torch.randint( + 0, num_pages, (num_pages_needed,), dtype=torch.int32, device=device + ) + + # Run baseline: separate torch_fast calls + k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices) + s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices) + + # Run fused Triton implementation + k_triton, s_triton = GetKAndS.triton(pool, buf, seq_len, page_indices) + + # Verify shapes + assert k_torch.shape == (seq_len, index_head_dim) + assert s_torch.shape == (seq_len, 4) + assert k_triton.shape == (seq_len, index_head_dim) + assert s_triton.shape == (seq_len, 4) + + # Verify dtypes + assert k_torch.dtype == torch.uint8 + assert s_torch.dtype == torch.uint8 + assert k_triton.dtype == torch.uint8 + assert s_triton.dtype == torch.uint8 + + # Compare K results + torch.testing.assert_close( + k_triton, k_torch, rtol=0, atol=0, msg="GetKAndS K outputs differ" + ) + + # Compare S results + torch.testing.assert_close( + s_triton, s_torch, rtol=0, atol=0, msg="GetKAndS S outputs differ" + ) + + def test_get_k_and_s_sequential_pages(self): + """Test GetKAndS with sequential page indices.""" + device = torch.device("cuda") + page_size = 64 + index_head_dim = 128 + num_pages = 10 + seq_len = 320 # 5 pages + + pool = MockNSATokenToKVPool( + page_size=page_size, index_head_dim=index_head_dim, device=device + ) + buf = create_test_buffer(num_pages, page_size, index_head_dim, device) + + # Sequential page indices [0, 1, 2, 3, 4] + page_indices = torch.arange(5, dtype=torch.int32, device=device) + + # Baseline + k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices) + s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices) + + # Fused + k_triton, s_triton = GetKAndS.triton(pool, buf, seq_len, page_indices) + + torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0) + torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0) + + def test_get_k_and_s_repeated_pages(self): + """Test GetKAndS with repeated page indices.""" + device = torch.device("cuda") + page_size = 64 + index_head_dim = 128 + num_pages = 5 + seq_len = 192 # 3 pages + + pool = MockNSATokenToKVPool( + page_size=page_size, index_head_dim=index_head_dim, device=device + ) + buf = create_test_buffer(num_pages, page_size, index_head_dim, device) + + # Repeated page indices [2, 2, 2] + page_indices = torch.full((3,), 2, dtype=torch.int32, device=device) + + # Baseline + k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices) + s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices) + + # Fused + k_triton, s_triton = GetKAndS.triton(pool, buf, seq_len, page_indices) + + torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0) + torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0) + + def test_get_k_and_s_partial_page(self): + """Test GetKAndS when seq_len is not a multiple of page_size.""" + device = torch.device("cuda") + page_size = 64 + index_head_dim = 128 + num_pages = 5 + seq_len = 100 # Not a multiple of 64 + + pool = MockNSATokenToKVPool( + page_size=page_size, index_head_dim=index_head_dim, device=device + ) + buf = create_test_buffer(num_pages, page_size, index_head_dim, device) + + num_pages_needed = (seq_len + page_size - 1) // page_size + page_indices = torch.arange(num_pages_needed, dtype=torch.int32, device=device) + + # Baseline + k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices) + s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices) + + # Fused + k_triton, s_triton = GetKAndS.triton(pool, buf, seq_len, page_indices) + + # Should handle partial pages correctly + torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0) + torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +class TestEdgeCases: + """Test edge cases and boundary conditions.""" + + def test_single_token(self): + """Test with seq_len=1 (single token).""" + device = torch.device("cuda") + page_size = 64 + index_head_dim = 128 + num_pages = 2 + seq_len = 1 + + pool = MockNSATokenToKVPool( + page_size=page_size, index_head_dim=index_head_dim, device=device + ) + buf = create_test_buffer(num_pages, page_size, index_head_dim, device) + page_indices = torch.tensor([0], dtype=torch.int32, device=device) + + # Test GetK + k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices) + k_triton = GetK.triton(pool, buf, seq_len, page_indices) + torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0) + + # Test GetS + s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices) + s_triton = GetS.triton(pool, buf, seq_len, page_indices) + torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0) + + # Test GetKAndS + k_triton2, s_triton2 = GetKAndS.triton(pool, buf, seq_len, page_indices) + torch.testing.assert_close(k_triton2, k_torch, rtol=0, atol=0) + torch.testing.assert_close(s_triton2, s_torch, rtol=0, atol=0) + + def test_exact_page_boundary(self): + """Test when seq_len exactly matches page boundaries.""" + device = torch.device("cuda") + page_size = 64 + index_head_dim = 128 + num_pages = 5 + seq_len = 192 # Exactly 3 pages + + pool = MockNSATokenToKVPool( + page_size=page_size, index_head_dim=index_head_dim, device=device + ) + buf = create_test_buffer(num_pages, page_size, index_head_dim, device) + page_indices = torch.arange(3, dtype=torch.int32, device=device) + + # Test GetK + k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices) + k_triton = GetK.triton(pool, buf, seq_len, page_indices) + torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0) + + # Test GetS + s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices) + s_triton = GetS.triton(pool, buf, seq_len, page_indices) + torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0) + + # Test GetKAndS + k_triton2, s_triton2 = GetKAndS.triton(pool, buf, seq_len, page_indices) + torch.testing.assert_close(k_triton2, k_torch, rtol=0, atol=0) + torch.testing.assert_close(s_triton2, s_torch, rtol=0, atol=0) + + def test_large_seq_len(self): + """Test with large sequence length.""" + device = torch.device("cuda") + page_size = 64 + index_head_dim = 128 + num_pages = 100 + seq_len = 4096 # 64 pages + + pool = MockNSATokenToKVPool( + page_size=page_size, index_head_dim=index_head_dim, device=device + ) + buf = create_test_buffer(num_pages, page_size, index_head_dim, device) + + num_pages_needed = (seq_len + page_size - 1) // page_size + page_indices = torch.randint( + 0, num_pages, (num_pages_needed,), dtype=torch.int32, device=device + ) + + # Test GetK + k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices) + k_triton = GetK.triton(pool, buf, seq_len, page_indices) + torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0) + + # Test GetS + s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices) + s_triton = GetS.triton(pool, buf, seq_len, page_indices) + torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0) + + # Test GetKAndS + k_triton2, s_triton2 = GetKAndS.triton(pool, buf, seq_len, page_indices) + torch.testing.assert_close(k_triton2, k_torch, rtol=0, atol=0) + torch.testing.assert_close(s_triton2, s_torch, rtol=0, atol=0) + + +def print_test_summary(): + """Print a summary message about the test suite.""" + print("\n" + "=" * 80) + print("NSA Indexer K/S Buffer Accessor Correctness Tests") + print("=" * 80) + print("Testing Triton implementations against torch_fast baseline:") + print(" - GetK.triton() vs GetK.torch_fast()") + print(" - GetS.triton() vs GetS.torch_fast()") + print(" - GetKAndS.triton() vs separate GetK/GetS torch_fast() calls") + print("=" * 80) + print() + + +if __name__ == "__main__": + # Run tests manually + if not torch.cuda.is_available(): + print("CUDA not available. Skipping tests.") + exit(0) + + print_test_summary() + + # Run a few sample tests + print("Running sample correctness tests...\n") + + # Test GetK + print("Testing GetK...") + test_getk = TestGetK() + test_getk.test_getk_correctness( + num_pages=4, seq_len=256, page_size=64, index_head_dim=128 + ) + test_getk.test_getk_sequential_pages() + print("✓ GetK tests passed\n") + + # Test GetS + print("Testing GetS...") + test_gets = TestGetS() + test_gets.test_gets_correctness( + num_pages=4, seq_len=256, page_size=64, index_head_dim=128 + ) + test_gets.test_gets_sequential_pages() + print("✓ GetS tests passed\n") + + # Test GetKAndS + print("Testing GetKAndS...") + test_get_k_and_s = TestGetKAndS() + test_get_k_and_s.test_get_k_and_s_correctness( + num_pages=4, seq_len=256, page_size=64, index_head_dim=128 + ) + test_get_k_and_s.test_get_k_and_s_sequential_pages() + test_get_k_and_s.test_get_k_and_s_partial_page() + print("✓ GetKAndS tests passed\n") + + # Test edge cases + print("Testing edge cases...") + test_edge = TestEdgeCases() + test_edge.test_single_token() + test_edge.test_exact_page_boundary() + test_edge.test_large_seq_len() + print("✓ Edge case tests passed\n") + + print("=" * 80) + print("All correctness tests passed successfully!") + print("=" * 80)