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(