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 12, 02:15 UTC
[ROCm] add fallback for aiter fp8 decode mla (#30005)
Signed-off-by: Ye (Charlotte) Qi <yeq@meta.com>
Ye (Charlotte) Qi committed
Dec 4, 2025 at 00:45 UTC
b8a6ae415859bf8c9a3509cbd714695866e35d66
1 file changed
+33
-4
vllm/_aiter_ops.py
+33
-4
@@ -283,6 +283,28 @@ def _rocm_aiter_grouped_topk_fake(
283
pass
284
285
286
+# Cache whether aiter supports FP8 MLA parameters
287
+_AITER_MLA_SUPPORTS_FP8: bool | None = None
288
+
289
+
290
+def _check_aiter_mla_fp8_support() -> bool:
291
+ """Check if aiter.mla.mla_decode_fwd supports q_scale and kv_scale parameters."""
292
+ global _AITER_MLA_SUPPORTS_FP8
293
+ if _AITER_MLA_SUPPORTS_FP8 is None:
294
+ try:
295
+ import inspect
296
+
297
+ from aiter.mla import mla_decode_fwd
298
+
299
+ sig = inspect.signature(mla_decode_fwd)
300
+ _AITER_MLA_SUPPORTS_FP8 = (
301
+ "q_scale" in sig.parameters and "kv_scale" in sig.parameters
302
+ )
303
+ except Exception:
304
+ _AITER_MLA_SUPPORTS_FP8 = False
305
+ return _AITER_MLA_SUPPORTS_FP8
306
+
307
+
308
def _rocm_aiter_mla_decode_fwd_impl(
309
q: torch.Tensor,
310
kv_buffer: torch.Tensor,
@@ -299,6 +321,16 @@ def _rocm_aiter_mla_decode_fwd_impl(
321
) -> None:
322
from aiter.mla import mla_decode_fwd
323
324
+ kwargs = {
325
+ "sm_scale": sm_scale,
326
+ "logit_cap": logit_cap,
327
+ }
328
+
329
+ # Only pass q_scale and kv_scale if the aiter library supports them
330
+ if _check_aiter_mla_fp8_support():
331
+ kwargs["q_scale"] = q_scale
332
+ kwargs["kv_scale"] = kv_scale
333
+
334
mla_decode_fwd(
335
q,
336
kv_buffer.view(-1, 1, 1, q.shape[-1]),
@@ -308,10 +340,7 @@ def _rocm_aiter_mla_decode_fwd_impl(
340
kv_indices,
341
kv_last_page_lens,
342
max_seqlen_qo,
311
- sm_scale=sm_scale,
312
- logit_cap=logit_cap,
313
- q_scale=q_scale,
314
- kv_scale=kv_scale,
343
+ **kwargs,
344
)
345
346