From b66f8e0b96c9bbc7aa51970d1fc65de1fa6c9ec1 Mon Sep 17 00:00:00 2001 From: Chunan Zeng Date: Tue, 26 May 2026 12:00:23 -0700 Subject: [PATCH] Sgl flashmla (#26132) --- sgl-kernel/cmake/flashmla.cmake | 4 +- sgl-kernel/csrc/flashmla_extension.cc | 74 ++++++++- sgl-kernel/include/sgl_kernel_ops.h | 19 ++- sgl-kernel/python/sgl_kernel/flash_mla.py | 194 +++++++++++++++++++++- 4 files changed, 274 insertions(+), 17 deletions(-) diff --git a/sgl-kernel/cmake/flashmla.cmake b/sgl-kernel/cmake/flashmla.cmake index b67ace6af..564d5a6f9 100644 --- a/sgl-kernel/cmake/flashmla.cmake +++ b/sgl-kernel/cmake/flashmla.cmake @@ -1,10 +1,8 @@ -include(FetchContent) - # flash_mla FetchContent_Declare( repo-flashmla GIT_REPOSITORY https://github.com/sgl-project/FlashMLA - GIT_TAG abb54777d4e08c8054c238f59889b52d4e9f0896 + GIT_TAG df022ebafb88578eab9f0300606ee765608d8b5c GIT_SHALLOW OFF ) FetchContent_Populate(repo-flashmla) diff --git a/sgl-kernel/csrc/flashmla_extension.cc b/sgl-kernel/csrc/flashmla_extension.cc index 12b09524a..b9f2fe003 100644 --- a/sgl-kernel/csrc/flashmla_extension.cc +++ b/sgl-kernel/csrc/flashmla_extension.cc @@ -16,8 +16,61 @@ limitations under the License. #include #include +#include "api/dense_decode.h" +#include "api/sparse_decode.h" +#include "api/sparse_fwd.h" #include "sgl_kernel_ops.h" +static std::tuple, std::optional> sgl_sparse_decode_fwd( + const at::Tensor& q, + const at::Tensor& kv, + const at::Tensor& indices, + const std::optional& topk_length, + const std::optional& attn_sink, + std::optional tile_scheduler_metadata, + std::optional num_splits, + const std::optional& extra_kv, + const std::optional& extra_indices, + const std::optional& extra_topk_length, + int64_t d_v, + double sm_scale) { + return sparse_attn_decode_interface( + q, + kv, + indices, + topk_length, + attn_sink, + tile_scheduler_metadata, + num_splits, + extra_kv, + extra_indices, + extra_topk_length, + static_cast(d_v), + static_cast(sm_scale)); +} + +static std::tuple, std::optional> sgl_dense_decode_fwd( + at::Tensor q, + const at::Tensor& kcache, + int64_t head_size_v, + const at::Tensor& seqlens_k, + const at::Tensor& block_table, + double softmax_scale, + bool is_causal, + std::optional tile_scheduler_metadata, + std::optional num_splits) { + return dense_attn_decode_interface( + q, + kcache, + static_cast(head_size_v), + seqlens_k, + block_table, + static_cast(softmax_scale), + is_causal, + tile_scheduler_metadata, + num_splits); +} + TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { /* * From FlashMLA @@ -32,7 +85,9 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { m.def( "fwd_kvcache_mla(Tensor q, Tensor kv_cache, int head_size_v, Tensor seqlens_k, Tensor block_table, float " - "softmax_scale, bool is_causal, Tensor tile_scheduler_metadata, Tensor num_splits, bool is_fp8, Tensor? indices) " + "softmax_scale, bool is_causal, Tensor tile_scheduler_metadata, Tensor num_splits, bool is_fp8, Tensor? indices, " + "Tensor? attn_sink, Tensor? extra_k_cache, Tensor? extra_indices_in_kvcache, Tensor? topk_length, Tensor? " + "extra_topk_length) " "-> Tensor[]"); m.impl("fwd_kvcache_mla", torch::kCUDA, &fwd_kvcache_mla); @@ -44,7 +99,22 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { m.impl("dense_prefill_fwd", torch::kCUDA, &FMHACutlassSM100FwdRun); #endif - m.def("sparse_prefill_fwd(Tensor q, Tensor kv, Tensor indices, float sm_scale, int d_v) -> Tensor[]"); + m.def( + "sparse_decode_fwd(Tensor q, Tensor kv, Tensor indices, Tensor? topk_length, Tensor? attn_sink, " + "Tensor? tile_scheduler_metadata, Tensor? num_splits, Tensor? extra_kv, Tensor? extra_indices, " + "Tensor? extra_topk_length, int d_v, float sm_scale) -> (Tensor, Tensor, Tensor?, Tensor?)"); + m.impl("sparse_decode_fwd", torch::kCUDA, &sgl_sparse_decode_fwd); + + m.def( + "dense_decode_fwd(Tensor q, Tensor kcache, int head_size_v, Tensor seqlens_k, Tensor block_table, float " + "softmax_scale, bool is_causal, Tensor? tile_scheduler_metadata, Tensor? num_splits) -> (Tensor, Tensor, " + "Tensor?, " + "Tensor?)"); + m.impl("dense_decode_fwd", torch::kCUDA, &sgl_dense_decode_fwd); + + m.def( + "sparse_prefill_fwd(Tensor q, Tensor kv, Tensor indices, float sm_scale, int d_v, Tensor? attn_sink=None, " + "Tensor? topk_length=None) -> Tensor[]"); m.impl("sparse_prefill_fwd", torch::kCUDA, &sparse_prefill_fwd); m.def( diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index d8833e80e..b16252007 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/sgl-kernel/include/sgl_kernel_ops.h @@ -22,6 +22,7 @@ limitations under the License. #include #include +#include #include #include @@ -877,8 +878,12 @@ std::vector fwd_kvcache_mla( const at::Tensor& tile_scheduler_metadata, // num_sm_parts x TileSchedulerMetaDataSize const at::Tensor& num_splits, // batch_size + 1 const bool& is_fp8, - const std::optional& indices // None, or batch_size x seqlen_q x topk -); + const std::optional& indices, // None, or batch_size x seqlen_q x topk + const std::optional& attn_sink, + const std::optional& extra_k_cache, + const std::optional& extra_indices_in_kvcache, + const std::optional& topk_length, + const std::optional& extra_topk_length); void FMHACutlassSM100FwdRun( at::Tensor workspace_buffer, @@ -895,8 +900,14 @@ void FMHACutlassSM100FwdRun( int64_t max_seqlen_kv, bool is_varlen); -std::vector -sparse_prefill_fwd(const at::Tensor& q, const at::Tensor& kv, const at::Tensor& indices, double sm_scale, int64_t d_v); +std::vector sparse_prefill_fwd( + const at::Tensor& q, + const at::Tensor& kv, + const at::Tensor& indices, + double sm_scale, + int64_t d_v, + const std::optional& attn_sink, + const std::optional& topk_length); std::vector fwd_kvcache_mla_fp8( at::Tensor& q, // batch_size x seqlen_q x num_heads x head_size diff --git a/sgl-kernel/python/sgl_kernel/flash_mla.py b/sgl-kernel/python/sgl_kernel/flash_mla.py index 3b4643cde..bd8d558a6 100644 --- a/sgl-kernel/python/sgl_kernel/flash_mla.py +++ b/sgl-kernel/python/sgl_kernel/flash_mla.py @@ -1,3 +1,4 @@ +import dataclasses from typing import Optional, Tuple import torch @@ -14,10 +15,33 @@ _IMPORT_ERROR = ImportError( ) +@dataclasses.dataclass +class FlashMLASchedMeta: + """Tile scheduler metadata for the newer FlashMLA Python API.""" + + @dataclasses.dataclass + class Config: + b: int + s_q: int + h_q: int + page_block_size: int + h_k: int + causal: bool + is_fp8_kvcache: bool + topk: Optional[int] + extra_page_block_size: Optional[int] + extra_topk: Optional[int] + + have_initialized: bool = False + config: Optional[Config] = None + tile_scheduler_metadata: Optional[torch.Tensor] = None + num_splits: Optional[torch.Tensor] = None + + def get_mla_metadata( - cache_seqlens: torch.Tensor, - num_q_tokens_per_head_k: int, - num_heads_k: int, + cache_seqlens: Optional[torch.Tensor] = None, + num_q_tokens_per_head_k: Optional[int] = None, + num_heads_k: Optional[int] = None, num_heads_q: Optional[int] = None, is_fp8_kvcache: bool = False, topk: Optional[int] = None, @@ -38,6 +62,12 @@ def get_mla_metadata( if _flashmla_import_error is not None: raise _IMPORT_ERROR from _flashmla_import_error + if cache_seqlens is None: + return FlashMLASchedMeta(), None + + assert num_q_tokens_per_head_k is not None + assert num_heads_k is not None + if is_fp8_kvcache and topk is None: return torch.ops.sgl_kernel.get_mla_decoding_metadata_dense_fp8.default( cache_seqlens, @@ -57,17 +87,22 @@ def get_mla_metadata( def flash_mla_with_kvcache( q: torch.Tensor, k_cache: torch.Tensor, - block_table: torch.Tensor, - cache_seqlens: torch.Tensor, + block_table: Optional[torch.Tensor], + cache_seqlens: Optional[torch.Tensor], head_dim_v: int, - tile_scheduler_metadata: torch.Tensor, - num_splits: torch.Tensor, + tile_scheduler_metadata: torch.Tensor | FlashMLASchedMeta, + num_splits: Optional[torch.Tensor] = None, softmax_scale: Optional[float] = None, causal: bool = False, descale_q: torch.Tensor | None = None, descale_k: torch.Tensor | None = None, is_fp8_kvcache: bool = False, indices: Optional[torch.Tensor] = None, + attn_sink: Optional[torch.Tensor] = None, + extra_k_cache: Optional[torch.Tensor] = None, + extra_indices_in_kvcache: Optional[torch.Tensor] = None, + topk_length: Optional[torch.Tensor] = None, + extra_topk_length: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Arguments: @@ -94,6 +129,34 @@ def flash_mla_with_kvcache( if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) + if isinstance(tile_scheduler_metadata, FlashMLASchedMeta): + return _flash_mla_with_kvcache_sched_meta( + q=q, + k_cache=k_cache, + block_table=block_table, + cache_seqlens=cache_seqlens, + head_dim_v=head_dim_v, + sched_meta=tile_scheduler_metadata, + num_splits=num_splits, + softmax_scale=softmax_scale, + causal=causal, + is_fp8_kvcache=is_fp8_kvcache, + indices=indices, + attn_sink=attn_sink, + extra_k_cache=extra_k_cache, + extra_indices_in_kvcache=extra_indices_in_kvcache, + topk_length=topk_length, + extra_topk_length=extra_topk_length, + ) + + assert num_splits is not None + assert block_table is not None + assert cache_seqlens is not None + assert attn_sink is None + assert extra_k_cache is None + assert extra_indices_in_kvcache is None + assert topk_length is None + assert extra_topk_length is None if indices is not None: assert causal == False, "causal must be `false` if sparse attention is enabled." assert (descale_q is None) == ( @@ -127,16 +190,131 @@ def flash_mla_with_kvcache( num_splits, is_fp8_kvcache, indices, + attn_sink, + extra_k_cache, + extra_indices_in_kvcache, + topk_length, + extra_topk_length, ) return out, softmax_lse +def _flash_mla_with_kvcache_sched_meta( + q: torch.Tensor, + k_cache: torch.Tensor, + block_table: Optional[torch.Tensor], + cache_seqlens: Optional[torch.Tensor], + head_dim_v: int, + sched_meta: FlashMLASchedMeta, + num_splits: Optional[torch.Tensor], + softmax_scale: float, + causal: bool, + is_fp8_kvcache: bool, + indices: Optional[torch.Tensor], + attn_sink: Optional[torch.Tensor], + extra_k_cache: Optional[torch.Tensor], + extra_indices_in_kvcache: Optional[torch.Tensor], + topk_length: Optional[torch.Tensor], + extra_topk_length: Optional[torch.Tensor], +) -> Tuple[torch.Tensor, torch.Tensor]: + assert num_splits is None, "num_splits must be None with FlashMLASchedMeta" + + topk = indices.shape[-1] if indices is not None else None + extra_page_block_size = ( + extra_k_cache.shape[1] if extra_k_cache is not None else None + ) + extra_topk = ( + extra_indices_in_kvcache.shape[-1] + if extra_indices_in_kvcache is not None + else None + ) + + if not sched_meta.have_initialized: + sched_meta.have_initialized = True + sched_meta.config = FlashMLASchedMeta.Config( + b=q.shape[0], + s_q=q.shape[1], + h_q=q.shape[2], + page_block_size=k_cache.shape[1], + h_k=k_cache.shape[2], + causal=causal, + is_fp8_kvcache=is_fp8_kvcache, + topk=topk, + extra_page_block_size=extra_page_block_size, + extra_topk=extra_topk, + ) + else: + helper_msg = ( + " Input arguments are inconsistent with FlashMLASchedMeta. Reuse a " + "scheduler only for matching tensor shapes and sparse settings." + ) + assert sched_meta.config is not None + assert sched_meta.config.b == q.shape[0], helper_msg + assert sched_meta.config.s_q == q.shape[1], helper_msg + assert sched_meta.config.h_q == q.shape[2], helper_msg + assert sched_meta.config.page_block_size == k_cache.shape[1], helper_msg + assert sched_meta.config.h_k == k_cache.shape[2], helper_msg + assert sched_meta.config.causal == causal, helper_msg + assert sched_meta.config.is_fp8_kvcache == is_fp8_kvcache, helper_msg + assert sched_meta.config.topk == topk, helper_msg + assert ( + sched_meta.config.extra_page_block_size == extra_page_block_size + ), helper_msg + assert sched_meta.config.extra_topk == extra_topk, helper_msg + + if topk is not None: + assert not causal, "causal must be False when sparse attention is enabled" + assert is_fp8_kvcache, "is_fp8_kvcache must be True for sparse attention" + out, lse, new_tile_scheduler_metadata, new_num_splits = ( + torch.ops.sgl_kernel.sparse_decode_fwd.default( + q, + k_cache, + indices, + topk_length, + attn_sink, + sched_meta.tile_scheduler_metadata, + sched_meta.num_splits, + extra_k_cache, + extra_indices_in_kvcache, + extra_topk_length, + head_dim_v, + softmax_scale, + ) + ) + else: + assert block_table is not None and cache_seqlens is not None + assert attn_sink is None + assert extra_k_cache is None + assert extra_indices_in_kvcache is None + assert topk_length is None + assert extra_topk_length is None + out, lse, new_tile_scheduler_metadata, new_num_splits = ( + torch.ops.sgl_kernel.dense_decode_fwd.default( + q, + k_cache, + head_dim_v, + cache_seqlens, + block_table, + softmax_scale, + causal, + sched_meta.tile_scheduler_metadata, + sched_meta.num_splits, + ) + ) + + sched_meta.tile_scheduler_metadata = new_tile_scheduler_metadata + sched_meta.num_splits = new_num_splits + return out, lse + + def flash_mla_sparse_fwd( q: torch.Tensor, kv: torch.Tensor, indices: torch.Tensor, sm_scale: float, d_v: int = 512, + attn_sink: Optional[torch.Tensor] = None, + topk_length: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Sparse attention prefill kernel @@ -159,6 +337,6 @@ def flash_mla_sparse_fwd( raise _IMPORT_ERROR from _flashmla_import_error results = torch.ops.sgl_kernel.sparse_prefill_fwd.default( - q, kv, indices, sm_scale, d_v + q, kv, indices, sm_scale, d_v, attn_sink, topk_length ) return results