Update fa3 interface and add unit test (#9150)

2025-08-13 20:05:02 +08:00
parent 3b3b3baf9f
commit 94f44b88d1
4 changed files with 54 additions and 12 deletions
--- a/sgl-kernel/include/sgl_flash_kernel_ops.h
+++ b/sgl-kernel/include/sgl_flash_kernel_ops.h
@@ -82,4 +82,5 @@ std::vector<at::Tensor> mha_fwd(
    std::optional<at::Tensor>& scheduler_metadata_,  // (b + 1)
    int num_splits,
    std::optional<bool> pack_gqa_,
-    int const sm_margin);
+    int const sm_margin,
+    std::optional<const at::Tensor>& sinks_);