From fc6a829e4b80f7511238c2d063781d71ed80fba5 Mon Sep 17 00:00:00 2001 From: Aarne Talman Date: Thu, 15 Oct 2026 12:25:54 +0000 Subject: [PATCH 1/4] [Perf] [AMD] deterministic: Allow AITER decode attention on gfx950 AITER's paged decode attention reduces each request's KV cache in fixed 256-token partitions, so a request's result does not depend on its batch. On gfx950, deterministic inference now accepts --decode-attention-backend aiter for non-MLA models; prefill stays on a deterministic backend. MLA decode, whose KV splits depend on the batch size, stays refused, and the AITER backend refuses the unified-attention and 5D KV cache decode paths in deterministic mode. --- python/sglang/srt/arg_groups/overrides.py | 9 ++++ .../srt/layers/attention/aiter_backend.py | 11 ++++- test/registered/unit/test_model_overrides.py | 48 ++++++++++++++++++- 3 files changed, 66 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index f9d50f8aeab8..1dc2239b9ae7 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -1078,6 +1078,12 @@ def _deterministic_allreduce_fusion_disable(view: Any) -> dict: return {} +def _deterministic_aiter_decode_supported(view: Any) -> bool: + # AITER paged decode reduces fixed 256-token KV partitions per request, so it is + # batch-invariant (verified on gfx950); MLA decode sizes its KV splits by batch size. + return is_gfx95_supported() and not use_mla_backend(view) + + @register_post_process def _deterministic_rocm_rccl_check(view: Any) -> dict: # Deterministic TP on ROCm relies on the one-stage custom all-reduce; above two ranks @@ -1106,6 +1112,9 @@ def _deterministic_attention_backend(view: Any) -> dict: # so it has to be one of the deterministic backends too. for field in ("prefill_attention_backend", "decode_attention_backend"): backend = getattr(view, field, None) + if field == "decode_attention_backend" and backend == "aiter": + if _deterministic_aiter_decode_supported(view): + continue if ( backend is not None and backend not in DETERMINISTIC_ATTENTION_BACKEND_CHOICES diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 85b9a9850925..df9af2fce73b 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -1,6 +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_exec, get_parallel, get_schedule, get_spec """ end to end attention solution with aiter kernels @@ -566,6 +566,15 @@ def __init__( self.use_triton_unified_attention = get_bool_env_var( "SGLANG_USE_AITER_UNIFIED_ATTN" ) + if get_exec().deterministic.enable_deterministic_inference and not self.use_mla: + # Deterministic mode allows aiter decode only for the paged_attention_ragged + # path; the unified-attention and 5D-pool kernels are not verified batch-invariant. + if self.use_triton_unified_attention or self.kv_cache_is_vectorized_5d: + raise ValueError( + "Deterministic inference supports aiter decode attention only with the " + "default paged KV cache: not with sliding-window models, " + "SGLANG_USE_AITER_UNIFIED_ATTN or a vectorized 5D KV cache layout." + ) # Decode metadata format. unified_attention and pa_decode_gluon (the # SHUFFLE 5D decode kernel) both read kv_indices as a 2-D diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index 7c6acdd14a2d..d87a298bd8c2 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -1983,7 +1983,13 @@ def test_deterministic_incompatible_phase_backend_raises(self): # A per-phase backend overrides --attention-backend for its phase, so a # non-deterministic one must be refused even when the main backend is fine. for field in ("prefill_attention_backend", "decode_attention_backend"): - with self.subTest(field=field): + with ( + self.subTest(field=field), + patch( + "sglang.srt.arg_groups.overrides.is_gfx95_supported", + return_value=False, + ), + ): view = ResolvedView( SimpleNamespace( enable_deterministic_inference=True, @@ -1994,6 +2000,46 @@ def test_deterministic_incompatible_phase_backend_raises(self): with self.assertRaisesRegex(ValueError, field.replace("_", "-")): _deterministic_attention_backend(view) + def test_deterministic_aiter_decode_only_on_gfx950_without_mla(self): + """AITER decode attention is batch-invariant only for non-MLA models on gfx950; + AITER prefill attention never is.""" + from sglang.srt.arg_groups.overrides import ( + ResolvedView, + _deterministic_attention_backend, + ) + + cases = [ + # (gfx950, MLA model, prefill backend, accepted) + (True, False, "triton", True), + (False, False, "triton", False), + (True, True, "triton", False), + (True, False, "aiter", False), + ] + for gfx950, mla, prefill, accepted in cases: + with ( + self.subTest(gfx950=gfx950, mla=mla, prefill=prefill), + patch( + "sglang.srt.arg_groups.overrides.is_gfx95_supported", + return_value=gfx950, + ), + patch( + "sglang.srt.arg_groups.overrides.use_mla_backend", return_value=mla + ), + ): + view = ResolvedView( + SimpleNamespace( + enable_deterministic_inference=True, + attention_backend="triton", + prefill_attention_backend=prefill, + decode_attention_backend="aiter", + ) + ) + if accepted: + self.assertEqual(_deterministic_attention_backend(view), {}) + else: + with self.assertRaises(ValueError): + _deterministic_attention_backend(view) + def test_deterministic_compatible_phase_backends_are_accepted(self): from sglang.srt.arg_groups.overrides import ( ResolvedView, From 7f42536993c67321df5c9bc246e08c11c8158f80 Mon Sep 17 00:00:00 2001 From: Aarne Talman Date: Thu, 15 Oct 2026 14:03:33 +0000 Subject: [PATCH 2/4] [Perf] [AMD] deterministic: Simplify the per-phase attention backend check Build the allowed backends for each phase first (AITER added for decode on gfx950 for non-MLA models), then check both phases with one loop instead of a special case inside it. Read the per-phase fields directly: they always exist on the server arguments; the unit tests now set both. The AITER backend guard no longer checks for MLA, which deterministic mode already refuses with AITER decode. --- python/sglang/srt/arg_groups/overrides.py | 42 ++++++++++--------- .../srt/layers/attention/aiter_backend.py | 19 +++++---- test/registered/unit/test_model_overrides.py | 12 +++++- 3 files changed, 43 insertions(+), 30 deletions(-) diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 1dc2239b9ae7..4612432cfbbc 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -1078,12 +1078,6 @@ def _deterministic_allreduce_fusion_disable(view: Any) -> dict: return {} -def _deterministic_aiter_decode_supported(view: Any) -> bool: - # AITER paged decode reduces fixed 256-token KV partitions per request, so it is - # batch-invariant (verified on gfx950); MLA decode sizes its KV splits by batch size. - return is_gfx95_supported() and not use_mla_backend(view) - - @register_post_process def _deterministic_rocm_rccl_check(view: Any) -> dict: # Deterministic TP on ROCm relies on the one-stage custom all-reduce; above two ranks @@ -1108,20 +1102,30 @@ def _deterministic_attention_backend(view: Any) -> dict: return {} from sglang.srt.server_args import DETERMINISTIC_ATTENTION_BACKEND_CHOICES - # A per-phase backend takes priority over --attention-backend for its phase, - # so it has to be one of the deterministic backends too. - for field in ("prefill_attention_backend", "decode_attention_backend"): - backend = getattr(view, field, None) - if field == "decode_attention_backend" and backend == "aiter": - if _deterministic_aiter_decode_supported(view): - continue - if ( - backend is not None - and backend not in DETERMINISTIC_ATTENTION_BACKEND_CHOICES - ): + decode_choices = list(DETERMINISTIC_ATTENTION_BACKEND_CHOICES) + # AITER paged decode reduces fixed 256-token KV partitions per request, so it is + # batch-invariant (verified on gfx950); MLA decode sizes its KV splits by batch size. + if ( + view.decode_attention_backend == "aiter" + and is_gfx95_supported() + and not use_mla_backend(view) + ): + decode_choices.append("aiter") + + # A per-phase backend overrides --attention-backend for its phase. Unset (None) + # means the phase uses --attention-backend, which is checked below. + for flag, backend, choices in ( + ( + "--prefill-attention-backend", + view.prefill_attention_backend, + DETERMINISTIC_ATTENTION_BACKEND_CHOICES, + ), + ("--decode-attention-backend", view.decode_attention_backend, decode_choices), + ): + if backend is not None and backend not in choices: raise ValueError( - f"Currently only {DETERMINISTIC_ATTENTION_BACKEND_CHOICES} attention backends are supported for deterministic inference, " - f"but --{field.replace('_', '-')} is '{backend}'." + f"Currently only {choices} attention backends are supported for deterministic inference, " + f"but {flag} is '{backend}'." ) if view.attention_backend is None: diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index df9af2fce73b..eec0e83fc3dc 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -566,15 +566,16 @@ def __init__( self.use_triton_unified_attention = get_bool_env_var( "SGLANG_USE_AITER_UNIFIED_ATTN" ) - if get_exec().deterministic.enable_deterministic_inference and not self.use_mla: - # Deterministic mode allows aiter decode only for the paged_attention_ragged - # path; the unified-attention and 5D-pool kernels are not verified batch-invariant. - if self.use_triton_unified_attention or self.kv_cache_is_vectorized_5d: - raise ValueError( - "Deterministic inference supports aiter decode attention only with the " - "default paged KV cache: not with sliding-window models, " - "SGLANG_USE_AITER_UNIFIED_ATTN or a vectorized 5D KV cache layout." - ) + # Deterministic mode allows aiter decode only for the paged_attention_ragged + # path; the unified-attention and 5D-pool kernels are not verified batch-invariant. + if get_exec().deterministic.enable_deterministic_inference and ( + self.use_triton_unified_attention or self.kv_cache_is_vectorized_5d + ): + raise ValueError( + "Deterministic inference supports aiter decode attention only with the " + "default paged KV cache: not with sliding-window models, " + "SGLANG_USE_AITER_UNIFIED_ATTN or a vectorized 5D KV cache layout." + ) # Decode metadata format. unified_attention and pa_decode_gluon (the # SHUFFLE 5D decode kernel) both read kv_indices as a 2-D diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index d87a298bd8c2..a654ac4df364 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -1968,7 +1968,10 @@ def test_deterministic_incompatible_backend_raises(self): view = ResolvedView( SimpleNamespace( - enable_deterministic_inference=True, attention_backend="flashmla" + enable_deterministic_inference=True, + attention_backend="flashmla", + prefill_attention_backend=None, + decode_attention_backend=None, ) ) with self.assertRaises(ValueError): @@ -1990,11 +1993,16 @@ def test_deterministic_incompatible_phase_backend_raises(self): return_value=False, ), ): + phases = { + "prefill_attention_backend": None, + "decode_attention_backend": None, + field: "aiter", + } view = ResolvedView( SimpleNamespace( enable_deterministic_inference=True, attention_backend="triton", - **{field: "aiter"}, + **phases, ) ) with self.assertRaisesRegex(ValueError, field.replace("_", "-")): From 8da4261f8e5b9bac3fd0bb564d310e2f36eb7299 Mon Sep 17 00:00:00 2001 From: Aarne Talman Date: Thu, 15 Oct 2026 18:10:18 +0000 Subject: [PATCH 3/4] [Perf] [AMD] deterministic: Refuse AITER decode with speculative decoding With speculative decoding, target verify runs on the decode backend with --speculative-attention-mode decode, and some draft workers (frozen-KV MTP) use the decode backend too. That would run AITER's verify and extend kernels, which are not verified batch-invariant. Refuse --decode-attention-backend aiter in deterministic mode when a speculative algorithm is set, and say in the error why AITER decode was refused. --- python/sglang/srt/arg_groups/overrides.py | 36 ++++++++++++++------ test/registered/unit/test_model_overrides.py | 20 ++++++----- 2 files changed, 36 insertions(+), 20 deletions(-) diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 4612432cfbbc..39e61947c71b 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -1103,29 +1103,43 @@ def _deterministic_attention_backend(view: Any) -> dict: from sglang.srt.server_args import DETERMINISTIC_ATTENTION_BACKEND_CHOICES decode_choices = list(DETERMINISTIC_ATTENTION_BACKEND_CHOICES) - # AITER paged decode reduces fixed 256-token KV partitions per request, so it is - # batch-invariant (verified on gfx950); MLA decode sizes its KV splits by batch size. - if ( - view.decode_attention_backend == "aiter" - and is_gfx95_supported() - and not use_mla_backend(view) - ): - decode_choices.append("aiter") + decode_note = "" + if view.decode_attention_backend == "aiter": + # AITER paged decode reduces fixed 256-token KV partitions per request, so it is + # batch-invariant (verified on gfx950); MLA decode sizes its KV splits by batch + # size, and speculative decoding would run AITER's unverified verify kernels. + if ( + is_gfx95_supported() + and not use_mla_backend(view) + and view.speculative_algorithm is None + ): + decode_choices.append("aiter") + else: + decode_note = ( + " (aiter decode is accepted only on gfx950, for non-MLA models and " + "without speculative decoding)" + ) # A per-phase backend overrides --attention-backend for its phase. Unset (None) # means the phase uses --attention-backend, which is checked below. - for flag, backend, choices in ( + for flag, backend, choices, note in ( ( "--prefill-attention-backend", view.prefill_attention_backend, DETERMINISTIC_ATTENTION_BACKEND_CHOICES, + "", + ), + ( + "--decode-attention-backend", + view.decode_attention_backend, + decode_choices, + decode_note, ), - ("--decode-attention-backend", view.decode_attention_backend, decode_choices), ): if backend is not None and backend not in choices: raise ValueError( f"Currently only {choices} attention backends are supported for deterministic inference, " - f"but {flag} is '{backend}'." + f"but {flag} is '{backend}'{note}." ) if view.attention_backend is None: diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index a654ac4df364..b61d0aeb6b07 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -2009,23 +2009,24 @@ def test_deterministic_incompatible_phase_backend_raises(self): _deterministic_attention_backend(view) def test_deterministic_aiter_decode_only_on_gfx950_without_mla(self): - """AITER decode attention is batch-invariant only for non-MLA models on gfx950; - AITER prefill attention never is.""" + """AITER decode attention is batch-invariant only for non-MLA models on gfx950, + and its speculative verify path is not; AITER prefill attention never is.""" from sglang.srt.arg_groups.overrides import ( ResolvedView, _deterministic_attention_backend, ) cases = [ - # (gfx950, MLA model, prefill backend, accepted) - (True, False, "triton", True), - (False, False, "triton", False), - (True, True, "triton", False), - (True, False, "aiter", False), + # (gfx950, MLA model, prefill backend, speculative algorithm, accepted) + (True, False, "triton", None, True), + (False, False, "triton", None, False), + (True, True, "triton", None, False), + (True, False, "aiter", None, False), + (True, False, "triton", "EAGLE", False), ] - for gfx950, mla, prefill, accepted in cases: + for gfx950, mla, prefill, spec, accepted in cases: with ( - self.subTest(gfx950=gfx950, mla=mla, prefill=prefill), + self.subTest(gfx950=gfx950, mla=mla, prefill=prefill, spec=spec), patch( "sglang.srt.arg_groups.overrides.is_gfx95_supported", return_value=gfx950, @@ -2040,6 +2041,7 @@ def test_deterministic_aiter_decode_only_on_gfx950_without_mla(self): attention_backend="triton", prefill_attention_backend=prefill, decode_attention_backend="aiter", + speculative_algorithm=spec, ) ) if accepted: From 5438fe5d4151029e44905cec5c9e747446da68ad Mon Sep 17 00:00:00 2001 From: Aarne Talman Date: Fri, 30 Oct 2026 17:41:19 +0000 Subject: [PATCH 4/4] [Perf] [AMD] deterministic: Treat missing per-phase backend fields as unset The attention-backend check read view.prefill_attention_backend and view.decode_attention_backend directly; views without those fields (as in TestDeterministicAttentionBackendOnXpu) raised AttributeError. Read them with getattr, as the check did before this PR. --- python/sglang/srt/arg_groups/overrides.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 39e61947c71b..2fffff7ad2c8 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -1102,9 +1102,13 @@ def _deterministic_attention_backend(view: Any) -> dict: return {} from sglang.srt.server_args import DETERMINISTIC_ATTENTION_BACKEND_CHOICES + # Views without the per-phase fields (e.g. partial views in tests) mean unset. + prefill_backend = getattr(view, "prefill_attention_backend", None) + decode_backend = getattr(view, "decode_attention_backend", None) + decode_choices = list(DETERMINISTIC_ATTENTION_BACKEND_CHOICES) decode_note = "" - if view.decode_attention_backend == "aiter": + if decode_backend == "aiter": # AITER paged decode reduces fixed 256-token KV partitions per request, so it is # batch-invariant (verified on gfx950); MLA decode sizes its KV splits by batch # size, and speculative decoding would run AITER's unverified verify kernels. @@ -1125,13 +1129,13 @@ def _deterministic_attention_backend(view: Any) -> dict: for flag, backend, choices, note in ( ( "--prefill-attention-backend", - view.prefill_attention_backend, + prefill_backend, DETERMINISTIC_ATTENTION_BACKEND_CHOICES, "", ), ( "--decode-attention-backend", - view.decode_attention_backend, + decode_backend, decode_choices, decode_note, ),