diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index f9d50f8aeab8..2fffff7ad2c8 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -1102,17 +1102,48 @@ 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) + # 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 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. if ( - backend is not None - and backend not in DETERMINISTIC_ATTENTION_BACKEND_CHOICES + 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, note in ( + ( + "--prefill-attention-backend", + prefill_backend, + DETERMINISTIC_ATTENTION_BACKEND_CHOICES, + "", + ), + ( + "--decode-attention-backend", + decode_backend, + decode_choices, + decode_note, + ), + ): + 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}'{note}." ) 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 85b9a9850925..eec0e83fc3dc 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,16 @@ def __init__( self.use_triton_unified_attention = get_bool_env_var( "SGLANG_USE_AITER_UNIFIED_ATTN" ) + # 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 7c6acdd14a2d..b61d0aeb6b07 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): @@ -1983,17 +1986,70 @@ 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, + ), + ): + 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("_", "-")): _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, + 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, 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, spec, accepted in cases: + with ( + self.subTest(gfx950=gfx950, mla=mla, prefill=prefill, spec=spec), + 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", + speculative_algorithm=spec, + ) + ) + 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,