makora-ai/gemma4-26b-a4b-dspark

Gemma-4 26B-A4B DSpark

DSpark speculative-decoding draft head for google/gemma-4-26B-A4B-it.

This is not a standalone model — it is a ~1.2B-parameter draft head (speculators format, DSparkDraftModel) that proposes tokens for the Gemma-4 26B-A4B target. It is loaded on top of the base model as a vLLM speculative-config and accelerates greedy/low-temperature decoding without changing the target's output distribution (lossless rejection sampling).

DSpark extends a DFlash-style multi-layer draft with two extra heads:

• a Markov logit-bias head (markov_rank=256) so each of the 7 block positions can condition on the tokens already sampled within the block, and
• a confidence head that predicts per-position acceptance probability.

The draft is a 5-layer transformer that reads the target's intermediate hidden states from layers [3, 10, 18, 25, 28] (vLLM auto-appends the final layer). It was trained on-policy — prompts from a diverse 37-dataset corpus, with responses regenerated by the Gemma-4 target itself — so the acceptance numbers below are honest, serve-time (autoregressive) measurements rather than teacher-forced.

Target (verifier)	google/gemma-4-26B-A4B-it
Draft params	~1.2B (bf16)
Draft layers	5 (block_size=7)
Target layer IDs	3, 10, 18, 25, 28
Draft vocab	32,000 (mapped to the target's 262,144)
Extra heads	Markov (rank 256, vanilla) + confidence
Default proposal	greedy, num_speculative_tokens=6

Benchmarks

Served, autoregressive mean acceptance length and end-to-end throughput as the on-policy training corpus is scaled up. All served numbers use the pinned nightly vLLM DSpark path on a single B200, Gemma sampling (temperature 1.0, top-p 0.95, top-k 64), thinking enabled. MATH-500 is the fixed deployable benchmark (OSL 1500) at concurrency C=1 and C=4; makora greedy-400 is an in-domain accept-length probe. A representative subset of training-data scales is shown.

train prompts	offline val accept-len @ckpt	served in-domain accept-len	served MATH-500 accept-len (C1 / C4)	MATH-500 tok/s (C1 / C4 agg)
50,000	3.684	2.825	1.60 / 1.60	231 / 730
100,000	3.964	2.366	1.86 / 1.90	268 / 886
200,000	4.098	2.857	3.11 / 3.14	442 / 1304
400,000	4.079	3.157	3.98 / 4.08	568 / 1875
600,000 (this checkpoint)	4.237	3.101	3.97 / 4.10	566 / 1866

Notes

• Speedup vs. the unaccelerated target. Vanilla Gemma-4 26B-A4B on the same rig runs MATH-500 at 233 tok/s (C1) and 721 tok/s (C4 aggregate). This checkpoint reaches 566 / 1866 tok/s — roughly +143% (C1) and +159% (C4) — while producing token-for-token identical output. Served MATH-500 acceptance saturates at ~4.0 from the 400k scale onward.
• Served, not teacher-forced. The acceptance lengths are measured autoregressively at serve time (vllm:spec_decode_num_accepted_tokens... /metrics deltas), i.e. the number you actually get in deployment. Mean accept-length ≈ 1 + accepted/drafts.
• Domain. The draft is tuned toward structured reasoning (math/coding), where served MATH-500 accept-length reaches ~4.0. Open-ended chat acceptance is lower.
• Offline validation accept-length is reported at the saved best-by-loss checkpoint (not the peak-accept epoch), so served numbers are, if anything, conservative.

Usage

The head is consumed by vLLM (a recent nightly with the dspark speculative method), not by transformers directly — it has no lm_head/embeddings of its own and only produces drafts for the target to verify. Load the base model normally and point vLLM's speculative config at this repo; vLLM reads the DSpark config (algorithm, num_speculative_tokens, and the target layer IDs to extract) from config.json automatically.

Recommended environment (Blackwell/Hopper): VLLM_USE_FLASHINFER_SAMPLER=0 (the FlashInfer sampler JIT needs nvcc on PATH; disabling it selects the precompiled FLASH_ATTN + Triton MoE path).

Serve (OpenAI-compatible)

export HF_TOKEN=hf_...            # both repos may be gated
export VLLM_USE_FLASHINFER_SAMPLER=0

vllm serve google/gemma-4-26B-A4B-it \
  --speculative-config '{"model": "makora-ai/gemma4-26b-a4b-dspark", "num_speculative_tokens": 6}' \
  --max-model-len 8192 \
  --port 8000

curl http://127.0.0.1:8000/v1/chat/completions \
  -H 'Content-Type: application/json' \
  -d '{
    "model": "google/gemma-4-26B-A4B-it",
    "messages": [{"role": "user", "content": "What is 17*23? Think step by step."}],
    "temperature": 1.0, "top_p": 0.95, "top_k": 64,
    "max_tokens": 512
  }'

Speculative decoding is transparent to the API: responses are identical to the target model's, only faster. Inspect the acceleration via curl http://127.0.0.1:8000/metrics | grep spec_decode.

Offline (Python)

from vllm import LLM, SamplingParams

# The DSpark head is attached to the base target as a speculative draft.
llm = LLM(
    model="google/gemma-4-26B-A4B-it",
    speculative_config={
        "model": "makora-ai/gemma4-26b-a4b-dspark",
        "num_speculative_tokens": 6,
    },
    max_model_len=8192,
)

messages = [{"role": "user", "content": "What is 17*23? Think step by step."}]
sampling = SamplingParams(temperature=1.0, top_p=0.95, top_k=64, max_tokens=512)

out = llm.chat(messages, sampling)
print(out[0].outputs[0].text)

Inspecting the head directly

The draft weights and config load with the speculators library (>= 0.7.0.dev69), e.g. for verification or re-serving:

from speculators import SpeculatorModel

draft = SpeculatorModel.from_pretrained("makora-ai/gemma4-26b-a4b-dspark")
print(draft.config)  # DSparkSpeculatorConfig: block_size=7, markov_rank=256, ...