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

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 Hopper DeepGEMM E8M0 for DeepSeekV3.1 scale_fmt (#23666)

Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: youkaichao <youkaichao@gmail.com> Co-authored-by: youkaichao <youkaichao@gmail.com>

Wentao Ye committed Aug 27, 2025 at 10:09 UTC 3af47c3cc693f432b59658019891393385aa0e2a

10 files changed +68 -53

tests/kernels/moe/test_block_fp8.py

+2 -3

@@ -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_e8m0_used

19 +from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used

21 dg_available = has_deep_gemm()

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

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

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

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

229 -@pytest.mark.skipif(is_blackwell_deep_gemm_e8m0_used(),

230 - reason="Not E8M0 scale MOE")

229 +@pytest.mark.skipif(is_deep_gemm_e8m0_used(), reason="Not E8M0 scale MOE")

230 @torch.inference_mode()

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

232 monkeypatch):

tests/kernels/moe/test_deepep_deepgemm_moe.py

+3 -4

@@ -20,8 +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_e8m0_used,

24 - is_deep_gemm_supported)

23 +from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used, is_deep_gemm_supported

25 from ...utils import multi_gpu_test

26 from .parallel_utils import ProcessGroupInfo, parallel_launch

@@ -374,7 +373,7 @@ NUM_EXPERTS = [32]

373 @multi_gpu_test(num_gpus=2)

374 @requires_deep_ep

375 @requires_deep_gemm

