[ROCm][Perf] gfx942: use FlyDSL fp8 MQA logits kernel (ROCm/aiter#3913) (#49544)
Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
akii96 committed
Aug 16, 2026 at 19:50 UTC
6914d60b1e1a2594f0066fe81e3684574413fbe7
1 file changed
+2
-9
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+2
-9
@@ -563,21 +563,14 @@ def rocm_fp8_mqa_logits(
563
Logits tensor of shape [M, N], dtype `torch.float32`.
564
"""
565
566
- # TODO(ganyi): Temporarily workaround, will remove the module check and reference
567
- # path after aiter merge this kernel into main
566
from vllm._aiter_ops import rocm_aiter_ops
567
568
k_fp8, scale = kv
569
572
- # Temporarily route gfx942 to the vendored ROCm/aiter#3257 workaround.
573
- # Remove this branch once vLLM bumps AITER to a version that includes
574
- # ROCm/aiter#3257.
570
if _ON_GFX942 and rocm_aiter_ops.is_enabled():
576
- from vllm.v1.attention.ops.triton_fp8_mqa_logits import (
577
- fp8_mqa_logits_gfx942,
578
- )
571
+ from aiter.ops.flydsl import flydsl_fp8_mqa_logits
572
580
- return fp8_mqa_logits_gfx942(
573
+ return flydsl_fp8_mqa_logits(
574
q, k_fp8, scale, weights, cu_seqlen_ks, cu_seqlen_ke
575
)
576