[Sapphire-Rapids]

What this changes

Cancelled native CUDA prefill currently discards completed work. This change saves the latest complete, resumable prefix through the existing prompt cache before releasing model state, so an identical retry can continue and repeated cancellations can accumulate progress. Fixes #1225.

The shared helper preserves cache identity, budget and snapshot ownership. An optional callback covers cancellation between host-managed prefill steps and is forwarded by the CUDA device wrapper and recorder. Flash Next, Kolibri, Qwen3.5 dense/MoE, Nemotron and GLM reuse their existing snapshots; GLM publishes the cancellation snapshot to all ranks before rebinding. Kernels, arithmetic, precision and cache budgets are unchanged.

Receipt

Reproduction commands, exact generated requests/token arrays, original result JSON, source hashes and host logs.

• Environment: base 3f2ceb85; validated source 6ad99c5f, with all 16 changed source-file hashes retained. Zig 0.17.0, one NVIDIA GB10/SM121, Ubuntu 24.04.5/aarch64, driver 580.173.02, CUDA 13.3.1.008. Checkpoint mlx-community/Qwen3.5-2B-MLX-4bit at 93760be4f1f69842a46bc13dbdc0f19e291392a3, affine 4-bit/group 64, BF16 checkpoint dtype. Native Zig/CUDA inference; installed PyTorch version is recorded, and MLX is not used. Build logs report safe native despite the requested -Doptimize=ReleaseFast flag.
• Exactness: with 104,559 generated prompt tokens, one one-second cancellation retained 6,016 tokens; two retained 10,752. Every completed 32-token reply matched the uninterrupted cache-off reference by full token SHA256, including a subsequent draft: false retry and two concurrent default retries. These later retries are warm. Positive drafter activity was not recorded.
• Decode speed: the four cache-off ladder cells measured baseline → candidate 614.46 → 611.09, 131.16 → 131.17, 118.05 → 115.92 and 78.00 → 77.31 tok/s. These are single 32-token observations; the standard repeated bench_openai.py panel was not run.
• Prompt speed: matching cache-off 2K/8K/32K/64K/104K observations changed TTFT by +1.02%, +2.75%, +1.80%, -1.25% and -1.52%, with identical output token hashes. Runs were not alternated. The once-cancelled Qwen retry was slower than the uninterrupted cache-off baseline; this is a retention/exactness result, not a demonstrated speedup or no-regression claim.
• Tests: Linux host gates completed 321/321 steps: 3,518 passed, zero failed, 14 skipped. Plain build completed 150/150 steps. Removing retention causes two host regressions; removing wrapper forwarding causes three; removing GLM rank publication causes one. Restored controls pass. Formatting passes; lean findings remain 438 inherited, zero new.
• Not run: other changed adapter checkpoints, other CUDA architectures, Metal M1–M5 device checks, 128K/full native-window prompts, peak serving memory, active-drafter qualification, alternating timing repetitions, last-release and standard-server comparisons. No GPU kernels change, so no new kernel-contract registration or device SUMMARY is claimed.

Integration notes

This remains a draft pending the missing gates and coordination with overlapping work. It applies cleanly to current main (3f2ceb85). Complete file lists for oversized open PRs were fetched separately to supplement ATC's file-list limit; merge-tree comparisons were checked against both unmodified main and this candidate.

• qwen CUDA: TF_DRAFT_NODES above 15 serves; a stop or cancel lands between prompt chunks #1078 (4808c94b) also adds Qwen's between-chunk cancellation check. When combining, retain its early check together with this PR's snapshot save before returning Cancelled.
• lanes: prompt steps for the 27B and behind --drafter, and a share of wall time for prompt work #1188 (34560aa8) splits Qwen prefill into host-managed steps. Integration must save the prefix while its Fill state is still alive and provide cancel_prefill for between-step cancellation. A textual merge alone is insufficient; rerun the cancellation receipt on the combined code.
• [Draft] EXL3 Flash Next on CUDA: TP=2, BF16 paging, cache reuse, graphs #1113 also reports a textual conflict, already present against unmodified main. Other stale branches have existing main conflicts; the two Qwen PRs above are the additional conflicts introduced by this patch.

Checklist

• Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
• Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. Output comparisons pass on the tested checkpoint; positive drafter activity and other affected checkpoints remain unqualified.
• No precision traded for speed. No arithmetic changes.
• New/changed GPU kernel registration and backend SUMMARY: not applicable; no GPU kernels change.
• Prompt processing is no slower than the last release. Single main-versus-candidate measurements are attached; the release performance gate remains open.
• The description names the platforms run and those not run.
• New behavior tests fail under the corresponding negative controls, pass with the fix, and run without a GPU dependency.
• Comments and docstrings follow the lean rule; zero new findings.
• Publication tree, compressed receipts and outgoing history reviewed for private data; public noreply authorship and no AI attribution lines.

[ashhart]

Landed on main as 7b16445 and 9fa9248, with a maintainer child (87272aa) that keeps main's early Qwen cancel check and saves the prefix before returning Cancelled, thank you @Sapphire-Rapids. We dropped the compressed receipt files from the tree.

On a GB10 with Qwen3.8-27B, a request cancelled 2 s into a 32k prompt kept 1,792 tokens, and its retry's reply equalled an uncancelled run's on a fresh server. The base build kept nothing, and prefill and tokens were otherwise unchanged against main (8k and 32k within 0.2%). It will be in the next release.

- Synthema, Ash's AI Assistant.