Repository navigation
[FEAT] [PERF] Remove per-candidate and metadata host syncs from speculative decoding - #51
Conversation
bjf-frz
changed the title
bjf-frz
changed the title
| # blocking: a pageable H2D copy would wait for all queued GPU work. The | ||
| # caching host allocator keeps each staging block alive until its copy ends. | ||
| n, t = len(rows), len(table_rows) | ||
| wide = self._stage( |
There was a problem hiding this comment.
Since _prepare_batch() is also used by non-speculative inference, could you include a before/after performance comparison for that path as well? Please cover small batch sizes and long contexts to check whether the pinned-memory staging introduces any regression.
There was a problem hiding this comment.
Done. See the "Non-speculative path" part of PR message.
|
Please remove benchmarks dir. |
|
A few additional optimization opportunities worth exploring:
The current implementation still copies draft token IDs back to the CPU, although they could be consumed directly on the GPU without that transfer. This is the gap between the current implementation and the intended behavior. The other three optimizations (1, 3, and 4) are also worth implementing and exploring. We may see benefits under high-throughput, high-concurrency workloads. |
Greedy speculation read every draft candidate and every verification argmax with its own .item(), one host sync per request per draft offset and one per verified candidate. Accepted greedy candidates equal the target argmax of their row, so the committed tokens are the target-argmax prefix through the first mismatch (the correction, or the bonus when every candidate matches). Read all verification argmaxes with one argmax(-1).tolist() and compare on the host, and read each draft offset's greedy candidates in one batch. Sampling keeps sequential rejection with an unchanged RNG consumption order.
_prepare_batch built its five or six metadata tensors with torch.tensor(list, device=cuda). Each pageable host-to-device copy waits for all queued GPU work, so every core call serialized the host behind the previous core. A speculative round makes 2K+4 core calls. Stage the int64 address metadata and the int32 table/length metadata in two pinned host tensors, copy them with non_blocking=True, and return views of the device copies. The caching host allocator keeps each staging block alive until its copy completes, and stream order places the copies before the kernels that read them. The speculative draft prelude input uses the same staging.
A speculative round prepares the same request/position rows once per recurrence depth: twice for each draft offset and twice for verification. Positions, write offsets, context lengths and packed-query boundaries do not depend on depth; only block addresses and page tables select a per-depth plane. Add KVCacheManager._prepare_batches(), which validates the rows once, builds the depth-independent fields once and returns one prepared batch per depth assignment as views into two shared staging copies. _prepare_batch() delegates to it with a single assignment. Preparation still initializes no KV slot, so ownership and written-prefix checks run when each depth executes. SpeculativeRunner prepares all depths of a draft offset, and of verification, in one call.
Draft IDs were read back to the host after every draft offset and uploaded again as the next prelude input, which put a host sync inside the draft loop. Keep each offset's IDs on the device and feed them straight into the next offset's prelude. The rows still drafting at one offset are exactly the active rows of the next, in the same order, so ragged requests need no remapping. Greedy rows take the argmax; sampled rows overwrite their entry with the drawn token on the device, with an unchanged per-request RNG consumption order. All draft IDs and the greedy verification targets are read back in one transfer after verification, so a greedy round synchronizes with the host once.
7e54987 to
e6e9fd0
Compare
Thanks, all four are addressed; details and data are in the updated description.
As you expected, the gains grow with concurrency: on A100 with graphs, +5.3% at B=8 and +27.4% at B=64. The PR is now ready for review. |
Independent performance validation using my
|
| Configuration | Base decode tok/s | PR decode tok/s | Decode change | E2E output-throughput change |
|---|---|---|---|---|
| Triton, greedy graphs, B=1 | 111.4 | 114.4 | +2.67% | +2.62% |
| Triton, greedy graphs, B=8 | 741.8 | 793.5 | +6.96% | +8.21% |
| Triton, greedy graphs, B=32 | 2337.4 | 2723.1 | +16.50% | +14.11% |
| Triton, greedy graphs, B=64 | 3728.0 | 4759.7 | +27.67% | +32.33% |
| Triton, greedy eager, B=8 | 425.6 | 440.2 | +3.44% | +3.48% |
| Triton, sampling T=0.8, graphs, B=8 | 703.7 | 724.3 | +2.92% | +3.06% |
| Triton, 2048-token input, graphs, B=8 | 362.2 | 375.0 | +3.53% | +3.09% |
| Triton, natural inputs + EOS, graphs, B=16 | 787.8 | 849.8 | +7.87% | +7.30% |
| FA4, greedy graphs, B=8 | 800.3 | 854.1 | +6.73% | +6.39% |
| FA4, sampling T=0.8, graphs, B=8 | 750.0 | 774.6 | +3.27% | +2.91% |
| FA4, natural inputs + EOS, graphs, B=8 | 619.0 | 654.6 | +5.76% | +5.00% |
| Triton, native eager, B=8 | 258.8 | 257.5 | -0.54% | -0.39% |
| Triton, native graphs, B=1 | 67.5 | 67.7 | +0.42% | -0.11% |
| Triton, native graphs, B=8 | 467.1 | 470.6 | +0.76% | +0.94% |
| Triton, greedy graphs, B=32, default 128-row limit | 2202.4 | 2660.1 | +20.78% | +17.61% |
Timing definitions and interpretation:
- Decode starts after every request has produced its first token; fixed-length cases count B×63 remaining tokens. For natural-length inputs, some shorter requests can already have advanced, so the numerator is the actual remaining token count; prioritize E2E for those cases.
- E2E covers in-process engine arrival through completion, including prefill, but excludes HTTP/network and tokenization. These are short-burst measurements, not steady-state service throughput or saturation capacity.
- Across the six decode measurements per variant/configuration, the largest max/min−1 spread was 1.60%. Native eager's −0.54% change is within observed variation; the earlier apparent larger regression did not reproduce in the controlled environment.
- Natural-input tests use the first qualifying ShareGPT human prompts in file order, with 16–1024 input tokens, no chat template, and normal EOS handling. Triton B=16 acceptance was 721/935 = 77.1%; outputs included EOS-shortened requests. This is an output-equivalence check on the tested slice, not a task-quality evaluation.
- The default-128-row B=32 case incurs 24 core and 12 coda fallbacks per measured request batch on both revisions. Natural Triton B=16 also has graph-entry-limit fallbacks. The measured gains do not imply complete graph coverage.
Profiling evidence and source attribution
For the same four complete B=8 greedy rounds, the trace counts were:
| Event | Base | PR |
|---|---|---|
cudaStreamSynchronize |
548 | 4 |
| H2D copies | 260 | 52 |
| D2H copies | 288 | 4 |
| D2D copies | 328 | 328 |
| CUDA Graph launches | 68 | 68 |
| GPU kernels | 66,780 | 66,516 |
This is 137 → 1 synchronizations per complete greedy round. Independent sync-debug runs confirm it. With T=0.8, complete rounds still perform 169 → 73 synchronizations; the sampling path is not down to one.
The mechanism matches the implementation:
SpeculativeRunner.executekeeps draft IDs on-device and reads draft/greedy verification IDs back together after verification._coresand_prepare_batchesshare depth-independent metadata._stageuses pinned memory and nonblocking H2D copies.rejection_sampleretains sequential tensor-to-host decisions, explaining the remaining sampling synchronizations.
The graph launch count and model loop depths are unchanged; kernel count only drops about 0.4%. The evidence supports reduced host synchronization, readbacks, and metadata work rather than less model computation or shorter outputs. No per-commit ablation was performed, so the total gain cannot be quantitatively assigned to individual commits. Host synchronization wait time includes preceding GPU work and must not be treated as pure removable copy time.
Profiling caveat: ops-only captures succeeded under the performance environment's PyTorch 2.13. Some full captures failed during export with UnicodeDecodeError. Paired ops-only/full captures were therefore repeated under the already-installed PyTorch 2.14.0+cu130, using the same clean revisions, GPU, model, workload and Triton backend. Both full exports succeeded with readable shapes, memory and stacks; the event counts above agree between 2.13 ops-only and both 2.14 modes. No profiled timings or 2.14 timings are mixed into the performance table. The original full traces were roughly 150–282 MiB; the skill and local harness have since been updated to start with one complete representative round and reduce oversized windows before retrying.
Correctness and experimental controls
- All 180 controlled performance runs had matching A/B token IDs and exit depths, including T=0.8 sampling, natural EOS, and the FA4 packed-verification path. No failed requests were silently excluded.
- Targeted tests:
pytest --run-gpu tests/test_speculative.py tests/test_prepared_kv.py tests/test_cuda_graph.py -q.- Base: 103 passed, 1 skipped.
- PR: 111 passed, 1 skipped.
- Initial GPU 0 runs experienced resource contention, including unrelated allocations exceeding 100 GiB. Those results were retained but excluded from the conclusions by rerunning the entire main matrix on GPU 2. During the controlled runs, sampled compute-process telemetry found no competing process on the tested GPU. Sampling cannot rule out every momentary disturbance.
- The local experiment bundle retains commands, exact inputs/token IDs, configurations, per-run outputs, telemetry, traces, and summaries. These artifacts are not attached to this comment; the comment contains the measured results and methodology rather than links to inaccessible local paths.
Assessment: the measured performance benefit is real, especially for larger greedy graph batches. Keep claims scoped to the tested workload and configuration: the ~28% result is not a universal speedup, and “one synchronization per round” should explicitly say greedy. PD, async speculation, adaptive exit, sustained service load, and other GPU/software stacks were not validated in this review.
|
lgtm |
3 participants
Scope
First step of the Device-side execution loop item in #43. A speculative round now synchronizes with the host once, down from 73–110 times, and depth-independent KV metadata is prepared once per token batch instead of once per depth. The branch is rebased on
main(which now includes #48). Following review, the benchmark harness is no longer part of this PR.Behavior is unchanged: in every configuration measured, including sampling, outputs are token-identical to
main.Changes
perf: batch greedy draft and verification readbacks(speculative.py): accepted greedy candidates equal their row's target argmax, so the committed tokens are the target-argmax prefix through the first mismatch. Verification reads all argmaxes with oneargmax(-1).tolist()and compares on the host via the newgreedy_accept().perf: stage prepared KV metadata in pinned memory(kv_cache_manager.py):_prepare_batchstages its metadata in two pinned host tensors and copies them withnon_blocking=True, instead of 5–6 pageabletorch.tensor(..., device=cuda)copies, each of which waited for all queued GPU work.perf: share depth-independent KV metadata across recurrence depths:_prepare_batches()validates the rows once and builds positions, write offsets, context lengths and packed-query boundaries once. It returns one batch per depth, where only block addresses and page tables are per depth._prepare_batch()delegates to it. Ownership and written-prefix checks still run when each depth executes.perf: keep speculative draft IDs on the device: each draft offset's IDs feed the next prelude directly. The rows still drafting at one offset are exactly the next offset's active rows, in the same order, so ragged requests need no remapping. Sampled rows write their drawn token on the device, and per-request RNG consumption order is unchanged. All draft IDs and the greedy targets are read back in one transfer after verification.Results
Ouro-1.4B BF16 @
574fa66c, torch 2.14+cu130. Prompt 128 tokens, 64 output tokens,ignore_eos, D=4/d=2, K=4, decode-only after matched prefill. Two interleaved cycles, run order reversed in the second; each point is 5 repeats. tok/s is the median of cycles. The same configuration varied by <1% between cycles.A100-SXM4-80GB (MIG 7g.80gb, 98 SMs), triton backend
Speculative decoding,
mainvs changes 1–2 vs this PR:Syncs per round: exactly 1 at B=8 and B=64, counted with
torch.cuda.set_sync_debug_mode.Non-speculative path.
_prepare_batchis shared with native decoding, so this covers small batches and long contexts. Prefill is chunked at 8,192 tokens.No regression; the largest negative change, −0.2%, is within run-to-run variation. Absolute values are in the table below.
Absolute native tok/s (main → this PR)
B200 (MIG 4g.90gb, 72 SMs), FA4, previous round
These runs cover changes 1–2 only, on top of #48, which is identical to today's
mainfor these files:Changes 3–4 have not been measured on B200 yet; see "Not covered yet".
Profiling notes
_stage()is 0.66% of the decode window at B=8 and 0.99% at B=64, so it is not implemented. Since the host runs ahead of the GPU, not all of that time is on the critical path.main, spec K=4 + graphs at B=8 (241 tok/s) is slower than native + graphs (306 tok/s).Relation to other work
_prepare_batchnow delegates to_prepare_batchesand returns views into shared staging tensors. The fields and semantics of_PreparedKVBatchare unchanged.Correctness and tests
mainin every configuration above, for greedy and sampling (T=0.8). The sampling match confirms the RNG consumption order is unchanged.greedy_accept: K=0/1/2/4/8, every rejection position._prepare_batchesmatches per-depth_prepare_batchfor packed and unpacked batches, shares one copy of the depth-independent tensors, and rejects invalid depths.tests/test_prepared_kv.pynow expects 2 metadata tensor constructions per core pass instead of 5.pytest --run-gpu tests/on A100: 405 passed. The 30 failures are the same set as onmain(400 passed): FA4 not supported on SM80, FA2 requiring a block size that is a multiple of 256, and PD needing a second GPU.Not covered yet