From 43d6a32045ec20902b0094f6be52a2aa0238a36d Mon Sep 17 00:00:00 2001 From: Fan Yin <1106310035@qq.com> Date: Sat, 7 Mar 2026 16:30:52 +0800 Subject: [PATCH] [sgl-kernel] rebase FlashMLA 0217 (#18902) Co-authored-by: Baizhou Zhang --- .../srt/layers/attention/flashmla_backend.py | 173 +++++++++++------- sgl-kernel/cmake/flashmla.cmake | 58 ++++-- sgl-kernel/python/sgl_kernel/flash_mla.py | 9 + 3 files changed, 164 insertions(+), 76 deletions(-) diff --git a/python/sglang/srt/layers/attention/flashmla_backend.py b/python/sglang/srt/layers/attention/flashmla_backend.py index c3d017583..0693d072d 100644 --- a/python/sglang/srt/layers/attention/flashmla_backend.py +++ b/python/sglang/srt/layers/attention/flashmla_backend.py @@ -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, diff --git a/sgl-kernel/cmake/flashmla.cmake b/sgl-kernel/cmake/flashmla.cmake index b1546b151..d52aadf3f 100644 --- a/sgl-kernel/cmake/flashmla.cmake +++ b/sgl-kernel/cmake/flashmla.cmake @@ -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 $<$:${FLASHMLA_CUDA_FLAGS}>) +target_compile_options(flashmla_ops PRIVATE + $<$:-std=c++20> + $<$:-std=c++20> + $<$:${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 diff --git a/sgl-kernel/python/sgl_kernel/flash_mla.py b/sgl-kernel/python/sgl_kernel/flash_mla.py index 144ddc31a..3b4643cde 100644 --- a/sgl-kernel/python/sgl_kernel/flash_mla.py +++ b/sgl-kernel/python/sgl_kernel/flash_mla.py @@ -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 )