Skip to content

perf(cuda): CUDA-graph-captured k-token spec verify — the ~10.6 ms fixed verify overhead is now DSpark's whole gap (#428 lever 2) #430

Description

@pekkah

Split out of #428 (its MMQ half is resolved-rejected; this is the remaining lever). With PR #429 merged, the DSpark draft round is ~11.8 ms flat and the verify pass is the entire remaining cost: verify(t) = 10.6 + 5.05·t ms on Qwen3-4B (t = verify batch tokens, k+1), i.e. the fixed ~10.6 ms/round is ~25–45% of the whole DSpark round depending on k, and the un-graphed 36-layer batched trunk (~300 launches) is where it goes.

Why not MMQ

Measured on the 4B pairing (#428 comment): SHARPI_BATCH_DECODE_MMQ=1 verify reaches 79.5 t/s but fails the 256-token parity oracle — output text diverges from plain greedy mid-sequence, so batched-MMQ argmax is not stable against the per-token decode path. The allowDecodeMmq:false pin in BatchVerify/BatchVerifyArgmax is load-bearing and stays. Any fixed-cost cut must keep the bit-exact WS matvec path — which is exactly what graph capture does (same kernels, fewer launch round-trips).

What capture needs

  1. Position-parameterized nodes: the batched verify trunk's positions/slots ride as by-value kernel params (ragged kernels' sharpi_seq_pos) and scalar args (RoPE base position, seq lens). A replayable graph needs cuGraphExecKernelNodeSetParams updates per round, or an indirection through a small device buffer the host writes before replay.
  2. Per-k graphs: the confidence trim varies k per round (1..block). Either one graph executable per k (block=7 → 8 graphs, built lazily) or pad-to-max-k with masked tails (wasted compute at small k — measure which wins; at 5.05 ms/token padding is expensive, per-k graphs look right).
  3. Tap capture inside the graph: this is the hard part and the same reason the plain decode graph bails when taps are on — CaptureTapRows does position-indexed copies, and a captured memcpy node can't retarget its destination (see the SnapKV Q-capture rationale). Design hint: copy tap rows keyed by batch row index (graph-stable) into a fixed staging area inside the graph, then scatter row n → positions[n] with one tiny post-replay kernel/host step outside the graph.
  4. On-device argmax (llm_argmax_rows) is already at the tail and graph-friendly; the k×8-byte D2H + sync stays outside the graph.

Projected gain (measured model, 4B laptop)

At today's best config (--dspark-verify-len 3, 67.9 t/s, ~42 ms/round = draft 12.6 + verify 29.3 + commit): cutting the fixed 10.6 → ~1.5 ms is ~9 ms/round → ~77–85 t/s, 1.25–1.4× plain (61.5). The 8B fixed cost is larger (verify_8B(t) ≈ 16.3 + 7.75·t), so the absolute cut there is bigger still.

Related: #401 (decode-trunk CUDA-graph follow-ups — shares the param-update machinery), #405#409 (batched-decode gap umbrella), #428 (desktop validation), PR #429 (draft-round work + SHARPI_DSPARK_TIMING=1 breakdown used for these numbers).

🤖 Generated with Claude Code

https://claude.ai/code/session_018EWo1niax1g8A5E79RSALJ

Metadata

Metadata

Assignees

No one assigned

    Labels

    perfPerformance optimization opportunity

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions