From 5ae0ac42446558aff7d035e283a3c9826928ad48 Mon Sep 17 00:00:00 2001 From: Kaixi Hou Date: Fri, 14 Nov 2025 12:51:11 -0800 Subject: [PATCH] [NVIDIA] Fix use case of SGLANG_ENABLE_FLASHINFER_GEMM (#13274) Co-authored-by: Baizhou Zhang --- docs/references/environment_variables.md | 2 +- python/sglang/srt/environ.py | 5 ++++- python/sglang/srt/layers/quantization/fp8_utils.py | 9 +++++---- python/sglang/srt/models/deepseek_v2.py | 4 +++- 4 files changed, 13 insertions(+), 7 deletions(-) diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 7e90f642b..fda66255c 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -68,7 +68,7 @@ SGLang supports various environment variables that can be used to configure its | `SGLANG_INT4_WEIGHT` | Enable INT4 weight quantization | `false` | | `SGLANG_MOE_PADDING` | Enable MoE padding (sets padding size to 128 if value is `1`, often set to `1` in Docker builds) | `0` | | `SGLANG_FORCE_FP8_MARLIN` | Force using FP8 MARLIN kernels even if other FP8 kernels are available | `false` | -| `SGLANG_ENABLE_FLASHINFER_GEMM` | Use flashinfer kernels when running blockwise fp8 GEMM on Blackwell GPUs | `false` | +| `SGLANG_ENABLE_FLASHINFER_FP8_GEMM` | Use flashinfer kernels when running blockwise fp8 GEMM on Blackwell GPUs | `false` | | `SGLANG_FLASHINFER_FP4_GEMM_BACKEND` | Select backend for `mm_fp4` on Blackwell GPUS | `` | | `SGLANG_SUPPORT_CUTLASS_BLOCK_FP8` | Use Cutlass kernels when running blockwise fp8 GEMM on Hopper or Blackwell GPUs | `false` | | `SGLANG_CUTLASS_MOE` (deprecated) | Use Cutlass FP8 MoE kernel on Blackwell GPUs (deprecated, use --moe-runner-backend=cutlass) | `false` | diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index a7afa81d3..a91cc8e77 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -206,7 +206,7 @@ class Envs: # Flashinfer SGLANG_IS_FLASHINFER_AVAILABLE = EnvBool(True) - SGLANG_ENABLE_FLASHINFER_GEMM = EnvBool(False) + SGLANG_ENABLE_FLASHINFER_FP8_GEMM = EnvBool(False) # Default to the pick from flashinfer SGLANG_FLASHINFER_FP4_GEMM_BACKEND = EnvStr("") SGLANG_FLASHINFER_WORKSPACE_SIZE = EnvInt(384 * 1024 * 1024) @@ -307,6 +307,9 @@ def _print_deprecated_env(new_name: str, old_name: str): def _convert_SGL_to_SGLANG(): _print_deprecated_env("SGLANG_LOG_GC", "SGLANG_GC_LOG") + _print_deprecated_env( + "SGLANG_ENABLE_FLASHINFER_FP8_GEMM", "SGLANG_ENABLE_FLASHINFER_GEMM" + ) for key, value in os.environ.items(): if key.startswith("SGL_"): diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index ba589cc74..747f123b3 100644 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -2,6 +2,7 @@ from typing import Callable, List, Optional, Tuple import torch +from sglang.srt.environ import envs from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.quantization.fp8_kernel import sglang_per_token_group_quant_fp8 from sglang.srt.layers.quantization.mxfp4_tensor import MXFP4QuantizeUtil @@ -127,17 +128,17 @@ def cutlass_block_fp8_supported() -> bool: CUTLASS_BLOCK_FP8_SUPPORTED = cutlass_block_fp8_supported() -ENABLE_FLASHINFER_GEMM = ( - get_bool_env_var("SGLANG_ENABLE_FLASHINFER_GEMM") +ENABLE_FLASHINFER_FP8_GEMM = ( + envs.SGLANG_ENABLE_FLASHINFER_FP8_GEMM.get() and is_blackwell_supported() and is_flashinfer_available() ) -if ENABLE_FLASHINFER_GEMM: +if ENABLE_FLASHINFER_FP8_GEMM: from flashinfer.gemm import gemm_fp8_nt_groupwise def dispatch_w8a8_block_fp8_linear() -> Callable: - if ENABLE_FLASHINFER_GEMM: + if ENABLE_FLASHINFER_FP8_GEMM: return flashinfer_gemm_w8a8_block_fp8_linear elif CUTLASS_BLOCK_FP8_SUPPORTED: return cutlass_w8a8_block_fp8_linear_with_fallback diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 3d5ea2c65..486e88805 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -89,6 +89,7 @@ from sglang.srt.layers.quantization.fp8_kernel import ( per_token_group_quant_mla_deep_gemm_masked_fp8, ) from sglang.srt.layers.quantization.fp8_utils import ( + ENABLE_FLASHINFER_FP8_GEMM, block_quant_dequant, block_quant_to_tensor_quant, channel_quant_to_tensor_quant, @@ -3420,7 +3421,8 @@ class DeepseekV2ForCausalLM(nn.Module): self_attn.use_deep_gemm_bmm = True if ( - deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM + not ENABLE_FLASHINFER_FP8_GEMM + and deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 and hasattr(self.quant_config, "weight_block_size") and self.quant_config.weight_block_size is not None