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", ]