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

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

[Feature] Integrate SM100 DeepGEMM support (#20087)

Wentao Ye committed Jul 10, 2025 at 23:18 UTC e2de455c349df8385b18fe447beb6325dcb6af9c

16 files changed +397 -114

benchmarks/kernels/benchmark_moe.py

@@ -86,6 +86,9 @@ def benchmark_config(

86 (num_experts, 2 * shard_intermediate_size), dtype=torch.float32

87 )

88 w2_scale = torch.randn((hidden_size, num_experts), dtype=torch.float32)

89 + if use_deep_gemm:

90 + # we use the default block shape for deepgemm

91 + block_quant_shape = [128, 128]

92 if use_fp8_w8a8:

93 if block_quant_shape:

94 block_n, block_k = block_quant_shape[0], block_quant_shape[1]

tests/kernels/moe/test_block_fp8.py

+8 -8

@@ -15,13 +15,13 @@ from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (

15 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 -dg_available = False

20 -try:

21 - import deep_gemm

22 - dg_available = True

23 -except ImportError:

24 - pass

21 +dg_available = has_deep_gemm()

22 +

23 +if dg_available:

24 + from deep_gemm import get_m_alignment_for_contiguous_layout

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

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

@@ -224,6 +224,7 @@ 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")

228 @torch.inference_mode()

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

230 monkeypatch):

@@ -238,8 +239,7 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed,

239 torch.manual_seed(seed)

240

241 monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", str(chunk_size))

241 -

242 - block_m = deep_gemm.get_m_alignment_for_contiguous_layout()

242 + block_m = get_m_alignment_for_contiguous_layout()

243 block_size = [block_m, block_m]

244 dtype = torch.bfloat16

245

tests/kernels/moe/test_deepep_deepgemm_moe.py

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

25 from .parallel_utils import ProcessGroupInfo, parallel_launch

26 from .utils import make_test_weights

@@ -368,6 +369,8 @@ NUM_EXPERTS = [32]

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

370 @requires_deep_ep

371 @requires_deep_gemm

372 +@pytest.mark.skipif(is_blackwell_deep_gemm_used(),

373 + reason="Skipping test for Blackwell DeepGEMM")

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

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

376 """

@@ -423,6 +426,8 @@ USE_FP8_DISPATCH = [False]

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

427 @requires_deep_ep

428 @requires_deep_gemm

429 +@pytest.mark.skipif(is_blackwell_deep_gemm_used(),

430 + reason="Skipping test for Blackwell DeepGEMM")

431 def test_ll_deepep_deepgemm_moe(

432 mnk: tuple[int, int, int],

433 num_experts: int,

tests/kernels/moe/test_deepgemm.py

+8 -47

@@ -13,48 +13,18 @@ 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 cdiv

16 +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)

20 -has_deep_gemm = importlib.util.find_spec("deep_gemm") is not None

21 -

22 -if has_deep_gemm:

23 - import deep_gemm

24 - BLOCK_M = deep_gemm.get_m_alignment_for_contiguous_layout()

25 - BLOCK_SIZE = [BLOCK_M, BLOCK_M]

20 +BLOCK_SIZE = [128, 128]

22 requires_deep_gemm = pytest.mark.skipif(

28 - not has_deep_gemm,

23 + not has_deep_gemm(),

24 reason="Requires deep_gemm kernels",

25 )

33 -def calc_diff(x: torch.Tensor, y: torch.Tensor):

34 - x, y = x.double(), y.double()

35 - denominator = (x * x + y * y).sum()

36 - sim = 2 * (x * y).sum() / denominator

37 - return 1 - sim

38 -

39 -

40 -def per_block_cast_to_fp8(

41 - x: torch.Tensor,

42 - block_size_n: int = 128) -> tuple[torch.Tensor, torch.Tensor]:

43 - assert x.dim() == 2

44 - m, n = x.shape

45 - x_padded = torch.zeros(

46 - (cdiv(m, 128) * 128, cdiv(n, block_size_n) * block_size_n),

47 - dtype=x.dtype,

48 - device=x.device)

49 - x_padded[:m, :n] = x

50 - x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, block_size_n)

51 - x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)

52 - x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn)

53 - x_scaled_sub = x_scaled.view_as(x_padded)[:m, :n].contiguous()

54 - scales = (x_amax / 448.0).view(x_view.size(0), x_view.size(2))

55 - return x_scaled_sub, scales

56 -

57 -

28 def make_block_quant_fp8_weights(

29 e: int,

30 n: int,

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

81 """

82 tokens_bf16 = torch.randn(

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

114 - _, a1_scale = per_token_group_quant_fp8(tokens_bf16, block_size[1])

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

86 # expert weight tensors

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

@@ -155,17 +125,8 @@ def run_single_case(m, n, k, topk, num_experts, block_size):

125 block_shape=block_size,

126 allow_deep_gemm=True,

127 )

158 -

159 - base = out_triton.abs().mean()

160 - atol = 0.1 * base.clamp(min=1e-2) # 10% of mean, but not lower than 1e-3

161 - rtol = 0.05

162 - # ----- Compare -----

163 - torch.testing.assert_close(

164 - out_deepgemm.to(torch.float32),

165 - out_triton.to(torch.float32),

166 - rtol=rtol,

167 - atol=float(atol),

168 - )

128 + diff = calc_diff(out_deepgemm, out_triton)

129 + assert diff < 0.001, f"Diff exceeded 1%: {diff}"

130

131

132 # Note: W1 has shape (E, 2N, K), so N = 512

tests/kernels/quantization/test_block_fp8.py

+11 -16

@@ -8,19 +8,15 @@ import pytest

8 import torch

10 from tests.kernels.quant_utils import (native_per_token_group_quant_fp8,

11 - native_w8a8_block_matmul,

12 - per_block_cast_to_fp8)

11 + native_w8a8_block_matmul)

