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

[AMD][AITER] Guard _use_mla_ps_kernel with self.use_mla in draft_extend_v2 paths (#20409)

Browse Source
This commit is contained in:
Chuan (Richard) Li
2026-03-18 00:45:22 -07:00
committed by GitHub
parent ead9d7aa43
commit 93422f27d6
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/layers/attention/aiter_backend.py
View File

@@ -1317,7 +1317,7 @@ class AiterAttnBackend(AttentionBackend):
kv_last_page_len = self.cuda_graph_kv_last_page_len[:bs]
max_q_len = num_tokens_per_bs
if _use_mla_ps_kernel:
if self.use_mla and _use_mla_ps_kernel:
num_kv_splits = self.max_split_per_batch
self.make_mla_meta_data(
@@ -1660,7 +1660,7 @@ class AiterAttnBackend(AttentionBackend):
kv_last_page_len = self.cuda_graph_kv_last_page_len[:bs]
max_q_len = num_tokens_per_bs
if _use_mla_ps_kernel:
if self.use_mla and _use_mla_ps_kernel:
num_kv_splits = self.max_split_per_batch