Make Q contiguous before TRT-LLM MHA decode (#31667)

2026-07-20 15:21:00 -07:00
parent 7fe9ad25ac
commit a82ead53bd
2 changed files with 4 additions and 9 deletions
@@ -973,7 +973,10 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
            and not use_fused_qkv
        ):
            q = q.to(torch.float8_e4m3fn)
-        q = q.reshape(-1, layer.tp_q_head_num, layer.head_dim)
+        if self.is_xqa_impl:
+            q = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim)
+        else:
+            q = q.reshape(-1, layer.tp_q_head_num, layer.head_dim)

        if self.is_nvfp4_kvcache:
            kv_cache, kv_cache_block_scales = self._get_nvfp4_decode_kv_cache(layer)