12 from vllm.config import VllmConfig

13 from vllm.model_executor.layers.quantization.utils.fp8_utils import (

15 - per_token_group_quant_fp8, w8a8_block_fp8_matmul)

14 + get_col_major_tma_aligned_tensor, per_token_group_quant_fp8,

15 + w8a8_block_fp8_matmul)

16 from vllm.platforms import current_platform

17 -

18 -dg_available = False

19 -try:

20 - import deep_gemm

21 - dg_available = True

22 -except ImportError:

23 - pass

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)

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

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

@@ -106,7 +102,8 @@ def test_w8a8_block_fp8_matmul(M, N, K, block_size, out_dtype, seed):

102 @pytest.mark.parametrize(

103 "M,N,K,block_size,out_dtype,seed",

104 itertools.product(M, N, K, BLOCK_SIZE, OUT_DTYPES, SEEDS))

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

105 +@pytest.mark.skipif(not has_deep_gemm(),

106 + reason="DeepGemm kernels not available.")

107 @torch.inference_mode()

108 def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):

109 # only aligned sizes

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

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

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

119

123 - _, block_k = block_size[0], block_size[1]

124 -

125 - A_fp8, As_fp8 = per_token_group_quant_fp8(A_fp32, block_k)

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

121 B_fp8, Bs_fp8 = per_block_cast_to_fp8(B_fp32)

122

123 As = As_fp8.to(torch.float32)

@@ -132,14 +127,14 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):

127 out_dtype)

128

129 # Transpose earlier so that the testing will not trigger transposing kernels

135 - As_fp8 = deep_gemm.get_col_major_tma_aligned_tensor(As_fp8)

130 + As_fp8 = get_col_major_tma_aligned_tensor(As_fp8)

131

132 out = torch.zeros((M, N), device='cuda', dtype=out_dtype)

133

134 assert As_fp8.shape == (M, (K + 127) //

135 128), f"{As_fp8.shape} != {(M, (K + 127) // 128)}"

136

142 - deep_gemm.gemm_fp8_fp8_bf16_nt((A_fp8, As_fp8), (B_fp8, Bs_fp8), out)

137 + fp8_gemm_nt((A_fp8, As_fp8), (B_fp8, Bs_fp8), out)

138

139 rel_diff = (torch.mean(

140 torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))) /

vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py

+9 -12

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

11 TopKWeightAndReduceDelegate)

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

16 logger = init_logger(__name__)

@@ -271,7 +272,6 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

272 assert expert_tokens_meta is not None

273 expert_num_tokens = expert_tokens_meta.expert_num_tokens

274

274 - import deep_gemm as dg

275 assert hidden_states.ndim == 3

276 assert self.block_shape is not None

277

@@ -289,18 +289,15 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

289 # for the M expectation of each batch, correctly setting this value

290 # may lead to better performance.

291 expected_m = max_num_tokens

292 -

293 - dg.m_grouped_gemm_fp8_fp8_bf16_nt_masked((a1q, a1q_scale),

294 - (w1, w1_scale),

295 - out=workspace1,

296 - masked_m=expert_num_tokens,

297 - expected_m=expected_m)

292 + fp8_m_grouped_gemm_nt_masked((a1q, a1q_scale), (w1, w1_scale),

293 + out=workspace1,

294 + masked_m=expert_num_tokens,

295 + expected_m=expected_m)

296

297 a2q, a2q_scale = silu_mul_fp8_quant_deep_gemm(workspace1,

298 expert_num_tokens)

299

302 - dg.m_grouped_gemm_fp8_fp8_bf16_nt_masked((a2q, a2q_scale),

303 - (w2, w2_scale),

304 - out=output,

305 - masked_m=expert_num_tokens,

306 - expected_m=expected_m)

300 + fp8_m_grouped_gemm_nt_masked((a2q, a2q_scale), (w2, w2_scale),

301 + out=output,

302 + masked_m=expert_num_tokens,

303 + expected_m=expected_m)

vllm/model_executor/layers/fused_moe/deep_gemm_moe.py

+11 -11

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

14 MoEPrepareAndFinalizeNoEP)

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 (

18 - _resize_cache, per_token_group_quant_fp8)

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

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

22 logger = init_logger(__name__)

@@ -127,7 +128,6 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

128 workspace2: torch.Tensor,

129 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

130 ):

130 - import deep_gemm as dg

131 assert self.block_shape is not None

132

133 a1q = hidden_states

@@ -164,19 +164,19 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

164 (M_sum, N // 2))

165 mm2_out = _resize_cache(workspace2, (M_sum, K))

166

167 - dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(

168 - (a1q, a1q_scale), (w1, w1_scale), mm1_out, expert_ids)

167 + m_grouped_fp8_gemm_nt_contiguous((a1q, a1q_scale), (w1, w1_scale),

168 + mm1_out, expert_ids)

169

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

171

172 a2q_scale: Optional[torch.Tensor] = None

173 - a2q, a2q_scale = per_token_group_quant_fp8(act_out,

174 - self.block_shape[1],

175 - column_major_scales=True,

176 - out_q=quant_out)

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)

177

178 - dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(

179 - (a2q, a2q_scale), (w2, w2_scale), mm2_out, expert_ids)

178 + m_grouped_fp8_gemm_nt_contiguous((a2q, a2q_scale), (w2, w2_scale),

179 + mm2_out, expert_ids)

180

181 torch.index_select(mm2_out, 0, inv_perm, out=output.view((-1, K)))

182

vllm/model_executor/layers/fused_moe/fused_moe.py

+9 -3

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

34 from vllm.platforms import current_platform

35 from vllm.triton_utils import tl, triton

36 from vllm.utils import direct_register_custom_op

37 +from vllm.utils.deep_gemm import is_blackwell_deep_gemm_used

39 from .rocm_aiter_fused_moe import is_rocm_aiter_moe_enabled

@@ -1171,9 +1172,15 @@ def fused_experts(

1172 allow_cutlass_block_scaled_grouped_gemm: bool = False) -> torch.Tensor:

1173 # For now, disable DeepGemm for small N (<= 512) until better

1174 # permute/unpermute ops are available.

1175 + # However, on B200, we use DeepGemm for all cases becuase they only support

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

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

1178 + # accuracy issue.

1179 N = w1.size(1)

1175 - if (allow_deep_gemm and use_fp8_w8a8 and N > 512

1176 - and _valid_deep_gemm(hidden_states, w1, w2)):

1180 + should_use_deep_gemm = ((N > 512

1181 + and _valid_deep_gemm(hidden_states, w1, w2))

1182 + or is_blackwell_deep_gemm_used())

1183 + if (allow_deep_gemm and use_fp8_w8a8 and should_use_deep_gemm):

1184 assert apply_router_weight_on_input is False

1185 return deep_gemm_moe_fp8(

1186 hidden_states=hidden_states,

@@ -1363,7 +1370,6 @@ def fused_experts_impl(

1370

1371 curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]

1372 curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]

1366 -

1373 qcurr_hidden_states, a1q_scale = moe_kernel_quantize_input(

1374 A=curr_hidden_states,

1375 A_scale=a1_scale,

vllm/model_executor/layers/fused_moe/prepare_finalize.py

-1

@@ -48,7 +48,6 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):

48 assert topk == 1, \

49 "apply_router_weight_on_input is only implemented for topk=1"

50 a1.mul_(topk_weights.to(a1.dtype))

51 -

51 a1q, a1q_scale = moe_kernel_quantize_input(

52 a1, a1_scale, quant_config.quant_dtype,

53 quant_config.per_act_token_quant, quant_config.block_shape)

vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py

+5 -2

@@ -9,6 +9,7 @@ from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

9 from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (

10 DeepGemmExperts, _valid_deep_gemm, _valid_deep_gemm_shape)

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

12 +from vllm.utils.deep_gemm import is_blackwell_deep_gemm_used

15 class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

@@ -102,7 +103,8 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

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

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

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

105 - if self.allow_deep_gemm and _valid_deep_gemm_shape(M, N, K):

106 + if self.allow_deep_gemm and (_valid_deep_gemm_shape(M, N, K)

107 + or is_blackwell_deep_gemm_used()):

108 assert self.deep_gemm_expert is not None

109 return self.deep_gemm_expert.workspace_shapes(

110 a, aq, M, N, K, topk, global_num_experts, local_num_experts)

@@ -132,7 +134,8 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

134 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

135 ):

136 use_deep_gemm = (self.allow_deep_gemm

135 - and _valid_deep_gemm(hidden_states, w1, w2))

137 + and (_valid_deep_gemm(hidden_states, w1, w2)

138 + or is_blackwell_deep_gemm_used()))

139

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

141 assert experts is not None

vllm/model_executor/layers/fused_moe/utils.py

+6 -1

@@ -15,6 +15,8 @@ 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)

22 @triton.jit

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

117 assert not per_act_token

118 assert len(block_shape) == 2

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

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

120 + if is_blackwell_deep_gemm_used():

121 + A, A_scale = per_token_group_cast_to_fp8(A, block_k)

122 + else:

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

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

125

126 return A, A_scale

vllm/model_executor/layers/mamba/ops/causal_conv1d.py

+1 -2

@@ -8,10 +8,9 @@ from typing import Optional, Union

9 import numpy as np

10 import torch

11 -import triton

12 -import triton.language as tl

12 from vllm.attention.backends.utils import PAD_SLOT_ID

13 +from vllm.triton_utils import tl, triton

16 @triton.jit()

vllm/model_executor/layers/quantization/deepgemm.py

+3 -5

@@ -6,10 +6,8 @@ import torch

7 from vllm.platforms import current_platform

8 from vllm.triton_utils import triton

9 -from vllm.utils import direct_register_custom_op, has_deep_gemm

10 -

11 -if has_deep_gemm():

12 - import deep_gemm

9 +from vllm.utils import direct_register_custom_op

10 +from vllm.utils.deep_gemm import fp8_gemm_nt

12 logger = logging.getLogger(__name__)

@@ -57,7 +55,7 @@ def w8a8_block_fp8_matmul_deepgemm(

55 output_dtype)

56 # Deepgemm only supports output tensor type as bfloat16

57 assert C.dtype == torch.bfloat16

60 - deep_gemm.gemm_fp8_fp8_bf16_nt((A, As), (B, Bs), C)

58 + fp8_gemm_nt((A, As), (B, Bs), C)

59 return C

vllm/model_executor/layers/quantization/fp8.py

+42 -4

@@ -23,6 +23,8 @@ from vllm.model_executor.layers.quantization import QuantizationMethods

23 from vllm.model_executor.layers.quantization.base_config import (

24 QuantizationConfig, QuantizeMethodBase)

25 from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod

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

27 + get_col_major_tma_aligned_tensor, requant_weight_ue8m0_inplace)

28 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (

29 apply_fp8_marlin_linear, prepare_fp8_layer_for_marlin,

30 prepare_moe_fp8_layer_for_marlin)

@@ -40,6 +42,7 @@ from vllm.model_executor.utils import set_weight_attrs

42 from vllm.platforms import current_platform

43 from vllm.scalar_type import scalar_types

44 from vllm.utils import has_deep_gemm

45 +from vllm.utils.deep_gemm import is_blackwell_deep_gemm_used

47 if TYPE_CHECKING:

48 from vllm.model_executor.models.utils import WeightsMapper

@@ -393,6 +396,19 @@ class Fp8LinearMethod(LinearMethodBase):

396 # Activations not quantized for marlin.

397 del layer.input_scale

398

399 + # On B200, DeepGemm only support E8M0 scale, which means we need to

400 + # requantize the weight and input to the specific scale

401 + # at the same time.

402 + if is_blackwell_deep_gemm_used():

403 + assert layer.weight_block_size is not None

404 + block_sz = tuple(layer.weight_block_size)

405 + requant_weight_ue8m0_inplace(

406 + layer.weight.data,

407 + layer.weight_scale_inv.data if hasattr(

408 + layer, "weight_scale_inv") else layer.weight_scale.data,

409 + block_sz,

410 + )

411 +

412 def apply(self,

413 layer: torch.nn.Module,

414 x: torch.Tensor,

@@ -670,15 +686,14 @@ class Fp8MoEMethod(FusedMoEMethodBase):

686

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

688 # it ahead of time for performance reasons.

673 - if self.allow_deep_gemm:

689 + if self.allow_deep_gemm and not is_blackwell_deep_gemm_used():

690 # Lazy import to avoid CUDA initialization problems.

675 - import deep_gemm as dg

691 if _is_col_major(layer.w13_weight_scale_inv):

692 layer.w13_weight_scale_inv = \

678 - dg.get_col_major_tma_aligned_tensor(layer.w13_weight_scale_inv).contiguous()

693 + get_col_major_tma_aligned_tensor(layer.w13_weight_scale_inv).contiguous()

694 if _is_col_major(layer.w2_weight_scale_inv):

695 layer.w2_weight_scale_inv = \

681 - dg.get_col_major_tma_aligned_tensor(layer.w2_weight_scale_inv).contiguous()

696 + get_col_major_tma_aligned_tensor(layer.w2_weight_scale_inv).contiguous()

697

698 # If checkpoint is fp16, quantize in place.

699 elif not self.quant_config.is_checkpoint_fp8_serialized:

@@ -797,6 +812,29 @@ class Fp8MoEMethod(FusedMoEMethodBase):

812 del layer.w13_input_scale

813 del layer.w2_input_scale

814

815 + if is_blackwell_deep_gemm_used():

816 + assert layer.weight_block_size is not None

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

818 + block_sz = tuple(layer.weight_block_size)

819 + requant_weight_ue8m0_inplace(

820 + layer.w13_weight.data,

821 + layer.w13_weight_scale_inv.data,

822 + block_sz,

823 + )

824 + requant_weight_ue8m0_inplace(

825 + layer.w2_weight.data,

826 + layer.w2_weight_scale_inv.data,

827 + block_sz,

828 + )

829 +

830 + # Ensure column-major TMA alignment expected by DeepGEMM.

831 + if _is_col_major(layer.w13_weight_scale_inv):

832 + layer.w13_weight_scale_inv = get_col_major_tma_aligned_tensor(

833 + layer.w13_weight_scale_inv).contiguous()

834 + if _is_col_major(layer.w2_weight_scale_inv):

835 + layer.w2_weight_scale_inv = get_col_major_tma_aligned_tensor(

836 + layer.w2_weight_scale_inv).contiguous()

837 +

838 def select_gemm_impl(

839 self,

840 prepare_finalize: FusedMoEPrepareAndFinalize,

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

+124 -2

@@ -5,6 +5,7 @@

5 import functools

6 import json

7 import os

8 +from collections.abc import Sequence

9 from typing import Any, Callable, Optional, Union

11 import torch

@@ -13,7 +14,7 @@ import vllm.envs as envs

14 from vllm import _custom_ops as ops

15 from vllm.logger import init_logger

16 from vllm.model_executor.layers.quantization.utils.quant_utils import (

16 - scaled_dequantize)

17 + group_broadcast)

18 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

19 CUTLASS_BLOCK_FP8_SUPPORTED)

