Sgl flashmla (#26132)

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)

sgl-kernel/csrc/flashmla_extension.cc

@@ -16,8 +16,61 @@ limitations under the License.
 #include <torch/all.h>
 #include <torch/library.h>
 
+#include "api/dense_decode.h"
+#include "api/sparse_decode.h"
+#include "api/sparse_fwd.h"
 #include "sgl_kernel_ops.h"
 
+static std::tuple<at::Tensor, at::Tensor, std::optional<at::Tensor>, std::optional<at::Tensor>> sgl_sparse_decode_fwd(
+    const at::Tensor& q,
+    const at::Tensor& kv,
+    const at::Tensor& indices,
+    const std::optional<at::Tensor>& topk_length,
+    const std::optional<at::Tensor>& attn_sink,
+    std::optional<at::Tensor> tile_scheduler_metadata,
+    std::optional<at::Tensor> num_splits,
+    const std::optional<at::Tensor>& extra_kv,
+    const std::optional<at::Tensor>& extra_indices,
+    const std::optional<at::Tensor>& 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<int>(d_v),
+      static_cast<float>(sm_scale));
+}
+
+static std::tuple<at::Tensor, at::Tensor, std::optional<at::Tensor>, std::optional<at::Tensor>> 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<at::Tensor> tile_scheduler_metadata,
+    std::optional<at::Tensor> num_splits) {
+  return dense_attn_decode_interface(
+      q,
+      kcache,
+      static_cast<int>(head_size_v),
+      seqlens_k,
+      block_table,
+      static_cast<float>(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(

sgl-kernel/include/sgl_kernel_ops.h

@@ -22,6 +22,7 @@ limitations under the License.
 #include <torch/library.h>
 #include <torch/torch.h>
 
+#include <optional>
 #include <tuple>
 #include <vector>
 
@@ -877,8 +878,12 @@ std::vector<at::Tensor> 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<at::Tensor>& indices  // None, or batch_size x seqlen_q x topk
-);
+    const std::optional<at::Tensor>& indices,  // None, or batch_size x seqlen_q x topk
+    const std::optional<at::Tensor>& attn_sink,
+    const std::optional<at::Tensor>& extra_k_cache,
+    const std::optional<at::Tensor>& extra_indices_in_kvcache,
+    const std::optional<at::Tensor>& topk_length,
+    const std::optional<at::Tensor>& 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<at::Tensor>
-sparse_prefill_fwd(const at::Tensor& q, const at::Tensor& kv, const at::Tensor& indices, double sm_scale, int64_t d_v);
+std::vector<at::Tensor> 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<at::Tensor>& attn_sink,
+    const std::optional<at::Tensor>& topk_length);
 
 std::vector<at::Tensor> fwd_kvcache_mla_fp8(
     at::Tensor& q,             // batch_size x seqlen_q x num_heads x head_size

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