DSpark_implementation_notes.md

# DSpark Implementation Notes

## Architecture Summary

DSpark introduces two key innovations over existing speculative decoding methods:

### 1. Semi-Autoregressive Generation
- **Parallel backbone** (DFlash-style): Processes all γ positions in a single forward pass, making drafting latency nearly independent of block size.
- **Sequential head**: Two variants:
  - *Markov head* (`VanillaMarkov`/`GatedMarkovHead`): Low-rank transition bias `B = W1[x_{k-1}] @ W2` (r=256).
  - *RNN head* (`RNNHead`): GRU-like recurrent state accumulating full prefix history.
- The sequential stage fixes the "multi-modal collision" problem where parallel drafters produce incoherent combinations like "of problem" instead of "of course".

### 2. Confidence-Scheduled Verification
- **Confidence head**: Predicts per-position conditional acceptance probability `c_k = σ(w^T[h_k; W1[x_{k-1}]])`.
- **Sequential Temperature Scaling (STS)**: Calibrates cumulative prefix survival probabilities left-to-right via 1D grid search minimizing ECE.
- **Hardware-Aware Prefix Scheduler (Algorithm 1)**: Maximizes `Θ = τ · SPS(B)` by greedily admitting highest-survival-probability tokens with early stopping.

## Algorithm Details

### Acceptance Rule (Lossless)
```
P(accept token x_k) = min(1, p^t_k(x_k) / p^d_k(x_k))
```
- First rejection at position k discards tokens k+1..γ
- One bonus token sampled from residual distribution at rejection position
- Preserves exact target distribution

### Draft Distribution
```
p_k(v | x_0, x_<k) = exp(U_k(v) + B_k(x_0, x_<k, v)) / Σ exp(U_k(u) + B_k(...))
```
where `U_k` are base logits from parallel backbone and `B_k` is the sequential transition bias.

### Training Objective (Eq. 12)
```
L = 0.1 · L_ce + 0.9 · L_tv + 1.0 · L_conf
```
- `L_ce`: Cross-entropy for next-token prediction
- `L_tv`: Total variation distance `||p_d - p_t||_1` (proxy for acceptance rate)
- `L_conf`: Binary cross-entropy on confidence predictions
- All position-weighted by `w_k = exp(-(k-1)/γ)`

### STS Calibration
For each position k = 1..γ:
- Compute cumulative product of calibrated scores up to k
- Find temperature T_k minimizing ECE of cumulative product
- Apply T_k to k-th position score
- Order-preserving: doesn't disrupt relative rankings

### Algorithm 1: Hardware-Aware Prefix Scheduler
```
1. Compute a_{r,j} = ∏_{i≤j} c_{r,i} for each request r, position j
2. Create candidate set E = {(r,j) | a_{r,j} > 0}
3. Sort E descending by a_{r,j}
4. Greedily add candidates:
   - Update ℓ_r = j, B += 1, τ += a_{r,j}
   - Compute Θ = τ * SPS(B)
   - If Θ > best, update best lengths; else break
5. Return per-request verification lengths
```

## What's Implemented Now

| Component | Status | Location |
|---|---|---|
| Markov head (vanilla, gated) | ✅ | `models/markov_head.py` |
| RNN head | ✅ | `models/rnn_head.py` |
| Confidence head | ✅ | `models/confidence_head.py` |
| Acceptance rule | ✅ | `decoding/acceptance.py` |
| Hardware-aware scheduler | ✅ | `decoding/scheduler.py` |
| Static scheduler (fallback) | ✅ | `decoding/scheduler.py` |
| Speculative decoding loop | ✅ | `decoding/speculative.py` |
| Training dataset | ✅ | `training/dataset.py` |
| Loss functions (CE + TV + Conf) | ✅ | `training/losses.py` |
| Training loop | ✅ | `training/train_drafter.py` |
| Target cache generation | ✅ | `training/cache_targets.py` |
| STS calibration | ✅ | `calibration/sts.py` |
| Acceptance evaluation | ✅ | `evaluation/eval_acceptance.py` |
| Latency benchmarking | ✅ | `evaluation/benchmark_latency.py` |
| Unit tests (80) | ✅ | `tests/` |
| Smoke train script | ✅ | `scripts/smoke_train.py` |
| Smoke eval script | ✅ | `scripts/smoke_eval.py` |
| Benchmark script | ✅ | `scripts/run_benchmark.py` |

## Deviations from Paper

1. **Parallel backbone**: Uses `nn.TransformerEncoder` instead of the full DFlash-style backbone with target model KV injection.
2. **Training efficiency**: No multi-GPU training, no 38TB target cache. Target logits computed on-the-fly.
3. **SPS profiling**: Uses a synthetic throughput profile instead of real engine profiling.
4. **Mask token**: Uses token 0 as placeholder; the paper uses learned mask embeddings.
5. **Anchor modification**: Paper treats anchor as first prediction position (γ inputs → γ outputs). We follow the original DFlash style (γ inputs → γ outputs) where anchor is part of the input.

## Future Work

1. Implement full DFlash backbone with target model KV injection
2. Multi-GPU distributed training (DeepSpeed/FSDP)
3. Real engine throughput profiling
4. Tree-based verification for autoregressive drafters
5. Integration with vLLM or similar serving frameworks
6. Qwen3-specific model classes with proper configuration