Skip to content

Model: Skip standalone load for drafters that load with their target - #515

Merged
turboderp merged 1 commit into
theroyallab:mainfrom
s1arsky:fix/assistant-drafter-load-with-target
Nov 3, 2026
Merged

turboderp merged 1 commit into
theroyallab:mainfrom
s1arsky:fix/assistant-drafter-load-with-target

Conversation

@s1arsky

@s1arsky s1arsky commented Nov 3, 2026

Copy link
Copy Markdown
Contributor

Note: this PR is AI-generated. My local AI agent wrote the commit and ran the tests below directly from a live debug session on my machine, under my supervision — every code reference and claim was verified by the agent against live runs and the cited commits during writing, including the negative control and the same-stack before/after measurement. I stand behind it. Follow-ups handled through me.

Closes #514 — root cause, minimal repro and the full environment write-up are in the issue.

Change

Gate the standalone drafter load on exllamav3's load_with_target cap and let the existing draft_model= hand-off to the target's load_gen do the work, mirroring exllamav3's own model_init.py:378-381 wiring. Seven lines in backends/exllamav3/model.py.

The draft_model is not None narrowing keeps the new condition from adding a fresh reportOptionalMemberAccess next to the preexisting one on the standalone line.

Tested (live; single RTX 3090; tabbyAPI @ 01f939c + this commit; exllamav3 1.6.6 release wheel +cu132.torch2.11.0, torch 2.11.0+cu130, Python 3.14.8)

  • Negative control: without this commit the same config dies at exllamav3/model/model.py:603 — the exact Gemma4 assistant drafter (google/gemma-4-*-it-assistant) can't load with exllamav3 1.6.6: "loads with its target model: pass it as draft_model to the target's load()" #514 traceback.
  • With it: Gemma4-31B exl3 4.00bpw-h6 target + google/gemma-4-31b-it-assistant drafter (draft_mode: model, draft_num_tokens: 3, cache 49,152) loads in 6.0–6.2 s (2.5 s warm relaunch); the guard logs the skip; full RP sessions run normally: 44.2 / 47.1 T/s decode (draft 175/567 = 31%, 286/828 = 35% at 17.8k–18.8k prompts), 930/938 T/s prefill (cold, none cached), a 28-token chat completion at 56.3 T/s (draft 17/33 = 52%), clean Model unloaded. on shutdown.
  • Same stack one hour earlier on exllamav3 1.6.5 (where the drafter loads standalone): 43.2/45.7 T/s, 927/926 T/s — target-forward rate 23.0/23.2 per s (1.6.6 + this guard) vs 23.5/22.7 per s (1.6.5): flat, no serving regression.
  • MTP / DFlash drafters: the elif path is byte-identical to before; grep -rl load_with_target across exl3 1.6.6 names architecture/gemma4_assistant.py as the only cap carrier.
  • CI gates simulated locally at upstream's pinned versions: ruff 0.11.10 check + format clean; the docker-smoke import line (import main; from common import model, sampling; from endpoints.OAI import router) imports fine against this tree.

Notes

  • Tensor-parallel targets keep exl3's own pre-allocation NotImplementedError from the target load_gen — this guard neither hides nor triggers it.
  • TabbyAPI's draft-cache allocation for this drafter matches exl3's own model_init (which builds a draft Cache for every drafter regardless of the cap) — no behavior divergence touched here.

exllamav3 1.6.6 (commit 07015a36) gave Gemma4AssistantForCausalLM the
capability `load_with_target`: its modules are loaded by the target's
load() on the devices of the K/V cache layers it reads, and a standalone
load now raises a RuntimeError. Gate the standalone draft load on that
cap and let the existing `draft_model=` hand-off to the target's
load_gen do the work, mirroring exllamav3's own model_init wiring.

Drafters without the cap (MTP component, DFlash) keep their standalone
load, and on exllamav3 < 1.6.6 no drafter carries the cap, so behavior
there is unchanged. Tensor-parallel targets are unaffected either way:
the drafter load path in the target's load_gen already raises a
NotImplementedError for them before any allocation.

Fixes theroyallab#514
@turboderp

Copy link
Copy Markdown
Collaborator

Thanks. 👍

@turboderp
turboderp merged commit 52c9e0c into theroyallab:main Nov 3, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

2 participants