[Attention] support softmax lse

include/sgl_flash_kernel_ops.h

@@ -46,7 +46,7 @@ limitations under the License.
 // Defined in the FMHA/MLA SYCL shared libraries (built with -fvisibility=hidden)
 // and called from common_ops; keep them exported with default visibility.
 #pragma GCC visibility push(default)
-std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> mha_fwd(
+void mha_fwd(
     const at::Tensor& q,  // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
     const at::Tensor& k,  // (b_k, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k or (num_pages, page_size,
                           // h_k, d) if there is page_table.
@@ -77,7 +77,8 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> mha_fwd(
     int num_kv_splits,
     std::optional<bool> pack_gqa_,
     int const sm_margin,
-    std::optional<at::Tensor>& out_);
+    at::Tensor& out,
+    std::optional<at::Tensor>& softmax_lse);
 
 void flash_mla_decode(
     torch::Tensor& out,

python/sgl_kernel/flash_attn.py

@@ -264,7 +264,22 @@ def flash_attn_with_kvcache(
     if cache_seqlens is not None:
         assert cache_seqlens.size(0) + 1 == cu_seqlens_q.size(0)
         cu_seqlens_k = cache_seqlens
-    out, softmax_lse, *rest = torch.ops.sgl_kernel.fwd.default(
+
+    # Pre-construct the output buffers and let the kernel write into them in
+    # place (passed by reference); nothing is returned by the op. Whether the
+    # logsumexp is computed is signalled by passing ``softmax_lse``: a real
+    # tensor requests it, ``None`` skips the LSE computation.
+    total_q = q.size(0)
+    num_heads = q.size(-2)
+    head_size_v = v_cache.size(-1)
+    if out is None:
+        out = q.new_empty(total_q, num_heads, head_size_v)
+    softmax_lse = (
+        q.new_empty(num_heads, total_q, dtype=torch.float32)
+        if return_softmax_lse
+        else None
+    )
+    torch.ops.sgl_kernel.fwd.default(
         q,
         k_cache,
         v_cache,
@@ -294,8 +309,9 @@ def flash_attn_with_kvcache(
         pack_gqa,
         sm_margin,
         out,
+        softmax_lse,
     )
-    return (out, softmax_lse, *rest) if return_softmax_lse else out
+    return (out, softmax_lse) if return_softmax_lse else out
 
 
 def flash_attn_varlen_func(
@@ -338,7 +354,20 @@ def flash_attn_varlen_func(
         max_seqlen_q = q.size(1)
         q = q.view(-1, q.size(-2), q.size(-1)).contiguous()
 
-    out, softmax_lse, *rest = torch.ops.sgl_kernel.fwd.default(
+    # Pre-construct the output buffers and let the kernel write into them in
+    # place (passed by reference); nothing is returned by the op. Whether the
+    # logsumexp is computed is signalled by passing ``softmax_lse``: a real
+    # tensor requests it, ``None`` skips the LSE computation.
+    total_q = q.size(0)
+    num_heads = q.size(-2)
+    head_size_v = v.size(-1)
+    out = q.new_empty(total_q, num_heads, head_size_v)
+    softmax_lse = (
+        q.new_empty(num_heads, total_q, dtype=torch.float32)
+        if return_softmax_lse
+        else None
+    )
+    torch.ops.sgl_kernel.fwd.default(
         q,
         k,
         v,
@@ -367,6 +396,8 @@ def flash_attn_varlen_func(
         num_splits,
         pack_gqa,
         sm_margin,
+        out,
+        softmax_lse,
     )
 
-    return (out, softmax_lse, *rest) if return_softmax_lse else out
+    return (out, softmax_lse) if return_softmax_lse else out