[Bugfix][ROCm] Skip FP8 MLA prefill PS-metadata build for chunked-context batches (#52356)
Signed-off-by: Shantipriya Parida <shantipriya.parida@amd.com>
Shantipriya Parida committed
Aug 16, 2026 at 16:21 UTC
fe1c317157d4478fdc0e02096447e61305b871e9
1 file changed
+5
-1
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+5
-1
@@ -777,7 +777,11 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
777
attn_metadata.reduce_indptr = self._mla_reduce_indptr
778
attn_metadata.reduce_final_map = self._mla_reduce_final_map
779
attn_metadata.reduce_partial_map = self._mla_reduce_partial_map
780
- if self._fp8_prefill_enabled and attn_metadata.prefill is not None:
780
+ if (
781
+ self._fp8_prefill_enabled
782
+ and attn_metadata.prefill is not None
783
+ and attn_metadata.prefill.chunked_context is None
784
+ ):
785
self._build_fp8_prefill_ps_metadata(attn_metadata, common_attn_metadata)
786
return attn_metadata
787