[ROCm][Kimi-K3] Fix non-contiguous state_indices crash and GPU-sync assert in fused KDA/MLA prefill (#56526)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Divakar Verma committed
Sep 11, 2026 at 21:17 UTC
06e57f622cf53ed1c3f8d8295cdd48ea2c133bd1
2 files changed
+5
-3
vllm/models/kimi_k3/amd/ops/kda_chunk.py
+2
-2
@@ -286,9 +286,9 @@ def fused_kda_chunk(
286
None if checkpoint_offsets is None else checkpoint_offsets.to(torch.int32),
287
None
288
if checkpoint_state_indices is None
289
- else checkpoint_state_indices.to(torch.int32),
289
+ else checkpoint_state_indices.to(torch.int32).contiguous(),
290
state_cache,
291
- None if state_indices is None else state_indices.to(torch.int32),
291
+ None if state_indices is None else state_indices.to(torch.int32).contiguous(),
292
has_initial_state,
293
)
294
return out, final_state
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+3
-1
@@ -22,6 +22,7 @@ from vllm.model_executor.layers.attention.mla_attention import (
22
QueryLenSupport,
23
)
24
from vllm.triton_utils import tl, triton
25
+from vllm.utils.gpu_sync_debug import gpu_sync_allowed
26
from vllm.utils.math_utils import cdiv, largest_power_of_2_divisor
27
from vllm.v1.attention.backend import (
28
AttentionCGSupport,
@@ -796,7 +797,8 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
797
# build) keeps it off the per-layer forward path where a sync would
798
# break CUDA Graph capture. Using the device-side reduce_indptr is
799
# acceptable since build is allowed to incur an occasional sync.
799
- num_partial_tiles = int(self.fp8_ps_reduce_indptr[-1].item())
800
+ with gpu_sync_allowed():
801
+ num_partial_tiles = int(self.fp8_ps_reduce_indptr[-1].item())
802
803
# Attach PS metadata to the metadata object so forward_mha can read it.
804
metadata.fp8_prefill_qo_indptr = qo_indptr