[API] Support caller-owned outputs for fused sparse attention

csrc/api/api.cpp

@@ -9,8 +9,8 @@ STABLE_TORCH_LIBRARY(_flashmla_C, m) {
     m.def("dense_decode_fwd(Tensor q, Tensor kcache, int head_size_v, Tensor seqlens_k, Tensor block_table, float softmax_scale, bool is_causal, Tensor(a)? tile_scheduler_metadata, Tensor(b)? num_splits, Tensor(c!)? out_) -> (Tensor(c!), Tensor, Tensor(a)?, Tensor(b)?)");
     m.def("sparse_prefill_fwd(Tensor q, Tensor kv, Tensor indices, float sm_scale, int d_v, Tensor? attn_sink, Tensor? topk_length, Tensor(a!)? out_) -> Tensor[]");
     m.def("dense_prefill_fwd(Tensor workspace_buffer, Tensor q, Tensor k, Tensor v, Tensor cumulative_seqlen_q, Tensor cumulative_seqlen_kv, Tensor(a!) o, Tensor(b!) lse, int mask_mode_code, float softmax_scale, int max_seqlen_q, int max_seqlen_kv, bool is_varlen) -> ()");
-    m.def("fused_norm_rope_attn_rope_cast_fwd(Tensor q, Tensor kv, Tensor indices, float sm_scale, int d_v, Tensor? attn_sink, Tensor? topk_length, bool enable_q_norm, float rms_norm_eps, Tensor token_positions, bool is_rope_neox_style, int rope_dim, Tensor cos_sin_cache, int n_wv_group, int num_per_channels, bool use_tma_aligned_col_major_sf, bool round_sf, bool use_packed_ue8m0) -> Tensor[]");
-    m.def("fused_norm_rope_attn_rope_cast_decode(Tensor q, Tensor kv, Tensor indices, float sm_scale, int d_v, Tensor? attn_sink, Tensor? topk_length, Tensor? extra_kv, Tensor? extra_indices, Tensor? extra_topk_length, bool enable_q_norm, float rms_norm_eps, Tensor token_positions, bool is_rope_neox_style, int rope_dim, Tensor cos_sin_cache, int n_wv_group, int num_per_channels, bool use_tma_aligned_col_major_sf, bool round_sf, bool use_packed_ue8m0) -> Tensor[]");
+    m.def("fused_norm_rope_attn_rope_cast_fwd(Tensor q, Tensor kv, Tensor indices, float sm_scale, int d_v, Tensor? attn_sink, Tensor? topk_length, bool enable_q_norm, float rms_norm_eps, Tensor token_positions, bool is_rope_neox_style, int rope_dim, Tensor cos_sin_cache, int n_wv_group, int num_per_channels, bool use_tma_aligned_col_major_sf, bool round_sf, bool use_packed_ue8m0, Tensor(a!)? out_fp8_=None, Tensor(b!)? out_sf_=None) -> Tensor[]");
+    m.def("fused_norm_rope_attn_rope_cast_decode(Tensor q, Tensor kv, Tensor indices, float sm_scale, int d_v, Tensor? attn_sink, Tensor? topk_length, Tensor? extra_kv, Tensor? extra_indices, Tensor? extra_topk_length, bool enable_q_norm, float rms_norm_eps, Tensor token_positions, bool is_rope_neox_style, int rope_dim, Tensor cos_sin_cache, int n_wv_group, int num_per_channels, bool use_tma_aligned_col_major_sf, bool round_sf, bool use_packed_ue8m0, Tensor(a!)? out_fp8_=None, Tensor(b!)? out_sf_=None) -> Tensor[]");
     m.def("permute_q_b_proj(Tensor q_b_proj, Tensor scale_factors, int h_q, int d_q) -> Tensor[]");
     m.def("permute_wv_proj(Tensor wv_proj, Tensor scale_factors, int wv_group_size, int d_o) -> Tensor[]");
 #ifdef FLASH_MLA_ENABLE_DENSE_BWD

csrc/api/fused_norm_rope_attn_rope_cast_fwd.cpp

@@ -39,6 +39,34 @@ static Tensor allocate_scale_factor(
     return sf;
 }
 
+static std::pair<Tensor, Tensor> prepare_fused_outputs(
+    const Tensor &q, int64_t s_q, int64_t n_wv_group, int64_t group_dim,
+    const std::optional<Tensor> &out_fp8_, const std::optional<Tensor> &out_sf_) {
+    Tensor out_fp8 = out_fp8_.has_value() ? out_fp8_.value() :
+        torch::stable::new_empty(q, {s_q, n_wv_group, group_dim}, ScalarType::Float8_e4m3fn);
+    Tensor out_sf = out_sf_.has_value() ? out_sf_.value() :
+        torch::stable::transpose(allocate_scale_factor(s_q, group_dim, 32, q, n_wv_group), 0, 1);
+    if (out_fp8_.has_value()) {
+        KU_CHECK_DEVICE(out_fp8);
+        STD_TORCH_CHECK(out_fp8.device() == q.device(), "out_fp8 must be on the same device as q");
+        KU_CHECK_DTYPE(out_fp8, ScalarType::Float8_e4m3fn);
+        KU_CHECK_SHAPE(out_fp8, s_q, n_wv_group, group_dim);
+        KU_CHECK_CONTIGUOUS(out_fp8);
+    }
+    if (out_sf_.has_value()) {
+        KU_CHECK_DEVICE(out_sf);
+        STD_TORCH_CHECK(out_sf.device() == q.device(), "out_sf must be on the same device as q");
+        KU_CHECK_DTYPE(out_sf, ScalarType::Int);
+        KU_CHECK_SHAPE(out_sf, s_q, n_wv_group, group_dim / 128);
+        STD_TORCH_CHECK(out_sf.stride(0) == 1, "out_sf must have stride(0) == 1");
+        STD_TORCH_CHECK(out_sf.stride(2) >= ((s_q + 3) / 4) * 4 && out_sf.stride(2) % 4 == 0,
+                        "out_sf columns must have a non-overlapping, 4-aligned stride");
+        STD_TORCH_CHECK(out_sf.stride(1) >= out_sf.stride(2) * (group_dim / 128),
+                        "out_sf groups must not overlap");
+    }
+    return {out_fp8, out_sf};
+}
+
 std::vector<Tensor> fused_norm_rope_attn_rope_cast_fwd(
     const Tensor &q,
     const Tensor &kv,
@@ -58,7 +86,9 @@ std::vector<Tensor> fused_norm_rope_attn_rope_cast_fwd(
     int64_t num_per_channels,
     bool use_tma_aligned_col_major_sf,
     bool round_sf,
-    bool use_packed_ue8m0
+    bool use_packed_ue8m0,
+    const std::optional<Tensor> &out_fp8_,
+    const std::optional<Tensor> &out_sf_
 ) {
     Arch arch = Arch();
     bool is_sm100f = arch.is_sm100f();
@@ -125,10 +155,7 @@ std::vector<Tensor> fused_norm_rope_attn_rope_cast_fwd(
     torch::stable::accelerator::DeviceGuard device_guard(q.get_device_index());
 
     STD_TORCH_CHECK(d_v % (num_per_channels * 4) == 0); // 4 is the number of uint8 in uint32, since `use_packed_ue8m0` is `True`
-    Tensor out_fp8 = torch::stable::new_empty(q, {s_q, n_wv_group, wv_group_size * d_v}, ScalarType::Float8_e4m3fn);
-    uint32_t out_sf_scale_gran = 32; // Since the weight is per-32 scaled and deep_gemm.einsum requires A and B to have the same scale granularity, the output sf is always stored in a per-32 scaled format, although it will be actually per-128 scaled when num_per_channels is 128
-    Tensor out_sf = allocate_scale_factor(s_q, wv_group_size * d_v, out_sf_scale_gran, q, n_wv_group);
-    out_sf = torch::stable::transpose(out_sf, 0, 1);   // [s_q, n_wv_group, wv_group_size * d_v / (out_sf_scale_gran*4)]
+    auto [out_fp8, out_sf] = prepare_fused_outputs(q, s_q, n_wv_group, wv_group_size * d_v, out_fp8_, out_sf_);
     Tensor max_logits = torch::stable::new_empty(q, {s_q, h_q}, ScalarType::Float);
     Tensor lse = torch::stable::new_empty(q, {s_q, h_q}, ScalarType::Float);
     KU_CHECK_CONTIGUOUS(out_fp8);
@@ -210,7 +237,9 @@ std::vector<Tensor> fused_norm_rope_attn_rope_cast_decode(
     int64_t num_per_channels,
     bool use_tma_aligned_col_major_sf,
     bool round_sf,
-    bool use_packed_ue8m0
+    bool use_packed_ue8m0,
+    const std::optional<Tensor> &out_fp8_,
+    const std::optional<Tensor> &out_sf_
 ) {
     Arch arch = Arch();
     STD_TORCH_CHECK(arch.is_sm100f(), "Fused Norm + RoPE + Core Attn + RoPE + Cast (fused_norm_rope_attn_rope_cast_decode) is only supported on SM100f architectures.");
@@ -334,10 +363,7 @@ std::vector<Tensor> fused_norm_rope_attn_rope_cast_decode(
     torch::stable::accelerator::DeviceGuard device_guard(q.get_device_index());
 
     STD_TORCH_CHECK(d_v % (num_per_channels * 4) == 0); // 4 is the number of uint8 in uint32, since `use_packed_ue8m0` is `True`
-    Tensor out_fp8 = torch::stable::new_empty(q, {s_q, n_wv_group, wv_group_size * d_v}, ScalarType::Float8_e4m3fn);
-    uint32_t out_sf_scale_gran = 32; // Since the weight is per-32 scaled and deep_gemm.einsum requires A and B to have the same scale granularity, the output sf is always stored in a per-32 scaled format, although it will be actually per-128 scaled when num_per_channels is 128
-    Tensor out_sf = allocate_scale_factor(s_q, wv_group_size * d_v, out_sf_scale_gran, q, n_wv_group);
-    out_sf = torch::stable::transpose(out_sf, 0, 1);   // [s_q, n_wv_group, wv_group_size * d_v / (out_sf_scale_gran*4)]
+    auto [out_fp8, out_sf] = prepare_fused_outputs(q, s_q, n_wv_group, wv_group_size * d_v, out_fp8_, out_sf_);
     Tensor lse = torch::stable::new_empty(q, {s_q, h_q}, ScalarType::Float);
     KU_CHECK_CONTIGUOUS(out_fp8);
     STD_TORCH_CHECK(out_sf.stride(0) == 1);

csrc/api/interfaces.h

@@ -64,7 +64,9 @@ std::vector<Tensor> fused_norm_rope_attn_rope_cast_fwd(
     int64_t num_per_channels,
     bool use_tma_aligned_col_major_sf,
     bool round_sf,
-    bool use_packed_ue8m0);
+    bool use_packed_ue8m0,
+    const std::optional<Tensor> &out_fp8_,
+    const std::optional<Tensor> &out_sf_);
 
 std::vector<Tensor> fused_norm_rope_attn_rope_cast_decode(
     const Tensor &q,
@@ -87,7 +89,9 @@ std::vector<Tensor> fused_norm_rope_attn_rope_cast_decode(
     int64_t num_per_channels,
     bool use_tma_aligned_col_major_sf,
     bool round_sf,
-    bool use_packed_ue8m0);
+    bool use_packed_ue8m0,
+    const std::optional<Tensor> &out_fp8_,
+    const std::optional<Tensor> &out_sf_);
 
 std::vector<Tensor> permute_q_b_proj(
     const Tensor &q_b_proj,

flash_mla/fused_norm_rope_attn_rope_cast.py

@@ -29,6 +29,9 @@ def prefill(
     d_v: int = 512,
     attn_sink: Optional[torch.Tensor] = None,
     topk_length: Optional[torch.Tensor] = None,
+    *,
+    out_fp8: Optional[torch.Tensor] = None,
+    out_sf: Optional[torch.Tensor] = None,
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     """
     A fused kernel for Q Norm + Q RoPE + Core Attn (sparse attention) + O RoPE + O cast to FP8, for DeepSeek-V4 & DeepSeek-V4.1
@@ -63,6 +66,15 @@ def prefill(
         topk_length: optional, [s_q], int32. If provided, the i-th q token will only attend to k tokens specified by indices[i, :, :topk_length[i]], ignoring later k tokens (even if provided in indices). This parameter is mainly used for variable-length topk attention scenarios, such as using sparse attention to simulate causal attention.
             In extremely rare cases (topk_length provided, there is a valid topk index between topk_length[i] ~ s_kv, and that topk index points to a k token containing NaN), operator output will contain NaN, so please avoid this situation.
 
+    Optional output buffers:
+        out_fp8: Caller-owned contiguous output with the returned shape/dtype.
+        out_sf: Caller-owned packed UE8M0 scales with the returned shape/dtype.
+            Its token stride must be 1, column stride a multiple of 4 at least
+            ceil(s_q / 4) * 4, and group stride must keep groups disjoint.
+            Token slices of a larger scale buffer are supported.
+        Supplied buffers must share q's device and not overlap inputs or each
+        other. Omitted outputs are allocated normally.
+
     Returns:
         - out_fp8: [s_q, n_wv_group, wv_group_size * d_v], fp8_e4m3, quantized attention result
         - out_sf: [s_q, n_wv_group, wv_group_size * d_v / (32*4)], int32_t, scaling factor. This scaling factor is ALWAYS stored in the per-32 scaled format, even if num_per_channels is 128
@@ -75,7 +87,8 @@ def prefill(
 
         enable_q_norm, rms_norm_eps, token_positions, is_rope_neox_style, rope_dim, cos_sin_cache,
 
-        n_wv_group, num_per_channels, use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0
+        n_wv_group, num_per_channels, use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0,
+        out_fp8, out_sf,
     )
     out_fp8, out_sf, max_logits, lse = results
     return out_fp8, out_sf, max_logits, lse
@@ -106,6 +119,9 @@ def decode(
     extra_k_cache: Optional[torch.Tensor] = None,
     extra_indices_in_kvcache: Optional[torch.Tensor] = None,
     extra_topk_length: Optional[torch.Tensor] = None,
+    *,
+    out_fp8: Optional[torch.Tensor] = None,
+    out_sf: Optional[torch.Tensor] = None,
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """
     Fused Decoding kernel: Q Norm + Q RoPE + Core Attn (decode, with paged FP8 KV cache) + O RoPE + O cast to FP8, for DeepSeek-V4 & DeepSeek-V4.1
@@ -143,6 +159,15 @@ def decode(
         extra_indices_in_kvcache: optional, [s_q, extra_topk], int32. Indices into the extra KV cache
         extra_topk_length: optional, [s_q], int32. Actual valid extra topk count of the request
 
+    Optional output buffers:
+        out_fp8: Caller-owned contiguous output with the returned shape/dtype.
+        out_sf: Caller-owned packed UE8M0 scales with the returned shape/dtype.
+            Its token stride must be 1, column stride a multiple of 4 at least
+            ceil(s_q / 4) * 4, and group stride must keep groups disjoint.
+            Token slices of a larger scale buffer are supported.
+        Supplied buffers must share q's device and not overlap inputs or each
+        other. Omitted outputs are allocated normally.
+
     Returns:
         - out_fp8: [s_q, n_wv_group, wv_group_size * d_v], fp8_e4m3, quantized attention result
         - out_sf: [s_q, n_wv_group, wv_group_size * d_v / (32*4)], int32, scaling factor
@@ -153,7 +178,8 @@ def decode(
         attn_sink, topk_length,
         extra_k_cache, extra_indices_in_kvcache, extra_topk_length,
         enable_q_norm, rms_norm_eps, token_positions, is_rope_neox_style, rope_dim, cos_sin_cache,
-        n_wv_group, num_per_channels, use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0
+        n_wv_group, num_per_channels, use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0,
+        out_fp8, out_sf,
     )
     return out_fp8, out_sf, lse