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

[hicache]: add mamba concurrency io transfer kernel (#30535)

Browse Source 
Co-authored-by: hzh0425 <hzh0425@apache.org>
This commit is contained in:
jojo
2026-07-16 18:13:02 +08:00
committed by GitHub
co-authored by hzh0425
parent e2d021d4ab
commit b296e1a503
4 changed files with 658 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/kvcacheio/transfer_mamba.cuh
+188
View File
@@ -0,0 +1,188 @@
#pragma once
#include "hicache.cuh"
#include <algorithm>
#include <cstdint>
namespace {
constexpr int kBlockSize = 1024;
constexpr int kBlockQuotaBackup = 2;
constexpr int kBlockQuotaLoad = 2;
constexpr int kBytesPerThreadPerStep = 16;
constexpr int kBytesPerBlockPerStep = kBlockSize * kBytesPerThreadPerStep;
struct MambaTransferParams {
const char* __restrict__ src_base;
char* __restrict__ dst_base;
const uintptr_t* __restrict__ layer_ptrs;
const int64_t* __restrict__ src_indices;
const int64_t* __restrict__ dst_indices;
int64_t item_size;
int64_t src_layout_dim;
int64_t dst_layout_dim;
int64_t layer_id;
int64_t num_items;
int64_t num_layers;
};
__global__
__launch_bounds__(kBlockSize, 1) void transfer_mamba_load_kernel(const __grid_constant__ MambaTransferParams params) {
const int tid = threadIdx.x;
for (int64_t item_id = static_cast<int64_t>(blockIdx.x); item_id < params.num_items; item_id += gridDim.x) {
const int64_t src_page = params.src_indices[item_id];
const int64_t dst_page = params.dst_indices[item_id];
const char* src = params.src_base + src_page * params.src_layout_dim + params.layer_id * params.item_size;
char* dst = params.dst_base + dst_page * params.item_size;
const int64_t base = static_cast<int64_t>(tid) * kBytesPerThreadPerStep;
if (base < params.item_size) {
uint4 v_cur = device::details::load_nc(reinterpret_cast<const uint4*>(src + base));
int64_t off;
for (off = base + kBytesPerBlockPerStep; off < params.item_size; off += kBytesPerBlockPerStep) {
uint4 v_next = device::details::load_nc(reinterpret_cast<const uint4*>(src + off));
device::details::store_nc(reinterpret_cast<uint4*>(dst + off - kBytesPerBlockPerStep), v_cur);
v_cur = v_next;
}
device::details::store_nc(reinterpret_cast<uint4*>(dst + off - kBytesPerBlockPerStep), v_cur);
}
}
}
__global__
__launch_bounds__(kBlockSize, 1) void transfer_mamba_backup_kernel(const __grid_constant__ MambaTransferParams params) {
const int tid = threadIdx.x;
const int64_t total_work = params.num_items * params.num_layers;
for (int64_t work_id = static_cast<int64_t>(blockIdx.x); work_id < total_work; work_id += gridDim.x) {
const int64_t layer_id = work_id % params.num_layers;
const int64_t item_id = work_id / params.num_layers;
const int64_t src_page = params.src_indices[item_id];
const int64_t dst_page = params.dst_indices[item_id];
const char* src = reinterpret_cast<const char*>(params.layer_ptrs[layer_id]) + src_page * params.item_size;
char* dst = params.dst_base + dst_page * params.dst_layout_dim + layer_id * params.item_size;
const int64_t base = static_cast<int64_t>(tid) * kBytesPerThreadPerStep;
if (base < params.item_size) {
uint4 v_cur = device::details::load_nc(reinterpret_cast<const uint4*>(src + base));
int64_t off;
for (off = base + kBytesPerBlockPerStep; off < params.item_size; off += kBytesPerBlockPerStep) {
uint4 v_next = device::details::load_nc(reinterpret_cast<const uint4*>(src + off));
device::details::store_nc(reinterpret_cast<uint4*>(dst + off - kBytesPerBlockPerStep), v_cur);
v_cur = v_next;
}
device::details::store_nc(reinterpret_cast<uint4*>(dst + off - kBytesPerBlockPerStep), v_cur);
}
}
}
struct TransferMambaKernel {
// Load: page_first -> layer_first (single layer at a time)
static void run_pf_lf(
const tvm::ffi::TensorView src,
const tvm::ffi::TensorView dst,
const tvm::ffi::TensorView src_indices,
const tvm::ffi::TensorView dst_indices,
const int64_t layer_id,
const int64_t item_size,
const int64_t src_layout_dim) {
using namespace host;
auto L = SymbolicSize{"num_indices"};
auto device_ = SymbolicDevice{};
TensorMatcher({L}) //
.with_dtype<int64_t>()
.with_device<kDLCUDA>(device_)
.verify(src_indices)
.verify(dst_indices);
RuntimeCheck(item_size > 0, "transfer_mamba: item_size must be positive");
RuntimeCheck(item_size % 16 == 0, "transfer_mamba: item_size must be 16-byte aligned (uint4)");
const auto num_items = L.unwrap();
if (num_items == 0) return;
const auto device = device_.unwrap();
const int grid_x = static_cast<int>(std::min(static_cast<int64_t>(kBlockQuotaLoad), num_items));
dim3 grid(grid_x);
const auto params = MambaTransferParams{
.src_base = static_cast<const char*>(src.data_ptr()),
.dst_base = static_cast<char*>(dst.data_ptr()),
.layer_ptrs = nullptr,
.src_indices = static_cast<const int64_t*>(src_indices.data_ptr()),
.dst_indices = static_cast<const int64_t*>(dst_indices.data_ptr()),
.item_size = item_size,
.src_layout_dim = src_layout_dim,
.dst_layout_dim = 0,
.layer_id = layer_id,
.num_items = num_items,
.num_layers = 1,
};
LaunchKernel(grid, kBlockSize, device)(transfer_mamba_load_kernel, params);
}
// Backup: layer_first -> page_first (all layers at once)
static void run_lf_pf(
const tvm::ffi::TensorView src_ptrs,
const tvm::ffi::TensorView dst,
const tvm::ffi::TensorView src_indices,
const tvm::ffi::TensorView dst_indices,
const int64_t item_size,
const int64_t dst_layout_dim,
const int64_t num_layers) {
using namespace host;
auto L = SymbolicSize{"num_indices"};
auto device_ = SymbolicDevice{};
TensorMatcher({L}) //
.with_dtype<int64_t>()
.with_device<kDLCUDA>(device_)
.verify(src_indices)
.verify(dst_indices);
// src_ptrs is a 1D tensor of device pointers (uint64) on CUDA
TensorMatcher({static_cast<int64_t>(num_layers)}) //
.with_dtype<uint64_t>()
.with_device<kDLCUDA>(device_)
.verify(src_ptrs);
RuntimeCheck(item_size > 0, "transfer_mamba: item_size must be positive");
RuntimeCheck(item_size % 16 == 0, "transfer_mamba: item_size must be 16-byte aligned (uint4)");
RuntimeCheck(num_layers > 0, "transfer_mamba: num_layers must be positive");
const auto num_items = L.unwrap();
if (num_items == 0) return;
const auto device = device_.unwrap();
const int64_t total_work = num_items * num_layers;
const int grid_x = static_cast<int>(std::min(static_cast<int64_t>(kBlockQuotaBackup), total_work));
dim3 grid(grid_x);
const auto params = MambaTransferParams{
.src_base = nullptr,
.dst_base = static_cast<char*>(dst.data_ptr()),
.layer_ptrs = static_cast<const uintptr_t*>(src_ptrs.data_ptr()),
.src_indices = static_cast<const int64_t*>(src_indices.data_ptr()),
.dst_indices = static_cast<const int64_t*>(dst_indices.data_ptr()),
.item_size = item_size,
.src_layout_dim = 0,
.dst_layout_dim = dst_layout_dim,
.layer_id = 0,
.num_items = num_items,
.num_layers = num_layers,
};
LaunchKernel(grid, kBlockSize, device)(transfer_mamba_backup_kernel, params);
}
};
} // namespace
python/sglang/jit_kernel/transfer_mamba.py
+83
View File
@@ -0,0 +1,83 @@
"""JIT-compiled Mamba KV cache transfer kernel.
Provides ``transfer_kv_mamba_pf_lf`` (load: page_first -> layer_first)
and ``transfer_kv_mamba_lf_pf`` (backup: layer_first -> page_first).
Uses the shared ``load_jit`` + ``cache_once`` infrastructure from
``sglang.jit_kernel.utils`` — the same mechanism used by ``hicache.py``
for MHA/MLA staged write-back kernels. This ensures consistent
content-addressed caching, CUDA arch detection, and multi-worker
JIT compilation behavior across all JIT kernels.
"""
from __future__ import annotations
import logging
from typing import TYPE_CHECKING
from sglang.jit_kernel.utils import cache_once, load_jit
from sglang.kernel_api_logging import debug_kernel_api
if TYPE_CHECKING:
import torch
from tvm_ffi.module import Module
logger = logging.getLogger(__name__)
@cache_once
def _jit_transfer_mamba_module() -> Module:
return load_jit(
"transfer_mamba",
cuda_files=["kvcacheio/transfer_mamba.cuh"],
cuda_wrappers=[
("transfer_kv_mamba_pf_lf", "&TransferMambaKernel::run_pf_lf"),
("transfer_kv_mamba_lf_pf", "&TransferMambaKernel::run_lf_pf"),
],
)
@debug_kernel_api
def transfer_kv_mamba_pf_lf(
src: torch.Tensor,
dst: torch.Tensor,
src_indices: torch.Tensor,
dst_indices: torch.Tensor,
layer_id: int,
item_size: int,
src_layout_dim: int,
num_warps_per_item: int = 32,
):
module = _jit_transfer_mamba_module()
module.transfer_kv_mamba_pf_lf(
src,
dst,
src_indices,
dst_indices,
layer_id,
item_size,
src_layout_dim,
)
@debug_kernel_api
def transfer_kv_mamba_lf_pf(
src_ptrs: torch.Tensor,
dst: torch.Tensor,
src_indices: torch.Tensor,
dst_indices: torch.Tensor,
item_size: int,
dst_layout_dim: int,
num_layers: int,
num_warps_per_item: int = 32,
):
module = _jit_transfer_mamba_module()
module.transfer_kv_mamba_lf_pf(
src_ptrs,
dst,
src_indices,
dst_indices,
item_size,
dst_layout_dim,
num_layers,
)
python/sglang/srt/mem_cache/memory_pool_host.py
+38 -85
View File
@@ -20,10 +20,7 @@ from sglang.jit_kernel.hicache import (
transfer_hicache_all_layer_mla_staged_lf_pf as jit_transfer_hicache_all_layer_mla_staged_lf_pf,
)
from sglang.jit_kernel.hisparse import transfer_cache_dsv4_mla
from sglang.srt.mem_cache.memory_pool import (
DSATokenToKVPool,
MambaPool,
)
from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool, MambaPool
from sglang.srt.utils import is_cuda, is_hip, is_mps, is_npu, is_xpu
_is_cuda = is_cuda()
@@ -41,6 +38,13 @@ if _is_cuda or _is_hip:
transfer_kv_per_layer_mla,
transfer_kv_per_layer_mla_pf_lf,
)
if _is_cuda:
from sglang.jit_kernel.transfer_mamba import (
transfer_kv_mamba_lf_pf,
transfer_kv_mamba_pf_lf,
)
if _is_npu:
pass
logger = logging.getLogger(__name__)
@@ -74,14 +78,10 @@ class MambaPoolHost(HostKVCache):
self.device_pool = device_pool
self.page_size = 1
# TODO: Mamba pool is currently incompatible with write-back staging
# kernel; only allow 'page_first_direct' + 'direct' for now.
# Relax this restriction once the staging bug is fixed.
if layout != "page_first_direct":
raise ValueError(
f"MambaPoolHost only supports layout='page_first_direct', "
f"got '{layout}'."
)
assert layout in [
"page_first",
"page_first_direct",
], f"Unsupported layout: {layout}"
self.layout = layout
self.pin_memory = pin_memory
@@ -217,50 +217,13 @@ class MambaPoolHost(HostKVCache):
def _init_write_back_staging_buffers(self):
self.temporal_staging_buffer = None
self.conv_staging_buffers = [None] * len(self.conv_buffer)
self.can_use_write_back_jit = False
# Must be True: HostPoolGroup computes can_use_write_back_jit as AND of
# all pools. When True, start_writing() keeps indices on CPU, which MLA's
# staged write-back kernel requires. MambaPoolHost's own backup path does
# not check this flag — it routes by layout + io_backend instead.
self.can_use_write_back_jit = True
self._temporal_can_use_jit = False
self._conv_can_use_jit = [False] * len(self.conv_buffer)
if self.layout != "page_first" or (_is_npu or _is_xpu or _is_mps):
return
self._temporal_can_use_jit = _is_cuda and can_use_write_back_jit_kernel(
element_size=self._item_size_per_index(self.temporal_buffer[0]),
)
self._conv_can_use_jit = [
_is_cuda
and can_use_write_back_jit_kernel(
element_size=self._item_size_per_index(buf[0]),
)
for buf in self.conv_buffer
]
self.can_use_write_back_jit = self._temporal_can_use_jit and all(
self._conv_can_use_jit
)
self.staging_page_capacity = min(self.page_num, _WRITE_BACK_STAGING_PAGE_CHUNK)
self.staging_token_capacity = self.staging_page_capacity * self.page_size
self.temporal_staging_buffer = torch.empty(
(
self.staging_token_capacity,
self.num_mamba_layers,
1,
*self.temporal_state_shape,
),
dtype=self.temporal_dtype,
device=self.device_pool.device,
)
self.conv_staging_buffers = [
torch.empty(
(
self.staging_token_capacity,
self.num_mamba_layers,
1,
*conv_shape,
),
dtype=self.conv_dtype,
device=self.device_pool.device,
)
for conv_shape in self.conv_state_shapes
]
def get_hybrid_pool_buffer(self):
# Expose all mamba host tensors that need Mooncake buffer registration.
@@ -369,7 +332,12 @@ class MambaPoolHost(HostKVCache):
return
if io_backend == "kernel":
item_size = MambaPoolHost._item_size_per_index(dst)
transfer_kv_per_layer_mla_pf_lf(
# Mamba JIT kernel expects all index tensors on CUDA.
# host_indices may be on CPU (kept there by start_writing when
# can_use_write_back_jit is True on the HostPoolGroup).
if src_indices.device.type != "cuda":
src_indices = src_indices.to(dst_indices.device, non_blocking=True)
transfer_kv_mamba_pf_lf(
src=src,
dst=dst,
src_indices=src_indices,
@@ -406,26 +374,21 @@ class MambaPoolHost(HostKVCache):
return
if io_backend == "kernel":
item_size = MambaPoolHost._item_size_per_index(src_layers[0])
if can_use_jit:
jit_transfer_hicache_all_layer_mla_staged_lf_pf(
ptr_src=src_ptrs,
src_indices=src_indices,
dst_indices=dst_indices,
staging=staging,
dst=dst,
page_size=1,
element_size=item_size,
)
else:
transfer_kv_all_layer_mla_lf_pf(
src_layers=src_ptrs,
dst=dst,
src_indices=src_indices,
dst_indices=dst_indices,
item_size=item_size,
dst_layout_dim=item_size * num_layers,
num_layers=num_layers,
)
# Mamba JIT kernel expects all index tensors on CUDA.
# When can_use_write_back_jit is True on the HostPoolGroup,
# start_writing() keeps host_indices on CPU (for MLA staged kernel).
# Move dst_indices to CUDA here to satisfy the kernel's requirement.
if dst_indices.device.type != "cuda":
dst_indices = dst_indices.to(src_indices.device, non_blocking=True)
transfer_kv_mamba_lf_pf(
src_ptrs=src_ptrs,
dst=dst,
src_indices=src_indices,
dst_indices=dst_indices,
item_size=item_size,
dst_layout_dim=item_size * num_layers,
num_layers=num_layers,
)
elif io_backend == "direct":
src_ptrs = [src_layers[i] for i in range(num_layers)]
transfer_kv_all_layer_direct_lf_pf(
@@ -446,11 +409,6 @@ class MambaPoolHost(HostKVCache):
layer_id,
io_backend="kernel",
):
if io_backend != "direct":
raise ValueError(
f"MambaPoolHost only supports io_backend='direct', "
f"got '{io_backend}'."
)
if self.layout in ["page_first", "page_first_direct"]:
self._copy_tensor_pf_lf(
src=self.temporal_buffer,
@@ -491,11 +449,6 @@ class MambaPoolHost(HostKVCache):
def backup_from_device_all_layer(
self, device_pool, host_indices, device_indices, io_backend="kernel"
):
if io_backend != "direct":
raise ValueError(
f"MambaPoolHost only supports io_backend='direct', "
f"got '{io_backend}'."
)
if self.layout in ["page_first", "page_first_direct"]:
self._copy_tensor_all_layers_lf_pf(
src_layers=device_pool.mamba_cache.temporal,