A high-throughput and memory-efficient inference and serving engine for LLMs
[ROCm][Bugfix]: Fallback GFX942 sparse MLA ops to Triton (#45782)
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
vllmellm committed
Jun 17, 2026 at 19:41 UTC
d537122398df9f8720c307c8b2fa2a8d66f0ff78
1 file changed
+15
-34
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+15
-34
@@ -639,7 +639,6 @@ def rocm_aiter_sparse_attn_indexer(
639
# careful! this will be None in dummy run
640
attn_metadata = get_forward_context().attn_metadata
641
fp8_dtype = current_platform.fp8_dtype()
642
- from vllm import _custom_ops as ops
642
from vllm.utils.torch_utils import _resolve_layer_name
643
644
k_cache_prefix = _resolve_layer_name(k_cache_prefix)
@@ -716,22 +715,13 @@ def rocm_aiter_sparse_attn_indexer(
715
raise ValueError("k must be provided when skip_k_cache_insert is False")
716
717
if not skip_k_cache_insert:
719
- if _ON_GFX942:
720
- ops.indexer_k_quant_and_cache(
721
- k,
722
- kv_cache,
723
- slot_mapping,
724
- quant_block_size,
725
- scale_fmt,
726
- )
727
- else:
728
- indexer_k_quant_and_cache_triton(
729
- k,
730
- kv_cache,
731
- slot_mapping,
732
- quant_block_size,
733
- scale_fmt,
734
- )
718
+ indexer_k_quant_and_cache_triton(
719
+ k,
720
+ kv_cache,
721
+ slot_mapping,
722
+ quant_block_size,
723
+ scale_fmt,
724
+ )
725
726
topk_indices_buffer[: hidden_states.shape[0]] = -1
727
if has_prefill:
@@ -746,23 +736,14 @@ def rocm_aiter_sparse_attn_indexer(
736
for chunk in prefill_metadata.chunks:
737
k_fp8 = k_fp8_full[: chunk.total_seq_lens]
738
k_scale = k_scale_full[: chunk.total_seq_lens]
749
- if _ON_GFX942:
750
- ops.cp_gather_indexer_k_quant_cache(
751
- kv_cache,
752
- k_fp8,
753
- k_scale,
754
- chunk.block_table,
755
- chunk.cu_seq_lens,
756
- )
757
- else:
758
- cp_gather_indexer_k_quant_cache_triton(
759
- kv_cache,
760
- k_fp8,
761
- k_scale,
762
- chunk.block_table,
763
- chunk.cu_seq_lens,
764
- token_to_seq=chunk.token_to_seq,
765
- )
739
+ cp_gather_indexer_k_quant_cache_triton(
740
+ kv_cache,
741
+ k_fp8,
742
+ k_scale,
743
+ chunk.block_table,
744
+ chunk.cu_seq_lens,
745
+ token_to_seq=chunk.token_to_seq,
746
+ )
747
logits = rocm_fp8_mqa_logits(
748
q_fp8[chunk.token_start : chunk.token_end],
749
(k_fp8, k_scale.view(torch.float32)),