From a12fea21ed08f1a5a808de2ed21bf8e43165e159 Mon Sep 17 00:00:00 2001 From: Minglei Zhu Date: Wed, 25 Mar 2026 13:17:27 -0700 Subject: [PATCH] perf(sgl-kernel): expose get_scheduler_metadata for FA3 decode optimization (#21103) --- sgl-kernel/csrc/flash_extension.cc | 34 ++++++++++++ sgl-kernel/include/sgl_flash_kernel_ops.h | 31 +++++++++++ sgl-kernel/python/sgl_kernel/flash_attn.py | 63 ++++++++++++++++++++++ 3 files changed, 128 insertions(+) diff --git a/sgl-kernel/csrc/flash_extension.cc b/sgl-kernel/csrc/flash_extension.cc index df6024dfa..376b69283 100644 --- a/sgl-kernel/csrc/flash_extension.cc +++ b/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) diff --git a/sgl-kernel/include/sgl_flash_kernel_ops.h b/sgl-kernel/include/sgl_flash_kernel_ops.h index b36af6b69..10ee93075 100644 --- a/sgl-kernel/include/sgl_flash_kernel_ops.h +++ b/sgl-kernel/include/sgl_flash_kernel_ops.h @@ -83,3 +83,34 @@ std::tuple mha_fwd( std::optional pack_gqa_, int64_t sm_margin, std::optional& 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 cu_seqlens_q_, + std::optional cu_seqlens_k_, + std::optional cu_seqlens_k_new_, + std::optional seqused_q_, + std::optional leftpad_k_, + std::optional 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 pack_gqa_, + int64_t sm_margin); diff --git a/sgl-kernel/python/sgl_kernel/flash_attn.py b/sgl-kernel/python/sgl_kernel/flash_attn.py index 7aff1af85..05227bad2 100644 --- a/sgl-kernel/python/sgl_kernel/flash_attn.py +++ b/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, + )