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

This commit is contained in:
Mohammad Miadh Angkad
2026-07-20 15:21:00 -07:00
committed by GitHub
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)
@@ -9,7 +9,6 @@ register_cuda_ci(
est_time=300,
stage="extra-a",
runner_config="1-gpu-small",
disabled="Temporarily disabled due to failing accuracy",
)
@@ -31,19 +30,12 @@ class TestLlama8BNVFP4KVCacheSM120(GSM8KMixin, DefaultServerBase):
other_args = [
"--quantization",
"modelopt_fp4",
"--fp4-gemm-backend",
"auto",
"--kv-cache-dtype",
"nvfp4",
"--prefill-attention-backend",
"flashinfer",
"--decode-attention-backend",
"trtllm_mha",
"--page-size",
"64",
"--mem-fraction-static",
"0.87",
"--cuda-graph-backend-prefill=disabled",
]