From a82ead53bde1dfae0ab485331982f08b5c132b56 Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Tue, 21 Jul 2026 06:21:00 +0800 Subject: [PATCH] Make Q contiguous before TRT-LLM MHA decode (#31667) --- python/sglang/srt/layers/attention/trtllm_mha_backend.py | 5 ++++- .../registered/quant/test_llama8b_nvfp4_kv_cache_sm120.py | 8 -------- 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 548c1c389..cefed569a 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -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) diff --git a/test/registered/quant/test_llama8b_nvfp4_kv_cache_sm120.py b/test/registered/quant/test_llama8b_nvfp4_kv_cache_sm120.py index cbbf6349e..62addea32 100644 --- a/test/registered/quant/test_llama8b_nvfp4_kv_cache_sm120.py +++ b/test/registered/quant/test_llama8b_nvfp4_kv_cache_sm120.py @@ -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", ]