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

[Kernel] Import flash_mla kernels from sglang kernel for deepseek v4 (#26499)

Browse Source
This commit is contained in:
Chunan Zeng
2026-05-27 14:32:44 -07:00
committed by GitHub
parent 19663aafcd
commit e06058ed62
3 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/deepseek_v4_backend.py
+3 -3
View File
@@ -58,7 +58,7 @@ from sglang.srt.speculative.spec_info import SpecInput
from sglang.srt.utils import ceil_align from sglang.srt.utils import ceil_align
if TYPE_CHECKING: if TYPE_CHECKING:
from flash_mla.flash_mla_interface import FlashMLASchedMeta from sgl_kernel.flash_mla import FlashMLASchedMeta
from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.model_executor.model_runner import ModelRunner
@@ -82,7 +82,7 @@ def _pad_last_dim(x: T, multiples_of: int = PAGE_INDEX_ALIGNED_SIZE) -> T:
def _create_flashmla_metadata(): def _create_flashmla_metadata():
import flash_mla import sgl_kernel.flash_mla as flash_mla
return flash_mla.get_mla_metadata()[0] return flash_mla.get_mla_metadata()[0]
@@ -1045,7 +1045,7 @@ class DeepseekV4AttnBackend(
extra_indices.shape[-1] % 64 == 0 extra_indices.shape[-1] % 64 == 0
), f"{extra_indices.shape=}'s last dimension is not aligned to 64" ), f"{extra_indices.shape=}'s last dimension is not aligned to 64"
import flash_mla import sgl_kernel.flash_mla as flash_mla
o = flash_mla.flash_mla_with_kvcache( o = flash_mla.flash_mla_with_kvcache(
q=q, q=q,
python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py
+2 -2
View File
@@ -55,7 +55,7 @@ from sglang.srt.speculative.spec_info import SpecInput
from sglang.srt.utils import ceil_align from sglang.srt.utils import ceil_align
if TYPE_CHECKING: if TYPE_CHECKING:
from flash_mla.flash_mla_interface import FlashMLASchedMeta from sgl_kernel.flash_mla import FlashMLASchedMeta
from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.model_executor.model_runner import ModelRunner
@@ -83,7 +83,7 @@ def _create_flashmla_metadata():
if is_hip(): if is_hip():
return None return None
import flash_mla import sgl_kernel.flash_mla as flash_mla
return flash_mla.get_mla_metadata()[0] return flash_mla.get_mla_metadata()[0]
python/sglang/srt/layers/attention/hip_flash_mla.py
+1 -1
View File
@@ -14,7 +14,7 @@ def flash_mla_with_kvcache_entrypoint(backend: str, **kwargs):
backend = os.environ.get("SGLANG_HACK_FLASHMLA_BACKEND", "tilelang") backend = os.environ.get("SGLANG_HACK_FLASHMLA_BACKEND", "tilelang")
else: else:
import flash_mla import sgl_kernel.flash_mla as flash_mla
if backend == "comparison": if backend == "comparison":
pack_ref, pack_fast_via_tester = flash_mla_with_kvcache_entrypoint( pack_ref, pack_fast_via_tester = flash_mla_with_kvcache_entrypoint(