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

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

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