fix(ds4): restore rejected compressor rows during speculative rollback

Merged
DU
1mo ago
+250-112

Summary

Fix the compact DS4 speculative rollback path so rejected tokens cannot leave
data in compressor rolling state. This is a correctness fix, independent of
decode optimizations and prefill PR #684.

Before verification, save both halves of each ratio-4 compressor and the
at-risk rows of each ratio-128 compressor. On rejection, restore the rejected
current/ring rows regardless of whether a ratio-4 boundary was crossed. Keep
accepted current rows and legitimate previous-window rotations. Q5 continues
to use the existing restore-to-start plus accepted-prefix replay path.

The old code saved only the ratio-4 previous half and skipped ratio-128 state.
Position-addressed writes alone are insufficient: a later compressor pool can
read a rejected future row before that row is overwritten by committed work.

Scope

  • Reuse existing vectors and pinned spans. A saved-mode boolean records whether
    pinned staging was actually selected (including allocation fallback).
  • Share the row-copy logic between synchronous, asynchronous, and pinned-host
    staging. Save accepts optional backend/pinned arguments; apply derives its
    transport from the saved state and cannot independently select the wrong mode.
  • Preserve raw-ring restoration, HC handling, compressed-row count truncation,
    the full-snapshot mode, verifier width policy, and wide rejection replay.
  • Retain the existing six-row staging capacity. This does not enable Q6
    verification, which remains explicitly unsupported.

Regression

The new checked-in test_deepseek4_unit case compares exact bytes against an
independent sequential accepted-prefix state transition. It covers all 128
starting phases, widths 1–6, every accepted count including zero/all, all four
KV/score tensors for each ratio, and F16 KV plus F32 score storage. Widths above
four exercise restore-to-start/replay, and width six tests staging capacity
only. There are 3,456 transactions per copy mode.

The regression fails against unmodified public main (b1c88230), first at
pos=0 q=1 accepted=0 tensor=0, and passes after this repair.

Validation

  • Base: b1c88230e3ab584ab90ff0b224ddc69e6f8f6b58.
  • Candidate: 2722c1c05d863c34809e60cf73bf99d053b86f99.
  • Review follow-up: reproduced the configured-save/default-apply pinned
    assertion on the preceding head. The regression now queues save, overwrite,
    and default apply without a post-save synchronization; async/pinned modes pass.
  • Current-head full CUDA build: pass. Full CTest: 525 tests, zero failures,
    five expected skips (61.87 seconds).
  • The following original qualification results were collected on predecessor
    aef40f182d50f33dedb8a4f72372a3dd77b3258d; AMD/sanitizer requalification of
    the follow-up is separate:
  • Public-branch CUDA builds of test_deepseek4_unit and dflash_server: pass,
    with build parallelism capped at four.
  • Complete ctest -R '^deepseek4_unit$': pass in 12.87 seconds on RTX 3090.
    The new CPU sweep and all three CUDA copy modes ran and passed. Six existing
    HIP-only/two-GPU checks skipped on this single-CUDA-device machine.
  • Focused sweep on AMD R9700 (gfx1201): all three copy modes pass.
  • Focused sweep on AMD 8060S (gfx1151): all three copy modes pass.
  • The AMD checks compile the exact candidate helper and checked-in regression
    against the existing staging HIP backend. They are not a fresh full HIP
    server build or a model-backed test. Pinned mode verifies the allocated
    buffer type, so pageable fallback cannot masquerade as pinned coverage.
  • CPU focused sweep with AddressSanitizer and UndefinedBehaviorSanitizer:
    pass (vptr instrumentation and leak detection disabled in the isolated
    test driver; GGML shared libraries are not sanitizer-instrumented).
  • git diff --check: pass.

Reproduce the checked-in regression on a configured build:

cmake --build server/build --target test_deepseek4_unit dflash_server -j4
ctest --test-dir server/build -R '^deepseek4_unit$' --output-on-failure

No throughput, full-model token parity, or general verifier-numerics claim is
made by this PR. In particular, it does not claim to deliver 51 tok/s on public
main. Upstream CI/review and maintainer merge are still required.