minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Khoa Pham <khoa.pham@radixark.ai>
This commit is contained in:
Tan Trinh
2026-08-25 22:42:23 -07:00
committed by GitHub
co-authored by Khoa Pham
parent d7baad0116
commit 04c1036bb3
6 changed files with 258 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/elementwise/set_mla_kv_buffer.cuh
+5 -2
View File
@@ -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
+36 -7
View File
@@ -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
+13 -1
View File
@@ -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
+1
View File
@@ -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
+199 -4
View File
@@ -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
+4 -3
View File
@@ -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(
torch.int64
)
loc = (
torch.randperm(MLA_PAGES - 1, generator=generator, device="cuda")[:batch_size]
+ 1
).to(torch.int64)
return (
pool,
loc,