frost(sdpa): THD execute host round-trip — D2H sync, per-total recompile, and CUDA-graph capture

[vedaanta]

Context

The FROST THD (varlen) execute paths derive everything from the graph's device-side SEQ_LEN tensors at execute time:

1. two .tolist() D2H syncs (Q and KV lengths),
2. a host-side cumsum,
3. one H2D upload of the combined [seq_kv | cu_q | cu_k] metadata buffer,
4. a per-execute compile() keyed on the exact packed totals (sq=t_q, skv=t_kv, plus max_sq on SM120) via lru_cache,
5. host computation of the exact unit/grid count (sum(ceil(len_i / tile))).

This buys an exact grid (no dead tiles) and per-shape kernel specialization, but costs three things that matter for real deployments:

• A hard D2H sync on every execute — a pipeline bubble right before the kernel.
• Per-total recompile: under continuous batching the packed totals change every step, so the lru_cache keyed on t_q/t_kv degenerates into a fresh cute.compile (seconds) per step.
• CUDA-graph capture is impossible: .tolist() during capture throws, and the pageable H2D doesn't capture either.

It also re-derives offsets the caller already has: the graph carries ragged-offset tensors on device, and framework callers (TE, vLLM, torch varlen) hold cu_seqlens on device natively — overlapping with #538.

How the other backends avoid this

backend	host values at execute	mechanism
FA2/3/4 varlen	max_seqlen, totals	pushed to the caller as host scalars in the API contract
cuDNN C++ backend	none	grid sized by the graph's declared s_max (host-known at plan time); tiles read device lengths and early-exit; pays dead-tile oversubscription + device setup kernels
FROST THD (today)	derived per execute	D2H sync + per-total compile

Proposal to evaluate

Compile and size the grid against plan-time constants (B, s_max from the tensor declarations — already available at check_support), read the exact cu_seqlens on device (ideally binding the caller's ragged-offset/cu tensor directly, per #538 / the PR #290 multiplier idea), and early-exit dead tiles like the C++ backend. That removes the sync, the recompile, and the capture blocker in one design move, at the price of dead-tile scheduling — the exact-grid vs oversubscription trade should be measured, not assumed.

Interim half-measures worth considering if the full move is too big:

• accept caller-provided host totals through the OSS API layer (FA-style) where the caller has them anyway;
• bucket t_q/t_kv in the compile key (pad to the next bucket) so continuous batching reuses kernels.

Related

• frost(sdpa): cu_seq_len (cu_ragged) form assumes packed strides #538 (cu_seq_len acceptance / packed-stride assumption)
• frost(sdpa): THD backward engines must adopt the declared-stride contract from day one #539 (THD backward contract)
• frost(sdpa): bind THD host prep to the launch stream; plan-time-only THD compile keys (#552) #543 (stream-binding of the same host prep — orthogonal correctness fix)

The two .tolist() syncs are documented in-code as "inherent to the lowering"; this issue is about making them not inherent.