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

[sgl-kernel] rebase FlashMLA 0217 (#18902)

Browse Source 
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
This commit is contained in:
Fan Yin
2026-03-07 16:30:52 +08:00
committed by GitHub
parent c267bdb805
commit 43d6a32045
3 changed files with 164 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,9 +1,9 @@
from __future__ import annotations
"""
Support attention backend for FlashMLA.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, Callable, Optional, Tuple, Union
@@ -23,11 +23,8 @@ if TYPE_CHECKING:
from sglang.srt.speculative.spec_info import SpecInput
# FlashMLA only supports pagesize=64
PAGE_SIZE = 64
# FlashMLA FP8 issue: https://github.com/deepseek-ai/FlashMLA/issues/56
@dataclass
class FlashMLADecodeMetadata:
@@ -47,8 +44,6 @@ class FlashMLADecodeMetadata:
class FlashMLABackend(FlashInferMLAAttnBackend):
"""Flashmla attention kernels."""
def __init__(
self,
model_runner: ModelRunner,
@@ -76,7 +71,6 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
self.data_type = model_runner.kv_cache_dtype
self.q_data_type = model_runner.dtype
self.kv_cache_dim = self.kv_lora_rank + self.qk_rope_head_dim
# Check if KV cache is FP8 (supports both e4m3 and e5m2)
self.is_fp8_kvcache = self.data_type in {
torch.float8_e4m3fn,
torch.float8_e5m2,
@@ -84,8 +78,13 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
self.num_draft_tokens = model_runner.server_args.speculative_num_draft_tokens
def init_forward_metadata(self, forward_batch: ForwardBatch):
self.cuda_graph_kv_indices = None
self.cuda_graph_mla_metadata = None
self.cuda_graph_num_splits = None
self.cuda_graph_mla_metadata_view = None
self.cuda_graph_num_splits_view = None
def init_forward_metadata(self, forward_batch: ForwardBatch):
bs = forward_batch.batch_size
if forward_batch.forward_mode.is_decode_or_idle():
max_seqlen_pad = triton.cdiv(
@@ -143,8 +142,6 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
1,
is_fp8_kvcache=self.is_fp8_kvcache,
)
# Use FlashMLADecodeMetadata which has the attributes forward_extend expects
self.forward_metadata = FlashMLADecodeMetadata(
mla_metadata,
num_splits,
@@ -160,34 +157,31 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
block_kv_indices: Optional[torch.Tensor] = None,
):
if block_kv_indices is None:
cuda_graph_kv_indices = torch.full(
self.cuda_graph_kv_indices = torch.full(
(max_bs, (self.max_context_len + PAGE_SIZE) // PAGE_SIZE),
1,
dtype=torch.int32,
device="cuda",
)
else:
cuda_graph_kv_indices = block_kv_indices
self.cuda_graph_kv_indices = block_kv_indices
if self.num_draft_tokens:
self.cuda_graph_mla_metadata, self.cuda_graph_num_splits = get_mla_metadata(
torch.ones(
max_bs, dtype=torch.int32, device=cuda_graph_kv_indices.device
),
self.num_draft_tokens * self.num_q_heads,
1,
is_fp8_kvcache=self.is_fp8_kvcache,
)
else:
self.cuda_graph_mla_metadata, self.cuda_graph_num_splits = get_mla_metadata(
torch.ones(
max_bs, dtype=torch.int32, device=cuda_graph_kv_indices.device
),
self.num_q_heads,
1,
is_fp8_kvcache=self.is_fp8_kvcache,
)
self.cuda_graph_kv_indices = cuda_graph_kv_indices
device_props = torch.cuda.get_device_properties(self.req_to_token.device)
max_num_sm_parts = device_props.multi_processor_count
self.cuda_graph_mla_metadata = torch.empty(
(max_num_sm_parts, 8),
dtype=torch.int32,
device="cuda",
)
self.cuda_graph_num_splits = torch.empty(
max_bs + 1,
dtype=torch.int32,
device="cuda",
)
self.cuda_graph_mla_metadata_view = None
self.cuda_graph_num_splits_view = None
def init_forward_metadata_capture_cuda_graph(
self,
@@ -211,20 +205,35 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
self.req_to_token.stride(0),
self.cuda_graph_kv_indices.stride(0),
)
num_q_heads = self.num_q_heads * (self.num_draft_tokens or 1)
num_q_heads = self.num_q_heads
mla_metadata, num_splits = get_mla_metadata(
seq_lens.to(torch.int32),
num_q_heads,
1,
is_fp8_kvcache=self.is_fp8_kvcache,
)
self.cuda_graph_mla_metadata.copy_(mla_metadata)
actual_num_sm_parts = mla_metadata.shape[0]
assert actual_num_sm_parts <= self.cuda_graph_mla_metadata.shape[0], (
f"num_sm_parts {actual_num_sm_parts} exceeds preallocated max "
f"{self.cuda_graph_mla_metadata.shape[0]}"
)
self.cuda_graph_mla_metadata[:actual_num_sm_parts].copy_(mla_metadata)
self.cuda_graph_num_splits[: bs + 1].copy_(num_splits)
self.cuda_graph_mla_metadata_view = self.cuda_graph_mla_metadata[
:actual_num_sm_parts
]
self.cuda_graph_num_splits_view = self.cuda_graph_num_splits[: bs + 1]
self.forward_metadata = FlashMLADecodeMetadata(
self.cuda_graph_mla_metadata,
self.cuda_graph_num_splits[: bs + 1],
self.cuda_graph_mla_metadata_view,
self.cuda_graph_num_splits_view,
self.cuda_graph_kv_indices[:bs, :max_seqlen_pad],
)
elif forward_mode.is_target_verify():
seq_lens = seq_lens + self.num_draft_tokens
max_seqlen_pad = triton.cdiv(seq_lens.max().item(), PAGE_SIZE)
@@ -238,17 +247,28 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
self.req_to_token.stride(0),
self.cuda_graph_kv_indices.stride(0),
)
mla_metadata, num_splits = get_mla_metadata(
seq_lens.to(torch.int32),
self.num_draft_tokens * self.num_q_heads,
1,
is_fp8_kvcache=self.is_fp8_kvcache,
)
self.cuda_graph_mla_metadata.copy_(mla_metadata)
actual_num_sm_parts = mla_metadata.shape[0]
assert actual_num_sm_parts <= self.cuda_graph_mla_metadata.shape[0]
self.cuda_graph_mla_metadata[:actual_num_sm_parts].copy_(mla_metadata)
self.cuda_graph_num_splits[: bs + 1].copy_(num_splits)
self.cuda_graph_mla_metadata_view = self.cuda_graph_mla_metadata[
:actual_num_sm_parts
]
self.cuda_graph_num_splits_view = self.cuda_graph_num_splits[: bs + 1]
self.forward_metadata = FlashMLADecodeMetadata(
self.cuda_graph_mla_metadata,
self.cuda_graph_num_splits[: bs + 1],
self.cuda_graph_mla_metadata_view,
self.cuda_graph_num_splits_view,
self.cuda_graph_kv_indices[:bs, :max_seqlen_pad],
)
else:
@@ -273,12 +293,12 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
spec_info: Optional[SpecInput],
seq_lens_cpu: Optional[torch.Tensor],
):
if forward_mode.is_decode_or_idle():
assert seq_lens_cpu is not None
seq_lens = seq_lens[:bs]
seq_lens_cpu = seq_lens_cpu[:bs]
max_seqlen_pad = triton.cdiv(seq_lens_cpu.max().item(), PAGE_SIZE)
create_flashmla_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices[:bs],
@@ -288,24 +308,46 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
self.req_to_token.stride(0),
self.cuda_graph_kv_indices.stride(0),
)
num_q_heads = self.num_q_heads * (self.num_draft_tokens or 1)
num_q_heads = self.num_q_heads
mla_metadata, num_splits = get_mla_metadata(
seq_lens.to(torch.int32),
num_q_heads,
1,
is_fp8_kvcache=self.is_fp8_kvcache,
)
self.cuda_graph_mla_metadata.copy_(mla_metadata)
actual_num_sm_parts = mla_metadata.shape[0]
if actual_num_sm_parts != self.cuda_graph_mla_metadata_view.shape[0]:
import logging
logger = logging.getLogger(__name__)
logger.warning(
f"num_sm_parts mismatch in CUDA Graph replay: "
f"capture={self.cuda_graph_mla_metadata_view.shape[0]}, "
f"replay={actual_num_sm_parts}. "
f"This may indicate batch size changed between capture and replay."
)
self.cuda_graph_mla_metadata_view = self.cuda_graph_mla_metadata[
:actual_num_sm_parts
]
self.cuda_graph_num_splits_view = self.cuda_graph_num_splits[: bs + 1]
self.cuda_graph_mla_metadata[:actual_num_sm_parts].copy_(mla_metadata)
self.cuda_graph_num_splits[: bs + 1].copy_(num_splits)
self.forward_metadata.mla_metadata = self.cuda_graph_mla_metadata
self.forward_metadata.num_splits = self.cuda_graph_num_splits[: bs + 1]
self.forward_metadata.mla_metadata = self.cuda_graph_mla_metadata_view
self.forward_metadata.num_splits = self.cuda_graph_num_splits_view
self.forward_metadata.block_kv_indices = self.cuda_graph_kv_indices[
:bs, :max_seqlen_pad
]
elif forward_mode.is_target_verify():
seq_lens = seq_lens[:bs] + self.num_draft_tokens
seq_lens_cpu = seq_lens_cpu[:bs] + self.num_draft_tokens
max_seqlen_pad = triton.cdiv(seq_lens_cpu.max().item(), PAGE_SIZE)
create_flashmla_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices[:bs],
@@ -315,16 +357,27 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
self.req_to_token.stride(0),
self.cuda_graph_kv_indices.stride(0),
)
mla_metadata, num_splits = get_mla_metadata(
seq_lens.to(torch.int32),
self.num_draft_tokens * self.num_q_heads,
1,
is_fp8_kvcache=self.is_fp8_kvcache,
)
self.cuda_graph_mla_metadata.copy_(mla_metadata)
actual_num_sm_parts = mla_metadata.shape[0]
if actual_num_sm_parts != self.cuda_graph_mla_metadata_view.shape[0]:
self.cuda_graph_mla_metadata_view = self.cuda_graph_mla_metadata[
:actual_num_sm_parts
]
self.cuda_graph_num_splits_view = self.cuda_graph_num_splits[: bs + 1]
self.cuda_graph_mla_metadata[:actual_num_sm_parts].copy_(mla_metadata)
self.cuda_graph_num_splits[: bs + 1].copy_(num_splits)
self.forward_metadata.mla_metadata = self.cuda_graph_mla_metadata
self.forward_metadata.num_splits = self.cuda_graph_num_splits[: bs + 1]
self.forward_metadata.mla_metadata = self.cuda_graph_mla_metadata_view
self.forward_metadata.num_splits = self.cuda_graph_num_splits_view
self.forward_metadata.block_kv_indices = self.cuda_graph_kv_indices[
:bs, :max_seqlen_pad
]
@@ -368,14 +421,11 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
reshape_q = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim)
if self.is_fp8_kvcache:
# For FP8 KV cache, Q needs to be converted to FP8 for FlashMLA kernel
# In SGLang, we use layer.k_scale for both q and k scales
if layer.k_scale is not None:
q_scale = layer.k_scale
descale_q = layer.k_scale.reshape(1)
descale_k = layer.k_scale.reshape(1)
else:
# Fallback to 1.0 if k_scale is not initialized
q_scale = torch.ones((1,), dtype=torch.float32, device=reshape_q.device)
descale_q = torch.ones(
(1,), dtype=torch.float32, device=reshape_q.device
@@ -384,7 +434,6 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
(1,), dtype=torch.float32, device=reshape_q.device
)
# Reshape to 2D for scaled_fp8_quant (which requires 2D input)
q_shape = reshape_q.shape
reshape_q_2d = reshape_q.reshape(-1, q_shape[-1])
reshape_q_fp8_2d, _ = scaled_fp8_quant(reshape_q_2d, q_scale)
@@ -394,7 +443,7 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
k_cache=k_cache.view(-1, PAGE_SIZE, 1, self.kv_cache_dim),
block_table=self.forward_metadata.block_kv_indices[:bs],
cache_seqlens=forward_batch.seq_lens.to(torch.int32),
head_dim_v=self.kv_lora_rank, # TODO Retrieve from config.
head_dim_v=self.kv_lora_rank,
tile_scheduler_metadata=self.forward_metadata.flashmla_metadata,
num_splits=self.forward_metadata.num_splits,
softmax_scale=layer.scaling,
@@ -405,13 +454,12 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
else:
# todo: need check all causal True or False?
o, _ = flash_mla_with_kvcache(
q=reshape_q,
k_cache=k_cache.view(-1, PAGE_SIZE, 1, self.kv_cache_dim),
block_table=self.forward_metadata.block_kv_indices[:bs],
cache_seqlens=forward_batch.seq_lens.to(torch.int32),
head_dim_v=self.kv_lora_rank, # TODO Retrieve from config.
head_dim_v=self.kv_lora_rank,
tile_scheduler_metadata=self.forward_metadata.flashmla_metadata,
num_splits=self.forward_metadata.num_splits,
softmax_scale=layer.scaling,
@@ -447,14 +495,11 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
reshape_q = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim)
if self.is_fp8_kvcache:
# For FP8 KV cache, Q needs to be converted to FP8 for FlashMLA kernel
# In SGLang, we use layer.k_scale for both q and k scales
if layer.k_scale is not None:
q_scale = layer.k_scale
descale_q = layer.k_scale.reshape(1)
descale_k = layer.k_scale.reshape(1)
else:
# Fallback to 1.0 if k_scale is not initialized
q_scale = torch.ones(
(1,), dtype=torch.float32, device=reshape_q.device
)
@@ -465,8 +510,6 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
(1,), dtype=torch.float32, device=reshape_q.device
)
# Quantize Q using scaled_fp8_quant (matching vLLM's approach)
# Reshape to 2D for scaled_fp8_quant (which requires 2D input)
q_shape = reshape_q.shape
reshape_q_2d = reshape_q.reshape(-1, q_shape[-1])
reshape_q_fp8_2d, _ = scaled_fp8_quant(reshape_q_2d, q_scale)
@@ -501,13 +544,7 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
# TODO: multi step kv indices optimization
class FlashMLAMultiStepDraftBackend:
"""
Wrap multiple flashmla attention backends as one for multiple consecutive
draft decoding steps.
"""
def __init__(
self,
model_runner: ModelRunner,
@@ -566,6 +603,10 @@ class FlashMLAMultiStepDraftBackend:
def init_forward_metadata_capture_cuda_graph(self, forward_batch: ForwardBatch):
def call_fn(i, forward_batch):
# EAGLE draft worker uses DECODE mode for draft steps
from sglang.srt.model_executor.forward_batch_info import ForwardMode
# Create a dummy forward_mode for draft step
self.attn_backends[i].init_forward_metadata_capture_cuda_graph(
forward_batch.batch_size,
forward_batch.batch_size * self.topk,
@@ -582,6 +623,8 @@ class FlashMLAMultiStepDraftBackend:
self, forward_batch: ForwardBatch, bs: int
):
def call_fn(i, forward_batch):
from sglang.srt.model_executor.forward_batch_info import ForwardMode
self.attn_backends[i].init_forward_metadata_replay_cuda_graph(
bs,
forward_batch.req_pool_indices,

58
sgl-kernel/cmake/flashmla.cmake
View File

@@ -4,7 +4,7 @@ include(FetchContent)
FetchContent_Declare(
repo-flashmla
GIT_REPOSITORY https://github.com/sgl-project/FlashMLA
GIT_TAG be055fb7df0090fde45f08e9cb5b8b4c0272da73
GIT_TAG 9804b12079e4c873514d3457aa588d3ccf40da28
GIT_SHALLOW OFF
)
FetchContent_Populate(repo-flashmla)
@@ -34,8 +34,9 @@ if(${CUDA_VERSION} VERSION_GREATER_EQUAL "13.0")
# Patch FlashMLA sources for SM103a support.
# These patches are only needed (and only valid) with CUDA 13+.
# Patch flashmla_utils.h: widen IS_SM100 to cover the full SM100 family
set(FLASHMLA_UTILS_FILE "${repo-flashmla_SOURCE_DIR}/csrc/flashmla_utils.h")
# Patch utils.h: widen IS_SM100 to cover the full SM100 family.
# Newer FlashMLA versions use csrc/utils.h.
set(FLASHMLA_UTILS_FILE "${repo-flashmla_SOURCE_DIR}/csrc/utils.h")
file(READ "${FLASHMLA_UTILS_FILE}" FLASHMLA_UTILS_CONTENT)
string(REPLACE
"#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000)
@@ -44,7 +45,7 @@ if(${CUDA_VERSION} VERSION_GREATER_EQUAL "13.0")
#define IS_SM100 1"
FLASHMLA_UTILS_CONTENT "${FLASHMLA_UTILS_CONTENT}")
file(WRITE "${FLASHMLA_UTILS_FILE}" "${FLASHMLA_UTILS_CONTENT}")
message(STATUS "Patched flashmla_utils.h for SM103a support")
message(STATUS "Patched utils.h for SM103a support")
# Patch cutlass/arch/config.h: add SM103 architecture defines.
# The new block is inserted right before the existing "// SM101 and SM101a"
@@ -87,16 +88,46 @@ endif()
set(FlashMLA_SOURCES
"csrc/flashmla_extension.cc"
# Compatibility shim for sgl-kernel torch.ops API.
${repo-flashmla_SOURCE_DIR}/csrc/python_api.cpp
${repo-flashmla_SOURCE_DIR}/csrc/smxx/get_mla_metadata.cu
${repo-flashmla_SOURCE_DIR}/csrc/smxx/mla_combine.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/splitkv_mla.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/splitkv_mla.cu
# Decode metadata/combine kernels.
${repo-flashmla_SOURCE_DIR}/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu
${repo-flashmla_SOURCE_DIR}/csrc/smxx/decode/combine/combine.cu
# sm90 dense decode.
${repo-flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/fp16.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/bf16.cu
# sm90 sparse decode.
${repo-flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h64.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h128.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h64.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h128.cu
# sm90 sparse prefill.
${repo-flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/fwd.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm100/decode/sparse_fp8/splitkv_mla.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu
# sm100 dense prefill/bwd.
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd.cu
# sm100 sparse prefill.
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k512.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k576.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_prefill_k512.cu
# sm100 sparse decode.
${repo-flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/model1.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu
${repo-flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/dense_fp8_python_api.cpp
${repo-flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/flash_fwd_mla_fp8_sm90.cu
@@ -104,9 +135,14 @@ set(FlashMLA_SOURCES
)
Python_add_library(flashmla_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI ${FlashMLA_SOURCES})
target_compile_options(flashmla_ops PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:${FLASHMLA_CUDA_FLAGS}>)
target_compile_options(flashmla_ops PRIVATE
$<$<COMPILE_LANGUAGE:CXX>:-std=c++20>
$<$<COMPILE_LANGUAGE:CUDA>:-std=c++20>
$<$<COMPILE_LANGUAGE:CUDA>:${FLASHMLA_CUDA_FLAGS}>
)
target_include_directories(flashmla_ops PRIVATE
${repo-flashmla_SOURCE_DIR}/csrc
${repo-flashmla_SOURCE_DIR}/csrc/kerutils/include
${repo-flashmla_SOURCE_DIR}/csrc/sm90
${repo-flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/
${repo-flashmla_SOURCE_DIR}/csrc/cutlass/include

9
sgl-kernel/python/sgl_kernel/flash_mla.py
View File

@@ -35,6 +35,9 @@ def get_mla_metadata(
tile_scheduler_metadata: (num_sm_parts, TileSchedulerMetaDataSize), dtype torch.int32.
num_splits: (batch_size + 1), dtype torch.int32.
"""
if _flashmla_import_error is not None:
raise _IMPORT_ERROR from _flashmla_import_error
if is_fp8_kvcache and topk is None:
return torch.ops.sgl_kernel.get_mla_decoding_metadata_dense_fp8.default(
cache_seqlens,
@@ -86,6 +89,9 @@ def flash_mla_with_kvcache(
out: (batch_size, seq_len_q, num_heads_q, head_dim_v).
softmax_lse: (batch_size, num_heads_q, seq_len_q), torch.float32.
"""
if _flashmla_import_error is not None:
raise _IMPORT_ERROR from _flashmla_import_error
if softmax_scale is None:
softmax_scale = q.shape[-1] ** (-0.5)
if indices is not None:
@@ -149,6 +155,9 @@ def flash_mla_sparse_fwd(
- max_logits: [s_q, h_q], float
- lse: [s_q, h_q], float, 2-based log-sum-exp
"""
if _flashmla_import_error is not None:
raise _IMPORT_ERROR from _flashmla_import_error
results = torch.ops.sgl_kernel.sparse_prefill_fwd.default(
q, kv, indices, sm_scale, d_v
)