diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index c987abdf1..47d63ea8f 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -28,6 +28,7 @@ SGLang supports various environment variables that can be used to configure its | `SGLANG_SKIP_P2P_CHECK` | Skip P2P (peer-to-peer) access check | `false` | | `SGL_CHUNKED_PREFIX_CACHE_THRESHOLD` | Sets the threshold for enabling chunked prefix caching | `8192` | | `SGLANG_FUSED_MLA_ENABLE_ROPE_FUSION` | Enable RoPE fusion in Fused Multi-Layer Attention | `1` | +| `SGLANG_DISABLE_FA4_WARMUP` | Disable Flash Attention 4 warmup passes (set to `1`, `true`, `yes`, or `on` to disable) | `false` | ## DeepGEMM Configuration (Advanced Optimization) diff --git a/sgl-kernel/python/sgl_kernel/_fa4_interface.py b/sgl-kernel/python/sgl_kernel/_fa4_interface.py index da1adeec4..5411cfe16 100644 --- a/sgl-kernel/python/sgl_kernel/_fa4_interface.py +++ b/sgl-kernel/python/sgl_kernel/_fa4_interface.py @@ -8,6 +8,7 @@ import copy import gc import logging import math +import os from typing import Callable, Optional, Tuple logger = logging.getLogger(__name__) @@ -416,6 +417,15 @@ def warmup_flash_attn(f): - Executes sequentially to minimize peak GPU mem - Does not modify user tensors (clones) """ + disable_warmup = os.getenv("SGLANG_DISABLE_FA4_WARMUP", "").lower() in ( + "1", + "true", + "yes", + "on", + ) + if disable_warmup: + return f + done = False def _clone_args(args, kwargs):