377 -@pytest.mark.skipif(is_blackwell_deep_gemm_e8m0_used(),

376 +@pytest.mark.skipif(is_deep_gemm_e8m0_used(),

377 reason="Skipping test for Blackwell DeepGEMM")

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

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

@@ -432,7 +431,7 @@ USE_FP8_DISPATCH = [False]

431 @multi_gpu_test(num_gpus=2)

432 @requires_deep_ep

433 @requires_deep_gemm

435 -@pytest.mark.skipif(is_blackwell_deep_gemm_e8m0_used(),

434 +@pytest.mark.skipif(is_deep_gemm_e8m0_used(),

435 reason="Skipping test for Blackwell DeepGEMM")

436 def test_ll_deepep_deepgemm_moe(

437 mnk: tuple[int, int, int],

vllm/envs.py

+7 -1

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

131 VLLM_TPU_USING_PATHWAYS: bool = False

132 VLLM_USE_DEEP_GEMM: bool = False

133 VLLM_USE_DEEP_GEMM_E8M0: bool = True

134 + VLLM_USE_DEEP_GEMM_E8M0_HOPPER: bool = False

135 VLLM_SKIP_DEEP_GEMM_WARMUP: bool = False

136 VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True

137 VLLM_USE_FLASHINFER_MOE_FP8: bool = False

@@ -954,9 +955,12 @@ environment_variables: dict[str, Callable[[], Any]] = {

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

956

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

957 - # E8M0 is faster on B200 but may reduce accuracy.

958 "VLLM_USE_DEEP_GEMM_E8M0":

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

960 + # TODO(wentao): unify the two E8M0 flags after verifying the correctness.

961 + # Whether to use E8M0 scaling when DeepGEMM is used on Hopper GPUs.

962 + "VLLM_USE_DEEP_GEMM_E8M0_HOPPER":

963 + lambda: bool(int(os.getenv("VLLM_USE_DEEP_GEMM_E8M0_HOPPER", "0"))),

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

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

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

@@ -1244,6 +1248,8 @@ def compute_hash() -> str:

1248 "VLLM_USE_FLASHINFER_SAMPLER",

1249 "VLLM_DISABLED_KERNELS",

1250 "VLLM_USE_DEEP_GEMM",

1251 + "VLLM_USE_DEEP_GEMM_E8M0",

1252 + "VLLM_USE_DEEP_GEMM_E8M0_HOPPER",

1253 "VLLM_USE_TRTLLM_FP4_GEMM",

1254 "VLLM_USE_FUSED_MOE_GROUPED_TOPK",

1255 "VLLM_USE_FLASHINFER_MOE_FP8",

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

15 + is_deep_gemm_e8m0_used)

17 logger = init_logger(__name__)

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

174 eps,

175 fp8_min,

176 fp8_max,

177 - is_blackwell_deep_gemm_e8m0_used(),

177 + is_deep_gemm_e8m0_used(),

178 BLOCK=group_size,

179 NUM_STAGES=4,

180 num_warps=1,

vllm/model_executor/layers/fused_moe/fused_moe.py

+3 -4

@@ -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_e8m0_used

43 +from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used

45 from .rocm_aiter_fused_moe import is_rocm_aiter_moe_enabled

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

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

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

1433 # accuracy issue.

1434 - if (allow_deep_gemm and use_fp8_w8a8

1435 - and (is_blackwell_deep_gemm_e8m0_used()

1436 - or _valid_deep_gemm(hidden_states, w1, w2))):

1434 + if (allow_deep_gemm and use_fp8_w8a8 and

1435 + (is_deep_gemm_e8m0_used() or _valid_deep_gemm(hidden_states, w1, w2))):

1436 assert apply_router_weight_on_input is False

1437 assert is_act_and_mul, (

1438 "DeepGemm only supports is_act_and_mul=True for now.")

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_e8m0_used

13 +from vllm.utils.deep_gemm import is_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_e8m0_used()

110 + if self.allow_deep_gemm and (is_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(

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

143 ):

144 use_deep_gemm = (self.allow_deep_gemm

145 and (_valid_deep_gemm(hidden_states, w1, w2)

146 - or is_blackwell_deep_gemm_e8m0_used()))

146 + or is_deep_gemm_e8m0_used()))

147

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

149 assert experts is not None

vllm/model_executor/layers/quantization/fp8.py

+4 -5

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

48 from vllm.platforms import current_platform

49 from vllm.scalar_type import scalar_types

50 from vllm.utils import has_deep_gemm

51 -from vllm.utils.deep_gemm import (is_blackwell_deep_gemm_e8m0_used,

52 - is_deep_gemm_supported)

51 +from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used, is_deep_gemm_supported

52 from vllm.utils.flashinfer import has_flashinfer_moe

54 if TYPE_CHECKING:

@@ -427,7 +426,7 @@ class Fp8LinearMethod(LinearMethodBase):

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

427 # requantize the weight and input to the specific scale

428 # at the same time.

430 - if is_blackwell_deep_gemm_e8m0_used():

429 + if is_deep_gemm_e8m0_used():

430 assert layer.weight_block_size is not None

431 block_sz = tuple(layer.weight_block_size)

432 requant_weight_ue8m0_inplace(

@@ -734,7 +733,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

733

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

735 # it ahead of time for performance reasons.

737 - if self.allow_deep_gemm and not is_blackwell_deep_gemm_e8m0_used():

736 + if self.allow_deep_gemm and not is_deep_gemm_e8m0_used():

737 # Lazy import to avoid CUDA initialization problems.

738 if _is_col_major(layer.w13_weight_scale_inv):

739 layer.w13_weight_scale_inv = \

@@ -871,7 +870,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

870 del layer.w13_input_scale

871 del layer.w2_input_scale

872

874 - if is_blackwell_deep_gemm_e8m0_used():

873 + if is_deep_gemm_e8m0_used():

874 assert layer.weight_block_size is not None

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

876 block_sz = tuple(layer.weight_block_size)

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

+2 -2

@@ -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

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

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

24 should_use_deepgemm_for_fp8_linear)

26 logger = init_logger(__name__)

@@ -385,7 +385,7 @@ def per_token_group_quant_fp8(

385 scaling factor.

386 """

387 if use_ue8m0 is None:

388 - use_ue8m0 = is_blackwell_deep_gemm_e8m0_used()

388 + use_ue8m0 = is_deep_gemm_e8m0_used()

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

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

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

vllm/transformers_utils/config.py

+18

@@ -501,6 +501,24 @@ def get_config(

501

502 if quantization_config is not None:

503 config.quantization_config = quantization_config

504 + # auto-enable DeepGEMM UE8M0 on Hopper if model config requests it

505 + scale_fmt = quantization_config.get("scale_fmt", None)

506 + if scale_fmt in ("ue8m0", ):

507 + if not envs.is_set("VLLM_USE_DEEP_GEMM_E8M0_HOPPER"):

508 + os.environ["VLLM_USE_DEEP_GEMM_E8M0_HOPPER"] = "1"

509 + logger.info_once(

510 + ("Detected quantization_config.scale_fmt=%s; "

511 + "enabling Hopper UE8M0."),

512 + scale_fmt,

513 + )

514 + elif not envs.VLLM_USE_DEEP_GEMM_E8M0_HOPPER:

515 + logger.warning_once(

516 + ("Model config requests UE8M0 "

517 + "(quantization_config.scale_fmt=%s), but "

518 + "VLLM_USE_DEEP_GEMM_E8M0_HOPPER=0 is set; "

519 + "Hopper UE8M0 disabled."),

520 + scale_fmt,

521 + )

522

523 if hf_overrides_kw:

524 logger.debug("Overriding HF config with %s", hf_overrides_kw)

vllm/utils/deep_gemm.py

+24 -29

@@ -31,34 +31,33 @@ def is_deep_gemm_supported() -> bool:

33 @functools.cache

34 -def is_blackwell_deep_gemm_e8m0_used() -> bool:

34 +def is_deep_gemm_e8m0_used() -> bool:

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

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

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

37 """

38 if not is_deep_gemm_supported():

39 - logger.debug_once(

39 + logger.info_once(

40 "DeepGEMM E8M0 disabled: DeepGEMM not supported on this system.")

41 return False

43 - if not envs.VLLM_USE_DEEP_GEMM_E8M0:

44 - logger.debug_once("DeepGEMM E8M0 disabled: VLLM_USE_DEEP_GEMM_E8M0=0.")

45 - return False

46 -

43 _lazy_init()

45 if _fp8_gemm_nt_impl is None:

50 - logger.debug_once(

51 - "DeepGEMM E8M0 disabled: _fp8_gemm_nt_impl not found")

46 + logger.info_once("DeepGEMM E8M0 disabled: _fp8_gemm_nt_impl not found")

47 return False

54 - enabled = (current_platform.is_cuda()

55 - and current_platform.has_device_capability(100))

56 - if enabled:

57 - logger.debug_once("DeepGEMM E8M0 enabled on Blackwell GPU.")

58 - else:

59 - logger.debug_once(

60 - "DeepGEMM E8M0 disabled: not running on Blackwell GPU.")

61 - return enabled

49 + if current_platform.is_device_capability(100) and \

50 + envs.VLLM_USE_DEEP_GEMM_E8M0:

51 + logger.info_once("DeepGEMM E8M0 enabled on Blackwell GPU.")

52 + return True

53 +

54 + if current_platform.is_device_capability(90) and \

55 + envs.VLLM_USE_DEEP_GEMM_E8M0_HOPPER:

56 + logger.info_once("DeepGEMM E8M0 enabled on Hopper GPU.")

57 + return True

58 +

59 + logger.info_once("DeepGEMM E8M0 disabled on current configuration.")

60 + return False

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

@@ -124,20 +123,18 @@ def fp8_gemm_nt(*args, **kwargs):

123 _lazy_init()

124 if _fp8_gemm_nt_impl is None:

125 return _missing(*args, **kwargs)

127 - return _fp8_gemm_nt_impl(

128 - *args,

129 - disable_ue8m0_cast=not is_blackwell_deep_gemm_e8m0_used(),

130 - **kwargs)

126 + return _fp8_gemm_nt_impl(*args,

127 + disable_ue8m0_cast=not is_deep_gemm_e8m0_used(),

128 + **kwargs)

129

130

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

132 _lazy_init()

133 if _grouped_impl is None:

134 return _missing(*args, **kwargs)

137 - return _grouped_impl(

138 - *args,

139 - disable_ue8m0_cast=not is_blackwell_deep_gemm_e8m0_used(),

140 - **kwargs)

135 + return _grouped_impl(*args,

136 + disable_ue8m0_cast=not is_deep_gemm_e8m0_used(),

137 + **kwargs)

138

139

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

@@ -145,9 +142,7 @@ def fp8_m_grouped_gemm_nt_masked(*args, **kwargs):

142 if _grouped_masked_impl is None:

143 return _missing(*args, **kwargs)

144 return _grouped_masked_impl(

148 - *args,

149 - disable_ue8m0_cast=not is_blackwell_deep_gemm_e8m0_used(),

150 - **kwargs)

145 + *args, disable_ue8m0_cast=not is_deep_gemm_e8m0_used(), **kwargs)

146

147

148 def _ceil_to_ue8m0(x: torch.Tensor):

@@ -211,7 +206,7 @@ __all__ = [

206 "m_grouped_fp8_gemm_nt_contiguous",

207 "fp8_m_grouped_gemm_nt_masked",

208 "per_block_cast_to_fp8",

214 - "is_blackwell_deep_gemm_e8m0_used",

209 + "is_deep_gemm_e8m0_used",

210 "is_deep_gemm_supported",

211 "should_use_deepgemm_for_fp8_linear",

212 ]