ITADN

[RFC]: LoRA adapter support for DFlash speculative decoding draft models

#52038Openanshkaggarwal22 创建于 6 天前
RFC
A
anshkaggarwal22commented
### Motivation. vLLM supports serving many LoRA adapters over a single base target model, but speculative decoding still requires a separate, fully-trained draft model per domain. When serving many domains, this is expensive: each domain needs its own drafter (e.g. 0.8B) held in memory, with corresponding I/O cost to load and swap. This RFC proposes applying LoRA adapters to the DFlash drafter, so a single base drafter can be specialized per-domain via small adapters instead of full drafters. In our experiments an R64 adapter is ~28x smaller than the 0.8B drafter it replaces, substantially reducing memory footprint and I/O latency when serving many domains, while keeping drafter quality within ~2% of a fully-trained per-domain drafter (measured by acceptance length behavior). This mirrors the efficiency win vLLM already provides for target-model LoRA, extended to the speculative-decoding drafter. ### Proposed Change. Allow a LoRA adapter to be applied to the DFlash drafter, independent of the target model: - The adapter wraps all drafter layers except `lm_head` and `embed_tokens`, loaded from a stage-2-trained LoRA checkpoint. - The drafter LoRA is specified via additional arguments in the speculative config, distinct from any target-model adapter. - The correct adapter is routed to the drafter and applied along its forward path. Initial support targets DFlash drafters specifically (tested with a GPT-OSS-120B target and a 0.8B DFlash drafter). This is distinct from target-side multi-adapter work (#49705, #48297) and from prior Eagle-style draft-model attempts. **Prior art / known issues** Closed PR #11966 (V0, Eagle-era) attempted LoRA + spec decode and hit a vocab-size mismatch (e.g. 768 vs 4096) from differing vocab padding on draft vs. target, and disabled draft-side LoRA "until the API interface is changed to inject the corresponding LoRA adapter for the draft model." In our DFlash testing we did not encounter this mismatch — likely because the adapter excludes `lm_head` and `embed_tokens`, avoiding the padded-vocab layers. We'd welcome maintainer input on whether this holds across all DFlash configurations. **Alternatives considered** - Full per-domain drafters (status quo): simple but costly — one full drafter per domain in memory. - Offline-merging the LoRA into drafter weights: avoids runtime adapter logic but still stores a full merged drafter per domain, forfeiting the 28x saving and per-request swappability. Serving the adapter dynamically preserves both the memory saving and per-request swappability. **Testing & benchmarks** Acceptance behavior and throughput comparing base drafter, LoRA'd drafter, and a fully-trained per-domain drafter; memory-footprint comparison (adapter vs. full drafter) across N domains. ### Feedback Period. One week. ### CC List. @LucasWilkinson @WoosukKwon ### Any Other Things. There are some details I can't share about the experiment results because that involved classified data and domains while my time at Capital One AI Foundations where I worked on this project, but I will share what I can based on your questions. ### Before submitting a new issue... - [x] Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the [documentation page](https://docs.vllm.ai/en/latest/), which can answer lots of frequently asked questions.
0 条评论