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

Co-authored-by: hzh0425 <hzh0425@apache.org>

python/sglang/jit_kernel/csrc/kvcacheio/transfer_mamba.cuh

@@ -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

@@ -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

@@ -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 (
+from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool, MambaPool
-    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
+        assert layout in [
-        # kernel; only allow 'page_first_direct' + 'direct' for now.
+            "page_first",
-        # Relax this restriction once the staging bug is fixed.
+            "page_first_direct",
-        if layout != "page_first_direct":
+        ], f"Unsupported layout: {layout}"
-            raise ValueError(
-                f"MambaPoolHost only supports layout='page_first_direct', "
-                f"got '{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:
+            # Mamba JIT kernel expects all index tensors on CUDA.
-                jit_transfer_hicache_all_layer_mla_staged_lf_pf(
+            # When can_use_write_back_jit is True on the HostPoolGroup,
-                    ptr_src=src_ptrs,
+            # start_writing() keeps host_indices on CPU (for MLA staged kernel).
-                    src_indices=src_indices,
+            # Move dst_indices to CUDA here to satisfy the kernel's requirement.
-                    dst_indices=dst_indices,
+            if dst_indices.device.type != "cuda":
-                    staging=staging,
+                dst_indices = dst_indices.to(src_indices.device, non_blocking=True)
-                    dst=dst,
+            transfer_kv_mamba_lf_pf(
-                    page_size=1,
+                src_ptrs=src_ptrs,
-                    element_size=item_size,
+                dst=dst,
-                )
+                src_indices=src_indices,
-            else:
+                dst_indices=dst_indices,
-                transfer_kv_all_layer_mla_lf_pf(
+                item_size=item_size,
-                    src_layers=src_ptrs,
+                dst_layout_dim=item_size * num_layers,
-                    dst=dst,
+                num_layers=num_layers,
-                    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

@@ -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"]))