perf(attention): copy triton metadata without blocking the host

[jason-fxz]

Summary

TritonAttentionBackend.prepare_metadata builds indptr, prefix_lens and, for an extend with a cached prefix, cu_seqlens_q with torch.tensor(list, device=cuda). That copy blocks the host: torch issues cudaMemcpyAsync and then cudaStreamSynchronize. In the overlap scheduler (Scheduler.overlap_loop), the scheduler stream already waits on the engine stream before the next batch is prepared, so this sync waits for the previous forward to finish. The host then builds and launches the next batch while the GPU sits idle.

This PR builds the three tensors in pinned host memory and copies them with non_blocking=True, as the fa and trtllm backends already do. Batch-1 decode gets 0.3 to 0.35 ms per token faster on Gemma 4, gpt-oss and Qwen3.6 with --attention-backend triton. The same code runs for every model on the triton backend, which includes all sliding-window models (no other backend serves them) and every model on ROCm.

Changes

attention/triton.py, TritonAttentionBackend.prepare_metadata: the three host lists become pinned int32 tensors (cumsum on the host) copied with .to(device, non_blocking=True). Memory is pinned only on CUDA devices, so the CPU-only tests still run.

Tested on

NVIDIA H100 80GB HBM3, Intel Xeon Platinum 8480+, driver 580.95.05, CUDA 13.0, torch 2.14.1+cu130, Triton 3.8.0. main is 9b68e00, the base of this branch.

End-to-end results

Models: nvidia/Gemma-4-26B-A4B-NVFP4, openai/gpt-oss-20b and nvidia/Qwen3.6-35B-A3B-NVFP4 with --attention-backend triton.

python -m freetoken --model-path <model> --port 8154 --served-model-name default --text-model-only \
  --max-running-requests 1 --cuda-graph-max-bs 1 --max-seq-len-override <40960 | 131072 | 200000> \
  --cache-type naive --moe-strategy fused [--attention-backend triton]

Long prompts: one streamed /v1/chat/completions request at a time, a prompt of about 1k to 192k tokens (source code, cut with each model's tokenizer), 256 output tokens with ignore_eos, temperature 1.0, top_p 0.95, top_k 64. Decode tok/s is (output tokens - 1) / (last token time - first token time). Median of 3 runs. TTFT is unchanged.

Model	1k	8k	32k	128k	192k
Gemma 4 26B-A4B	192.2 -> 203.8	175.9 -> 186.1	138.7 -> 145.1	-	-
gpt-oss-20b	258.9 -> 284.6	234.4 -> 252.9	173.8 -> 183.8	87.0 -> 89.5	-
Qwen3.6-35B-A3B (triton)	226.4 -> 243.1	200.8 -> 212.6	142.3 -> 148.2	65.7 -> 66.9	48.3 -> 49.0

For reference, Qwen3.6 with the flashinfer backend decodes at 244.9 tok/s at 1k on the same machine. At long context the triton backend is still far behind it (49.0 vs 183.2 tok/s at 192k). That gap is in the decode kernel, which splits every request's keys into a fixed 8 parts, and a separate PR fixes it.

AIME25 problem 1: the model's recommended sampling (Gemma 4: temperature 1.0, top_p 0.95, top_k 64; gpt-oss: temperature 1.0, top_p 1.0, no top_k), natural end, 3 runs per side.

Model	main decode tok/s	this PR decode tok/s	answers (main / this PR)
Gemma 4 26B-A4B	194.6, 195.6, 195.6	207.4, 207.5, 207.8	70, 70, 70 / 70, 70, 70
gpt-oss-20b	269.0, 267.8, 266.6	286.8, 290.6, 290.0	70, 70, 70 / 70, 70, 70

Other checks

• PYTHONPATH=python pytest tests/kernels/test_triton_attention.py: 38 passed; with CUDA_VISIBLE_DEVICES=, 11 passed and 27 skipped, the same as main.