[Performance] Optimize NSA Indexer K/S Buffer Access with Fused Triton Kernels (#13812)

Co-authored-by: Johnsonms <johnson@together.ai>

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)

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(
+            # Use fused Triton kernel to get both K and scale in a single call
-                layer_id,
+            k_fp8, k_scale = forward_batch.token_to_kv_pool.get_index_k_scale_buffer(
-                seq_len,
-                block_tables[i],
-            )
-            k_scale = forward_batch.token_to_kv_pool.get_index_k_scale_continuous(
                 layer_id,
                 seq_len,
                 block_tables[i],

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,

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)