EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Cutlass MLA decode - fix dtype error (#5868)

Browse Source
This commit is contained in:
Trevor Morris
2025-04-28 21:12:58 -07:00
committed by GitHub
parent 26fc32d168
commit 8d463fe351
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/srt/layers/attention/cutlass_mla_backend.py
View File

@@ -268,7 +268,7 @@ class CutlassMLABackend(FlashInferMLAAttnBackend):
reshape_q = q.view(-1, layer.tp_q_head_num, layer.head_dim) reshape_q = q.view(-1, layer.tp_q_head_num, layer.head_dim)
o = cutlass_mla_decode( o = cutlass_mla_decode(
q_nope_and_q_pe=reshape_q, q_nope_and_q_pe=reshape_q.to(self.q_data_type),
kv_c_and_k_pe_cache=k_cache.view(-1, PAGE_SIZE, self.kv_cache_dim), kv_c_and_k_pe_cache=k_cache.view(-1, PAGE_SIZE, self.kv_cache_dim),
seq_lens=forward_batch.seq_lens.to(torch.int32), seq_lens=forward_batch.seq_lens.to(torch.int32),
page_table=self.forward_metadata.block_kv_indices, page_table=self.forward_metadata.block_kv_indices,