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<int64_t>()
-        .with_device<kDLCUDA>(device_)
+        .with_device<kDLGPU>(device_)
         .verify(src_indices)
         .verify(dst_indices);
 
@@ -147,13 +147,13 @@ struct TransferMambaKernel {
 
     TensorMatcher({L})  //
         .with_dtype<int64_t>()
-        .with_device<kDLCUDA>(device_)
+        .with_device<kDLGPU>(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_)
+        .with_device<kDLGPU>(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,