From 12eadf86f12aec2e6f81a6e38b61b964a4c6b529 Mon Sep 17 00:00:00 2001 From: kangwangamd <100359556+kangwangamd@users.noreply.github.com> Date: Mon, 3 Aug 2026 07:40:17 +0800 Subject: [PATCH] [AMD] Enable mamba JIT transfer kernel on ROCm (fix transfer_kv_mamba NameError) (#31741) --- python/sglang/kernels/jit/csrc/kvcacheio/transfer_mamba.cuh | 6 +++--- python/sglang/srt/mem_cache/memory_pool_host.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/python/sglang/kernels/jit/csrc/kvcacheio/transfer_mamba.cuh b/python/sglang/kernels/jit/csrc/kvcacheio/transfer_mamba.cuh index 04a8633a6..795ff2f28 100644 --- a/python/sglang/kernels/jit/csrc/kvcacheio/transfer_mamba.cuh +++ b/python/sglang/kernels/jit/csrc/kvcacheio/transfer_mamba.cuh @@ -101,7 +101,7 @@ struct TransferMambaKernel { TensorMatcher({L}) // .with_dtype() - .with_device(device_) + .with_device(device_) .verify(src_indices) .verify(dst_indices); @@ -147,13 +147,13 @@ struct TransferMambaKernel { TensorMatcher({L}) // .with_dtype() - .with_device(device_) + .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_) + .with_device(device_) .verify(src_ptrs); RuntimeCheck(item_size > 0, "transfer_mamba: item_size must be positive"); diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 041f34a32..2670de52c 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -38,7 +38,7 @@ if _is_cuda or _is_hip: transfer_kv_per_layer_mla, transfer_kv_per_layer_mla_pf_lf, ) -if _is_cuda: +if _is_cuda or _is_hip: from sglang.kernels.ops.mamba.transfer_mamba import ( transfer_kv_mamba_lf_pf, transfer_kv_mamba_pf_lf,