From 06e57f622cf53ed1c3f8d8295cdd48ea2c133bd1 Mon Sep 17 00:00:00 2001 From: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Date: Fri, 11 Sep 2026 21:17:39 -0400 Subject: [PATCH] [ROCm][Kimi-K3] Fix non-contiguous state_indices crash and GPU-sync assert in fused KDA/MLA prefill (#56526) Signed-off-by: Divakar Verma --- vllm/models/kimi_k3/amd/ops/kda_chunk.py | 4 ++-- vllm/v1/attention/backends/mla/rocm_aiter_mla.py | 4 +++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/vllm/models/kimi_k3/amd/ops/kda_chunk.py b/vllm/models/kimi_k3/amd/ops/kda_chunk.py index efcdb1bc75..748581771b 100644 --- a/vllm/models/kimi_k3/amd/ops/kda_chunk.py +++ b/vllm/models/kimi_k3/amd/ops/kda_chunk.py @@ -286,9 +286,9 @@ def fused_kda_chunk( None if checkpoint_offsets is None else checkpoint_offsets.to(torch.int32), None if checkpoint_state_indices is None - else checkpoint_state_indices.to(torch.int32), + else checkpoint_state_indices.to(torch.int32).contiguous(), state_cache, - None if state_indices is None else state_indices.to(torch.int32), + None if state_indices is None else state_indices.to(torch.int32).contiguous(), has_initial_state, ) return out, final_state diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index c341bfc348..b850deb070 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -22,6 +22,7 @@ from vllm.model_executor.layers.attention.mla_attention import ( QueryLenSupport, ) from vllm.triton_utils import tl, triton +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.math_utils import cdiv, largest_power_of_2_divisor from vllm.v1.attention.backend import ( AttentionCGSupport, @@ -796,7 +797,8 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): # build) keeps it off the per-layer forward path where a sync would # break CUDA Graph capture. Using the device-side reduce_indptr is # acceptable since build is allowed to incur an occasional sync. - num_partial_tiles = int(self.fp8_ps_reduce_indptr[-1].item()) + with gpu_sync_allowed(): + num_partial_tiles = int(self.fp8_ps_reduce_indptr[-1].item()) # Attach PS metadata to the metadata object so forward_mha can read it. metadata.fp8_prefill_qo_indptr = qo_indptr