CURE: Local Uncertainty Repair for Block-Parallel Speculative Decoding

Let $x_{<t}$ denote the generated context sequence at decoding step $t$. A block-parallel draft model $q$ generates a length-$K$ candidate sequence $\mathbf{y}_{t:t+K-1}=(y_{t},y_{t+1},\dots,y_{t+K-1})$ alongside its corresponding logit distributions in a single forward pass. A target autoregressive model $p$ subsequently verifies candidate trajectories and accepts a prefix. In standard single-path speculative decoding, candidate generation yields a single path $\pi_{0}$. When an error occurs at position $k$, standard prefix verification rejects all trailing tokens $y_{j}$ ($j>k$), leading to severe truncation of the accepted length.

To recover valid suffixes beyond early rejection points without full autoregressive drafting, we formulate local repair as a constrained candidate-tree optimization problem. Let $\Pi=\{\pi_{0},\pi_{1},\dots,\pi_{R}\}$ represent a tree of candidate execution paths rooted at $x_{<t}$, where $\pi_{0}$ denotes the default block-parallel path and $\{\pi_{r}\}_{r=1}^{R}$ denote candidate repair branches spawned at identified error points. Let $a(\pi)$ be the accepted token length along path $\pi$, and let $V(\Pi)$ denote the total token budget (i.e., tree node count) presented to the target verifier. The optimal candidate trajectory $\pi^{\star}$ maximizes the accepted sequence length under a strict verification compute budget $V_{\max}$:

From an efficiency perspective, the effective token generation rate scales with $\frac{\mathbb{E}[a(\pi^{\star})]+1}{\mathcal{C}_{d}+\mathcal{C}_{v}(\Pi)}$, where $\mathcal{C}_{d}$ is the constant draft overhead and $\mathcal{C}_{v}(\Pi)$ is the memory/compute cost of target tree verification. Therefore, candidate tree expansion yields a positive wall-clock trade-off if and only if the gain in accepted tokens outpaces the increased target verification complexity $\mathcal{C}_{v}(\Pi)$.

To solve the constrained candidate-tree optimization problem under budget $V_{\max}$, CURE introduces an end-to-end, inference-only pipeline that operates directly atop any pretrained block-parallel backend without model retraining. As illustrated in Figure 1, each decoding iteration of CURE proceeds through three integrated phases:

1. 1.

   Uncertainty Focal Point (UFP) Detection: The draft model generates a single-pass block proposal $\pi_{0}$. CURE inspects the per-position top-1/top-2 logit margins to dynamically locate fragile positions (UFPs) where draft confidence drops precipitously.

2. 2.

   Budget-Aware Dynamic Tree Construction: Instead of uniform or unconstrained tree expansion, CURE dynamically calculates candidate branch allocations strictly at identified UFPs. It attaches alternative candidate tokens to $\pi_{0}$, constructing a sparse repair tree $\mathcal{T}$ that satisfies $V(\mathcal{T})\leq V_{\max}$.

3. 3.

   Single-Pass Verification and Cache Resynchronization: The target model $p$ verifies all branches in $\mathcal{T}$ concurrently using Tree Attention. If a repair branch $\pi^{\star}\neq\pi_{0}$ yields a longer accepted prefix, CURE accepts the repaired path and resynchronizes the draft model’s internal KV-cache to realign subsequent drafting states.

For each position $i\in\{0,\dots,K-1\}$ within a generated draft block, let $\ell_{i,1}$ and $\ell_{i,2}$ denote the top-1 and top-2 unnormalized log-probabilities emitted by the draft model $q$, respectively. We define the predictive confidence margin $m_{i}$ as:

A smaller predictive margin $m_{i}$ indicates that the draft model exhibits lower confidence in distinguishing the leading candidate tokens. Positions with $m_{i}<\tau_{\text{margin}}$ are identified as Uncertainty Focal Points (UFPs). We select positions in ascending order of $m_{i}$ to prioritize candidate expansion at the most fragile nodes while capping total tree size to remain within the verifier budget $V_{\max}$.

In our baseline setup, each selected position expands at most $B=5$ candidate branches, with a maximum tree depth of 15, a target verification beam width of 5, and a repair margin threshold $\tau_{\text{margin}}=1.0$. To optimize tree efficiency, we further introduce a dynamic branching mechanism that dynamically maps the normalized uncertainty $u_{i}\in[0,1]$ to a local branch budget $b_{i}$:

where $s$ denotes the margin scale, $b_{\min}=1$, and $b_{\max}=5$. This dynamic allocation adjusts only the local tree topology without modifying the underlying draft checkpoint or requiring retraining.

For each identified UFP, CURE constructs candidate repair paths by preserving the original prefix while substituting the fragile token with high-scoring alternatives. The remaining suffix tokens are then resampled under the substituted token. Paths sharing a common prefix naturally merge into shared parent nodes within the verification tree, enabling the target verifier to process all candidate paths in a single pass via Tree Attention without independent batch computation. Crucially, the original block-parallel path ($\pi_{0}$) is anchored as branch 0 and is protected from pruning by the uncertainty gate.

To ensure computational efficiency, candidate construction is strictly bounded by three complementary constraints: the position gate limits where repair branches may originate, the branch budget caps alternatives per selected position, and the verification beam constrains the active tree frontier. In our primary configuration, at most one low-margin position is repaired per block. This localized strategy concentrates the extra verification budget on critical ambiguous nodes rather than enumerating combinatorial paths across the full sequence block.

The constructed candidate paths and the default path $\pi_{0}$ are unified into a sparse verification tree $\mathcal{T}$. A single forward pass of the target model $p$ evaluates all tree positions concurrently via Tree Attention and computes the accepted prefix length $a(\pi)$ for each candidate trajectory $\pi\in\mathcal{T}$. The framework identifies the optimal trajectory that yields the longest accepted prefix:

If branch 0 wins ($\pi^{\star}=\pi_{0}$), the current draft internal states are preserved. If a repair branch wins ($\pi^{\star}\neq\pi_{0}$), the emitted sequence deviates from the original draft proposal. To maintain state consistency for subsequent drafting steps, CURE replays the selected repair tokens into the draft model’s Key-Value (KV) cache, realigning the draft model’s internal context representation with the verified output sequence.

While the pretrained block-parallel draft checkpoint requires specialized multi-token training, all components of CURE (including the dynamic repair tree, margin gating, dynamic branch allocation, and KV-cache resynchronization) operate entirely at inference time. These mechanisms introduce zero additional parameters to either the draft or target models and require no retraining or fine-tuning, establishing CURE as a flexible, plug-and-play module for existing parallel drafting backends.

Table 1 presents our controlled, same-system evaluation, where CURE is configured with the fast_margin100 variant. To contextualize our results within the broader literature, Table 2 separately summarizes reported speedups from existing literature on Qwen3-8B. We utilize the speedup ratio relative to each benchmark’s target-only AR baseline as the primary unifying metric, since raw TPOT is inherently dependent on specific hardware and execution environments, while accepted-length metrics vary by candidate tree structure. While the target model (Qwen3-8B), decoding temperature ($0$), and evaluation benchmarks are aligned, we explicitly acknowledge that hardware platforms, generation sequence lengths, and underlying framework implementations differ across published works.

Across the three primary code generation benchmarks, CURE achieves an end-to-end speedup of $2.66\times$ to $3.49\times$ over target-only AR decoding, and an acceleration of $4.38\times$ to $5.43\times$ over Naive SD. Relative to the single-path parallel draft backend, local uncertainty repair consistently increases the accepted token sequence length by 4.2% to 7.5%, while raw TPOT increases by $1.40\times$ to $1.61\times$. These findings establish that the core mechanism of CURE lies in recovering valid candidate trajectories at fragile nodes rather than improving raw latency over the parallel draft baseline. On the out-of-domain GSM8K benchmark, CURE demonstrates a notable 32.0% increase in accepted length, which is analyzed separately as an empirical stress test. Table 2 positions our speedup metrics alongside reported figures for established speculative decoding backends on Qwen3-8B. On shared evaluation benchmarks, the end-to-end speedups achieved by CURE surpass those reported for EAGLE-3 and DART, while remaining lower than DFlash and Domino. Given the variations in execution hardware, generation length, and framework optimizations across these implementations, we interpret this comparison as a system-level context positioning rather than a direct head-to-head benchmark.

To verify that leveraging token error patterns preserves target model alignment without degrading downstream task performance, we evaluate the functional correctness of generated code across standard benchmarks (HumanEval, MBPP, and LCB-lite). Table 3 summarizes the pass rates alongside latency and average accepted sequence length.

Task Performance Parity. Our method matches target-only decoding on all three reported pass@1 counts. Compared with the parallel drafting baseline, it differs by two HumanEval problems and matches MBPP and LCB-lite. These single-run results support comparable test pass counts under the evaluated suites; they do not establish functional equivalence beyond the provided tests.

