perf(sgl-kernel): expose get_scheduler_metadata for FA3 decode optimization (#21103)

sgl-kernel/csrc/flash_extension.cc

@@ -61,6 +61,40 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
       ") -> (Tensor, Tensor, Tensor, Tensor)");  // NEW return type: tuple of 4 tensors
 
   m.impl("fwd", torch::kCUDA, make_pytorch_shim(&mha_fwd));
+
+  /*
+   * From flash-attention: get_scheduler_metadata
+   * Precomputes tile scheduling for FA3 to avoid per-layer prepare_varlen_num_blocks calls.
+   */
+  m.def(
+      "get_scheduler_metadata("
+      "    int      batch_size,"
+      "    int      max_seqlen_q,"
+      "    int      max_seqlen_k,"
+      "    int      num_heads,"
+      "    int      num_heads_k,"
+      "    int      headdim,"
+      "    int      headdim_v,"
+      "    ScalarType qkv_dtype,"
+      "    Tensor   seqused_k,"         // b
+      "    Tensor?  cu_seqlens_q,"      // b+1
+      "    Tensor?  cu_seqlens_k,"      // b+1
+      "    Tensor?  cu_seqlens_k_new,"  // b+1
+      "    Tensor?  seqused_q,"         // b
+      "    Tensor?  leftpad_k,"         // b
+      "    int?     page_size,"
+      "    int      max_seqlen_k_new = 0,"
+      "    bool     is_causal = False,"
+      "    int      window_size_left = -1,"
+      "    int      window_size_right = -1,"
+      "    int      attention_chunk = 0,"
+      "    bool     has_softcap = False,"
+      "    int      num_splits = 0,"
+      "    bool?    pack_gqa = None,"
+      "    int      sm_margin = 0"
+      ") -> Tensor");
+
+  m.impl("get_scheduler_metadata", torch::kCUDA, make_pytorch_shim(&mha_fwd_get_scheduler_metadata));
 }
 
 REGISTER_EXTENSION(flash_ops)

sgl-kernel/include/sgl_flash_kernel_ops.h

@@ -83,3 +83,34 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> mha_fwd(
     std::optional<bool> pack_gqa_,
     int64_t sm_margin,
     std::optional<const at::Tensor>& sinks_);  // (h)
+
+/*
+ * From flash-attention: get_scheduler_metadata
+ * Precomputes tile scheduling metadata for FA3 so that the prepare_varlen_num_blocks
+ * kernel does not need to run per-layer.
+ */
+at::Tensor mha_fwd_get_scheduler_metadata(
+    int64_t batch_size,
+    int64_t max_seqlen_q,
+    int64_t max_seqlen_k,
+    int64_t num_heads,
+    int64_t num_heads_k,
+    int64_t headdim,
+    int64_t headdim_v,
+    at::ScalarType qkv_dtype,
+    at::Tensor seqused_k,
+    std::optional<at::Tensor> cu_seqlens_q_,
+    std::optional<at::Tensor> cu_seqlens_k_,
+    std::optional<at::Tensor> cu_seqlens_k_new_,
+    std::optional<at::Tensor> seqused_q_,
+    std::optional<at::Tensor> leftpad_k_,
+    std::optional<int64_t> page_size,
+    int64_t max_seqlen_k_new,
+    bool is_causal,
+    int64_t window_size_left,
+    int64_t window_size_right,
+    int64_t attention_chunk,
+    bool has_softcap,
+    int64_t num_splits,
+    std::optional<bool> pack_gqa_,
+    int64_t sm_margin);

sgl-kernel/python/sgl_kernel/flash_attn.py

@@ -312,3 +312,66 @@ def flash_attn_varlen_func(
     )
 
     return (out, softmax_lse, *rest) if return_softmax_lse else out
+
+
+def get_scheduler_metadata(
+    batch_size: int,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    num_heads: int,
+    num_heads_k: int,
+    headdim: int,
+    cache_seqlens: torch.Tensor,
+    qkv_dtype=torch.bfloat16,
+    headdim_v: Optional[int] = None,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    cu_seqlens_k_new: Optional[torch.Tensor] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    page_size: Optional[int] = None,
+    max_seqlen_k_new: int = 0,
+    causal: bool = False,
+    window_size=(-1, -1),
+    attention_chunk: int = 0,
+    has_softcap: bool = False,
+    num_splits: int = 0,
+    pack_gqa: Optional[bool] = None,
+    sm_margin: int = 0,
+):
+    """Precompute FA3 tile scheduling metadata.
+
+    Call this once per batch (not per layer) and pass the result as
+    scheduler_metadata to flash_attn_with_kvcache / flash_attn_varlen_func.
+    This avoids the prepare_varlen_num_blocks kernel running on every layer.
+    """
+    cache_seqlens = maybe_contiguous(cache_seqlens)
+    if headdim_v is None:
+        headdim_v = headdim
+
+    return torch.ops.sgl_kernel.get_scheduler_metadata(
+        batch_size,
+        max_seqlen_q,
+        max_seqlen_k,
+        num_heads,
+        num_heads_k,
+        headdim,
+        headdim_v,
+        qkv_dtype,
+        cache_seqlens,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        cu_seqlens_k_new,
+        seqused_q,
+        leftpad_k,
+        page_size,
+        max_seqlen_k_new,
+        causal,
+        window_size[0],
+        window_size[1],
+        attention_chunk,
+        has_softcap,
+        num_splits,
+        pack_gqa,
+        sm_margin,
+    )