chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Import flash_mla from sgl-kernel (#12135)

Browse Source
This commit is contained in:
Baizhou Zhang
2025-10-29 23:54:21 -07:00
committed by GitHub
parent fb52d35f63
commit 621dfb8886
8 changed files with 9 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Callable, Optional, Tuple, Union
import torch
import triton
from flash_mla import flash_mla_with_kvcache, get_mla_metadata
from sgl_kernel.flash_mla import flash_mla_with_kvcache, get_mla_metadata
from sglang.srt.layers.attention.flashinfer_mla_backend import FlashInferMLAAttnBackend
from sglang.srt.layers.attention.utils import create_flashmla_kv_indices_triton

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

@@ -1098,7 +1098,7 @@ class NativeSparseAttnBackend(AttentionBackend):
page_table_1: torch.Tensor,
sm_scale: float,
) -> torch.Tensor:
from flash_mla import flash_mla_sparse_fwd
from sgl_kernel.flash_mla import flash_mla_sparse_fwd
o, _, _ = flash_mla_sparse_fwd(
q=q_all,
@@ -1119,7 +1119,7 @@ class NativeSparseAttnBackend(AttentionBackend):
metadata: NSAMetadata,
page_table_1,
) -> torch.Tensor:
from flash_mla import flash_mla_with_kvcache
from sgl_kernel.flash_mla import flash_mla_with_kvcache
cache_seqlens = metadata.nsa_cache_seqlens_int32
@@ -1261,7 +1261,7 @@ class NativeSparseAttnBackend(AttentionBackend):
)
def _compute_flashmla_metadata(self, cache_seqlens: torch.Tensor, seq_len_q: int):
from flash_mla import get_mla_metadata
from sgl_kernel.flash_mla import get_mla_metadata
flashmla_metadata, num_splits = get_mla_metadata(
cache_seqlens=cache_seqlens,