Scope and Numerical Precision. We note that execution-based test suites evaluate behavioral correctness rather than strict token-level identity. The optimized bfloat16/FlashAttention path exhibits token-level mismatches, while a strict replay path restores target consistency in the tested path at the cost of additional latency.

To isolate the system-level contributions of the parallel fallback branch (Branch 0) and the KV cache resynchronization mechanism, we evaluate their individual and combined removals. Experiments are conducted on a representative 64-example subset sampled from HumanEval, MBPP, and LCB-lite. Table 4 summarizes the trade-offs in mean accepted length, end-to-end speedup, and TPOT ratio.

Branch 0 as a Low-Cost Safety Net. Removing Branch 0 results in a 5.7% drop in the average accepted length and reduces the end-to-end speedup from $2.708\times$ to $2.616\times$. This degradation demonstrates that retaining the vanilla parallel proposal path serves as a highly efficient fallback baseline; it ensures candidate diversity when error-pattern repair candidates fail to provide a valid token sequence.

Cache Resynchronization for State Alignment. Omitting the cache resynchronization mechanism leads to a catastrophic performance collapse. The average accepted length plummets by 40.1%, while the TPOT ratio deteriorates sharply from 1.776 to 2.642. From a system perspective, this severe overhead stems from stale draft states accumulating across consecutive decoding blocks, which breaks KV cache alignment and invalidates subsequent token verifications.

To understand whether token error repairs can be selectively prioritized before verification, we analyze the ranking capacity of a lightweight scoring mechanism. We define a useful repair as a block where the tree candidate yields strictly more accepted tokens than the fallback Branch 0, quantified by the net token gain:

A lightweight scorer ranks repair candidates immediately after block construction using draft confidence margins, block positions, and candidate statistics. We evaluate the ranking quality under an admitted-block budget $r$ across three metrics: Precision (fraction of selected blocks with $\Delta a>0$), Recall (fraction of useful blocks covered), and Extra-Token Recall (fraction of recovered $\Delta a$ covered).

High Predictability of Repair Utility. As reported in Table 5, the offline scorer provides a strong ranking signal for repair effectiveness. Prioritizing the top 10% of blocks ($r=0.10$) captures $39.6\%$ of all extra accepted tokens with a precision of $35.9\%$. Expanding the budget to 30% ($r=0.30$) recovers $71.6\%$ of the total extra tokens while maintaining a precision of $24.6\%$. This non-linear coverage confirms that useful error repairs are highly concentrated and predictable based on early draft features.

System Trade-offs and Runtime Implications. Because the scorer operates post-candidate construction, it does not bypass the initial candidate generation overhead. In our runtime smoke tests, incorporating a dynamic learned gate yielded negligible end-to-end TPOT improvements due to feature extraction and inference latency overheads. Consequently, we treat Table 5 not as a runtime acceleration claim, but as an insightful offline diagnostic. It demonstrates that a small fraction of candidate blocks drives the majority of acceptance gains, laying the foundation for future zero-overhead, static pattern-pruning strategies.

Table 1 reports the latency (TPOT) and average accepted sequence length across code generation and reasoning benchmarks.

HumanEval & Code Synthesis. On HumanEval, CURE achieves a TPOT of $8.718$ ms/token, delivering a $4.84\times$ speedup over Naive SD ($42.203$ ms/token). Compared with the vanilla parallel draft, our repair tree elevates the accepted length per step from $7.165$ to $7.641$ (+6.6%). While tree verification adds runtime compute, local repair recovers additional candidates at a measurable cost relative to the parallel backend.

Trade-off on General/Lighter Workloads. On MBPP and LCB-lite, CURE maintains consistent acceptance gains (e.g., boosting accepted length by 7.5% on LCB-lite). On MBPP, raw TPOT reaches $11.378$ ms/token versus $49.853$ ms/token in Naive SD. We observe that on simpler token distributions, the extra verification latency of candidate trees can marginally exceed the wall-clock savings from higher accepted lengths. This reveals an inherent latency-acceptance trade-off: tree-based repair is most beneficial in high-entropy or error-prone sub-sequences where target-draft misalignment is severe.

Generalization to Non-Code Reasoning (GSM8K). To evaluate cross-domain behavior under the legacy verification protocol, we evaluate on GSM8K. CURE achieves $11.231$ accepted tokens per step, compared with $8.514$ for the parallel draft and $9.496$ for Naive SD, yielding a $3.634\times$ speedup over target-only AR. This stress test suggests that local repair can recover rejected tokens outside code, but it does not establish protocol-level generality.