Add reusable workspace support to the Python API

[0z5a]

Summary

• make out optional in the Python wrapper and consistently return the output tensor
• accept a caller-owned reusable workspace, with allocate_workspace() for automatic sizing; the existing per-call allocation remains the default
• validate workspace dtype, contiguity, device, and capacity at the C++ boundary
• document ownership/concurrency rules and cover fixed-length, varlen, oversized reuse, invalid buffers, and CUDA Graph replay
• add a reproducible fixed/batched/varlen benchmark for allocation, latency, peak memory, output equality, and graph replay

Closes #23.

RTX 5080 validation

Built for SM120 with PyTorch 2.10.0+cu130 and CUDA Toolkit 13.2.

python -m pytest -q tests/test_python_api.py
10 passed in 0.33s

python -m pytest -q tests/test_fwd.py::test_fwd tests/test_fwd.py::test_fwd_varlen
2 passed in 61.20s

python -m pytest -q \
  'tests/test_fwd_full.py::test_fwd_fixed[in+out-bf16-T37-H4]' \
  'tests/test_fwd_full.py::test_fwd_varlen[in+out-bf16-seqs17_33_65-H4]' \
  'tests/test_fwd_full.py::test_fwd_batched[in+out-bf16-B2_T64-H4]'
3 passed in 1.12s

Workspace benchmark

The benchmark preallocates out in both modes, so the only API difference is automatic workspace allocation versus caller-owned reuse.

python benchmarks/bench_workspace.py \
  --warmup 30 --host-iterations 500 \
  --e2e-repeats 50 --calls-per-repeat 20 \
  --cuda-iterations 500 --profile-calls 50 --memory-calls 10 \
  --json-out workspace_benchmark_5080.json

Times are microseconds per call. Arrows are automatic allocation -> reused workspace.

Case	Workspace	Host p50/p95	Batched E2E p50/p95	Stream mean	aten::empty / call	Incremental peak allocated
fixed T=256 H=1	0.22 MiB	15.40/16.40 -> 17.80/18.30	35.27/35.86 -> 35.16/36.59	32.95 -> 32.86	1 -> 0	0.22 -> 0.00 MiB
fixed T=2048 H=8	13.61 MiB	25.90/215.30 -> 24.00/226.10	218.10/218.84 -> 217.69/218.31	214.37 -> 214.36	3 -> 2	14.03 -> 0.03 MiB
batched B=4 T=512 H=8	13.92 MiB	25.80/47.30 -> 23.60/24.30	78.03/78.89 -> 77.55/78.27	74.67 -> 75.00	3 -> 2	14.03 -> 0.03 MiB
varlen [17,33,65,257] H=4	1.48 MiB	28.80/58.30 -> 26.40/52.80	46.42/48.61 -> 45.31/46.62	42.90 -> 42.26	3 -> 2	1.48 -> 0.00 MiB
fixed T=8192 H=32	216.42 MiB	26.00/2078.10 -> 23.80/2073.50	2076.14/2084.07 -> 2071.98/2089.55	2074.23 -> 2070.01	3 -> 2	216.92 -> 0.50 MiB

All five cases produced exact equality between automatic and reused output, and all five captured/replayed exactly with torch.cuda.CUDAGraph. Reuse removes exactly one profiler-visible allocation per call and nearly all incremental workspace memory. Batched E2E and stream time are mostly neutral; this is evidence for explicit lifetime control and allocation stability, not a kernel-speedup claim. Host p95 includes queue backpressure for the longer-running shapes, which is why the E2E and stream columns are included separately.

Compatibility

Existing positional calls and calls that omit workspace retain their behavior. Callers that reuse storage must provide one workspace per overlapping call; sequential calls may share the same buffer, and a buffer sized for a larger input may serve a smaller one on the same device.