[API] Accept caller-provided FP8 output buffers in the fused sparse attention API

[zyongye]

Summary

Add an optional out=(out_fp8, out_sf) argument to flash_mla.fused_norm_rope_attn_rope_cast.prefill and .decode. When given, the fused norm + RoPE + attn + RoPE + cast kernel writes the FP8 activation and the packed ue8m0 scales into the caller's tensors instead of allocating them. Omitting it keeps the current allocating behaviour. Stable-ABI schema: Tensor(d!)? out_fp8=None, Tensor(e!)? out_sf=None appended to both operators, so existing torch.ops._flashmla_C callers are unaffected.

Motivation

In a mixed prefill + decode step the two fused launches (non-paged bf16 KV for prefill, paged FP8/FP4 cache for decode) can now write disjoint token ranges of one shared buffer pair sized for the total N tokens, and a single deep_gemm.fp8_einsum("bhr,hdr->bhd", ...) consumes the whole step. Without this, the two outputs have to be concatenated (about h_q * 512 bytes per token) before the Wv projection. Caller-owned outputs are also what a CUDA-graph capture of the downstream einsum needs.

Interface and checks

Both tensors must be given together. The checks are exactly what the kernel needs, deliberately not what DeepGEMM needs:

• out_fp8: [s_q, n_wv_group, wv_group_size*d_v], float8_e4m3fn, contiguous, 32-byte aligned (256-bit stores). A token range of a larger contiguous buffer qualifies, since the kernel hard-codes the token stride.
• out_sf: [s_q, n_wv_group, wv_group_size*d_v/128], int32, stride(0) == 1 (MN-major, skipped for s_q == 1 because PyTorch normalizes size-1 strides). The group and head-dim strides are passed through to the kernel and only have to keep token ranges from overlapping (stride(2) >= s_q, stride(1) >= size(2)*stride(2)).

The head-dim stride is not required to equal ceil4(s_q). DeepGEMM's check_sf_layout requires sf.stride(-1) == get_tma_aligned_size(mn, 4) as an exact equality for the mn it is finally called with, so a shared buffer must be allocated for the total N and consumed by one einsum over all N tokens. Per-segment einsums on views of a shared scale buffer are rejected by DeepGEMM, which is why the fused API leaves that constraint to the caller:

N_al = (N + 3) // 4 * 4
sf_buf = torch.empty((n_g, 8 * 512 // 128, N_al), dtype=torch.int32, device="cuda")
out_sf_all = sf_buf.permute(2, 0, 1)[:N]                      # [N, n_g, 32], strides (1, 32*N_al, N_al)
out_all = torch.empty((N, n_g, 8 * 512), dtype=torch.float8_e4m3fn, device="cuda")

prefill(..., out=(out_all[:n_p], out_sf_all[:n_p]))
decode(...,  out=(out_all[n_p:], out_sf_all[n_p:]))
deep_gemm.fp8_einsum("bhr,hdr->bhd", (out_all, out_sf_all), (wv_fp8, wv_sf), wv_out, recipe=(1, 1, 32))

Tests

tests/test_fused_optional_out.py is self-contained (torch + flash_mla only, no DeepGEMM / TileKernels / kernelkit) and checks:

• bit-identical out_fp8 / out_sf / lse / max_logits between provided and allocated outputs, for prefill and decode, V4 and V4.1 fp8 caches, h_q 64 and 128, s_q 1 / 5 / 7 / 128 / 200, with and without Q norm, including a query with no valid KV token
• the shared-buffer scenario above: a prefill launch and a decode launch writing disjoint ranges of one buffer pair sized for N, with sentinels proving neither launch touches the other's range or the scale padding columns, for boundaries that are not multiples of 4
• rejection of a lone out_fp8, wrong dtypes, wrong shape, non-contiguous activation, misaligned activation, non-unit token stride, and overlapping head-dim columns

On GB200 (Torch 2.13 / CUDA 13, extension built for sm_100a + sm_103a):

$ .venv/bin/python -m pytest tests/test_fused_optional_out.py -q
31 passed in 2.36s

Not run here: the upstream suites (tests/test_fused_norm_rope_attn_rope_cast.py, tests/test_flash_mla_sparse_decoding.py) need kernelkit / deep_gemm / tile_kernels, which are not installed on this machine. Those code paths are untouched by this change.

Relation to #22

#22 by @foraxe proposes the same capability with separate keyword arguments. This PR takes the single out=(out_fp8, out_sf) pair, adds the shared-buffer segment scenario and the DeepGEMM stride analysis above; @foraxe is credited as co-author on the commit.

Files

• csrc/api/api.cpp: schema for both fused operators
• csrc/api/interfaces.h, csrc/api/fused_norm_rope_attn_rope_cast_fwd.cpp: resolve_fp8_outputs validate-or-allocate helper, used by both entry points
• flash_mla/fused_norm_rope_attn_rope_cast.py: out parameter and docs
• tests/test_fused_optional_out.py: new

🤖 Generated with Claude Code