Sgl flashmla (#26132)

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(