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

[foraxe]

Summary

Add independently optional out_fp8 and out_sf destinations to the fused sparse prefill/decode wrappers and stable torch operator ABI. Existing callers retain allocating behavior. Supplied buffers are checked for device, dtype, shape, contiguous FP8 values, and the packed scale strides needed by the kernel; token slices of larger scale buffers are supported.

ASIS
fused attention -> allocate FP8 values + scales -> caller consumes

PR
caller-owned FP8 values + scales
             |
             v
fused attention writes into them -> caller consumes the same storage

This supplies the output-ownership contract needed when fused attention runs in an eager portion of a breakable CUDA graph and its FP8 consumer is captured afterwards. The kernel math and the other returned statistics are unchanged.

Related work

Extends the output-buffer support previously added for ordinary sparse/dense attention in #13 to the fused V4.1 APIs introduced by #20. No open fused-output-buffer PR was found in this repository or the DeepSeek upstream search. The vLLM consumer prototype builds on vllm-project/vllm#56344; this PR contains only the native API and its tests.

Validation

On GB200, Torch 2.13 / CUDA 13, with the inference extension built for SM100a:

CUDA_VISIBLE_DEVICES=0 .venv/bin/python -m pytest -q tests/test_output_buffer_api.py

27 passed. Coverage includes allocating/provided output equality, independently omitted destinations, head counts 64/128, sliced scale buffers and neighboring-row guards, invalid device/dtype/shape/strides, CUDA graph replay with changing inputs, and the FP4 secondary KV cache.

A separate vLLM integration prototype captures wo_a after the eager boundary using these destinations. Its TP4 A/B/B/A check preserved every generated token and mixed-request output; the 17-token piecewise-prefill TTFT median changed from 41.85 to 36.92 ms (11.8%). Default FULL decode already captures wo_a, and TPOT stayed within 0.21%; no decode or long-context speedup is claimed. This result belongs to the combined integration, not to the native API alone.

Report, raw results, test log, reproducible drivers, and checksums. Measured vLLM prototype diff. The final follow-up also preserves the original direct-output split-KV fallback to avoid copy overhead.

AI assistance: implementation, tests, and investigation used OpenAI Codex.

[zyongye]

Hi @foraxe I only notice your PR after my counterpart. I tagged you as the co-author of #23 and close this. Thank you for your contribution.

[foraxe]

Thanks @zyongye , appreciate it