[ROCm] Credit ROCm/aiter for the block32 GEMM's packed kernel and in-launch split-K (#58659)
Signed-off-by: Lingpeng Jin <103567126+valarLip@users.noreply.github.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Lingpeng Jin committed
Sep 25, 2026 at 12:37 UTC
35f92d2abcea5e8529846b4953a1060da1f88756
1 file changed
+26
vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py
+26
@@ -1,5 +1,27 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+# Portions adapted from ROCm/aiter (https://github.com/ROCm/aiter/pull/5750):
5
+# SPDX-License-Identifier: MIT
6
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
7
+#
8
+# Permission is hereby granted, free of charge, to any person obtaining a copy
9
+# of this software and associated documentation files (the "Software"), to deal
10
+# in the Software without restriction, including without limitation the rights
11
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
12
+# copies of the Software, and to permit persons to whom the Software is
13
+# furnished to do so, subject to the following conditions:
14
+#
15
+# The above copyright notice and this permission notice shall be included in
16
+# all copies or substantial portions of the Software.
17
+#
18
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
19
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
20
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
21
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
22
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
23
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
24
+# SOFTWARE.
25
"""MXFP8 GEMM on 32x32 block-scaled weights for gfx950 (``tl.dot_scaled``).
26
27
``y = x @ w.T`` with an MXFP8 activation (e4m3 values, one E8M0 scale per
@@ -18,6 +40,10 @@ Two kernels, picked per shape from a table tuned on MI355X:
40
Either can split K. The partials are normally summed in the same launch by
41
the last program of each output tile to finish, in split order, so the result
42
does not depend on scheduling. Otherwise a second launch reduces them.
43
+
44
+The packed kernel and the in-launch split-K reduction on one XCD
45
+(``_split_tile``, ``_sum_splits``, ``_split_counters``, ``_k_partition``) are
46
+adapted from the group32 GEMM in ROCm/aiter#5750.
47
"""
48
49
from typing import NamedTuple