chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[DeepseekV32]: use _concat_mla_absorb_q_general to replace torch.cat (#12215)

Browse Source 
Co-authored-by: Guangda Liu <bingps@users.noreply.github.com>
This commit is contained in:
bingps
2025-11-13 07:20:01 +08:00
committed by GitHub
parent 03a7e6f4db
commit 4eda9969e8
1 changed files with 7 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
python/sglang/srt/layers/attention/nsa_backend.py
View File

@@ -21,6 +21,7 @@ from sglang.srt.layers.attention.nsa.utils import (
NSA_FUSE_TOPK,
compute_nsa_seqlens,
)
from sglang.srt.layers.attention.trtllm_mla_backend import _concat_mla_absorb_q_general
from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.utils import is_hip
@@ -911,7 +912,7 @@ class NativeSparseAttnBackend(AttentionBackend):
if NSA_PREFILL_IMPL == "tilelang":
if q_rope is not None:
q_all = torch.cat([q_nope, q_rope], dim=-1)
q_all = _concat_mla_absorb_q_general(q_nope, q_rope)
return self._forward_tilelang(
q_all=q_all,
kv_cache=kv_cache,
@@ -921,7 +922,7 @@ class NativeSparseAttnBackend(AttentionBackend):
)
elif NSA_PREFILL_IMPL == "flashmla_sparse":
if q_rope is not None:
q_all = torch.cat([q_nope, q_rope], dim=-1)
q_all = _concat_mla_absorb_q_general(q_nope, q_rope)
# NSA_FLASHMLA_BACKEND_DECODE_COMPUTE_FP8 has no effect here,
# because the flashmla_sparse kernel doesn't support fp8 compute
@@ -947,7 +948,7 @@ class NativeSparseAttnBackend(AttentionBackend):
)
elif NSA_PREFILL_IMPL == "flashmla_kv":
if q_rope is not None:
q_all = torch.cat([q_nope, q_rope], dim=-1)
q_all = _concat_mla_absorb_q_general(q_nope, q_rope)
return self._forward_flashmla_kv(
q_all=q_all,
kv_cache=kv_cache,
@@ -1031,7 +1032,7 @@ class NativeSparseAttnBackend(AttentionBackend):
if NSA_DECODE_IMPL == "flashmla_sparse":
if q_rope is not None:
q_all = torch.cat([q_nope, q_rope], dim=-1)
q_all = _concat_mla_absorb_q_general(q_nope, q_rope)
return self._forward_flashmla_sparse(
q_all=q_all,
kv_cache=kv_cache,
@@ -1041,7 +1042,7 @@ class NativeSparseAttnBackend(AttentionBackend):
)
elif NSA_DECODE_IMPL == "flashmla_kv":
if q_rope is not None:
q_all = torch.cat([q_nope, q_rope], dim=-1)
q_all = _concat_mla_absorb_q_general(q_nope, q_rope)
return self._forward_flashmla_kv(
q_all=q_all,
kv_cache=kv_cache,
@@ -1054,7 +1055,7 @@ class NativeSparseAttnBackend(AttentionBackend):
)
elif NSA_DECODE_IMPL == "tilelang":
if q_rope is not None:
q_all = torch.cat([q_nope, q_rope], dim=-1)
q_all = _concat_mla_absorb_q_general(q_nope, q_rope)
return self._forward_tilelang(
q_all=q_all,
kv_cache=kv_cache,