ta serve: fast kernels in the CUDA-graphed decode step; Space uploads FLAC

[alexkroman]

Summary

Decode speed (scripts/inference.py). The CUDA-graphed decode step never ran the fast linear-attention kernels. In torch 2.8, dynamo answers torch.compiler.is_exporting() with True inside any torch.compile. transformers' kernel-or-reference wrappers read that as "exporting" and run the torch reference, so every graphed step ran the reference gated delta rule and depthwise conv. Only eager prefill used fla and causal-conv1d. Nothing warns: the fallback warning is suppressed while compiling. Four fixes, applied when serving on CUDA (the lm_head one applies everywhere):

• use_fast_kernels_under_compile: points the wrappers straight at the kernels they resolved.
• gated_delta_step_: a custom op that updates the recurrent state in place, so the cache no longer copies 18 layers of fp32 state every step.
• sdpa_gqa_decode: one-token queries fold their query heads into groups instead of repeat_kv copying the KV cache 4× each step (attention 1.45 → 0.24 ms).
• contiguous_lm_head_input: prefill ran lm_head as 32 GEMVs over the 1 GB embedding; one GEMM now.

Measured on an RTX 4090 (Qwen3.5-2B decoder, batch 32):

	Before	After
Decode step	9.7 ms	6.8 ms
1,200 eval clips, 12 datasets	20.5 s	15.6 s (1.31×)
ta bench RTFx at 32 / 64 / 128 concurrent	319 / 386 / 398	372 / 452 / 463
p50 latency, 1 request	0.28 s	0.27 s

WER changes stay inside run-to-run noise: running the same code twice changes about 100 of the 1,200 transcripts, and changing only the batch size changes 117. These changes moved 86, for a net of +7 word errors out of 20,616.

Space uploads (demo/app.py). RunPod's HTTP proxy (Cloudflare) stops accepting request bodies over 500 MiB. Long recordings failed with 413, or with a 502–504 that the Space reported as "starting up or offline". The Space now converts uploads to 16 kHz mono FLAC (a 476 MB 44.1 kHz stereo WAV became 42 MB) and sends the raw bytes with options in the query string. ffmpeg is restricted to the file protocol and plain media demuxers, so an uploaded playlist can't read Space files or fetch URLs. Conversion is capped at 6 hours of audio with a 10-minute timeout. packages.txt installs ffmpeg.

Docs. The README documents running and calling the HTTP API on RunPod, including request formats, options, errors, and the proxy's limits. The README and model card publish WER from a ta eval run against ta serve: mean of 12 datasets 8.71, pooled over all 11,822 samples 7.42.

Review note

LibriSpeech test-other moved from 3.09 to 6.38 in the new published numbers. In my 1,200-clip check of these kernel changes, test-other stayed at 2.33 before and after, so the jump likely comes from that eval run's path rather than the model. It's worth diffing those transcripts before this card is pushed to the Hub.

Test plan

• ta dev precommit passes.
• New tests: GQA attention matches stock sdpa, the export-check bypass, the lm_head hook, the in-place delta step bit-identical to fla (on CUDA), and Space upload conversion, fallback, query encoding, 413 message, and playlist rejection.
• tests/test_inference.py passes on the RTX 4090 pod (18/18, including the CUDA test).
• Live pod: ta bench before/after; 45-minute 44.1 kHz WAV through the new client (10 s plain, 21 s with speakers).
• Gradio API /transcribe via gradio_client: plain, timestamps, speakers, and a 20-minute file.

🤖 Generated with Claude Code