From f97c09dac16192d1578d05ecb340b4dca1923e0a Mon Sep 17 00:00:00 2001 From: YC Yen-Ching Tseng Date: Sat, 21 Mar 2026 03:07:41 +0800 Subject: [PATCH] [AMD] Enable aiter unified attention for non-SWA models (Qwen3-VL) (#20897) Co-authored-by: wunhuang --- .../srt/layers/attention/aiter_backend.py | 33 +++++++++++++++++-- python/sglang/srt/server_args.py | 11 +++++++ 2 files changed, 41 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index eb7006c39..0b8bfbd51 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -204,7 +204,9 @@ class AiterAttnBackend(AttentionBackend): self.token_to_kv_pool = model_runner.token_to_kv_pool self.use_triton_unified_attention = True else: - self.use_triton_unified_attention = False + self.use_triton_unified_attention = get_bool_env_var( + "SGLANG_USE_AITER_UNIFIED_ATTN" + ) # aiter kernel related initialization self.max_num_partitions = ( @@ -682,6 +684,8 @@ class AiterAttnBackend(AttentionBackend): kv_indices = self._transform_table_1_to_real(kv_indices) swa_page_table = self._transform_table_1_to_real(swa_page_table) + elif self.page_size > 1: + kv_indices = self._transform_table_1_to_real(kv_indices) qo_indptr = self.qo_indptr[: bs + 1] qo_indptr[1 : bs + 1] = torch.cumsum( @@ -1280,6 +1284,11 @@ class AiterAttnBackend(AttentionBackend): kv_indices[:new_rows, :new_cols].copy_(page_indices) swa_page_table = self.cuda_graph_swa_page_table swa_page_table[:new_rows, :new_cols].copy_(swa_page_indices) + elif self.page_size > 1: + page_indices = self._transform_table_1_to_real(page_indices) + new_rows = page_indices.shape[0] + new_cols = page_indices.shape[1] + kv_indices[:new_rows, :new_cols].copy_(page_indices) qo_indptr = self.qo_indptr[: bs + 1] qo_indptr[1 : bs + 1] = torch.cumsum( @@ -1657,6 +1666,11 @@ class AiterAttnBackend(AttentionBackend): kv_indices[:new_rows, :new_cols].copy_(page_indices) swa_page_table = self.cuda_graph_swa_page_table swa_page_table[:new_rows, :new_cols].copy_(swa_page_indices) + elif self.page_size > 1: + page_indices = self._transform_table_1_to_real(page_indices) + new_rows = page_indices.shape[0] + new_cols = page_indices.shape[1] + kv_indices[:new_rows, :new_cols].copy_(page_indices) qo_indptr = self.qo_indptr[: bs + 1] qo_indptr[1 : bs + 1] = torch.cumsum( @@ -1983,7 +1997,15 @@ class AiterAttnBackend(AttentionBackend): if k is not None: assert v is not None if save_kv_cache: - if self.use_triton_unified_attention: + # Only use SWA-specific kv cache write (reshape_and_cache_flash) when + # both unified attention and sliding window kv pool are active. + # Non-SWA models (e.g. Qwen3-VL) enabled via SGLANG_USE_AITER_UNIFIED_ATTN + # use standard set_kv_buffer, as they lack SWA-specific attributes + # like full_to_swa_index_mapping. + if ( + self.use_triton_unified_attention + and self.use_sliding_window_kv_pool + ): token_to_kv_pool = forward_batch.token_to_kv_pool k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer( layer.layer_id @@ -2385,7 +2407,12 @@ class AiterAttnBackend(AttentionBackend): o = torch.empty_like(q, dtype=self.input_dtype) if save_kv_cache: - if self.use_triton_unified_attention: + # Only use SWA-specific kv cache write (reshape_and_cache_flash) when + # both unified attention and sliding window kv pool are active. + # Non-SWA models (e.g. Qwen3-VL) enabled via SGLANG_USE_AITER_UNIFIED_ATTN + # use standard set_kv_buffer, as they lack SWA-specific attributes + # like full_to_swa_index_mapping. + if self.use_triton_unified_attention and self.use_sliding_window_kv_pool: token_to_kv_pool = forward_batch.token_to_kv_pool k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer( layer.layer_id diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 5faa2e2cc..dcf5c5683 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2021,6 +2021,17 @@ class ServerArgs: "as the first layer might not be an attention layer" ) + if ( + model_arch in ["Qwen3VLForConditionalGeneration"] + and is_hip() + and get_bool_env_var("SGLANG_USE_AITER_UNIFIED_ATTN") + and self.page_size is None + ): + self.page_size = 16 + logger.info( + "Setting page_size=16 for aiter unified attention on Qwen3VLForConditionalGeneration." + ) + if envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set(): self.disable_overlap_schedule = True logger.warning(