biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Feature] Migrate DeepGEMM API from get_m_alignment_for_contiguous_layout to get_mk_alignment_for_contiguous_layout (#26935)

Browse Source 
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>
This commit is contained in:
Wentao Ye
2025-10-16 16:46:48 -04:00
committed by GitHub
parent fb0571b077
commit b3dda72c23
8 changed files with 57 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
vllm/utils/deep_gemm.py
View File

@@ -75,6 +75,7 @@ _fp8_mqa_logits_impl: Callable[..., Any] | None = None
_fp8_paged_mqa_logits_impl: Callable[..., Any] | None = None
_get_paged_mqa_logits_metadata_impl: Callable[..., Any] | None = None
_get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None
_get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None
def _lazy_init() -> None:
@@ -83,7 +84,7 @@ def _lazy_init() -> None:
global _fp8_mqa_logits_impl, _fp8_paged_mqa_logits_impl
global _get_paged_mqa_logits_metadata_impl
global _get_mn_major_tma_aligned_tensor_impl
global _get_mk_alignment_for_contiguous_layout_impl
# fast path
if (
_fp8_gemm_nt_impl is not None
@@ -92,6 +93,7 @@ def _lazy_init() -> None:
or _fp8_mqa_logits_impl is not None
or _fp8_paged_mqa_logits_impl is not None
or _get_paged_mqa_logits_metadata_impl is not None
or _get_mk_alignment_for_contiguous_layout_impl is not None
):
return
@@ -118,6 +120,9 @@ def _lazy_init() -> None:
_get_mn_major_tma_aligned_tensor_impl = getattr(
_dg, "get_mn_major_tma_aligned_tensor", None
)
_get_mk_alignment_for_contiguous_layout_impl = getattr(
_dg, "get_mk_alignment_for_contiguous_layout", None
)
def get_num_sms() -> int:
@@ -126,6 +131,15 @@ def get_num_sms() -> int:
return int(_dg.get_num_sms())
@functools.cache
def get_mk_alignment_for_contiguous_layout() -> list[int]:
_lazy_init()
if _get_mk_alignment_for_contiguous_layout_impl is None:
return _missing()
mk_align_size = _get_mk_alignment_for_contiguous_layout_impl()
return [mk_align_size, mk_align_size]
def get_col_major_tma_aligned_tensor(x: torch.Tensor) -> torch.Tensor:
"""Wrapper for DeepGEMM's get_mn_major_tma_aligned_tensor"""
_lazy_init()
@@ -338,4 +352,5 @@ __all__ = [
"get_num_sms",
"should_use_deepgemm_for_fp8_linear",
"get_col_major_tma_aligned_tensor",
"get_mk_alignment_for_contiguous_layout",
]