diff --git a/python/sglang/jit_kernel/csrc/kvcacheio/transfer_mamba.cuh b/python/sglang/jit_kernel/csrc/kvcacheio/transfer_mamba.cuh new file mode 100644 index 000000000..04a8633a6 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/kvcacheio/transfer_mamba.cuh @@ -0,0 +1,188 @@ +#pragma once + +#include "hicache.cuh" +#include +#include + +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(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(tid) * kBytesPerThreadPerStep; + if (base < params.item_size) { + uint4 v_cur = device::details::load_nc(reinterpret_cast(src + base)); + int64_t off; + for (off = base + kBytesPerBlockPerStep; off < params.item_size; off += kBytesPerBlockPerStep) { + uint4 v_next = device::details::load_nc(reinterpret_cast(src + off)); + device::details::store_nc(reinterpret_cast(dst + off - kBytesPerBlockPerStep), v_cur); + v_cur = v_next; + } + device::details::store_nc(reinterpret_cast(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(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(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(tid) * kBytesPerThreadPerStep; + if (base < params.item_size) { + uint4 v_cur = device::details::load_nc(reinterpret_cast(src + base)); + int64_t off; + for (off = base + kBytesPerBlockPerStep; off < params.item_size; off += kBytesPerBlockPerStep) { + uint4 v_next = device::details::load_nc(reinterpret_cast(src + off)); + device::details::store_nc(reinterpret_cast(dst + off - kBytesPerBlockPerStep), v_cur); + v_cur = v_next; + } + device::details::store_nc(reinterpret_cast(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() + .with_device(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(std::min(static_cast(kBlockQuotaLoad), num_items)); + dim3 grid(grid_x); + + const auto params = MambaTransferParams{ + .src_base = static_cast(src.data_ptr()), + .dst_base = static_cast(dst.data_ptr()), + .layer_ptrs = nullptr, + .src_indices = static_cast(src_indices.data_ptr()), + .dst_indices = static_cast(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() + .with_device(device_) + .verify(src_indices) + .verify(dst_indices); + // src_ptrs is a 1D tensor of device pointers (uint64) on CUDA + TensorMatcher({static_cast(num_layers)}) // + .with_dtype() + .with_device(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(std::min(static_cast(kBlockQuotaBackup), total_work)); + dim3 grid(grid_x); + + const auto params = MambaTransferParams{ + .src_base = nullptr, + .dst_base = static_cast(dst.data_ptr()), + .layer_ptrs = static_cast(src_ptrs.data_ptr()), + .src_indices = static_cast(src_indices.data_ptr()), + .dst_indices = static_cast(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 diff --git a/python/sglang/jit_kernel/transfer_mamba.py b/python/sglang/jit_kernel/transfer_mamba.py new file mode 100644 index 000000000..a64afa060 --- /dev/null +++ b/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, + ) diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index d17ac17b2..729aa8d42 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/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 ( - 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, diff --git a/test/registered/jit/test_transfer_mamba.py b/test/registered/jit/test_transfer_mamba.py new file mode 100644 index 000000000..8f7644c23 --- /dev/null +++ b/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"]))