Upgrade fa3 hash (#28394)

Co-authored-by: Shijin <dovis.zhang02@gmail.com>
2026-06-17 13:32:11 -07:00
parent 753aa89a83
commit 4b817f5d7f
6 changed files with 82 additions and 16 deletions
@@ -57,7 +57,9 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
      "    int      num_splits,"
      "    bool?    pack_gqa,"
      "    int      sm_margin,"
-      "    Tensor?  sinks"
+      "    Tensor?  sinks,"
+      "    Tensor?  sparse_mask_fine,"  // [total_q, max_k_blocks, num_int32_per_block]
+      "    bool     only_qv"
      ") -> (Tensor(a!), Tensor, Tensor, Tensor)");  // first return aliases out

  m.impl("fwd", torch::kCUDA, make_pytorch_shim(&mha_fwd));