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

[AMD] Enable fused GDN QKV split Triton kernel on HIP (#27583)

Browse Source
This commit is contained in:
jacky.cheng
2026-06-11 02:01:27 -07:00
committed by GitHub
parent 493bb6ae0d
commit f4b3b99413
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/linear/gdn_backend.py
+3 -3
View File
@@ -18,7 +18,7 @@ from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
from sglang.srt.mem_cache.memory_pool import MambaPool from sglang.srt.mem_cache.memory_pool import MambaPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.utils import is_cpu, is_cuda, is_npu from sglang.srt.utils import is_cpu, is_cuda, is_hip, is_npu
from sglang.srt.utils.common import rank0_log from sglang.srt.utils.common import rank0_log
if not is_cpu(): if not is_cpu():
@@ -26,7 +26,7 @@ if not is_cpu():
CHUNK_SIZE as FLA_CHUNK_SIZE, CHUNK_SIZE as FLA_CHUNK_SIZE,
) )
if is_cuda(): if is_cuda() or is_hip():
from sglang.jit_kernel.triton.gdn_fused_proj import fused_qkv_split_gdn_prefill from sglang.jit_kernel.triton.gdn_fused_proj import fused_qkv_split_gdn_prefill
MAX_FUSED_QKV_SPLIT_DIM = 8192 MAX_FUSED_QKV_SPLIT_DIM = 8192
@@ -451,7 +451,7 @@ class GDNAttnBackend(MambaAttnBackendBase):
actual_seq_len = mixed_qkv.shape[0] actual_seq_len = mixed_qkv.shape[0]
qkv_dim = layer.q_dim + layer.k_dim + layer.v_dim qkv_dim = layer.q_dim + layer.k_dim + layer.v_dim
if is_cuda() and qkv_dim <= MAX_FUSED_QKV_SPLIT_DIM: if (is_cuda() or is_hip()) and qkv_dim <= MAX_FUSED_QKV_SPLIT_DIM:
query, key, value = fused_qkv_split_gdn_prefill( query, key, value = fused_qkv_split_gdn_prefill(
mixed_qkv, mixed_qkv,
layer.num_q_heads, layer.num_q_heads,