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