Speculative Decoding (DFlash)
Relevant source files
The following files were used as context for generating this wiki page:
Speculative decoding is an optimization technique that uses a small, fast draft model (DFlash) to accelerate inference for a large target model. The draft model generates candidate tokens that the target model verifies in parallel. This approach reduces the total number of expensive target model forward passes, increasing overall throughput.
Sources: oxidize-core/src/model/speculative.rs:1-8, PRD.md:214-215
System Architecture
The speculative decoding system consists of a SpeculativeDecoder that orchestrates interactions between a target Model and a DFlashDraftModel. It maintains generation state, manages KV caches for both models, and handles acceptance logic.
The diagram shows the iterative cycle where the draft model proposes tokens and the target model verifies them.
Sources: oxidize-core/src/model/speculative.rs:11-18, oxidize-core/src/model/speculative.rs:105-132
Core Components
- SpeculativeDecoder: The primary controller that manages the generation loop and session state. Sources: oxidize-core/src/model/speculative.rs:105-117
- DFlashDraftModel: A specialized, lightweight model optimized for fast token proposal. Sources: oxidize-core/src/model/speculative.rs:107, oxidize-core/src/model/AGENTS.md:14
- Session: Tracks token consumption and allows the target model to rewind its KV cache during rejections. Sources: oxidize-core/src/model/speculative.rs:114-116
Generation Workflow
Generation transitions through four distinct states: Prefill, Speculating, Fallback, and Done.
Prefill Stage
The target model processes the initial prompt to initialize its KV cache. The decoder then samples the first token and initializes the draft model context to match the target.
Sources: oxidize-core/src/model/speculative.rs:168-175, oxidize-core/src/model/speculative.rs:197-202
Speculation and Verification
In the speculation phase, the draft model generates K tokens autoregressively. The target model then performs a batched forward pass to verify these tokens. If the acceptance rate falls below a configured threshold, the system switches to Fallback mode, sampling directly from the target model to maintain quality.
This sequence illustrates the parallel verification process and the subsequent synchronization of KV caches.
Sources: oxidize-core/src/model/speculative.rs:205-240, oxidize-core/src/model/speculative.rs:260-278
Configuration and Performance
The SpeculativeConfig struct controls the behavior of the decoder. Developers can choose between predefined conservative or aggressive profiles or use a builder pattern for custom tuning.
Configuration Parameters
| Parameter | Type | Description | Default |
|---|---|---|---|
draft_tokens_per_step | usize | Number of tokens the draft model proposes at once. | 4 |
max_new_tokens | usize | Limit for generated output length. | 128 |
min_acceptance_rate | f32 | Threshold to switch to fallback mode. | 0.3 |
strict_mode | bool | If true, rejects all subsequent tokens on first mismatch. | false |
Sources: oxidize-core/src/model/speculative.rs:32-47, oxidize-core/src/model/speculative.rs:49-58
Performance Metrics
The system tracks SpeculativeStats to monitor efficiency:
- Acceptance Rate: The ratio of accepted draft tokens to total draft tokens generated.
- Estimated Speedup: Calculated as
(accepted + fallback tokens) / target_forward_passes.
Sources: oxidize-core/src/model/speculative.rs:81-99
CLI Usage
The oxidize-cli provides flags to enable and tune speculative decoding during local inference.
--draft-model <PATH>: Specifies the path to the DFlash draft model file.--draft-tokens <N>: Sets the number of draft tokens per step.--force-dflash: Forces DFlash use even if the draft model was trained for a different target hidden size.--no-mtp: Disables native in-GGUF multi-token prediction if present.
Sources: oxidize-cli/src/main.rs:125-135
Implementation Example
let config = SpeculativeConfigBuilder::new()
.draft_tokens_per_step(6)
.min_acceptance_rate(0.4)
.build();
let mut decoder = SpeculativeDecoder::new(target_model, draft_model, config);
let tokens = decoder.generate(prompt_tokens)?;
Sources: oxidize-core/src/model/speculative.rs:356-368
Summary
Speculative decoding via DFlash provides a significant performance boost for local LLM inference by offloading token proposals to a smaller model. By leveraging parallel verification and dynamic fallback mechanisms, the system achieves higher tokens-per-second while maintaining the output quality of the original target model. The integration supports various draft model architectures, including DFlash and EAGLE3, providing flexibility for different hardware profiles.
Sources: oxidize-core/src/model/speculative.rs:1-10, oxidize-cli/src/main.rs:565-580