Make Q contiguous before TRT-LLM MHA decode (#31667)
This commit is contained in:
@@ -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",
|
||||
]
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user