diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 722c598b6..0fe16379e 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -728,7 +728,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): # For XQA, q_dtype should be bf16 if self.data_type == torch.float8_e4m3fn and (not self.is_xqa_impl): q = q.to(torch.float8_e4m3fn) - q = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) + q = q.reshape(-1, layer.tp_q_head_num, layer.head_dim) k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id) # shape conversion: # [num_pages, page_size, num_kv_heads, head_dim] -> [num_pages, num_kv_heads, page_size, head_dim] @@ -813,7 +813,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): if self.data_type == torch.float8_e4m3fn: q = q.to(torch.float8_e4m3fn) - q = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) + q = q.reshape(-1, layer.tp_q_head_num, layer.head_dim) # [num_pages, page_size, num_kv_heads, head_dim] -> [num_pages, num_kv_heads, page_size, head_dim] k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id) k_cache = k_cache.view(