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

Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>

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:
+        device_props = torch.cuda.get_device_properties(self.req_to_token.device)
-            self.cuda_graph_mla_metadata, self.cuda_graph_num_splits = get_mla_metadata(
+        max_num_sm_parts = device_props.multi_processor_count
-                torch.ones(
+
-                    max_bs, dtype=torch.int32, device=cuda_graph_kv_indices.device
+        self.cuda_graph_mla_metadata = torch.empty(
-                ),
+            (max_num_sm_parts, 8),
-                self.num_draft_tokens * self.num_q_heads,
+            dtype=torch.int32,
-                1,
+            device="cuda",
-                is_fp8_kvcache=self.is_fp8_kvcache,
+        )
-            )
+        self.cuda_graph_num_splits = torch.empty(
-        else:
+            max_bs + 1,
-            self.cuda_graph_mla_metadata, self.cuda_graph_num_splits = get_mla_metadata(
+            dtype=torch.int32,
-                torch.ones(
+            device="cuda",
-                    max_bs, dtype=torch.int32, device=cuda_graph_kv_indices.device
+        )
-                ),
+
-                self.num_q_heads,
+        self.cuda_graph_mla_metadata_view = None
-                1,
+        self.cuda_graph_num_splits_view = None
-                is_fp8_kvcache=self.is_fp8_kvcache,
-            )
-        self.cuda_graph_kv_indices = cuda_graph_kv_indices
 
     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_mla_metadata_view,
-                self.cuda_graph_num_splits[: bs + 1],
+                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_mla_metadata_view,
-                self.cuda_graph_num_splits[: bs + 1],
+                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,

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
+    # Patch utils.h: widen IS_SM100 to cover the full SM100 family.
-    set(FLASHMLA_UTILS_FILE "${repo-flashmla_SOURCE_DIR}/csrc/flashmla_utils.h")
+    # 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
+    # Decode metadata/combine kernels.
-    ${repo-flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/splitkv_mla.cu
+    ${repo-flashmla_SOURCE_DIR}/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu
-    ${repo-flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/splitkv_mla.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

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
     )