[AMD][AITER] Guard _use_mla_ps_kernel with self.use_mla in draft_extend_v2 paths (#20409)
This commit is contained in:
@@ -1317,7 +1317,7 @@ class AiterAttnBackend(AttentionBackend):
|
|||||||
kv_last_page_len = self.cuda_graph_kv_last_page_len[:bs]
|
kv_last_page_len = self.cuda_graph_kv_last_page_len[:bs]
|
||||||
max_q_len = num_tokens_per_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
|
num_kv_splits = self.max_split_per_batch
|
||||||
|
|
||||||
self.make_mla_meta_data(
|
self.make_mla_meta_data(
|
||||||
@@ -1660,7 +1660,7 @@ class AiterAttnBackend(AttentionBackend):
|
|||||||
kv_last_page_len = self.cuda_graph_kv_last_page_len[:bs]
|
kv_last_page_len = self.cuda_graph_kv_last_page_len[:bs]
|
||||||
max_q_len = num_tokens_per_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
|
num_kv_splits = self.max_split_per_batch
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user