| 1 | # LiLiCorr |
| 2 | |
| 3 | LiLiCorr uses a DFlash backbone to produce per-position candidates, then scores |
| 4 | candidate combinations with a learned correlator. It requires a compatible |
| 5 | trained checkpoint and Model Runner V2, which is selected automatically. |
| 6 | Both plain and grouped-convolution checkpoints are supported. |
| 7 | |
| 8 | ## Usage |
| 9 | |
| 10 | For a checkpoint trained with block size 16: |
| 11 | |
| 12 | ```bash |
| 13 | vllm serve /path/to/target --dtype bfloat16 \ |
| 14 | --speculative-config '{"method":"dflash","model":"/path/to/lilicorr","num_speculative_tokens":15,"draft_sample_method":"probabilistic"}' |
| 15 | ``` |
| 16 | |
| 17 | We strongly recommend explicitly setting `num_speculative_tokens` to the |
| 18 | checkpoint's trained `block_size - 1` (for example, `15` for `block_size=16`). |
| 19 | The trained block includes one anchor position. Other supported lengths are |
| 20 | available, but may degrade acceptance length unpredictably; benchmark them with |
| 21 | your checkpoint and workload before using them. |
| 22 | |
| 23 | Shorter drafts, down to one token, use the corresponding prefix of the learned |
| 24 | slot embeddings and attention bias without changing checkpoint parameter shapes. |
| 25 | Lengths above `block_size - 1` are unsupported because the checkpoint has no |
| 26 | learned slot embeddings for those positions. |
| 27 | |
| 28 | Fixed-length drafting with shorter verification is a separate option provided by |
| 29 | adaptive verification (`enable_adaptive_verification`). LiLiCorr shares the |
| 30 | candidate-sampling and acceptance-estimation path, but this combination has not |
| 31 | been validated end-to-end. Adaptive verification also requires compatible target |
| 32 | attention backends and variable-length CUDA graph support; do not assume it works |
| 33 | with Mamba2 targets. |
| 34 | |
| 35 | ## Proposal sampling |
| 36 | |
| 37 | - `"draft_sample_method":"greedy"` selects the highest-scoring candidate at each |
| 38 | step of the correlator's conditional walk. |
| 39 | - `"draft_sample_method":"probabilistic"` samples from those conditional scores |
| 40 | using the request's temperature and supplies the realized proposal |
| 41 | distributions to the rejection sampler. |
| 42 | |
| 43 | These settings control proposals; they do not replace the target model's |
| 44 | sampling settings. With request temperature zero, probabilistic proposals also |
| 45 | use argmax. For nonzero temperatures, benchmark both modes with your checkpoint |
| 46 | and workload: compare acceptance length and end-to-end throughput or latency. |
| 47 | Neither mode is universally faster. |
| 48 | |
| 49 | The default rejection method is `"rejection_sample_method":"standard"`. |
| 50 | The shared Model Runner V2 rejection sampler also exposes `"block"`, but |
| 51 | LiLiCorr-specific GPU correctness and performance validation for that combination |
| 52 | is still pending. Do not assume that changing the rejection method improves |
| 53 | throughput; evaluate it separately from the proposal sampling mode. |
| 54 | |
| 55 | ## Checkpoint requirements |
| 56 | |
| 57 | The checkpoint must declare `LiLiCorrDraftModel` and include all `lilicorr_*` |
| 58 | geometry fields and trained head weights. Geometry is read from `dflash_config`. |
| 59 | Convolution tensors must match the configured `conv_kernel_size` and |
| 60 | `conv_group_size`. Target input embeddings must be available on the draft rank. |
| 61 | By default, candidates use the target LM head, which must also be available there. |
| 62 | A checkpoint with top-level `"has_own_lm_head": true` instead uses its own |
| 63 | `lm_head` with the draft quantization configuration and exclusions. Its weights |
| 64 | must be present; for ModelOpt NVFP4 this includes `lm_head.weight`, |
| 65 | `lm_head.weight_scale`, and `lm_head.weight_scale_2`. The owned head is preserved |
| 66 | and is not replaced by the target head. Candidate token embeddings still come |
| 67 | from the target. |
| 68 | |
| 69 | Correlator linear layers and convolution kernel projections use the draft |
| 70 | quantization configuration and its module exclusions. The following parameters |
| 71 | must remain in floating-point model dtype: |
| 72 | |
| 73 | - Lattice QKV `in_proj_weight` and `in_proj_bias`, which retain the exported |
| 74 | parameter layout. |
| 75 | - `lilicorr.factor_input_proj`, `lilicorr.out_head`, and `lilicorr.in_head`, which |
| 76 | are used by the split factor and fused edge projections. Exclude these modules |
| 77 | when exporting a quantized checkpoint. |