20 from vllm.platforms import current_platform

@@ -235,7 +236,7 @@ def block_quant_to_tensor_quant(

236 The outputs are tensor-wise quantization tensor and tensor-wise

237 quantization scale. Note only float8 is supported for now.

238 """

238 - x_dq_block = scaled_dequantize(x_q_block, x_s)

239 + x_dq_block = group_broadcast(x_q_block, x_s)

240 x_q_tensor, scale = input_to_float8(x_dq_block, dtype=x_q_block.dtype)

241 return x_q_tensor, scale

242

@@ -651,3 +652,124 @@ def w8a8_block_fp8_matmul(

652 )

653

654 return C

655 +

656 +

657 +# Taken from https://github.com/deepseek-ai/DeepGEMM/blob/0c88cd01392c1073c7049a97d6328c7bba9b3947

658 +# TODO(wentao): remove this function when DeepGEMM exposes this function

659 +def get_tma_aligned_size(x: int, element_size: int) -> int:

660 + """

661 + Global memory address of TMA must be 16-byte aligned.

662 + Since we use column-major layout for the LHS scaling tensor,

663 + the M-axis of the LHS scaling tensor needs to be padded to a multiple of

664 + 16 bytes.

665 +

666 + Arguments:

667 + x: original M-axis shape of the LHS scaling tensor.

668 + element_size: element size of the LHS scaling tensor.

669 +

670 + Returns:

671 + M-axis shape of the LHS scaling tensor after padding.

672 + """

673 + tma_alignment_bytes = 16

674 + assert tma_alignment_bytes % element_size == 0

675 + alignment = tma_alignment_bytes // element_size

676 + return cdiv(x, alignment) * alignment

677 +

678 +

679 +# Taken from https://github.com/deepseek-ai/DeepGEMM/blob/0c88cd01392c1073c7049a97d6328c7bba9b3947

680 +# TODO(wentao): remove this function when DeepGEMM exposes this function

681 +def get_col_major_tma_aligned_tensor(x: torch.Tensor) -> torch.Tensor:

682 + """

683 + Returns TMA-aligned transposed format of the input tensor. `torch.transpose`

684 + will be called if necessary.

685 + If the input tensor is already column-major layout and 16-byte aligned along

686 + the M axis (thus meets the requirement of LHS scaling tensor in

687 + DeepGEMM), this function will do nothing.

688 +

689 + Arguments:

690 + x: usually the LHS scaling tensor in GEMM.

691 +

692 + Returns:

693 + The LHS scaling tensor of TMA-aligned transposed format.

694 + """

695 + # NOTES: for the extreme performance, you may rewrite/fuse this function in

696 + # CUDA

697 + assert x.dim() in (2, 3)

698 + remove_dim = False

699 + m, n = x.shape[-2], x.shape[-1]

700 + aligned_m = get_tma_aligned_size(m, x.element_size())

701 + if x.dim() == 2:

702 + if x.stride(0) == 1 and x.stride(1) == aligned_m:

703 + return x

704 + x, remove_dim = x.unsqueeze(0), True

705 +

706 + b = x.shape[0]

707 +

708 + # The last kernel gives a column-major TMA aligned layout

709 + if x.stride(0) == aligned_m * n and x.stride(1) == 1 and x.stride(

710 + 2) == aligned_m:

711 + return x.squeeze(0) if remove_dim else x

712 +

713 + # Normal layout requires transposing

714 + aligned_x = torch.transpose(

715 + torch.empty((b, n, aligned_m), device=x.device, dtype=x.dtype), 1, 2)

716 + aligned_x[:, :m, :] = x

717 + aligned_x = aligned_x[:, :m, :]

718 + return aligned_x.squeeze(0) if remove_dim else aligned_x

719 +

720 +

721 +def requant_weight_ue8m0_inplace(

722 + weight: torch.Tensor,

723 + weight_scale: torch.Tensor,

724 + block_size: Sequence[int] = (128, 128),

725 +) -> None:

726 + """Re-quantise *weight* so that its per-block scaling factors are in the

727 + UE8M0 (power-of-two) format expected by the new DeepGEMM kernels inplace.

728 +

729 + Args:

730 + weight: Block-quantised weight tensor stored in ``torch.float8_e4m3fn``.

731 + Expected shape ``(..., M, K)``.

732 + weight_scale: Corresponding per-block scale tensor (``torch.float32``)

733 + with shape ``(..., M // block_size[0], K // block_size[1])``.

734 + block_size: 2-element iterable ``[block_m, block_k]`` describing the

735 + block quantisation granularity.

736 + """

737 + if weight.numel() == 0:

738 + return

739 +

740 + if weight.dtype != torch.float8_e4m3fn:

741 + raise ValueError("Expected *weight* to be torch.float8_e4m3fn, got "

742 + f"{weight.dtype} instead.")

743 +

744 + from vllm.utils.deep_gemm import per_block_cast_to_fp8

745 +

746 + block_m, block_k = int(block_size[0]), int(block_size[1])

747 +

748 + # Flatten leading dimensions so we can iterate over the last two dims.

749 + leading_shape = weight.shape[:-2]

750 + if len(leading_shape) == 0:

751 + w_view = weight.unsqueeze(0)

752 + s_view = weight_scale.unsqueeze(0)

753 + else:

754 + w_view = weight.reshape(-1, weight.shape[-2], weight.shape[-1])

755 + s_view = weight_scale.reshape(-1, *weight_scale.shape[-2:])

756 +

757 + num_mats = w_view.size(0)

758 + for idx in range(num_mats):

759 + w_q = w_view[idx]

760 + s_old = s_view[idx]

761 +

762 + # De-quantise with the *old* scaling factors (float32).

763 + m_cur, k_cur = w_q.shape

764 + s_float = s_old.to(torch.float32)

765 + # Expand scales along rows and cols by block size, then crop.

766 + s_exp_r = torch.repeat_interleave(s_float, block_m, dim=0)

767 + s_exp = torch.repeat_interleave(s_exp_r, block_k, dim=1)

768 + s_exp = s_exp[:m_cur, :k_cur]

769 + w_dq = w_q.to(torch.float32) * s_exp

770 + # Re-quantise using power-of-two scaling (UE8M0).

771 + w_requant, s_requant = per_block_cast_to_fp8(w_dq, [block_m, block_k])

772 +

773 + # Write back the results in-place.

774 + w_q.copy_(w_requant)

775 + s_old.copy_(s_requant)

vllm/utils/deep_gemm.py new

+152

@@ -0,0 +1,152 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Compatibility wrapper for DeepGEMM API changes.

4 +

5 +Users of vLLM should always import **only** these wrappers.

6 +"""

7 +from __future__ import annotations

8 +

9 +import functools

10 +import importlib

11 +from typing import Any, Callable, NoReturn

12 +

13 +import torch

14 +

15 +import vllm.envs as envs

16 +from vllm.utils import cuda_get_device_properties, has_deep_gemm

17 +

18 +

19 +@functools.cache

20 +def is_blackwell_deep_gemm_used() -> bool:

21 + """Return ``True`` if vLLM is configured to use DeepGEMM on a

22 + Blackwell-class GPU.

23 + """

24 +

25 + if not (envs.VLLM_USE_DEEP_GEMM and has_deep_gemm()

26 + and _per_block_cast_impl is not None):

27 + return False

28 +

29 + return cuda_get_device_properties(0, ("major", ))[0] == 10

30 +

31 +

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

33 + """Placeholder for unavailable DeepGEMM backend."""

34 + raise RuntimeError(

35 + "DeepGEMM backend is not available. Please install the `deep_gemm` "

36 + "package to enable FP8 kernels.")

37 +

38 +

39 +def _resolve_symbol(module, new: str, old: str) -> Callable[..., Any] | None:

40 + """Return the *new* symbol if it exists, otherwise the *old* one."""

41 + if hasattr(module, new):

42 + return getattr(module, new)

43 + if hasattr(module, old):

44 + return getattr(module, old)

45 + return None

46 +

47 +

48 +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

53 + _per_block_cast_impl: Callable[..., Any] | None = None

54 +else:

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

56 +

57 + _fp8_gemm_nt_impl = _resolve_symbol(

58 + _dg,

59 + "fp8_gemm_nt",

60 + "gemm_fp8_fp8_bf16_nt",

61 + )

62 + _grouped_impl = _resolve_symbol(

63 + _dg,

64 + "m_grouped_fp8_gemm_nt_contiguous",

65 + "m_grouped_gemm_fp8_fp8_bf16_nt_contiguous",

66 + )

67 + _grouped_masked_impl = _resolve_symbol(

68 + _dg,

69 + "fp8_m_grouped_gemm_nt_masked",

70 + "m_grouped_gemm_fp8_fp8_bf16_nt_masked",

71 + )

72 +

73 + # Try to get per_token_cast_to_fp8 from DeepGEMM math utils.

74 + try:

75 + _math_mod = importlib.import_module(

76 + "deep_gemm.utils.math") # type: ignore

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

78 + None)

79 + _per_block_cast_impl = getattr(_math_mod, "per_block_cast_to_fp8",

80 + None)

81 + except ModuleNotFoundError:

82 + _per_token_cast_impl = None

83 + _per_block_cast_impl = None

84 +

85 +

86 +def fp8_gemm_nt(*args, **kwargs):

87 + if _fp8_gemm_nt_impl is None:

88 + return _missing(*args, **kwargs)

89 + return _fp8_gemm_nt_impl(*args, **kwargs)

90 +

91 +

92 +def m_grouped_fp8_gemm_nt_contiguous(*args, **kwargs):

93 + if _grouped_impl is None:

94 + return _missing(*args, **kwargs)

95 + return _grouped_impl(*args, **kwargs)

96 +

97 +

98 +def fp8_m_grouped_gemm_nt_masked(*args, **kwargs):

99 + if _grouped_masked_impl is None:

100 + return _missing(*args, **kwargs)

101 + return _grouped_masked_impl(*args, **kwargs)

102 +

103 +

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 +

120 +def per_block_cast_to_fp8(x, *args, **kwargs):

121 + if _per_block_cast_impl is not None and is_blackwell_deep_gemm_used():

122 + return _per_block_cast_impl(x)

123 + # TODO: refactor the `per_block_cast_to_fp8` from tests to vllm utils

124 + from tests.kernels.quant_utils import per_block_cast_to_fp8 as _pbcf

125 + return _pbcf(x, *args, **kwargs)

126 +

127 +

128 +def calc_diff(x: torch.Tensor, y: torch.Tensor):

129 + """Return a global difference metric for unit tests.

130 +

131 + DeepGEMM kernels on Blackwell/B200 currently exhibit noticeable per-element

132 + error, causing ``torch.testing.assert_close`` to fail. Instead of checking

133 + every element, we compute a cosine-style similarity over the whole tensor

134 + and report ``1 - sim``. Once kernel accuracy improves this helper can be

135 + removed.

136 + """

137 +

138 + x, y = x.double(), y.double()

139 + denominator = (x * x + y * y).sum()

140 + sim = 2 * (x * y).sum() / denominator

141 + return 1 - sim

142 +

143 +

144 +__all__ = [

145 + "calc_diff",

146 + "fp8_gemm_nt",

147 + "m_grouped_fp8_gemm_nt_contiguous",

148 + "fp8_m_grouped_gemm_nt_masked",

149 + "per_token_group_cast_to_fp8",

150 + "per_block_cast_to_fp8",

151 + "is_blackwell_deep_gemm_used",

152 +]