[Kernel] Skip reserved writes in MLA KV cache (#36003)

Co-authored-by: Khoa Pham <khoa.pham@radixark.ai>

python/sglang/kernels/jit/csrc/elementwise/set_mla_kv_buffer.cuh

@@ -46,6 +46,7 @@ struct SetMlaKVBufferParams {
   int64_t stride_rope_bytes;
   int64_t stride_buffer_bytes;
   uint32_t batch_size;
+  int64_t reserved_skip_index;
 };
 
 template <int64_t kNopeBytes, int64_t kRopeBytes, int kNumWarps, bool kUsePDL, typename TLoc>
@@ -81,7 +82,7 @@ __global__ void set_mla_kv_buffer_kernel(const __grid_constant__ SetMlaKVBufferP
   asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
 
   // Lane 0 issues one bulk store from the smem slot to the scattered gmem row.
-  if (threadIdx.x % kWarpThreads == 0) {
+  if (threadIdx.x % kWarpThreads == 0 && loc != params.reserved_skip_index) {
     cuda::ptx::cp_async_bulk(
         cuda::ptx::space_global,
         cuda::ptx::space_shared,
@@ -114,7 +115,8 @@ struct SetMlaKVBufferKernel {
       tvm::ffi::TensorView loc,
       tvm::ffi::TensorView k_nope,
       tvm::ffi::TensorView k_rope,
-      int64_t num_warps_per_block) {
+      int64_t num_warps_per_block,
+      int64_t reserved_skip_index) {
     using namespace host;
 
     auto B = SymbolicSize{"batch_size"};
@@ -182,6 +184,7 @@ struct SetMlaKVBufferKernel {
         .stride_rope_bytes = S_rope.unwrap() * dtype_size,
         .stride_buffer_bytes = S_buf.unwrap() * dtype_size,
         .batch_size = batch,
+        .reserved_skip_index = reserved_skip_index,
     };
 
     const auto use_int32 = loc_dtype.is_type<int32_t>();

python/sglang/kernels/ops/kvcache/mla_buffer.py

@@ -14,6 +14,7 @@ def set_mla_kv_buffer_kernel(
     cache_k_nope_ptr,
     cache_k_rope_ptr,
     loc_ptr,
+    reserved_skip_index,
     buffer_stride: tl.constexpr,
     nope_stride: tl.constexpr,
     rope_stride: tl.constexpr,
@@ -36,7 +37,7 @@ def set_mla_kv_buffer_kernel(
         tl.extra.cuda.gdc_wait()
 
     loc = tl.load(loc_ptr + pid_loc).to(tl.int64)
-    is_valid = loc % DCP_WORLD_SIZE == DCP_RANK
+    is_valid = (loc != reserved_skip_index) & (loc % DCP_WORLD_SIZE == DCP_RANK)
     safe_loc = tl.where(is_valid, loc, 0)
     safe_loc = safe_loc // DCP_WORLD_SIZE
     dst_ptr = kv_buffer_ptr + safe_loc * buffer_stride + offs
@@ -92,6 +93,8 @@ def set_mla_kv_buffer_triton(
     loc: torch.Tensor,
     cache_k_nope: torch.Tensor,
     cache_k_rope: torch.Tensor,
+    *,
+    reserved_skip_index: int = 0,
 ):
     """Dispatch MLA paged-KV scatter writes to the fastest available path.
 
@@ -115,6 +118,9 @@ def set_mla_kv_buffer_triton(
 
     Name retained for caller compatibility; the implementation is no longer
     Triton-only.
+
+    Writes targeting ``reserved_skip_index`` are skipped. Slot 0 is reserved
+    for CUDA-graph padding by default; pass -1 to disable skipping.
     """
     from sglang.kernels.ops.kvcache.set_mla_kv_buffer import (
         can_use_set_mla_kv_buffer,
@@ -132,7 +138,13 @@ def set_mla_kv_buffer_triton(
         and can_use_set_mla_kv_buffer(nope_bytes, rope_bytes)
         and not get_parallel().dcp_enabled
     ):
-        jit_set_mla_kv_buffer(kv_buffer, loc, cache_k_nope, cache_k_rope)
+        jit_set_mla_kv_buffer(
+            kv_buffer,
+            loc,
+            cache_k_nope,
+            cache_k_rope,
+            reserved_skip_index=reserved_skip_index,
+        )
         return
 
     # Fallback: Triton with BLOCK = next_pow2(total_dim). One CTA per loc; the
@@ -151,6 +163,7 @@ def set_mla_kv_buffer_triton(
         cache_k_nope,
         cache_k_rope,
         loc,
+        reserved_skip_index,
         kv_buffer.stride(0),
         cache_k_nope.stride(0),
         cache_k_rope.stride(0),
@@ -169,6 +182,7 @@ def set_mla_kv_buffer_fp8_quant_kernel(
     cache_k_nope_ptr,
     cache_k_rope_ptr,
     loc_ptr,
+    reserved_skip_index,
     buffer_stride: tl.constexpr,
     nope_stride: tl.constexpr,
     rope_stride: tl.constexpr,
@@ -190,7 +204,9 @@ def set_mla_kv_buffer_fp8_quant_kernel(
         tl.extra.cuda.gdc_wait()
 
     loc = tl.load(loc_ptr + pid_loc).to(tl.int64)
-    dst_ptr = kv_buffer_fp8_ptr + loc * buffer_stride + offs
+    is_valid = loc != reserved_skip_index
+    safe_loc = tl.where(is_valid, loc, 0)
+    dst_ptr = kv_buffer_fp8_ptr + safe_loc * buffer_stride + offs
 
     if base + BLOCK <= nope_dim:
         src = tl.load(
@@ -220,7 +236,7 @@ def set_mla_kv_buffer_fp8_quant_kernel(
         src = tl.where(is_nope, src_nope, src_rope)
 
     # Destination pointer is FP8-typed view; tl.store performs downcast.
-    tl.store(dst_ptr, src, mask=mask)
+    tl.store(dst_ptr, src, mask=mask & is_valid)
 
     if USE_GDC:
         tl.extra.cuda.gdc_launch_dependents()
@@ -232,8 +248,13 @@ def set_mla_kv_buffer_triton_fp8_quant(
     cache_k_nope: torch.Tensor,
     cache_k_rope: torch.Tensor,
     fp8_dtype: torch.dtype,
+    *,
+    reserved_skip_index: int = 0,
 ):
-    """Fuse BF16/FP16 MLA K quantization with paged KV write."""
+    """Fuse BF16/FP16 MLA K quantization with paged KV write.
+
+    Writes targeting ``reserved_skip_index`` are skipped. Pass -1 to disable.
+    """
     kv_buffer_fp8 = kv_buffer.view(fp8_dtype)
 
     nope_dim = cache_k_nope.shape[-1]
@@ -250,6 +271,7 @@ def set_mla_kv_buffer_triton_fp8_quant(
         cache_k_nope,
         cache_k_rope,
         loc,
+        reserved_skip_index,
         kv_buffer_fp8.stride(0),
         cache_k_nope.stride(0),
         cache_k_rope.stride(0),
@@ -266,6 +288,7 @@ def set_mla_kv_scale_buffer_kernel(
     cache_k_nope_ptr,
     cache_k_rope_ptr,
     loc_ptr,
+    reserved_skip_index,
     buffer_stride: tl.constexpr,
     nope_stride: tl.constexpr,
     rope_stride: tl.constexpr,
@@ -282,7 +305,9 @@ def set_mla_kv_scale_buffer_kernel(
     mask = offs < total_dim  # Make sure don't cross the boundary
 
     loc = tl.load(loc_ptr + pid_loc)
-    dst_ptr = kv_buffer_ptr + loc * buffer_stride + offs
+    is_valid = loc != reserved_skip_index
+    safe_loc = tl.where(is_valid, loc, 0)
+    dst_ptr = kv_buffer_ptr + safe_loc * buffer_stride + offs
 
     # Check each offs should read 'nope' or 'rope'
     is_nope = offs < nope_dim
@@ -297,7 +322,7 @@ def set_mla_kv_scale_buffer_kernel(
 
     # Combine nope + rope
     src = src_nope + src_rope
-    tl.store(dst_ptr, src, mask=mask)
+    tl.store(dst_ptr, src, mask=mask & is_valid)
 
 
 def set_mla_kv_scale_buffer_triton(
@@ -305,7 +330,10 @@ def set_mla_kv_scale_buffer_triton(
     loc: torch.Tensor,
     cache_k_nope: torch.Tensor,
     cache_k_rope: torch.Tensor,
+    *,
+    reserved_skip_index: int = 0,
 ):
+    """Write MLA scale rows while preserving the reserved padding slot."""
     nope_dim = cache_k_nope.shape[-1]
     rope_dim = cache_k_rope.shape[-1]
     total_dim = nope_dim + rope_dim
@@ -318,6 +346,7 @@ def set_mla_kv_scale_buffer_triton(
         cache_k_nope,
         cache_k_rope,
         loc,
+        reserved_skip_index,
         kv_buffer.stride(0),
         cache_k_nope.stride(0),
         cache_k_rope.stride(0),

python/sglang/kernels/ops/kvcache/set_mla_kv_buffer.py

@@ -89,6 +89,8 @@ def set_mla_kv_buffer(
     cache_k_nope: torch.Tensor,
     cache_k_rope: torch.Tensor,
     num_warps: int = 0,
+    *,
+    reserved_skip_index: int = 0,
 ) -> None:
     """Write packed [k_nope | k_rope] rows into ``kv_buffer`` at ``loc`` indices
     via a TMA bulk-store. SM90+ only — the caller is expected to gate.
@@ -99,6 +101,9 @@ def set_mla_kv_buffer(
         cache_k_nope: [n_loc, nope_dim] or [n_loc, 1, nope_dim]
         cache_k_rope: [n_loc, rope_dim] or [n_loc, 1, rope_dim]
         loc:          [n_loc]
+
+    Writes targeting ``reserved_skip_index`` are skipped. Slot 0 is reserved
+    for CUDA-graph padding by default; pass -1 to disable skipping.
     """
     n_loc = loc.shape[0]
     if n_loc == 0:
@@ -114,4 +119,11 @@ def set_mla_kv_buffer(
         num_warps = _pick_num_warps(n_loc)
 
     module = set_mla_kv_buffer_module(nope_bytes, rope_bytes, is_arch_support_pdl())
-    module.set_mla_kv_buffer(buf, loc, src_nope, src_rope, num_warps)
+    module.set_mla_kv_buffer(
+        buf,
+        loc,
+        src_nope,
+        src_rope,
+        num_warps,
+        reserved_skip_index,
+    )

test/registered/kernels/benchmark/kvcache/bench_set_mla_kv_buffer.py

@@ -44,6 +44,7 @@ def _triton_baseline(kv_buffer, loc, cache_k_nope, cache_k_rope):
         cache_k_nope,
         cache_k_rope,
         loc,
+        0,
         kv_buffer.stride(0),
         cache_k_nope.stride(0),
         cache_k_rope.stride(0),

test/registered/kernels/ops/kvcache/test_set_mla_kv_buffer.py

@@ -4,16 +4,28 @@ import pytest
 import torch
 
 from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.kvcache.mla_buffer import (
+    set_mla_kv_buffer_triton,
+    set_mla_kv_buffer_triton_fp8_quant,
+    set_mla_kv_scale_buffer_triton,
+)
 from sglang.kernels.ops.kvcache.set_mla_kv_buffer import (
     can_use_set_mla_kv_buffer,
     set_mla_kv_buffer,
 )
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_amd_ci(est_time=15, stage="jit-kernel-unit", runner_config="amd")
 
 DEVICE = "cuda"
 CACHE_SIZE = 4096
+TRITON_NOPE_DIM = 128
+TRITON_ROPE_DIM = 64
+CUDA_TMA_ONLY = pytest.mark.skipif(
+    torch.version.hip is not None,
+    reason="The TMA bulk-store kernel requires CUDA SM90+",
+)
 
 # (nope_dim, rope_dim) pairs: standard MLA, MLA scale buffer, FP8 nope-extended layout.
 SHAPES = get_ci_test_range(
@@ -36,6 +48,7 @@ def _ref(kv_buffer, loc, cache_k_nope, cache_k_rope):
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 @pytest.mark.parametrize("shape", SHAPES)
 @pytest.mark.parametrize("batch_size", BATCH_SIZES)
+@CUDA_TMA_ONLY
 def test_set_mla_kv_buffer_correctness(dtype, shape, batch_size):
     nope_dim, rope_dim = shape
     total_dim = nope_dim + rope_dim
@@ -45,7 +58,7 @@ def test_set_mla_kv_buffer_correctness(dtype, shape, batch_size):
     kv_buffer = torch.randn((CACHE_SIZE, 1, total_dim), dtype=dtype, device=DEVICE)
     kv_ref = kv_buffer.clone()
 
-    loc = torch.randperm(CACHE_SIZE, device=DEVICE)[:batch_size]
+    loc = torch.randperm(CACHE_SIZE - 1, device=DEVICE)[:batch_size] + 1
 
     set_mla_kv_buffer(kv_buffer, loc, cache_k_nope, cache_k_rope)
     _ref(kv_ref, loc, cache_k_nope, cache_k_rope)
@@ -54,6 +67,7 @@ def test_set_mla_kv_buffer_correctness(dtype, shape, batch_size):
 
 
 @pytest.mark.parametrize("loc_dtype", [torch.int32, torch.int64])
+@CUDA_TMA_ONLY
 def test_set_mla_kv_buffer_loc_dtypes(loc_dtype):
     nope_dim, rope_dim = 512, 64
     batch_size = 128
@@ -66,7 +80,7 @@ def test_set_mla_kv_buffer_loc_dtypes(loc_dtype):
     )
     kv_ref = kv_buffer.clone()
 
-    loc = torch.randperm(CACHE_SIZE, device=DEVICE)[:batch_size].to(loc_dtype)
+    loc = (torch.randperm(CACHE_SIZE - 1, device=DEVICE)[:batch_size] + 1).to(loc_dtype)
 
     set_mla_kv_buffer(kv_buffer, loc, cache_k_nope, cache_k_rope)
     _ref(kv_ref, loc, cache_k_nope, cache_k_rope)
@@ -74,6 +88,7 @@ def test_set_mla_kv_buffer_loc_dtypes(loc_dtype):
     assert torch.equal(kv_buffer, kv_ref)
 
 
+@CUDA_TMA_ONLY
 def test_set_mla_kv_buffer_uint8_byte_layout():
     """FP8 DSA byte-layout: cache_k_nope is uint8 with [fp8(512) | scales(16)] = 528,
     cache_k_rope is uint8 [128]; total payload = 656 bytes."""
@@ -92,7 +107,7 @@ def test_set_mla_kv_buffer_uint8_byte_layout():
     )
     kv_ref = kv_buffer.clone()
 
-    loc = torch.randperm(CACHE_SIZE, device=DEVICE)[:batch_size]
+    loc = torch.randperm(CACHE_SIZE - 1, device=DEVICE)[:batch_size] + 1
 
     set_mla_kv_buffer(kv_buffer, loc, cache_k_nope, cache_k_rope)
     _ref(kv_ref, loc, cache_k_nope, cache_k_rope)
@@ -100,6 +115,7 @@ def test_set_mla_kv_buffer_uint8_byte_layout():
     assert torch.equal(kv_buffer, kv_ref)
 
 
+@CUDA_TMA_ONLY
 def test_set_mla_kv_buffer_empty_loc():
     nope_dim, rope_dim = 512, 64
     dtype = torch.bfloat16
@@ -116,6 +132,185 @@ def test_set_mla_kv_buffer_empty_loc():
     assert torch.equal(kv_buffer, kv_before)
 
 
+@pytest.mark.parametrize("loc_dtype", [torch.int32, torch.int64])
+@CUDA_TMA_ONLY
+def test_set_mla_kv_buffer_reserved_skip_index(loc_dtype):
+    nope_dim, rope_dim = 512, 64
+    dtype = torch.bfloat16
+    cache_k_nope = torch.randn((4, 1, nope_dim), dtype=dtype, device=DEVICE)
+    cache_k_rope = torch.randn((4, 1, rope_dim), dtype=dtype, device=DEVICE)
+    cache_k_nope[[0, 2]] = torch.nan
+    cache_k_rope[[0, 2]] = torch.nan
+    kv_buffer = torch.randn(
+        (CACHE_SIZE, 1, nope_dim + rope_dim), dtype=dtype, device=DEVICE
+    )
+    reserved_before = kv_buffer[0].clone()
+    loc = torch.tensor([0, 7, 0, 9], dtype=loc_dtype, device=DEVICE)
+
+    set_mla_kv_buffer(kv_buffer, loc, cache_k_nope, cache_k_rope)
+
+    torch.testing.assert_close(kv_buffer[0], reserved_before, rtol=0.0, atol=0.0)
+    torch.testing.assert_close(
+        kv_buffer[7, 0, :nope_dim], cache_k_nope[1, 0], rtol=0.0, atol=0.0
+    )
+    torch.testing.assert_close(
+        kv_buffer[7, 0, nope_dim:], cache_k_rope[1, 0], rtol=0.0, atol=0.0
+    )
+    torch.testing.assert_close(
+        kv_buffer[9, 0, :nope_dim], cache_k_nope[3, 0], rtol=0.0, atol=0.0
+    )
+    torch.testing.assert_close(
+        kv_buffer[9, 0, nope_dim:], cache_k_rope[3, 0], rtol=0.0, atol=0.0
+    )
+
+
+@CUDA_TMA_ONLY
+def test_set_mla_kv_buffer_zero_index_can_be_written_when_skip_disabled():
+    nope_dim, rope_dim = 512, 64
+    dtype = torch.bfloat16
+    cache_k_nope = torch.randn((1, 1, nope_dim), dtype=dtype, device=DEVICE)
+    cache_k_rope = torch.randn((1, 1, rope_dim), dtype=dtype, device=DEVICE)
+    kv_buffer = torch.randn(
+        (CACHE_SIZE, 1, nope_dim + rope_dim), dtype=dtype, device=DEVICE
+    )
+    loc = torch.zeros(1, dtype=torch.int64, device=DEVICE)
+
+    set_mla_kv_buffer(
+        kv_buffer,
+        loc,
+        cache_k_nope,
+        cache_k_rope,
+        reserved_skip_index=-1,
+    )
+
+    torch.testing.assert_close(
+        kv_buffer[0, 0, :nope_dim], cache_k_nope[0, 0], rtol=0.0, atol=0.0
+    )
+    torch.testing.assert_close(
+        kv_buffer[0, 0, nope_dim:], cache_k_rope[0, 0], rtol=0.0, atol=0.0
+    )
+
+
+@pytest.mark.parametrize("loc_dtype", [torch.int32, torch.int64])
+def test_set_mla_kv_buffer_triton_reserved_skip_index(loc_dtype):
+    dtype = torch.bfloat16
+    cache_k_nope = torch.randn((4, 1, TRITON_NOPE_DIM), dtype=dtype, device=DEVICE)
+    cache_k_rope = torch.randn((4, 1, TRITON_ROPE_DIM), dtype=dtype, device=DEVICE)
+    cache_k_nope[[0, 2]] = torch.nan
+    cache_k_rope[[0, 2]] = torch.nan
+    kv_buffer = torch.randn(
+        (CACHE_SIZE, 1, TRITON_NOPE_DIM + TRITON_ROPE_DIM),
+        dtype=dtype,
+        device=DEVICE,
+    )
+    reserved_before = kv_buffer[0].clone()
+    loc = torch.tensor([0, 7, 0, 9], dtype=loc_dtype, device=DEVICE)
+
+    set_mla_kv_buffer_triton(kv_buffer, loc, cache_k_nope, cache_k_rope)
+
+    torch.testing.assert_close(kv_buffer[0], reserved_before, rtol=0.0, atol=0.0)
+    torch.testing.assert_close(
+        kv_buffer[7, 0],
+        torch.cat((cache_k_nope[1, 0], cache_k_rope[1, 0])),
+        rtol=0.0,
+        atol=0.0,
+    )
+    torch.testing.assert_close(
+        kv_buffer[9, 0],
+        torch.cat((cache_k_nope[3, 0], cache_k_rope[3, 0])),
+        rtol=0.0,
+        atol=0.0,
+    )
+
+
+def test_set_mla_kv_buffer_triton_zero_index_can_be_written_when_skip_disabled():
+    dtype = torch.bfloat16
+    cache_k_nope = torch.randn((1, 1, TRITON_NOPE_DIM), dtype=dtype, device=DEVICE)
+    cache_k_rope = torch.randn((1, 1, TRITON_ROPE_DIM), dtype=dtype, device=DEVICE)
+    kv_buffer = torch.randn(
+        (CACHE_SIZE, 1, TRITON_NOPE_DIM + TRITON_ROPE_DIM),
+        dtype=dtype,
+        device=DEVICE,
+    )
+    loc = torch.zeros(1, dtype=torch.int64, device=DEVICE)
+
+    set_mla_kv_buffer_triton(
+        kv_buffer,
+        loc,
+        cache_k_nope,
+        cache_k_rope,
+        reserved_skip_index=-1,
+    )
+
+    torch.testing.assert_close(
+        kv_buffer[0, 0],
+        torch.cat((cache_k_nope[0, 0], cache_k_rope[0, 0])),
+        rtol=0.0,
+        atol=0.0,
+    )
+
+
+def test_set_mla_kv_buffer_triton_fp8_quant_reserved_skip_index():
+    fp8_dtype = torch.float8_e4m3fnuz if torch.version.hip else torch.float8_e4m3fn
+    cache_k_nope = torch.randn(
+        (4, 1, TRITON_NOPE_DIM), dtype=torch.bfloat16, device=DEVICE
+    )
+    cache_k_rope = torch.randn(
+        (4, 1, TRITON_ROPE_DIM), dtype=torch.bfloat16, device=DEVICE
+    )
+    cache_k_nope[[0, 2]] = torch.nan
+    cache_k_rope[[0, 2]] = torch.nan
+    kv_buffer = torch.randint(
+        0,
+        256,
+        (CACHE_SIZE, 1, TRITON_NOPE_DIM + TRITON_ROPE_DIM),
+        dtype=torch.uint8,
+        device=DEVICE,
+    )
+    reserved_before = kv_buffer[0].clone()
+    loc = torch.tensor([0, 7, 0, 9], dtype=torch.int64, device=DEVICE)
+
+    set_mla_kv_buffer_triton_fp8_quant(
+        kv_buffer,
+        loc,
+        cache_k_nope,
+        cache_k_rope,
+        fp8_dtype,
+    )
+
+    torch.testing.assert_close(kv_buffer[0], reserved_before, rtol=0.0, atol=0.0)
+    expected = torch.cat((cache_k_nope[1, 0], cache_k_rope[1, 0])).to(fp8_dtype)
+    torch.testing.assert_close(
+        kv_buffer[7, 0], expected.view(torch.uint8), rtol=0.0, atol=0.0
+    )
+
+
+def test_set_mla_kv_scale_buffer_triton_reserved_skip_index():
+    cache_k_nope = torch.randn((4, 1, 16), dtype=torch.float32, device=DEVICE)
+    cache_k_rope = torch.randn((4, 1, 4), dtype=torch.float32, device=DEVICE)
+    cache_k_nope[[0, 2]] = torch.nan
+    cache_k_rope[[0, 2]] = torch.nan
+    kv_buffer = torch.randn((CACHE_SIZE, 1, 20), dtype=torch.float32, device=DEVICE)
+    reserved_before = kv_buffer[0].clone()
+    loc = torch.tensor([0, 7, 0, 9], dtype=torch.int64, device=DEVICE)
+
+    set_mla_kv_scale_buffer_triton(
+        kv_buffer,
+        loc,
+        cache_k_nope,
+        cache_k_rope,
+    )
+
+    torch.testing.assert_close(kv_buffer[0], reserved_before, rtol=0.0, atol=0.0)
+    torch.testing.assert_close(
+        kv_buffer[7, 0],
+        torch.cat((cache_k_nope[1, 0], cache_k_rope[1, 0])),
+        rtol=0.0,
+        atol=0.0,
+    )
+
+
+@CUDA_TMA_ONLY
 def test_can_use_set_mla_kv_buffer():
     assert can_use_set_mla_kv_buffer(1024, 128)  # bf16 (512,64)
     assert can_use_set_mla_kv_buffer(528, 128)  # fp8 byte layout

test/registered/kernels/ops/test_kimi_k3_prerequisite_ops.py

@@ -103,9 +103,10 @@ def _make_mla_inputs(batch_size, num_heads, seed):
     pool = randn(MLA_PAGES, MLA_DIM)
     latent = randn(batch_size, MLA_DIM)
     query = randn(batch_size, num_heads, MLA_DIM)
-    loc = torch.randperm(MLA_PAGES, generator=generator, device="cuda")[:batch_size].to(
+    loc = (
-        torch.int64
+        torch.randperm(MLA_PAGES - 1, generator=generator, device="cuda")[:batch_size]
-    )
+        + 1
+    ).to(torch.int64)
     return (
         pool,
         loc,