minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Use reshape instead of contiguous().view() in TRTLLMHAAttnBackend (#22517)

Browse Source
This commit is contained in:
Lianmin Zheng
2026-04-13 14:29:12 -07:00
committed by GitHub
parent ff13dfee45
commit ba7bcca6b3
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/trtllm_mha_backend.py
+2 -2
View File
@@ -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(