chore: vendor sglang v0.5.10 snapshot

third_party/sglang/test/manual/layers/attention/nsa/test_act_quant_triton.py

@@ -0,0 +1,281 @@
+"""
+Unit tests comparing TileLang and Triton implementations of activation quantization.
+Tests both accuracy and performance.
+"""
+
+import time
+from typing import Tuple
+
+import pytest
+import torch
+
+from sglang.srt.layers.attention.nsa.tilelang_kernel import act_quant
+from sglang.srt.layers.attention.nsa.triton_kernel import act_quant as act_quant_triton
+
+
+def benchmark_kernel(
+    fn,
+    x: torch.Tensor,
+    block_size: int,
+    scale_fmt,
+    warmup: int = 10,
+    repeat: int = 100,
+    use_cuda_graph: bool = True,
+) -> Tuple[float, torch.Tensor, torch.Tensor]:
+    """
+    Benchmark a kernel function.
+
+    Args:
+        fn: Function to benchmark
+        x: Input tensor
+        block_size: Block size for quantization
+        scale_fmt: Scale format
+        warmup: Number of warmup iterations
+        repeat: Number of repeat iterations
+        use_cuda_graph: Whether to use CUDA graphs for more accurate timing
+
+    Returns:
+        Tuple of (avg_time_ms, quantized_output, scales)
+    """
+    # Warmup
+    for _ in range(warmup):
+        y, s = fn(x, block_size=block_size, scale_fmt=scale_fmt)
+
+    if not x.is_cuda or not use_cuda_graph:
+        # Fallback to regular timing
+        if x.is_cuda:
+            torch.cuda.synchronize()
+
+        start = time.perf_counter()
+        for _ in range(repeat):
+            y, s = fn(x, block_size=block_size, scale_fmt=scale_fmt)
+
+        if x.is_cuda:
+            torch.cuda.synchronize()
+
+        end = time.perf_counter()
+        avg_time_ms = (end - start) / repeat * 1000
+
+        return avg_time_ms, y, s
+
+    # Use CUDA graph for more accurate timing
+    torch.cuda.synchronize()
+
+    # Allocate output buffers
+    N = x.size(-1)
+    y = torch.empty_like(x, dtype=torch.float8_e4m3fn)
+    s = x.new_empty(*x.size()[:-1], N // block_size, dtype=torch.float32)
+
+    # Capture CUDA graph
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        y_cap, s_cap = fn(x, block_size=block_size, scale_fmt=scale_fmt)
+
+    # Warmup with graph
+    for _ in range(warmup):
+        graph.replay()
+
+    torch.cuda.synchronize()
+
+    # Timing with CUDA graph
+    start_event = torch.cuda.Event(enable_timing=True)
+    end_event = torch.cuda.Event(enable_timing=True)
+
+    start_event.record()
+    for _ in range(repeat):
+        graph.replay()
+    end_event.record()
+
+    torch.cuda.synchronize()
+
+    avg_time_ms = start_event.elapsed_time(end_event) / repeat
+
+    return avg_time_ms, y_cap, s_cap
+
+
+def check_accuracy(
+    y_ref: torch.Tensor,
+    s_ref: torch.Tensor,
+    y_test: torch.Tensor,
+    s_test: torch.Tensor,
+    rtol: float = 1e-2,
+    atol: float = 1e-2,
+) -> Tuple[bool, dict]:
+    """
+    Check accuracy between reference and test outputs.
+
+    Args:
+        y_ref: Reference quantized output
+        s_ref: Reference scales
+        y_test: Test quantized output
+        s_test: Test scales
+        rtol: Relative tolerance
+        atol: Absolute tolerance
+
+    Returns:
+        Tuple of (passed, metrics_dict)
+    """
+    # Convert FP8 to float for comparison
+    y_ref_float = y_ref.float()
+    y_test_float = y_test.float()
+
+    # Compute differences
+    y_diff = torch.abs(y_ref_float - y_test_float)
+    s_diff = torch.abs(s_ref - s_test)
+
+    # Compute metrics
+    y_max_diff = y_diff.max().item()
+    y_mean_diff = y_diff.mean().item()
+    s_max_diff = s_diff.max().item()
+    s_mean_diff = s_diff.mean().item()
+
+    # Check relative and absolute tolerance
+    y_close = torch.allclose(y_ref_float, y_test_float, rtol=rtol, atol=atol)
+    s_close = torch.allclose(s_ref, s_test, rtol=rtol, atol=atol)
+
+    # Compute percentage of matching elements
+    y_match_pct = (y_ref_float == y_test_float).float().mean().item() * 100
+
+    metrics = {
+        "y_max_diff": y_max_diff,
+        "y_mean_diff": y_mean_diff,
+        "y_match_pct": y_match_pct,
+        "s_max_diff": s_max_diff,
+        "s_mean_diff": s_mean_diff,
+        "y_close": y_close,
+        "s_close": s_close,
+    }
+
+    passed = y_close and s_close
+
+    return passed, metrics
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_act_quant_comprehensive_benchmark(scale_fmt=None):
+    """Comprehensive benchmark across multiple sizes with CUDA graphs."""
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+    block_size = 128
+
+    shapes = [
+        (128, 512),
+        (256, 1024),
+        (512, 2048),
+        (1024, 4096),
+        (2048, 8192),
+        (4096, 16384),
+    ]
+
+    print("\n" + "=" * 100)
+    print("Comprehensive Performance Benchmark with CUDA Graphs")
+    print("=" * 100)
+    print(
+        f"{'Shape':<20} {'TileLang (ms)':<15} {'Triton (ms)':<15} {'Speedup':<10} {'Status'}"
+    )
+    print("-" * 100)
+
+    for shape in shapes:
+        torch.manual_seed(42)
+        x = torch.randn(shape, dtype=dtype, device=device)
+
+        try:
+            # Benchmark both with CUDA graphs
+            time_tilelang, y_ref, s_ref = benchmark_kernel(
+                act_quant,
+                x,
+                block_size,
+                scale_fmt,
+                warmup=5,
+                repeat=50,
+                use_cuda_graph=True,
+            )
+            time_triton, y_triton, s_triton = benchmark_kernel(
+                act_quant_triton,
+                x,
+                block_size,
+                scale_fmt,
+                warmup=5,
+                repeat=50,
+                use_cuda_graph=True,
+            )
+
+            # Check accuracy
+            passed, _ = check_accuracy(y_ref, s_ref, y_triton, s_triton)
+
+            speedup = time_tilelang / time_triton if time_triton > 0 else 0
+            status = "✓ PASS" if passed else "✗ FAIL"
+
+            print(
+                f"{str(shape):<20} {time_tilelang:<15.4f} {time_triton:<15.4f} "
+                f"{speedup:<10.2f} {status}"
+            )
+        except Exception as e:
+            print(f"{str(shape):<20} ERROR: {str(e)}")
+
+    print("=" * 100)
+
+    # Also run without CUDA graphs for comparison
+    print("\n" + "=" * 100)
+    print("Performance Benchmark WITHOUT CUDA Graphs (for comparison)")
+    print("=" * 100)
+    print(
+        f"{'Shape':<20} {'TileLang (ms)':<15} {'Triton (ms)':<15} {'Speedup':<10} {'Status'}"
+    )
+    print("-" * 100)
+
+    for shape in shapes:
+        torch.manual_seed(42)
+        x = torch.randn(shape, dtype=dtype, device=device)
+
+        try:
+            # Benchmark both without CUDA graphs
+            time_tilelang, y_ref, s_ref = benchmark_kernel(
+                act_quant,
+                x,
+                block_size,
+                scale_fmt,
+                warmup=5,
+                repeat=50,
+                use_cuda_graph=False,
+            )
+            time_triton, y_triton, s_triton = benchmark_kernel(
+                act_quant_triton,
+                x,
+                block_size,
+                scale_fmt,
+                warmup=5,
+                repeat=50,
+                use_cuda_graph=False,
+            )
+
+            # Check accuracy
+            passed, _ = check_accuracy(y_ref, s_ref, y_triton, s_triton)
+
+            speedup = time_tilelang / time_triton if time_triton > 0 else 0
+            status = "✓ PASS" if passed else "✗ FAIL"
+
+            print(
+                f"{str(shape):<20} {time_tilelang:<15.4f} {time_triton:<15.4f} "
+                f"{speedup:<10.2f} {status}"
+            )
+        except Exception as e:
+            print(f"{str(shape):<20} ERROR: {str(e)}")
+
+    print("=" * 100)
+
+
+if __name__ == "__main__":
+    # Run comprehensive benchmark
+    if torch.cuda.is_available():
+        print("\n" + "=" * 80)
+        print("Running Comprehensive Benchmark with scale_fmt=None")
+        print("=" * 80)
+        test_act_quant_comprehensive_benchmark(scale_fmt=None)
+
+        print("\n" + "=" * 80)
+        print("Running Comprehensive Benchmark with scale_fmt!=None")
+        print("=" * 80)
+        test_act_quant_comprehensive_benchmark(scale_fmt="any")
+    else:
+        print("CUDA not available. Skipping tests.")

third_party/sglang/test/manual/layers/attention/nsa/test_get_k_scale_triton_kernel.py

@@ -0,0 +1,191 @@
+import torch
+
+from sglang.srt.layers.attention.nsa.index_buf_accessor import (
+    _get_k_and_s_triton_kernel,
+)
+
+
+def golden_torch_gen(
+    seq_len_tensor: torch.Tensor,
+    buffer_indexer: torch.Tensor,
+    buffer: torch.Tensor,
+    index_head_dim,
+    page_size,
+):
+    dim_split = page_size * index_head_dim
+    torch_k_out = buffer[:, 0:dim_split]
+    torch_s_out = buffer[:, dim_split:]
+
+    torch_k_out = torch_k_out.reshape(-1, 128)
+    torch_s_out = torch_s_out.reshape(-1, 4)
+
+    batch = seq_len_tensor.shape[0]
+    index_list = []
+    for i in range(batch):
+        seq_len = seq_len_tensor[i].item()
+        buffer_index_ = buffer_indexer[i]
+        align_seq_len = ((seq_len + page_size - 1) / page_size) * page_size
+        needed_block_num = int((seq_len + page_size - 1) / page_size)
+        for j in range(needed_block_num):
+            block_idx = buffer_index_[j].item()
+            start_idx = block_idx * page_size
+            end_idx = 0
+            if j == (needed_block_num - 1):
+                end_idx = block_idx * page_size + (
+                    seq_len - (needed_block_num - 1) * page_size
+                )
+            else:
+                end_idx = (block_idx + 1) * page_size
+
+            index_tensor = (
+                torch.arange(start=start_idx, end=end_idx, step=1)
+                .type(torch.int32)
+                .cuda()
+            )
+            index_list.append(index_tensor)
+
+    index_list_ = torch.cat(index_list, dim=0)
+    torch_k_out = torch.index_select(torch_k_out, dim=0, index=index_list_)
+    torch_s_out = torch.index_select(torch_s_out, dim=0, index=index_list_)
+
+    return torch_k_out, torch_s_out
+
+
+def get_k_and_s_triton():
+    index_head_dim = 128
+    page_size = 64
+    num_page = 128
+    s_offset_in_page = page_size * index_head_dim
+
+    seq_len_tensor = torch.tensor(
+        [256, 267, 215, 32, 129], dtype=torch.int64, device="cuda"
+    )  # 4 + 5 + 3 + 1 + 3 block
+    buffer_indexer = torch.tensor(
+        [
+            [1, 2, 3, 4, 0],
+            [7, 6, 5, 8, 9],
+            [10, 11, 12, 0, 0],
+            [13, 0, 0, 0, 0],
+            [14, 15, 16, 0, 0],
+        ],
+        dtype=torch.int32,
+        device="cuda",
+    )
+    seq_len_sum = seq_len_tensor.sum()
+    batch = seq_len_tensor.shape[0]
+
+    triton_k_out = torch.empty(
+        (seq_len_sum, index_head_dim), dtype=torch.uint8, device="cuda"
+    )
+    triton_s_out = torch.empty((seq_len_sum, 4), dtype=torch.uint8, device="cuda")
+    buffer = torch.randint(
+        0,
+        num_page,
+        (num_page, page_size * index_head_dim + page_size * 4),
+        device="cuda",
+    ).type(torch.uint8)
+
+    _, buf_numel_per_page = buffer.shape
+    _, page_indice_batch_offset = buffer_indexer.shape
+    max_seq_len = seq_len_tensor.max().item()
+
+    BLOCK_SIZE = 256
+    BLOCK_SIZE_K = 128
+
+    num_token_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
+    num_k_threads = (index_head_dim + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
+
+    grid = (batch, num_token_blocks, num_k_threads)
+    seq_num_pow2 = 1
+    while seq_num_pow2 < batch:
+        seq_num_pow2 *= 2
+
+    # acc test =====================
+    _get_k_and_s_triton_kernel[grid](
+        buf_ptr=buffer,
+        page_indices_ptr=buffer_indexer,
+        k_out_ptr=triton_k_out,
+        s_out_ptr=triton_s_out,
+        seq_len_ptr=seq_len_tensor,
+        seq_len_num_pow=seq_num_pow2,
+        page_size=page_size,
+        buf_numel_per_page=buf_numel_per_page,
+        index_head_dim=index_head_dim,
+        s_offset_in_page=s_offset_in_page,
+        page_indice_batch_offset=page_indice_batch_offset,
+        BLOCK_SIZE=BLOCK_SIZE,
+        BLOCK_SIZE_K=BLOCK_SIZE_K,
+    )
+
+    torch_k_out, torch_s_out = golden_torch_gen(
+        seq_len_tensor=seq_len_tensor,
+        buffer_indexer=buffer_indexer,
+        buffer=buffer,
+        index_head_dim=index_head_dim,
+        page_size=page_size,
+    )
+
+    torch.testing.assert_close(
+        triton_k_out, torch_k_out, rtol=0, atol=0, msg="k outputs differ!"
+    )
+    torch.testing.assert_close(
+        triton_s_out, torch_s_out, rtol=0, atol=0, msg="s outputs differ!"
+    )
+    print("_get_k_and_s_triton_kernel test pass")
+
+    # perf test =====================
+    import time
+
+    torch.cuda.synchronize()
+    for _ in range(10):
+        _get_k_and_s_triton_kernel[grid](
+            buf_ptr=buffer,
+            page_indices_ptr=buffer_indexer,
+            k_out_ptr=triton_k_out,
+            s_out_ptr=triton_s_out,
+            seq_len_ptr=seq_len_tensor,
+            seq_len_num_pow=seq_num_pow2,
+            page_size=page_size,
+            buf_numel_per_page=buf_numel_per_page,
+            index_head_dim=index_head_dim,
+            s_offset_in_page=s_offset_in_page,
+            page_indice_batch_offset=page_indice_batch_offset,
+            BLOCK_SIZE=BLOCK_SIZE,
+            BLOCK_SIZE_K=BLOCK_SIZE_K,
+        )
+
+    torch.cuda.synchronize()
+    start_time = time.perf_counter()
+
+    _get_k_and_s_triton_kernel[grid](
+        buf_ptr=buffer,
+        page_indices_ptr=buffer_indexer,
+        k_out_ptr=triton_k_out,
+        s_out_ptr=triton_s_out,
+        seq_len_ptr=seq_len_tensor,
+        seq_len_num_pow=seq_num_pow2,
+        page_size=page_size,
+        buf_numel_per_page=buf_numel_per_page,
+        index_head_dim=index_head_dim,
+        s_offset_in_page=s_offset_in_page,
+        page_indice_batch_offset=page_indice_batch_offset,
+        BLOCK_SIZE=BLOCK_SIZE,
+        BLOCK_SIZE_K=BLOCK_SIZE_K,
+    )
+
+    end_time = time.perf_counter()
+    print(
+        f"_get_k_and_s_triton_kernel triton kernel infer time is {((end_time-start_time)*1000):.4f} ms\n"
+    )
+
+
+if __name__ == "__main__":
+    if not torch.cuda.is_available():
+        print("CUDA not available. Skipping tests.")
+        exit(0)
+
+    print("Start test cases...\n")
+
+    get_k_and_s_triton()
+
+    print("End test cases...\n")

third_party/sglang/test/manual/layers/attention/nsa/test_index_buf_accessor.py

@@ -0,0 +1,591 @@
+"""
+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)
+        seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
+
+        # 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
+        )
+        page_indices_ = page_indices.unsqueeze(0)
+
+        # 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, page_indices_, seq_len_tensor, seq_len, seq_len
+        )
+
+        # 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
+        seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
+
+        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)
+        page_indices_ = page_indices.unsqueeze(0)
+
+        # 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, page_indices_, seq_len_tensor, seq_len, seq_len
+        )
+
+        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
+        seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
+
+        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)
+        page_indices_ = page_indices.unsqueeze(0)
+
+        # 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, page_indices_, seq_len_tensor, seq_len, seq_len
+        )
+
+        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
+        seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
+
+        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)
+        page_indices_ = page_indices.unsqueeze(0)
+
+        # 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, page_indices_, seq_len_tensor, seq_len, seq_len
+        )
+
+        # 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
+        seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
+
+        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)
+        page_indices_ = page_indices.unsqueeze(0)
+
+        # 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, page_indices_, seq_len_tensor, seq_len, seq_len
+        )
+        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
+        seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
+
+        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)
+        page_indices_ = page_indices.unsqueeze(0)
+
+        # 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, page_indices_, seq_len_tensor, seq_len, seq_len
+        )
+        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
+        seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
+
+        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
+        )
+        page_indices_ = page_indices.unsqueeze(0)
+
+        # 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, page_indices_, seq_len_tensor, seq_len, seq_len
+        )
+        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 SeqLen=256...")
+    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 SeqLen=256 tests passed\n")
+
+    print("Testing GetKAndS SeqLen=128K...")
+    test_get_k_and_s = TestGetKAndS()
+    test_get_k_and_s.test_get_k_and_s_correctness(
+        num_pages=2048, seq_len=131072, 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 SeqLen=128K 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)