minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] Enable mamba JIT transfer kernel on ROCm (fix transfer_kv_mamba NameError) (#31741)

Browse Source
This commit is contained in:
kangwangamd
2026-08-02 16:40:17 -07:00
committed by GitHub
parent 88e5a0f635
commit 12eadf86f1
2 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/kvcacheio/transfer_mamba.cuh
+3 -3
View File
@@ -101,7 +101,7 @@ struct TransferMambaKernel {
TensorMatcher({L}) // TensorMatcher({L}) //
.with_dtype<int64_t>() .with_dtype<int64_t>()
.with_device<kDLCUDA>(device_) .with_device<kDLGPU>(device_)
.verify(src_indices) .verify(src_indices)
.verify(dst_indices); .verify(dst_indices);
@@ -147,13 +147,13 @@ struct TransferMambaKernel {
TensorMatcher({L}) // TensorMatcher({L}) //
.with_dtype<int64_t>() .with_dtype<int64_t>()
.with_device<kDLCUDA>(device_) .with_device<kDLGPU>(device_)
.verify(src_indices) .verify(src_indices)
.verify(dst_indices); .verify(dst_indices);
// src_ptrs is a 1D tensor of device pointers (uint64) on CUDA // src_ptrs is a 1D tensor of device pointers (uint64) on CUDA
TensorMatcher({static_cast<int64_t>(num_layers)}) // TensorMatcher({static_cast<int64_t>(num_layers)}) //
.with_dtype<uint64_t>() .with_dtype<uint64_t>()
.with_device<kDLCUDA>(device_) .with_device<kDLGPU>(device_)
.verify(src_ptrs); .verify(src_ptrs);
RuntimeCheck(item_size > 0, "transfer_mamba: item_size must be positive"); RuntimeCheck(item_size > 0, "transfer_mamba: item_size must be positive");
python/sglang/srt/mem_cache/memory_pool_host.py
+1 -1
View File
@@ -38,7 +38,7 @@ if _is_cuda or _is_hip:
transfer_kv_per_layer_mla, transfer_kv_per_layer_mla,
transfer_kv_per_layer_mla_pf_lf, transfer_kv_per_layer_mla_pf_lf,
) )
if _is_cuda: if _is_cuda or _is_hip:
from sglang.kernels.ops.mamba.transfer_mamba import ( from sglang.kernels.ops.mamba.transfer_mamba import (
transfer_kv_mamba_lf_pf, transfer_kv_mamba_lf_pf,
transfer_kv_mamba_pf_lf, transfer_kv_mamba_pf_lf,