From 73792629d42d8dca96bbee5cf9dd9b35c64e3dab Mon Sep 17 00:00:00 2001 From: HAI Date: Fri, 8 May 2026 01:31:00 -0700 Subject: [PATCH] [AMD] Intro SGLANG_DIFFUSION_AITER_FP8_ATTN (#24677) --- .../runtime/layers/attention/backends/aiter.py | 6 +++--- test/registered/amd/test_wan22_fp8_mla.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py index 49b02b031..7f69be843 100755 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py @@ -20,7 +20,7 @@ from sglang.srt.models.deepseek_common.utils import _use_aiter_gfx95 logger = logging.getLogger(__name__) -_use_fp8_attn = os.environ.get("SGLANG_AITER_FP8_ATTN", "0") == "1" +_use_fp8_attn = os.environ.get("SGLANG_DIFFUSION_AITER_FP8_ATTN", "0") == "1" _fp8_dtype = torch.float8_e4m3fn # ── MLA prefill ASM kernel constraints ────────────────────────────── @@ -47,7 +47,7 @@ _MLA_PREFILL_HEAD_TILE = 8 if _use_fp8_attn: - logger.info("DiT FP8 attention enabled via SGLANG_AITER_FP8_ATTN=1") + logger.info("DiT FP8 attention enabled via SGLANG_DIFFUSION_AITER_FP8_ATTN=1") def _can_use_mla_prefill(v_head_dim: int, num_heads: int) -> bool: @@ -335,7 +335,7 @@ class AITerImpl(AttentionImpl): ) -> torch.Tensor: """ Performs attention using one of: - - _mla_prefill_ps_attention (FP8, SGLANG_AITER_FP8_ATTN=1) + - _mla_prefill_ps_attention (FP8, SGLANG_DIFFUSION_AITER_FP8_ATTN=1) - flash_attn_func (BF16, default or FP8 fallback for unsupported shapes) Args: diff --git a/test/registered/amd/test_wan22_fp8_mla.py b/test/registered/amd/test_wan22_fp8_mla.py index 74f1b3f70..7bbfd82a4 100644 --- a/test/registered/amd/test_wan22_fp8_mla.py +++ b/test/registered/amd/test_wan22_fp8_mla.py @@ -28,7 +28,7 @@ logger = logging.getLogger(__name__) register_amd_ci(est_time=3600, suite="nightly-amd-fp8-mla-diffusion", nightly=True) MODEL = "Wan-AI/Wan2.2-T2V-A14B-Diffusers" -FP8_ENV = {"SGLANG_AITER_FP8_ATTN": "1"} +FP8_ENV = {"SGLANG_DIFFUSION_AITER_FP8_ATTN": "1"} FP8_MLA_CASES = [ DiffusionTestCase(