From 7e3ae15f73b4339070282aa44a33c1d8c9ace72d Mon Sep 17 00:00:00 2001 From: Vignesh Sethuraman Date: Thu, 10 Sep 2026 19:27:09 -0700 Subject: [PATCH] [AMD] aiter: honor per-layer softmax scale (#38754) --- .../srt/layers/attention/aiter_backend.py | 26 +++++-------------- 1 file changed, 6 insertions(+), 20 deletions(-) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index d88ee7589..04cdd2487 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -1,10 +1,6 @@ from __future__ import annotations -from sglang.srt.runtime_context import ( - get_parallel, - get_schedule, - get_spec, -) +from sglang.srt.runtime_context import get_parallel, get_schedule, get_spec """ end to end attention solution with aiter kernels @@ -30,9 +26,7 @@ from sglang.kernels.ops.kvcache.aiter_unified_attention import ( scatter_req_to_token_to_page_table_kernel, ) from sglang.srt.layers.attention.base_attn_backend import AttentionBackend -from sglang.srt.layers.dp_attention import ( - is_dp_attention_enabled, -) +from sglang.srt.layers.dp_attention import is_dp_attention_enabled from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.speculative.spec_utils import ( draft_kv_indices_buffer_width, @@ -77,10 +71,7 @@ from sglang.kernels.ops.attention.utils import ( launch_reshape_and_cache_flash, pad_sequence_with_mask, ) -from sglang.kernels.ops.quantization.fp8_kernel import ( - fp8_dtype, - scaled_fp8_quant, -) +from sglang.kernels.ops.quantization.fp8_kernel import fp8_dtype, scaled_fp8_quant from sglang.srt.configs.model_config import AttentionArch from sglang.srt.environ import envs from sglang.srt.layers.attention.aiter_mla_gluon import ( @@ -247,9 +238,7 @@ class AiterAttnBackend(AttentionBackend): ): super().__init__() # Lazy import to avoid the initialization of cuda context - from sglang.kernels.ops.attention.extend_attention import ( - extend_attention_fwd, - ) + from sglang.kernels.ops.attention.extend_attention import extend_attention_fwd self.input_dtype = model_runner.model_config.dtype @@ -636,7 +625,6 @@ class AiterAttnBackend(AttentionBackend): max_split_per_batch, intra_batch_mode, ): - nhead_kv = 1 page_size = self.page_size dtype = self.kv_cache_dtype @@ -1853,7 +1841,6 @@ class AiterAttnBackend(AttentionBackend): seq_lens_cpu: Optional[torch.Tensor], verify_tokens_per_req: Optional[int], ): - num_kv_splits = None # num_kv_splits_indptr = None @@ -3266,7 +3253,7 @@ class AiterAttnBackend(AttentionBackend): out=o.view(-1, layer.tp_q_head_num, layer.v_head_dim), seqused_k=forward_batch.seq_lens, max_seqlen_k=max_kv_len, - softmax_scale=self.scale, + softmax_scale=layer.scaling, block_table=page_table, k_descale=k_descale, v_descale=v_descale, @@ -3293,7 +3280,7 @@ class AiterAttnBackend(AttentionBackend): seqused_k=forward_batch.seq_lens, max_seqlen_q=self.forward_metadata.max_q_len, max_seqlen_k=max_kv_len, - softmax_scale=self.scale, + softmax_scale=layer.scaling, causal=True, window_size=window_size, block_table=page_table, @@ -3381,7 +3368,6 @@ class AiterIndicesUpdaterPrefill: encoder_lens: Optional[torch.Tensor], spec_info: Optional[SpecInput], ): - kv_start_idx = None kv_indptr = self.kv_indptr qo_indptr = self.qo_indptr