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] Migrate DeepGEMM API from `get_m_alignment_for_contiguous_layout` to `get_mk_alignment_for_contiguous_layout` (#26935)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Wentao Ye committed
Oct 16, 2025 at 16:46 UTC
b3dda72c236f42fc8a2a7bd2003e0d394533bccd
8 files changed
+57
-46
tests/kernels/moe/test_block_fp8.py
+5
-6
@@ -22,13 +22,13 @@ from vllm.model_executor.layers.fused_moe.fused_moe import (
22
)
23
from vllm.platforms import current_platform
24
from vllm.utils import has_deep_gemm
25
-from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used
25
+from vllm.utils.deep_gemm import (
26
+ get_mk_alignment_for_contiguous_layout,
27
+ is_deep_gemm_e8m0_used,
28
+)
29
30
dg_available = has_deep_gemm()
31
29
-if dg_available:
30
- from deep_gemm import get_m_alignment_for_contiguous_layout
31
-
32
if current_platform.get_device_capability() < (9, 0):
33
pytest.skip("FP8 Triton requires CUDA 9.0 or higher", allow_module_level=True)
34
@@ -218,8 +218,7 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed, monkeypatch)
218
torch.manual_seed(seed)
219
220
monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", str(chunk_size))
221
- block_m = get_m_alignment_for_contiguous_layout()
222
- block_size = [block_m, block_m]
221
+ block_size = get_mk_alignment_for_contiguous_layout()
222
dtype = torch.bfloat16
223
224
a = torch.randn((M, K), dtype=dtype) / 10
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+6
-3
@@ -6,14 +6,17 @@ import torch
6
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7
from vllm.logger import init_logger
8
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
9
-from vllm.model_executor.layers.fused_moe.deep_gemm_utils import deep_gemm_block_shape
9
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
10
TopKWeightAndReduceDelegate,
11
)
12
from vllm.model_executor.layers.fused_moe.utils import _resize_cache
13
from vllm.platforms import current_platform
14
from vllm.triton_utils import tl, triton
16
-from vllm.utils.deep_gemm import fp8_m_grouped_gemm_nt_masked, is_deep_gemm_e8m0_used
15
+from vllm.utils.deep_gemm import (
16
+ fp8_m_grouped_gemm_nt_masked,
17
+ get_mk_alignment_for_contiguous_layout,
18
+ is_deep_gemm_e8m0_used,
19
+)
20
21
logger = init_logger(__name__)
22
@@ -227,7 +230,7 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
230
quant_config: Quantization configuration
231
"""
232
super().__init__(quant_config)
230
- assert self.block_shape == deep_gemm_block_shape()
233
+ assert self.block_shape == get_mk_alignment_for_contiguous_layout()
234
self.max_num_tokens = max_num_tokens
235
self.num_dispatchers = num_dispatchers
236
vllm/model_executor/layers/fused_moe/batched_triton_or_deep_gemm_moe.py
+2
-2
@@ -8,8 +8,8 @@ from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
8
BatchedDeepGemmExperts,
9
)
10
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
11
-from vllm.model_executor.layers.fused_moe.deep_gemm_utils import deep_gemm_block_shape
11
from vllm.model_executor.layers.fused_moe.fused_batched_moe import BatchedTritonExperts
12
+from vllm.utils.deep_gemm import get_mk_alignment_for_contiguous_layout
13
14
15
class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
@@ -31,7 +31,7 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
31
self.allow_deep_gemm = (
32
allow_deep_gemm
33
and self.quant_config.use_fp8_w8a8
34
- and self.block_shape == deep_gemm_block_shape()
34
+ and self.block_shape == get_mk_alignment_for_contiguous_layout()
35
)
36
37
self.batched_deep_gemm_experts = (
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
+10
-8
@@ -13,7 +13,6 @@ from vllm.model_executor.layers.fused_moe.config import (
13
)
14
from vllm.model_executor.layers.fused_moe.deep_gemm_utils import (
15
compute_aligned_M,
16
- deep_gemm_block_shape,
16
deepgemm_moe_permute,
17
deepgemm_unpermute_and_reduce,
18
)
@@ -28,14 +27,17 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
27
per_token_group_quant_fp8,
28
)
29
from vllm.utils import has_deep_gemm
31
-from vllm.utils.deep_gemm import m_grouped_fp8_gemm_nt_contiguous
30
+from vllm.utils.deep_gemm import (
31
+ get_mk_alignment_for_contiguous_layout,
32
+ m_grouped_fp8_gemm_nt_contiguous,
33
+)
34
from vllm.utils.functools import run_once
35
36
logger = init_logger(__name__)
37
38
39
def _valid_deep_gemm_shape(M: int, N: int, K: int) -> bool:
38
- align = deep_gemm_block_shape()[0]
40
+ align = get_mk_alignment_for_contiguous_layout()[0]
41
return align <= M and N % align == 0 and K % align == 0
42
43
@@ -54,7 +56,7 @@ def _valid_deep_gemm(
56
M = hidden_states.size(0)
57
_, K, N = w2.size()
58
57
- align = deep_gemm_block_shape()[0]
59
+ align = get_mk_alignment_for_contiguous_layout()[0]
60
61
if not _valid_deep_gemm_shape(M, N, K):
62
logger.debug_once(
@@ -124,7 +126,7 @@ def warmup_deepgemm_gg_contiguous_kernels(
126
127
assert w1.size(0) == w2.size(0), "w1 and w2 must have the same number of experts"
128
127
- block_m = deep_gemm_block_shape()[0]
129
+ block_m = get_mk_alignment_for_contiguous_layout()[0]
130
num_experts = w1.size(0)
131
device = w1.device
132
@@ -173,7 +175,7 @@ def warmup_deepgemm_gg_contiguous_kernels(
175
class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
176
def __init__(self, quant_config: FusedMoEQuantConfig):
177
super().__init__(quant_config)
176
- assert quant_config.block_shape == deep_gemm_block_shape()
178
+ assert quant_config.block_shape == get_mk_alignment_for_contiguous_layout()
179
assert quant_config.quant_dtype == torch.float8_e4m3fn
180
assert not quant_config.per_act_token_quant
181
assert not quant_config.per_out_ch_quant
@@ -255,7 +257,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
257
M=topk_ids.size(0),
258
num_topk=topk_ids.size(1),
259
local_num_experts=local_num_experts,
258
- alignment=deep_gemm_block_shape()[0],
260
+ alignment=get_mk_alignment_for_contiguous_layout()[0],
261
expert_tokens_meta=expert_tokens_meta,
262
)
263
@@ -364,7 +366,7 @@ def deep_gemm_moe_fp8(
366
w2_scale=w2_scale,
367
a1_scale=a1_scale,
368
a2_scale=a2_scale,
367
- block_shape=deep_gemm_block_shape(),
369
+ block_shape=get_mk_alignment_for_contiguous_layout(),
370
)
371
372
fn = mk.FusedMoEModularKernel(
vllm/model_executor/layers/fused_moe/deep_gemm_utils.py
+2
-13
@@ -5,23 +5,13 @@ Taken from https://github.com/ModelTC/LightLLM/blob/8ed97c74c18f11505b048b1ba00b
5
and updated to fit vllm needs and terminology.
6
"""
7
8
-import functools
9
-
8
import torch
9
10
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
from vllm.model_executor.layers.fused_moe.utils import count_expert_num_tokens
12
from vllm.triton_utils import tl, triton
13
from vllm.utils import round_up
16
-
17
-
18
-@functools.cache
19
-def deep_gemm_block_shape() -> list[int]:
20
- # Lazy import to avoid CUDA initialization problems.
21
- import deep_gemm as dg
22
-
23
- block = dg.get_m_alignment_for_contiguous_layout()
24
- return [block, block]
14
+from vllm.utils.deep_gemm import get_mk_alignment_for_contiguous_layout
15
16
17
def expert_num_tokens_round_up_and_sum(
@@ -354,8 +344,7 @@ def deepgemm_moe_permute(
344
H = aq.size(1)
345
device = aq.device
346
357
- block_m = deep_gemm_block_shape()[0]
358
- block_k = deep_gemm_block_shape()[1]
347
+ block_m, block_k = get_mk_alignment_for_contiguous_layout()
348
349
M_sum = compute_aligned_M(
350
M=topk_ids.size(0),
vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py
+5
-3
@@ -10,9 +10,11 @@ from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
10
_valid_deep_gemm,
11
_valid_deep_gemm_shape,
12
)
13
-from vllm.model_executor.layers.fused_moe.deep_gemm_utils import deep_gemm_block_shape
13
from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
15
-from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used
14
+from vllm.utils.deep_gemm import (
15
+ get_mk_alignment_for_contiguous_layout,
16
+ is_deep_gemm_e8m0_used,
17
+)
18
19
20
class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
@@ -28,7 +30,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
30
self.allow_deep_gemm = (
31
allow_deep_gemm
32
and self.quant_config.use_fp8_w8a8
31
- and self.block_shape == deep_gemm_block_shape()
33
+ and self.block_shape == get_mk_alignment_for_contiguous_layout()
34
)
35
36
self.deep_gemm_expert = (
vllm/model_executor/warmup/deep_gemm_warmup.py
+11
-10
@@ -12,10 +12,7 @@ from tqdm import tqdm
12
import vllm.envs as envs
13
from vllm.distributed.parallel_state import get_dp_group
14
from vllm.model_executor.layers.fused_moe.deep_gemm_moe import DeepGemmExperts
15
-from vllm.model_executor.layers.fused_moe.deep_gemm_utils import (
16
- compute_aligned_M,
17
- deep_gemm_block_shape,
18
-)
15
+from vllm.model_executor.layers.fused_moe.deep_gemm_utils import compute_aligned_M
16
from vllm.model_executor.layers.fused_moe.layer import FusedMoE
17
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel
18
from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
@@ -23,7 +20,11 @@ from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
20
)
21
from vllm.model_executor.layers.linear import LinearBase
22
from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
26
-from vllm.utils.deep_gemm import fp8_gemm_nt, m_grouped_fp8_gemm_nt_contiguous
23
+from vllm.utils.deep_gemm import (
24
+ fp8_gemm_nt,
25
+ get_mk_alignment_for_contiguous_layout,
26
+ m_grouped_fp8_gemm_nt_contiguous,
27
+)
28
29
30
def _generate_optimal_warmup_m_values(
@@ -129,7 +130,7 @@ def _fp8_linear_may_use_deep_gemm(module: torch.nn.Module) -> bool:
130
"""
131
Return True if the input module/layer could be processed with DeepGEMM.
132
"""
132
- block_size = deep_gemm_block_shape()[0]
133
+ block_size = get_mk_alignment_for_contiguous_layout()[0]
134
if not (
135
isinstance(module, LinearBase)
136
and isinstance(module.quant_method, Fp8LinearMethod)
@@ -139,7 +140,7 @@ def _fp8_linear_may_use_deep_gemm(module: torch.nn.Module) -> bool:
140
141
w, _, block_sizes = _extract_data_from_linear_base_module(module)
142
return (
142
- block_sizes == deep_gemm_block_shape()
143
+ block_sizes == get_mk_alignment_for_contiguous_layout()
144
and w.ndim == 2
145
and w.shape[0] % block_size == 0
146
and w.shape[1] % block_size == 0
@@ -155,7 +156,7 @@ def _fused_moe_grouped_gemm_may_use_deep_gemm(module: torch.nn.Module) -> bool:
156
if (
157
moe_quant_config is None
158
or moe_quant_config.quant_dtype != torch.float8_e4m3fn
158
- or moe_quant_config.block_shape != deep_gemm_block_shape()
159
+ or moe_quant_config.block_shape != get_mk_alignment_for_contiguous_layout()
160
):
161
return False
162
@@ -176,7 +177,7 @@ def _deepgemm_fp8_gemm_nt_warmup(w: torch.Tensor, ws: torch.Tensor, max_tokens:
177
return
178
179
n, k = w.size()
179
- block_m = deep_gemm_block_shape()[0]
180
+ block_m = get_mk_alignment_for_contiguous_layout()[0]
181
182
device = w.device
183
a1q = torch.empty((max_tokens, k), device=device, dtype=torch.float8_e4m3fn)
@@ -229,7 +230,7 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
230
231
assert w1.size(0) == w2.size(0), "w1 and w2 must have the same number of experts"
232
232
- block_m = deep_gemm_block_shape()[0]
233
+ block_m = get_mk_alignment_for_contiguous_layout()[0]
234
num_experts = w1.size(0)
235
device = w1.device
236
vllm/utils/deep_gemm.py
+16
-1
@@ -75,6 +75,7 @@ _fp8_mqa_logits_impl: Callable[..., Any] | None = None
75
_fp8_paged_mqa_logits_impl: Callable[..., Any] | None = None
76
_get_paged_mqa_logits_metadata_impl: Callable[..., Any] | None = None
77
_get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None
78
+_get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None
79
80
81
def _lazy_init() -> None:
@@ -83,7 +84,7 @@ def _lazy_init() -> None:
84
global _fp8_mqa_logits_impl, _fp8_paged_mqa_logits_impl
85
global _get_paged_mqa_logits_metadata_impl
86
global _get_mn_major_tma_aligned_tensor_impl
86
-
87
+ global _get_mk_alignment_for_contiguous_layout_impl
88
# fast path
89
if (
90
_fp8_gemm_nt_impl is not None
@@ -92,6 +93,7 @@ def _lazy_init() -> None:
93
or _fp8_mqa_logits_impl is not None
94
or _fp8_paged_mqa_logits_impl is not None
95
or _get_paged_mqa_logits_metadata_impl is not None
96
+ or _get_mk_alignment_for_contiguous_layout_impl is not None
97
):
98
return
99
@@ -118,6 +120,9 @@ def _lazy_init() -> None:
120
_get_mn_major_tma_aligned_tensor_impl = getattr(
121
_dg, "get_mn_major_tma_aligned_tensor", None
122
)
123
+ _get_mk_alignment_for_contiguous_layout_impl = getattr(
124
+ _dg, "get_mk_alignment_for_contiguous_layout", None
125
+ )
126
127
128
def get_num_sms() -> int:
@@ -126,6 +131,15 @@ def get_num_sms() -> int:
131
return int(_dg.get_num_sms())
132
133
134
+@functools.cache
135
+def get_mk_alignment_for_contiguous_layout() -> list[int]:
136
+ _lazy_init()
137
+ if _get_mk_alignment_for_contiguous_layout_impl is None:
138
+ return _missing()
139
+ mk_align_size = _get_mk_alignment_for_contiguous_layout_impl()
140
+ return [mk_align_size, mk_align_size]
141
+
142
+
143
def get_col_major_tma_aligned_tensor(x: torch.Tensor) -> torch.Tensor:
144
"""Wrapper for DeepGEMM's get_mn_major_tma_aligned_tensor"""
145
_lazy_init()
@@ -338,4 +352,5 @@ __all__ = [
352
"get_num_sms",
353
"should_use_deepgemm_for_fp8_linear",
354
"get_col_major_tma_aligned_tensor",
355
+ "get_mk_alignment_for_contiguous_layout",
356
]