[ROCm][Perf] Enable fused KDA decode on gfx942 (MI325X) (#52293)
Signed-off-by: Matvei Pashkovskii <mpashkov@amd.com> Co-authored-by: Claude <noreply@anthropic.com>
Matvei Pashkovskii committed
Aug 18, 2026 at 21:01 UTC
ad5e71b276f6a8d9642560205736bee800464937
3 files changed
+15
-14
CMakeLists.txt
+6
-6
@@ -1180,14 +1180,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
1180
endif()
1181
1182
# HIP counterpart of the FUSED_KDA_DECODE block above: same `fused_kda_decode`
1183
- # op, same VLLM_ENABLE_FUSED_KDA_DECODE guard in torch_bindings.cpp. Only
1184
- # measured on gfx950, so it is only built when gfx950 is in the arch list.
1185
- # A multi-arch build that includes gfx950 still emits this source for every
1186
- # arch in that list. The runtime gate in kda_decode.py keeps the kernel off
1187
- # non-gfx950 devices.
1183
+ # op, same VLLM_ENABLE_FUSED_KDA_DECODE guard in torch_bindings.cpp. Built for
1184
+ # the CDNA archs the kernel supports (gfx942 / gfx950), so it is only emitted
1185
+ # when one of those is in the arch list. A multi-arch build that includes them
1186
+ # still emits this source for every arch in that list. The runtime gate in
1187
+ # kda_decode.py keeps the kernel off unsupported devices.
1188
if(VLLM_GPU_LANG STREQUAL "HIP")
1189
set(FUSED_KDA_DECODE_HIP_ARCHS ${VLLM_GPU_ARCHES})
1190
- list(FILTER FUSED_KDA_DECODE_HIP_ARCHS INCLUDE REGEX "gfx950")
1190
+ list(FILTER FUSED_KDA_DECODE_HIP_ARCHS INCLUDE REGEX "gfx942|gfx950")
1191
if(FUSED_KDA_DECODE_HIP_ARCHS)
1192
set(FUSED_KDA_DECODE_HIP_SRC
1193
"csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu")
tests/models/kimi_k3/test_amd_kda_decode.py
+5
-5
@@ -14,17 +14,17 @@ import torch
14
from vllm.platforms import current_platform
15
16
17
-def _on_gfx950() -> bool:
17
+def _on_supported_arch() -> bool:
18
if not current_platform.is_rocm():
19
return False
20
- from vllm.platforms.rocm import on_gfx950
20
+ from vllm.platforms.rocm import on_gfx942, on_gfx950
21
22
- return on_gfx950()
22
+ return on_gfx950() or on_gfx942()
23
24
25
pytestmark = pytest.mark.skipif(
26
- not _on_gfx950(),
27
- reason="The fused KDA decode kernel is only built for gfx950",
26
+ not _on_supported_arch(),
27
+ reason="The fused KDA decode kernel is only built for gfx942 / gfx950",
28
)
29
30
# Kimi-K3 KDA: 96 heads x 128, conv width 4, gate_lower_bound -5.0.
vllm/models/kimi_k3/amd/ops/kda_decode.py
+4
-3
@@ -35,7 +35,7 @@ def is_fused_kda_decode_supported(
35
conv_state_dtype: torch.dtype,
36
) -> bool:
37
"""Whether the fused decode kernel can serve this layer on this device."""
38
- from vllm.platforms.rocm import on_gfx950
38
+ from vllm.platforms.rocm import on_gfx942, on_gfx950
39
40
if (
41
num_heads not in SUPPORTED_NUM_HEADS
@@ -48,8 +48,9 @@ def is_fused_kda_decode_supported(
48
or not hasattr(torch.ops._C, "fused_kda_decode")
49
):
50
return False
51
- # TODO: Verify on other archs; only measured on gfx950 for now
52
- return on_gfx950()
51
+ # gfx950 (MI355X) and gfx942 (MI325X): both CDNA, sharing the wave64 / DPP /
52
+ # bf16 primitives the kernel relies on.
53
+ return on_gfx950() or on_gfx942()
54
55
56
def make_decode_conv1d_weight_loader(