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,
test/registered/jit/test_transfer_mamba.py
+349
View File
@@ -0,0 +1,349 @@
"""Unit tests for the Mamba JIT transfer kernel.
Verifies kernel backup (D2H) and load (H2D) correctness for
``MambaPoolHost`` via the ``io_backend='kernel'`` path, across both
supported layouts and multiple index scenarios.
"""
import sys
import threading
from types import SimpleNamespace
import pytest
import torch
from sglang.srt.mem_cache.memory_pool_host import MambaPoolHost
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
register_amd_ci(est_time=10, suite="nightly-amd-kernel-1-gpu", nightly=True)
pytestmark = pytest.mark.skipif(
not torch.cuda.is_available(), reason="Mamba transfer kernel tests require CUDA."
)
DEVICE = "cuda"
NUM_LAYERS = 3
SIZE = 16
TEMPORAL_SHAPE = (16, 128) # 16*128*2 = 4096 bytes (fp16), 16-byte aligned
CONV_SHAPE = (4, 128) # 4*128*2 = 1024 bytes (fp16), 16-byte aligned
DTYPES = [torch.float16, torch.bfloat16]
LAYOUTS = ["page_first", "page_first_direct"]
def make_device_pool(dtype, device=DEVICE):
"""Create a minimal mock device pool that MambaPoolHost can use."""
temporal = torch.zeros(
(NUM_LAYERS, SIZE) + TEMPORAL_SHAPE, dtype=dtype, device=device
)
conv = [torch.zeros((NUM_LAYERS, SIZE) + CONV_SHAPE, dtype=dtype, device=device)]
mamba_cache = SimpleNamespace(temporal=temporal, conv=conv)
return SimpleNamespace(
mamba_cache=mamba_cache,
size=SIZE,
device=device,
)
def make_host_pool(dtype, layout):
"""Create a MambaPoolHost bypassing __init__, manually setting attributes.
NOTE: If MambaPoolHost adds/renames attributes accessed by
backup_from_device_all_layer or load_to_device_per_layer, this mock
must be updated to match. See assert_host_mock_complete() below.
"""
host = MambaPoolHost.__new__(MambaPoolHost)
host.layout = layout
host.page_size = 1
host.page_num = SIZE
host.size = SIZE
host.pin_memory = True
host.device = "cpu"
host.num_mamba_layers = NUM_LAYERS
host.conv_state_shapes = [CONV_SHAPE]
host.temporal_state_shape = TEMPORAL_SHAPE
host.temporal_state_elem_size = int(torch.prod(torch.tensor(TEMPORAL_SHAPE)).item())
host.conv_state_elem_sizes = [int(torch.prod(torch.tensor(CONV_SHAPE)).item())]
host.conv_dtype = dtype
host.temporal_dtype = dtype
host.dtype = dtype
host.size_per_token = host.get_size_per_token()
# Allocate host buffers (page_first layout)
temporal_dims = (SIZE, NUM_LAYERS, 1) + TEMPORAL_SHAPE
host.temporal_buffer = torch.zeros(temporal_dims, dtype=dtype).pin_memory()
host.conv_buffer = []
conv_dims = (SIZE, NUM_LAYERS, 1) + CONV_SHAPE
host.conv_buffer.append(torch.zeros(conv_dims, dtype=dtype).pin_memory())
# Staging buffers and JIT flags
host.temporal_staging_buffer = None
host.conv_staging_buffers = [None]
host.can_use_write_back_jit = True
host._temporal_can_use_jit = False
host._conv_can_use_jit = [False]
# Device pointers (needed for backup kernel path)
device_pool = make_device_pool(dtype)
host.device_pool = device_pool
host.temporal_device_ptrs = torch.tensor(
[device_pool.mamba_cache.temporal[i].data_ptr() for i in range(NUM_LAYERS)],
dtype=torch.uint64,
device=DEVICE,
)
host.conv_device_ptrs = [
torch.tensor(
[conv_state[i].data_ptr() for i in range(NUM_LAYERS)],
dtype=torch.uint64,
device=DEVICE,
)
for conv_state in device_pool.mamba_cache.conv
]
host.lock = threading.RLock()
host.clear()
return host
def assert_host_mock_complete(host):
"""Sanity check: ensure mock covers attributes used by backup/load paths."""
required = [
"layout",
"page_size",
"page_num",
"size",
"pin_memory",
"device",
"num_mamba_layers",
"conv_state_shapes",
"temporal_state_shape",
"temporal_state_elem_size",
"conv_state_elem_sizes",
"conv_dtype",
"temporal_dtype",
"dtype",
"size_per_token",
"temporal_buffer",
"conv_buffer",
"temporal_staging_buffer",
"conv_staging_buffers",
"can_use_write_back_jit",
"_temporal_can_use_jit",
"_conv_can_use_jit",
"device_pool",
"temporal_device_ptrs",
"conv_device_ptrs",
"lock",
]
missing = [attr for attr in required if not hasattr(host, attr)]
assert not missing, f"Mock MambaPoolHost missing attributes: {missing}"
def fill_device_data(device_pool, dtype):
"""Fill device temporal and conv states with deterministic data."""
for layer_id in range(NUM_LAYERS):
offset = layer_id * 1000
data = torch.arange(
device_pool.mamba_cache.temporal[layer_id].numel(),
device=DEVICE,
dtype=dtype,
)
device_pool.mamba_cache.temporal[layer_id].copy_(
(data + offset).view_as(device_pool.mamba_cache.temporal[layer_id])
)
for conv_idx in range(len(device_pool.mamba_cache.conv)):
conv_data = torch.arange(
device_pool.mamba_cache.conv[conv_idx][layer_id].numel(),
device=DEVICE,
dtype=dtype,
)
device_pool.mamba_cache.conv[conv_idx][layer_id].copy_(
(conv_data + offset + conv_idx * 500).view_as(
device_pool.mamba_cache.conv[conv_idx][layer_id]
)
)
def assert_host_matches_device(host, device_pool, host_indices, device_indices):
"""Verify host backup data matches device source data."""
for layer_id in range(NUM_LAYERS):
# Temporal
host_temporal = host.temporal_buffer[host_indices, layer_id, 0].cpu()
dev_temporal = device_pool.mamba_cache.temporal[layer_id][device_indices].cpu()
torch.testing.assert_close(host_temporal, dev_temporal)
# Conv
for conv_idx in range(len(host.conv_buffer)):
host_conv = host.conv_buffer[conv_idx][host_indices, layer_id, 0].cpu()
dev_conv = device_pool.mamba_cache.conv[conv_idx][layer_id][
device_indices
].cpu()
torch.testing.assert_close(host_conv, dev_conv)
def assert_device_matches_host(host, device_pool, host_indices, device_indices):
"""Verify device load data matches host source data."""
for layer_id in range(NUM_LAYERS):
# Temporal
host_temporal = host.temporal_buffer[host_indices, layer_id, 0].to(DEVICE)
dev_temporal = device_pool.mamba_cache.temporal[layer_id][device_indices]
torch.testing.assert_close(dev_temporal, host_temporal)
# Conv
for conv_idx in range(len(host.conv_buffer)):
host_conv = host.conv_buffer[conv_idx][host_indices, layer_id, 0].to(DEVICE)
dev_conv = device_pool.mamba_cache.conv[conv_idx][layer_id][device_indices]
torch.testing.assert_close(dev_conv, host_conv)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("layout", LAYOUTS)
def test_mamba_kernel_backup_load_roundtrip(dtype, layout):
"""Test D2H backup + H2D load roundtrip with io_backend='kernel'."""
host = make_host_pool(dtype, layout)
assert_host_mock_complete(host)
device_pool = host.device_pool
# Fill device with known data
fill_device_data(device_pool, dtype)
# Use a few indices for the test
device_indices = torch.tensor([1, 5, 10], dtype=torch.int64, device=DEVICE)
host_indices = torch.tensor([0, 1, 2], dtype=torch.int64)
load_indices = torch.tensor([3, 7, 12], dtype=torch.int64, device=DEVICE)
# --- Backup: device -> host (kernel) ---
host.backup_from_device_all_layer(
device_pool, host_indices, device_indices, io_backend="kernel"
)
torch.cuda.synchronize()
assert_host_matches_device(host, device_pool, host_indices, device_indices)
# --- Clear device buffers ---
for layer_id in range(NUM_LAYERS):
device_pool.mamba_cache.temporal[layer_id].zero_()
for conv_idx in range(len(device_pool.mamba_cache.conv)):
device_pool.mamba_cache.conv[conv_idx][layer_id].zero_()
# --- Load: host -> device (kernel), per layer ---
for layer_id in range(NUM_LAYERS):
host.load_to_device_per_layer(
device_pool,
host_indices,
load_indices,
layer_id,
io_backend="kernel",
)
torch.cuda.synchronize()
assert_device_matches_host(host, device_pool, host_indices, load_indices)
# Verify non-target positions remain zero (catch kernel writing wrong indices)
all_indices = set(range(SIZE))
target_set = set(load_indices.tolist())
untouched = sorted(all_indices - target_set)
if untouched:
untouched_t = torch.tensor(untouched, dtype=torch.int64, device=DEVICE)
for layer_id in range(NUM_LAYERS):
assert (
device_pool.mamba_cache.temporal[layer_id][untouched_t].abs().max() == 0
)
for conv_idx in range(len(device_pool.mamba_cache.conv)):
assert (
device_pool.mamba_cache.conv[conv_idx][layer_id][untouched_t]
.abs()
.max()
== 0
)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("layout", LAYOUTS)
def test_mamba_kernel_empty_indices(dtype, layout):
"""Test that empty indices are handled gracefully (no crash)."""
host = make_host_pool(dtype, layout)
device_pool = host.device_pool
fill_device_data(device_pool, dtype)
empty_device = torch.tensor([], dtype=torch.int64, device=DEVICE)
empty_host = torch.tensor([], dtype=torch.int64)
host.backup_from_device_all_layer(
device_pool, empty_host, empty_device, io_backend="kernel"
)
torch.cuda.synchronize()
# Host buffers should remain all zeros
assert host.temporal_buffer.abs().max() == 0
for layer_id in range(NUM_LAYERS):
host.load_to_device_per_layer(
device_pool, empty_host, empty_device, layer_id, io_backend="kernel"
)
torch.cuda.synchronize()
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("layout", LAYOUTS)
def test_mamba_kernel_single_item(dtype, layout):
"""Test single item backup + load."""
host = make_host_pool(dtype, layout)
device_pool = host.device_pool
fill_device_data(device_pool, dtype)
device_indices = torch.tensor([7], dtype=torch.int64, device=DEVICE)
host_indices = torch.tensor([3], dtype=torch.int64)
load_indices = torch.tensor([9], dtype=torch.int64, device=DEVICE)
host.backup_from_device_all_layer(
device_pool, host_indices, device_indices, io_backend="kernel"
)
torch.cuda.synchronize()
assert_host_matches_device(host, device_pool, host_indices, device_indices)
for layer_id in range(NUM_LAYERS):
device_pool.mamba_cache.temporal[layer_id].zero_()
for conv_idx in range(len(device_pool.mamba_cache.conv)):
device_pool.mamba_cache.conv[conv_idx][layer_id].zero_()
for layer_id in range(NUM_LAYERS):
host.load_to_device_per_layer(
device_pool, host_indices, load_indices, layer_id, io_backend="kernel"
)
torch.cuda.synchronize()
assert_device_matches_host(host, device_pool, host_indices, load_indices)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("layout", LAYOUTS)
def test_mamba_kernel_full_indices(dtype, layout):
"""Test full-size backup + load (all SIZE items)."""
host = make_host_pool(dtype, layout)
device_pool = host.device_pool
fill_device_data(device_pool, dtype)
device_indices = torch.arange(SIZE, dtype=torch.int64, device=DEVICE)
host_indices = torch.arange(SIZE, dtype=torch.int64)
load_indices = torch.arange(SIZE, dtype=torch.int64, device=DEVICE)
host.backup_from_device_all_layer(
device_pool, host_indices, device_indices, io_backend="kernel"
)
torch.cuda.synchronize()
assert_host_matches_device(host, device_pool, host_indices, device_indices)
for layer_id in range(NUM_LAYERS):
device_pool.mamba_cache.temporal[layer_id].zero_()
for conv_idx in range(len(device_pool.mamba_cache.conv)):
device_pool.mamba_cache.conv[conv_idx][layer_id].zero_()
for layer_id in range(NUM_LAYERS):
host.load_to_device_per_layer(
device_pool, host_indices, load_indices, layer_id, io_backend="kernel"
)
torch.cuda.synchronize()
assert_device_matches_host(host, device_pool, host_indices, load_indices)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))