build_flattened_request_row_ids ran once per layer in the per-layer
prefill attention loop and rebuilt the indexer-seq-len-derived row ids
on the GPU via repeat_interleave with a *device* repeats tensor, which
forces a device sync (it reads sum(repeats) to size the output) and
serializes the launch thread every layer.
The row ids depend only on indexer_seq_lens_cpu, which is identical
across every layer of a forward pass. Add
get_cp_shared_kv_flattened_request_row_ids, which builds them once and
caches on the ForwardBatch (keyed by expected flattened length + device
so a different batch shape or the draft-vs-target pass never reuses a
stale tensor), and build on CPU then move once to drop the sync.
In a TP-5 warm cachebench trace this collapsed 1817 wrapper calls to 46
actual builds (97.5% fewer) and dropped ipc_materialize visibly.
GSM8K 200q x2 = 0.955 / 0.950, 0 invalid.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
leavelet
34763d92b1