From 93422f27d6f41c2cfb0315c9569c98fc3a55a322 Mon Sep 17 00:00:00 2001 From: "Chuan (Richard) Li" Date: Wed, 18 Mar 2026 00:45:22 -0700 Subject: [PATCH] [AMD][AITER] Guard _use_mla_ps_kernel with self.use_mla in draft_extend_v2 paths (#20409) --- python/sglang/srt/layers/attention/aiter_backend.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 56db2f601..44c868eb9 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -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