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

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 15, 02:15 UTC

[Perf] Use Triton instead of Torch for DeepGEMM Per Token Group Quant (#20841)

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

Wentao Ye committed Jul 12, 2025 at 22:38 UTC 42d440c22bc85a95dc457d4a0970507f8e66f0da

6 files changed +26 -42

tests/kernels/moe/test_deepgemm.py

+4 -3

@@ -13,9 +13,10 @@ import torch

14 # vLLM fused-expert reference (Triton fallback + DeepGEMM option)

15 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts

16 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (

17 + per_token_group_quant_fp8)

18 from vllm.utils import has_deep_gemm

17 -from vllm.utils.deep_gemm import (calc_diff, per_block_cast_to_fp8,

18 - per_token_group_cast_to_fp8)

19 +from vllm.utils.deep_gemm import calc_diff, per_block_cast_to_fp8

21 BLOCK_SIZE = [128, 128]

@@ -81,7 +82,7 @@ def run_single_case(m, n, k, topk, num_experts, block_size):

82 """

83 tokens_bf16 = torch.randn(

84 m, k, device="cuda", dtype=torch.bfloat16).clamp_min_(-1).clamp_max_(1)

84 - _, a1_scale = per_token_group_cast_to_fp8(tokens_bf16, block_size[1])

85 + _, a1_scale = per_token_group_quant_fp8(tokens_bf16, block_size[1])

87 # expert weight tensors

88 w1, w2, w1_s, w2_s = make_block_quant_fp8_weights(num_experts, n, k,

tests/kernels/quantization/test_block_fp8.py

+2 -3

@@ -15,8 +15,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (

15 w8a8_block_fp8_matmul)

16 from vllm.platforms import current_platform

17 from vllm.utils import has_deep_gemm

18 -from vllm.utils.deep_gemm import (fp8_gemm_nt, per_block_cast_to_fp8,

19 - per_token_group_cast_to_fp8)

18 +from vllm.utils.deep_gemm import fp8_gemm_nt, per_block_cast_to_fp8

20 if current_platform.get_device_capability() < (9, 0):

21 pytest.skip("FP8 Triton requires CUDA 9.0 or higher",

@@ -117,7 +116,7 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):

116 A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * fp8_max

117 B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * fp8_max

118

120 - A_fp8, As_fp8 = per_token_group_cast_to_fp8(A_fp32, block_size[1])

119 + A_fp8, As_fp8 = per_token_group_quant_fp8(A_fp32, block_size[1])

120 B_fp8, Bs_fp8 = per_block_cast_to_fp8(B_fp32)

121

122 As = As_fp8.to(torch.float32)

vllm/model_executor/layers/fused_moe/deep_gemm_moe.py

+7 -6

@@ -15,9 +15,10 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize import (

15 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

16 TopKWeightAndReduceDelegate)

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

18 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (

19 + per_token_group_quant_fp8)

20 from vllm.utils import has_deep_gemm, round_up

19 -from vllm.utils.deep_gemm import (m_grouped_fp8_gemm_nt_contiguous,

20 - per_token_group_cast_to_fp8)

21 +from vllm.utils.deep_gemm import m_grouped_fp8_gemm_nt_contiguous

23 logger = init_logger(__name__)

@@ -170,10 +171,10 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

171 self.activation(activation, act_out, mm1_out.view(-1, N))

172

173 a2q_scale: Optional[torch.Tensor] = None

173 - a2q, a2q_scale = per_token_group_cast_to_fp8(act_out,

174 - self.block_shape[1],

175 - column_major_scales=True,

176 - out_q=quant_out)

174 + a2q, a2q_scale = per_token_group_quant_fp8(act_out,

175 + self.block_shape[1],

176 + column_major_scales=True,

177 + out_q=quant_out)

178

179 m_grouped_fp8_gemm_nt_contiguous((a2q, a2q_scale), (w2, w2_scale),

180 mm2_out, expert_ids)

vllm/model_executor/layers/fused_moe/utils.py

+1 -6

@@ -15,8 +15,6 @@ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (

15 from vllm.platforms import current_platform

16 from vllm.triton_utils import tl, triton

17 from vllm.utils import cdiv

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

19 - per_token_group_cast_to_fp8)

20 @triton.jit

@@ -119,10 +117,7 @@ def _fp8_quantize(

117 assert not per_act_token

118 assert len(block_shape) == 2

119 _, block_k = block_shape[0], block_shape[1]

122 - if is_blackwell_deep_gemm_used():

123 - A, A_scale = per_token_group_cast_to_fp8(A, block_k)

124 - else:

125 - A, A_scale = per_token_group_quant_fp8(A, block_k)

120 + A, A_scale = per_token_group_quant_fp8(A, block_k)

121 assert cdiv(A.size(-1), block_k) == A_scale.size(-1)

122

123 return A, A_scale

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

+12 -3

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

25 logger = init_logger(__name__)

@@ -256,6 +257,7 @@ def _per_token_group_quant_fp8(

257 # Information for float8

258 fp8_min,

259 fp8_max,

260 + use_ue8m0: tl.constexpr,

261 # Meta-parameters

262 BLOCK: tl.constexpr,

263 ):

@@ -285,7 +287,8 @@ def _per_token_group_quant_fp8(

287 y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32)

288 # Quant

289 _absmax = tl.maximum(tl.max(tl.abs(y)), eps)

288 - y_s = _absmax / fp8_max

290 + scale_raw = _absmax / fp8_max

291 + y_s = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) if use_ue8m0 else scale_raw

292 y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty)

293

294 tl.store(y_q_ptr + cols, y_q, mask=mask)

@@ -309,6 +312,7 @@ def _per_token_group_quant_fp8_colmajor(

312 # Information for float8

313 fp8_min,

314 fp8_max,

315 + use_ue8m0: tl.constexpr,

316 # Meta-parameters

317 BLOCK: tl.constexpr,

318 ):

@@ -347,7 +351,8 @@ def _per_token_group_quant_fp8_colmajor(

351 y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32)

352 # Quant

353 _absmax = tl.maximum(tl.max(tl.abs(y)), eps)

350 - y_s = _absmax / fp8_max

354 + scale_raw = _absmax / fp8_max

355 + y_s = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) if use_ue8m0 else scale_raw

356 y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty)

357

358 tl.store(y_q_ptr + cols, y_q, mask=mask)

@@ -373,9 +378,11 @@ def per_token_group_quant_fp8(

378 is supported for now.

379 column_major_scales: Outputs scales in column major.

380 out_q: Optional output tensor. If not provided, function will create.

376 - Returns:

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

382 scaling factor for quantization.

383 + Returns:

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

385 + scaling factor.

386 """

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

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

@@ -418,6 +425,7 @@ def per_token_group_quant_fp8(

425 eps,

426 fp8_min=fp8_min,

427 fp8_max=fp8_max,

428 + use_ue8m0=is_blackwell_deep_gemm_used(),

429 BLOCK=BLOCK,

430 num_warps=num_warps,

431 num_stages=num_stages,

@@ -433,6 +441,7 @@ def per_token_group_quant_fp8(

441 eps,

442 fp8_min=fp8_min,

443 fp8_max=fp8_max,

444 + use_ue8m0=is_blackwell_deep_gemm_used(),

445 BLOCK=BLOCK,

446 num_warps=num_warps,

447 num_stages=num_stages,

vllm/utils/deep_gemm.py

-21

@@ -49,7 +49,6 @@ if not has_deep_gemm():

49 _fp8_gemm_nt_impl: Callable[..., Any] | None = None

50 _grouped_impl: Callable[..., Any] | None = None

51 _grouped_masked_impl: Callable[..., Any] | None = None

52 - _per_token_cast_impl: Callable[..., Any] | None = None

52 _per_block_cast_impl: Callable[..., Any] | None = None

53 else:

54 _dg = importlib.import_module("deep_gemm") # type: ignore

@@ -74,12 +73,9 @@ else:

73 try:

74 _math_mod = importlib.import_module(

75 "deep_gemm.utils.math") # type: ignore

77 - _per_token_cast_impl = getattr(_math_mod, "per_token_cast_to_fp8",

78 - None)

76 _per_block_cast_impl = getattr(_math_mod, "per_block_cast_to_fp8",

77 None)

78 except ModuleNotFoundError:

82 - _per_token_cast_impl = None

79 _per_block_cast_impl = None

@@ -101,22 +97,6 @@ def fp8_m_grouped_gemm_nt_masked(*args, **kwargs):

97 return _grouped_masked_impl(*args, **kwargs)

104 -def per_token_group_cast_to_fp8(x, group_size, *args, **kwargs):

105 - """Wrapper for token-wise FP8 quantisation.

106 -

107 - • If DeepGEMM provides ``per_token_cast_to_fp8`` (new API), use it.

108 - • Otherwise, fall back to vLLM's ``per_token_group_quant_fp8``

109 - """

110 -

111 - if _per_token_cast_impl is not None and is_blackwell_deep_gemm_used():

112 - assert group_size == 128, "group_size must be 128 for deepgemm"

113 - return _per_token_cast_impl(x)

114 -

115 - from vllm.model_executor.layers.quantization.utils.fp8_utils import (

116 - per_token_group_quant_fp8 as _ptg)

117 - return _ptg(x, group_size, *args, **kwargs)

118 -

119 -

100 def per_block_cast_to_fp8(x, *args, **kwargs):

101 if _per_block_cast_impl is not None and is_blackwell_deep_gemm_used():

102 return _per_block_cast_impl(x)

@@ -146,7 +126,6 @@ __all__ = [

126 "fp8_gemm_nt",

127 "m_grouped_fp8_gemm_nt_contiguous",

128 "fp8_m_grouped_gemm_nt_masked",

149 - "per_token_group_cast_to_fp8",

129 "per_block_cast_to_fp8",

130 "is_blackwell_deep_gemm_used",

131 ]