Skip to content

perf(qwen35): add opt-in F32 rollback checkpoints for single-GPU inference#515

Open
Amidwestnoob wants to merge 2 commits into
Luce-Org:mainfrom
Amidwestnoob:perf/single3090-f32-rollback-clean
Open

perf(qwen35): add opt-in F32 rollback checkpoints for single-GPU inference#515
Amidwestnoob wants to merge 2 commits into
Luce-Org:mainfrom
Amidwestnoob:perf/single3090-f32-rollback-clean

Conversation

@Amidwestnoob

@Amidwestnoob Amidwestnoob commented Jul 12, 2026

Copy link
Copy Markdown

Summary

Adds an experimental, default-off F32 checkpoint mode for single-chain speculative rollback on memory-constrained single-GPU Qwen deployments.

The existing behavior remains unchanged unless explicitly enabled:

DFLASH_SINGLE_CHAIN_CHECKPOINT_F32=1
DFLASH_FAST_ROLLBACK_THRESHOLD=1   # valid range: 1-5

Optional diagnostics:

DFLASH_SINGLE_CHAIN_ROLLBACK_DIAG=1

Without the F32 opt-in, checkpoints remain F16 and the rollback threshold remains 5.

Motivation

On a single RTX 3090 running Qwen3.6-27B Q4, replaying rejected speculative tokens can cost more than restoring a saved recurrent-state checkpoint. F32 checkpoints permit direct restoration and avoid the conversion/replay path for sufficiently deep rollbacks.

This is intentionally separate from the split/layer-sharded work in #506. It targets a different deployment: one 24 GiB GPU, with no layer split.

Implementation

  • Centralizes checkpoint dtype, threshold, and diagnostic policy resolution.
  • Allocates F32 single-chain rollback checkpoints only when opted in.
  • Uses direct F32 checkpoint copy/restore when available.
  • Preserves the existing F16 conversion path by default.
  • Adds unit coverage for defaults, opt-in behavior, threshold validation, and diagnostics.
  • Leaves split/multi-GPU checkpoint allocation unchanged.

Validation

Final standalone branch, based on current main:

  • CUDA release build: passed
  • check: 15/15 tests passed
  • Single RTX 3090 context stress: 8K through 65K active tokens
  • OOMs: 0
  • Segfaults: 0
  • Second GPU usage: 15 MiB (idle)

Fresh A/B/A validation on the standalone build, six prompt families and three repetitions per arm:

  • F16 threshold 5, first baseline: 53.299 tok/s
  • F32 threshold 1: 56.764 tok/s
  • F16 threshold 5, repeated baseline: 53.338 tok/s
  • F32 threshold 1 versus bracketed baseline: +6.46% weighted decode throughput

The broader development campaign previously found threshold 2 to be the best pooled policy, but this PR deliberately uses the 6.46% threshold-1 result from the exact standalone candidate as its directly reproducible performance claim.

Limitations

  • Experimental and default-off.
  • F32 checkpoints require approximately 1.15 GiB more VRAM than F16 checkpoints in the tested configuration.
  • Generated text can differ from the legacy replay trajectory; this is not a claim of bit-identical output.
  • Performance evidence is specific to Qwen3.6-27B Q4 on an RTX 3090 and should not be generalized to other models or GPUs without measurement.
  • Some workloads may see smaller gains or no gain.
  • Split/multi-GPU behavior is out of scope and remains unchanged.

Suggested test configuration

cmake -S server -B server/build -DDFLASH27B_TESTS=ON
cmake --build server/build --target check

Review in cubic

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1 issue found across 11 files

Prompt for AI agents (unresolved issues)

Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="server/src/qwen35/qwen35_backend.cpp">

<violation number="1" location="server/src/qwen35/qwen35_backend.cpp:1932">
P2: The rollback diagnostic counters and print logic are duplicated across both decode loops (`qwen35_backend.cpp` and `dflash_spec_decode.cpp`): five counter variables, the accept_n histogram indexing, three counter-update sites, and the diagnostics print block. Approximately 26 lines are verbatim copies. This makes the two paths prone to drift when someone updates one side but forgets the other. Consider extracting the counters and the print block into a small RAII accumulator struct (e.g. `RollbackDiag`) shared from `chain_rollback_policy.h`, then use it from both decode loops. The struct could hold the counters, expose `record_fast_rollback(accept_n)`, `record_legacy_replay()`, `record_fallback()`, and a `print()` method — keeping the two decode loops focused on decode logic rather than diagnostic bookkeeping.</violation>
</file>

Reply with feedback, questions, or to request a fix.

Re-trigger cubic

int n_hint_accepted = 0;
int target_forwards = 0;
const ChainRollbackPolicy rollback_policy = resolve_chain_rollback_policy();
int rollback_accept_hist[17] = {};

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2: The rollback diagnostic counters and print logic are duplicated across both decode loops (qwen35_backend.cpp and dflash_spec_decode.cpp): five counter variables, the accept_n histogram indexing, three counter-update sites, and the diagnostics print block. Approximately 26 lines are verbatim copies. This makes the two paths prone to drift when someone updates one side but forgets the other. Consider extracting the counters and the print block into a small RAII accumulator struct (e.g. RollbackDiag) shared from chain_rollback_policy.h, then use it from both decode loops. The struct could hold the counters, expose record_fast_rollback(accept_n), record_legacy_replay(), record_fallback(), and a print() method — keeping the two decode loops focused on decode logic rather than diagnostic bookkeeping.

Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At server/src/qwen35/qwen35_backend.cpp, line 1932:

<comment>The rollback diagnostic counters and print logic are duplicated across both decode loops (`qwen35_backend.cpp` and `dflash_spec_decode.cpp`): five counter variables, the accept_n histogram indexing, three counter-update sites, and the diagnostics print block. Approximately 26 lines are verbatim copies. This makes the two paths prone to drift when someone updates one side but forgets the other. Consider extracting the counters and the print block into a small RAII accumulator struct (e.g. `RollbackDiag`) shared from `chain_rollback_policy.h`, then use it from both decode loops. The struct could hold the counters, expose `record_fast_rollback(accept_n)`, `record_legacy_replay()`, `record_fallback()`, and a `print()` method — keeping the two decode loops focused on decode logic rather than diagnostic bookkeeping.</comment>

<file context>
@@ -1927,12 +1928,33 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen,
     int n_hint_accepted = 0;
     int target_forwards = 0;
+    const ChainRollbackPolicy rollback_policy = resolve_chain_rollback_policy();
+    int rollback_accept_hist[17] = {};
+    int rollback_fast_low = 0;
+    int rollback_fast_high = 0;
</file context>

Comment thread server/src/common/chain_rollback_policy.h
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant