@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f7dcce7a4a

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Feature] Add `VLLM_USE_DEEP_GEMM_E8M0` Env to Control E8M0 Scale (#21968)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

Wentao Ye committed Aug 11, 2025 at 12:39 UTC f7dcce7a4aabb1445c2827ac5d978a9c5e18be30

9 files changed +65 -39

tests/kernels/moe/test_block_fp8.py

+3 -2

@@ -16,7 +16,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe import (

16 fused_topk, modular_triton_fused_moe)

17 from vllm.platforms import current_platform

18 from vllm.utils import has_deep_gemm

19 -from vllm.utils.deep_gemm import is_blackwell_deep_gemm_used

19 +from vllm.utils.deep_gemm import is_blackwell_deep_gemm_e8m0_used

21 dg_available = has_deep_gemm()

@@ -224,7 +224,8 @@ def test_w8a8_block_fp8_fused_moe(M, N, K, E, topk, block_size, dtype, seed,

224 @pytest.mark.parametrize("topk", TOP_KS)

225 @pytest.mark.parametrize("seed", SEEDS)

226 @pytest.mark.skipif(not dg_available, reason="DeepGemm kernels not available.")

227 -@pytest.mark.skipif(is_blackwell_deep_gemm_used(), reason="Not E8M0 scale MOE")

227 +@pytest.mark.skipif(is_blackwell_deep_gemm_e8m0_used(),

228 + reason="Not E8M0 scale MOE")

229 @torch.inference_mode()

230 def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed,

231 monkeypatch):

tests/kernels/moe/test_deepep_deepgemm_moe.py

+3 -3

@@ -20,7 +20,7 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (

20 FusedMoEModularKernel)

21 from vllm.platforms import current_platform

22 from vllm.utils import has_deep_ep, has_deep_gemm

23 -from vllm.utils.deep_gemm import (is_blackwell_deep_gemm_used,

23 +from vllm.utils.deep_gemm import (is_blackwell_deep_gemm_e8m0_used,

24 is_deep_gemm_supported)

26 from .parallel_utils import ProcessGroupInfo, parallel_launch

@@ -370,7 +370,7 @@ NUM_EXPERTS = [32]

370 @pytest.mark.parametrize("world_dp_size", [(2, 1)])

371 @requires_deep_ep

372 @requires_deep_gemm

373 -@pytest.mark.skipif(is_blackwell_deep_gemm_used(),

373 +@pytest.mark.skipif(is_blackwell_deep_gemm_e8m0_used(),

374 reason="Skipping test for Blackwell DeepGEMM")

375 def test_ht_deepep_deepgemm_moe(mnk: tuple[int, int, int], num_experts: int,

376 topk: int, world_dp_size: tuple[int, int]):

@@ -427,7 +427,7 @@ USE_FP8_DISPATCH = [False]

427 @pytest.mark.parametrize("world_dp_size", [(2, 1)])

428 @requires_deep_ep

429 @requires_deep_gemm

430 -@pytest.mark.skipif(is_blackwell_deep_gemm_used(),

430 +@pytest.mark.skipif(is_blackwell_deep_gemm_e8m0_used(),

431 reason="Skipping test for Blackwell DeepGEMM")

432 def test_ll_deepep_deepgemm_moe(

433 mnk: tuple[int, int, int],

vllm/envs.py

@@ -127,6 +127,7 @@ if TYPE_CHECKING:

127 VLLM_TPU_MOST_MODEL_LEN: Optional[int] = None

128 VLLM_TPU_USING_PATHWAYS: bool = False

129 VLLM_USE_DEEP_GEMM: bool = False

130 + VLLM_USE_DEEP_GEMM_E8M0: bool = True

131 VLLM_SKIP_DEEP_GEMM_WARMUP: bool = False

132 VLLM_USE_FLASHINFER_MOE_FP8: bool = False

133 VLLM_USE_FLASHINFER_MOE_FP4: bool = False

@@ -925,6 +926,10 @@ environment_variables: dict[str, Callable[[], Any]] = {

926 "VLLM_USE_DEEP_GEMM":

927 lambda: bool(int(os.getenv("VLLM_USE_DEEP_GEMM", "0"))),

928

929 + # Whether to use E8M0 scaling when DeepGEMM is used on Blackwell GPUs.

930 + # E8M0 is faster on B200 but may reduce accuracy.

931 + "VLLM_USE_DEEP_GEMM_E8M0":

932 + lambda: bool(int(os.getenv("VLLM_USE_DEEP_GEMM_E8M0", "1"))),

933 # DeepGemm JITs the kernels on-demand. The warmup attempts to make DeepGemm

934 # JIT all the required kernels before model execution so there is no

935 # JIT'ing in the hot-path. However, this warmup increases the engine

vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py

+2 -2

@@ -12,7 +12,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

12 from vllm.model_executor.layers.fused_moe.utils import _resize_cache

13 from vllm.triton_utils import tl, triton

14 from vllm.utils.deep_gemm import (fp8_m_grouped_gemm_nt_masked,

15 - is_blackwell_deep_gemm_used)

15 + is_blackwell_deep_gemm_e8m0_used)

17 logger = init_logger(__name__)

@@ -176,7 +176,7 @@ def silu_mul_fp8_quant_deep_gemm(

176 eps,

177 fp8_min,

178 fp8_max,

179 - is_blackwell_deep_gemm_used(),

179 + is_blackwell_deep_gemm_e8m0_used(),

180 BLOCK=group_size,

181 NUM_STAGES=8,

182 num_warps=1,

vllm/model_executor/layers/fused_moe/fused_moe.py

+3 -3

@@ -40,7 +40,7 @@ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (

40 from vllm.platforms import current_platform

41 from vllm.triton_utils import tl, triton

42 from vllm.utils import direct_register_custom_op, is_torch_equal_or_newer

43 -from vllm.utils.deep_gemm import is_blackwell_deep_gemm_used

43 +from vllm.utils.deep_gemm import is_blackwell_deep_gemm_e8m0_used

45 from .rocm_aiter_fused_moe import is_rocm_aiter_moe_enabled

@@ -1387,8 +1387,8 @@ def fused_experts(hidden_states: torch.Tensor,

1387 # E8M0 scale, which means we requantize the weight and input to the specific

1388 # scale. Fallen back to cutlass or triton for some cases would cause

1389 # accuracy issue.

1390 - should_use_deep_gemm = is_blackwell_deep_gemm_used() or _valid_deep_gemm(

1391 - hidden_states, w1, w2)

1390 + should_use_deep_gemm = is_blackwell_deep_gemm_e8m0_used(

1391 + ) or _valid_deep_gemm(hidden_states, w1, w2)

1392 if (allow_deep_gemm and use_fp8_w8a8 and should_use_deep_gemm):

1393 assert apply_router_weight_on_input is False

1394 assert is_act_and_mul, (

vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py

+3 -3

@@ -10,7 +10,7 @@ from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (

10 DeepGemmExperts, _valid_deep_gemm, _valid_deep_gemm_shape,

11 deep_gemm_block_shape)

12 from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts

13 -from vllm.utils.deep_gemm import is_blackwell_deep_gemm_used

13 +from vllm.utils.deep_gemm import is_blackwell_deep_gemm_e8m0_used

16 class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

@@ -107,7 +107,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

107 # Note: the deep gemm workspaces are strictly larger than the triton

108 # workspaces so we can be pessimistic here and allocate for DeepGemm

109 # even if we fall back to triton later, e.g. if expert maps are set.

110 - if self.allow_deep_gemm and (is_blackwell_deep_gemm_used()

110 + if self.allow_deep_gemm and (is_blackwell_deep_gemm_e8m0_used()

111 or _valid_deep_gemm_shape(M, N, K)):

112 assert self.deep_gemm_expert is not None

113 return self.deep_gemm_expert.workspace_shapes(

@@ -133,7 +133,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

133 extra_expert_args: Optional[dict[str, Any]]):

134 use_deep_gemm = (self.allow_deep_gemm

135 and (_valid_deep_gemm(hidden_states, w1, w2)

136 - or is_blackwell_deep_gemm_used()))

136 + or is_blackwell_deep_gemm_e8m0_used()))

137

138 experts = self.deep_gemm_expert if use_deep_gemm else self.triton_expert

139 assert experts is not None

vllm/model_executor/layers/quantization/fp8.py

+7 -12

@@ -45,7 +45,8 @@ from vllm.model_executor.utils import set_weight_attrs

45 from vllm.platforms import current_platform

46 from vllm.scalar_type import scalar_types

47 from vllm.utils import has_deep_gemm

48 -from vllm.utils.deep_gemm import is_blackwell_deep_gemm_used

48 +from vllm.utils.deep_gemm import (is_blackwell_deep_gemm_e8m0_used,

49 + is_deep_gemm_supported)

50 from vllm.utils.flashinfer import has_flashinfer_moe

52 if TYPE_CHECKING:

@@ -415,10 +416,10 @@ class Fp8LinearMethod(LinearMethodBase):

416 # Activations not quantized for marlin.

417 del layer.input_scale

418

418 - # On B200, DeepGemm only support E8M0 scale, which means we need to

419 + # On B200, if E8M0 for DeepGemm is used, we need to

420 # requantize the weight and input to the specific scale

421 # at the same time.

421 - if is_blackwell_deep_gemm_used():

422 + if is_blackwell_deep_gemm_e8m0_used():

423 assert layer.weight_block_size is not None

424 block_sz = tuple(layer.weight_block_size)

425 requant_weight_ue8m0_inplace(

@@ -505,15 +506,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):

506 elif not self.block_quant:

507 logger.warning_once("Model is not block quantized. Not using "

508 "DeepGemm kernels")

508 - elif (current_platform.is_cuda()

509 - and current_platform.is_device_capability(90)):

509 + elif (is_deep_gemm_supported()):

510 logger.info_once("Using DeepGemm kernels for Fp8MoEMethod.")

511 self.allow_deep_gemm = True

512 - elif (current_platform.is_cuda()

513 - and is_blackwell_deep_gemm_used()):

514 - logger.info_once("Using DeepGemm SM100 kernels for "

515 - "Fp8MoEMethod.")

516 - self.allow_deep_gemm = True

512 else:

513 logger.warning_once(

514 "DeepGemm not supported on the current platform.")

@@ -725,7 +720,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

720

721 # DeepGemm scales need to be transposed and aligned. We try to do

722 # it ahead of time for performance reasons.

728 - if self.allow_deep_gemm and not is_blackwell_deep_gemm_used():

723 + if self.allow_deep_gemm and not is_blackwell_deep_gemm_e8m0_used():

724 # Lazy import to avoid CUDA initialization problems.

725 if _is_col_major(layer.w13_weight_scale_inv):

726 layer.w13_weight_scale_inv = \

@@ -851,7 +846,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

846 del layer.w13_input_scale

847 del layer.w2_input_scale

848

854 - if is_blackwell_deep_gemm_used():

849 + if is_blackwell_deep_gemm_e8m0_used():

850 assert layer.weight_block_size is not None

851 # Re-quantise the expert weights so their scales are UE8M0.

852 block_sz = tuple(layer.weight_block_size)

vllm/model_executor/layers/quantization/utils/fp8_utils.py

+2 -4

@@ -20,7 +20,7 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

20 from vllm.platforms import current_platform

21 from vllm.triton_utils import tl, triton

22 from vllm.utils import cdiv, direct_register_custom_op, has_deep_gemm

23 -from vllm.utils.deep_gemm import is_blackwell_deep_gemm_used

23 +from vllm.utils.deep_gemm import is_blackwell_deep_gemm_e8m0_used

25 logger = init_logger(__name__)

@@ -394,10 +394,8 @@ def per_token_group_quant_fp8(

394 tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the

395 scaling factor.

396 """

397 - # TODO(wentao): refactor this

398 - # use_ue8m0 should be a global flag that could be set by user

397 if use_ue8m0 is None:

400 - use_ue8m0 = is_blackwell_deep_gemm_used()

398 + use_ue8m0 = is_blackwell_deep_gemm_e8m0_used()

399 dtype = current_platform.fp8_dtype() if dtype is None else dtype

400 assert (x.shape[-1] % group_size == 0), (

401 f"the last dimension of `x` {x.shape[-1]} must be divisible "

vllm/utils/deep_gemm.py

+37 -10

@@ -31,19 +31,37 @@ def is_deep_gemm_supported() -> bool:

33 @functools.cache

34 -def is_blackwell_deep_gemm_used() -> bool:

35 - """Return ``True`` if vLLM is configured to use DeepGEMM on a

36 - Blackwell-class GPU.

34 +def is_blackwell_deep_gemm_e8m0_used() -> bool:

35 + """Return ``True`` if vLLM is configured to use DeepGEMM "

36 + "E8M0 scale on a Blackwell-class GPU.

37 """

38 - if not (envs.VLLM_USE_DEEP_GEMM and has_deep_gemm()):

38 + if not (envs.VLLM_USE_DEEP_GEMM):

39 + logger.debug_once("DeepGEMM E8M0 disabled: VLLM_USE_DEEP_GEMM=0.")

40 + return False

41 +

42 + if not has_deep_gemm():

43 + logger.debug_once("DeepGEMM E8M0 disabled: DeepGEMM backend missing.")

44 + return False

45 +

46 + if not envs.VLLM_USE_DEEP_GEMM_E8M0:

47 + logger.debug_once("DeepGEMM E8M0 disabled: VLLM_USE_DEEP_GEMM_E8M0=0.")

48 return False

50 _lazy_init()

51 +

52 if _fp8_gemm_nt_impl is None:

53 + logger.debug_once(

54 + "DeepGEMM E8M0 disabled: _fp8_gemm_nt_impl not found")

55 return False

45 - return (current_platform.is_cuda()

46 - and current_platform.is_device_capability(100))

57 + enabled = (current_platform.is_cuda()

58 + and current_platform.has_device_capability(100))

59 + if enabled:

60 + logger.debug_once("DeepGEMM E8M0 enabled on Blackwell GPU.")

61 + else:

62 + logger.debug_once(

63 + "DeepGEMM E8M0 disabled: not running on Blackwell GPU.")

64 + return enabled

67 def _missing(*_: Any, **__: Any) -> NoReturn:

@@ -109,21 +127,30 @@ def fp8_gemm_nt(*args, **kwargs):

127 _lazy_init()

128 if _fp8_gemm_nt_impl is None:

129 return _missing(*args, **kwargs)

112 - return _fp8_gemm_nt_impl(*args, **kwargs)

130 + return _fp8_gemm_nt_impl(

131 + *args,

132 + disable_ue8m0_cast=not is_blackwell_deep_gemm_e8m0_used(),

133 + **kwargs)

134

135

136 def m_grouped_fp8_gemm_nt_contiguous(*args, **kwargs):

137 _lazy_init()

138 if _grouped_impl is None:

139 return _missing(*args, **kwargs)

119 - return _grouped_impl(*args, **kwargs)

140 + return _grouped_impl(

141 + *args,

142 + disable_ue8m0_cast=not is_blackwell_deep_gemm_e8m0_used(),

143 + **kwargs)

144

145

146 def fp8_m_grouped_gemm_nt_masked(*args, **kwargs):

147 _lazy_init()

148 if _grouped_masked_impl is None:

149 return _missing(*args, **kwargs)

126 - return _grouped_masked_impl(*args, **kwargs)

150 + return _grouped_masked_impl(

151 + *args,

152 + disable_ue8m0_cast=not is_blackwell_deep_gemm_e8m0_used(),

153 + **kwargs)

154

155

156 def _ceil_to_ue8m0(x: torch.Tensor):

@@ -181,6 +208,6 @@ __all__ = [

208 "m_grouped_fp8_gemm_nt_contiguous",

209 "fp8_m_grouped_gemm_nt_masked",

210 "per_block_cast_to_fp8",

184 - "is_blackwell_deep_gemm_used",

211 + "is_blackwell_deep_gemm_e8m0_used",

212 "is_deep_gemm_supported",

213 ]