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

[bugfix] Fix page size for create_flashmla_kv_indices_triton() for cutlass mla (#8685)

Browse Source
This commit is contained in:
Trevor Morris
2025-08-01 14:28:04 -07:00
committed by GitHub
parent 2ae95d17e8
commit 6a7528e623
2 changed files with 9 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -102,7 +102,7 @@ class CutlassMLABackend(FlashInferMLAAttnBackend):
block_kv_indices,
self.req_to_token.stride(0),
max_seqlen_pad,
PAGE_SIZE,
PAGED_SIZE=PAGE_SIZE,
)
workspace_size = cutlass_mla_get_workspace_size(
max_seqlen_pad * PAGE_SIZE, bs, num_kv_splits=1
@@ -165,7 +165,7 @@ class CutlassMLABackend(FlashInferMLAAttnBackend):
self.cuda_graph_kv_indices,
self.req_to_token.stride(0),
self.cuda_graph_kv_indices.stride(0),
PAGE_SIZE,
PAGED_SIZE=PAGE_SIZE,
)
self.forward_metadata = CutlassMLADecodeMetadata(
self.cuda_graph_mla_workspace,
@@ -206,7 +206,7 @@ class CutlassMLABackend(FlashInferMLAAttnBackend):
self.cuda_graph_kv_indices,
self.req_to_token.stride(0),
self.cuda_graph_kv_indices.stride(0),
PAGE_SIZE,
PAGED_SIZE=PAGE_SIZE,
)
else:
super().init_forward_metadata_replay_cuda_graph(

12
python/sglang/srt/layers/attention/trtllm_mla_backend.py
View File

@@ -147,8 +147,8 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
block_kv_indices,
self.req_to_token.stride(0),
max_blocks,
TRITON_PAD_NUM_PAGE_PER_BLOCK,
self.page_size,
NUM_PAGE_PER_BLOCK=TRITON_PAD_NUM_PAGE_PER_BLOCK,
PAGED_SIZE=self.page_size,
)
return block_kv_indices
@@ -204,8 +204,8 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
block_kv_indices,
self.req_to_token.stride(0),
max_seqlen_pad,
TRITON_PAD_NUM_PAGE_PER_BLOCK,
self.page_size,
NUM_PAGE_PER_BLOCK=TRITON_PAD_NUM_PAGE_PER_BLOCK,
PAGED_SIZE=self.page_size,
)
metadata = TRTLLMMLADecodeMetadata(self.cuda_graph_workspace, block_kv_indices)
@@ -248,8 +248,8 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
metadata.block_kv_indices,
self.req_to_token.stride(0),
metadata.block_kv_indices.shape[1],
TRITON_PAD_NUM_PAGE_PER_BLOCK,
self.page_size,
NUM_PAGE_PER_BLOCK=TRITON_PAD_NUM_PAGE_PER_BLOCK,
PAGED_SIZE=self.page_size,
)
def get_cuda_graph_seq_len_fill_value(self) -> int: