From a80bcb5a68494cb542800af6ef14469615d1b3c6 Mon Sep 17 00:00:00 2001
From: yinghui <32845984+cicirori@users.noreply.github.com>
Date: Fri, 31 Oct 2025 12:25:00 -0700
Subject: [PATCH] Add env var to disable FA4 warmup (#12430)

---
 docs/references/environment_variables.md       |  1 +
 sgl-kernel/python/sgl_kernel/_fa4_interface.py | 10 ++++++++++
 2 files changed, 11 insertions(+)

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):