naveenrajk/Qwen3.5-4B-DFlash-W8A16

Qwen3.5-4B-DFlash · W8A16 (int8 weight-only)

A weight-only int8 version of z-lab/Qwen3.5-4B-DFlash, the block-diffusion DFlash drafter for Qwen3.5-4B. Architecture, vocab, and behavior are unchanged from the upstream checkpoint; only the Linear layer weights are quantized to int8.

Paper: DFlash: Block Diffusion for Flash Speculative Decoding · Upstream model: z-lab/Qwen3.5-4B-DFlash · Upstream code: z-lab/dflash

What changed vs upstream

	upstream bf16	this checkpoint
Linear weights	bfloat16	int8 (W8A16, RTN, symmetric, group=128)
RMSNorm / scales	bfloat16	bfloat16 (unchanged)
fc / lm_head	inherited	fc quantized, lm_head tied to verifier (unchanged)
Format	dense safetensors	compressed-tensors pack-quantized
On-disk size	1.00 GB	521 MB
block_size / arch / vocab	16 / 5-layer Qwen3 / 248320	identical

Quantization recipe (also saved to recipe.yaml):

DEFAULT_stage:
  DEFAULT_modifiers:
    QuantizationModifier:
      targets: [Linear]
      ignore: [lm_head]
      scheme: W8A16

Produced via llm-compressor 0.11.0 (QuantizationModifier oneshot, no calibration data needed — pure RTN on the weights).

Validation: this checkpoint vs upstream bf16

Tested as the DFlash drafter for a paroquant w4a16-quantized Qwen3.5-4B target in vLLM 0.22.1 (num_speculative_tokens=15, greedy, 20 prompts / 60% no-think ultrachat-style + 40% gsm8k-style; total ~425 draft steps each).

	upstream bf16	W8A16 (this)	Δ
Mean accept length	3.60 tok / step	3.56 tok / step	-0.04
Per-token accept rate	24.03%	23.74%	-0.29 pp
End-to-end throughput	1.41 req/s	1.39 req/s	-1.4%

Per-position acceptance (positions 0–14 of the 15 spec tokens):

pos	bf16	W8A16	Δ
0	86.8%	87.1%	+0.3
1	67.2%	67.3%	+0.1
2	47.6%	48.1%	+0.5
3	32.1%	32.5%	+0.4
4	25.2%	25.7%	+0.5
5	22.2%	22.4%	+0.2
6	16.5%	16.8%	+0.3
7	11.8%	12.1%	+0.3
8	10.8%	11.2%	+0.4
9	9.9%	10.0%	+0.1
10	9.0%	9.1%	+0.1
11	7.5%	7.7%	+0.2
12	7.1%	7.0%	-0.1
13	6.1%	6.1%	0
14	5.7%	5.6%	-0.1

All differences are within sampling noise for 425 trials. Treat the int8 drafter as lossless for practical purposes.

The chain-throughput gain on this hardware (RTX 3090) was negligible because the drafter is only 30% of the spec-decode chain wall-clock — the target verifier dominates. The win here is memory (480 MB of GPU returned), which matters mostly when stacking multiple speculators on the same device or pushing longer max_model_len / larger batch.

Usage

vLLM (recommended)

The drafter pairs with any Qwen3.5-4B (or paroquant'd Qwen3.5-4B) target.

pip install "vllm>=0.20"

vllm serve Qwen/Qwen3.5-4B \
  --dtype bfloat16 \
  --speculative-config '{"method":"dflash","model":"naveenrajk/Qwen3.5-4B-DFlash-W8A16","num_speculative_tokens":15}'

For the paroquant w4a16 variant of Qwen3.5-4B used in our measurements:

pip install "vllm>=0.20" paroquant

VLLM_PLUGINS=paroquant vllm serve <paroquant-qwen3.5-4b> \
  --quantization paroquant --dtype float16 --trust-remote-code \
  --speculative-config '{"method":"dflash","model":"naveenrajk/Qwen3.5-4B-DFlash-W8A16","num_speculative_tokens":15}'

Required vLLM patches (until upstream lands them)

vLLM 0.22.1's qwen3_dflash.py reads qkv_proj.weight / fc.weight directly via F.linear. This bypasses the quant-aware Linear.forward and trips on any quantized drafter (compressed-tensors, bnb, paroquant, …). The patch covers four sites and keeps the fused-KV fast path active even for quantized weights — no per-layer fallback at runtime:

1. Defer _build_fused_kv_buffers() out of load_weights. The eager call runs before process_weights_after_loading and so before Marlin / etc. have set up their internals (e.g. g_idx_sort_indices). Move the build to the lazy path in precompute_and_store_context_kv, which fires after warmup.
2. _build_fused_kv_buffers — detect quantized qkv_proj. For the FP case, keep the original weight[q_size:] slice + cat. For the quantized case, recover the effective KV-projection weight by running an identity matrix through each layer's qkv_proj(...), slicing the K+V output rows, and concatenating into one bf16 fused tensor. This works for any quant scheme without needing knowledge of its internal storage. ~50 MB extra runtime memory (K+V slice across 5 layers).
3. DFlashQwen3Attention.forward — call self.qkv_proj(hidden_states) when .weight isn't a plain FP tensor, instead of F.linear(qkv_proj.weight, ...). This is the regular draft-pass attention; uses Marlin int8 GEMM directly.
4. combine_hidden_states — read fc.weight_scale.dtype instead of fc.weight.dtype to pick the activation cast dtype when fc is quantized.

The first time you load this checkpoint with the patch in place, vLLM logs:

DFlash drafter: qkv_proj is quantized; extracting fused KV weight via identity
probe (dtype=torch.float16, hidden=2560). Restores the fast fused-GEMM path in
context KV precompute.

A reference implementation of the three-site patch lives in the discussion thread on this repo, and we plan to submit it to vLLM upstream.

Reproducing the quantization

import torch
from transformers import AutoModel
from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier

model = AutoModel.from_pretrained(
    "z-lab/Qwen3.5-4B-DFlash", trust_remote_code=True, dtype=torch.bfloat16
)
oneshot(
    model=model,
    recipe=QuantizationModifier(
        targets=["Linear"], ignore=["lm_head"], scheme="W8A16"
    ),
    output_dir="Qwen3.5-4B-DFlash-W8A16",
    save_compressed=True,
)

Requires transformers==4.57.x for the llmcompressor pass; restore your normal transformers after.

Credits

• All training and architectural work belongs to Z-Lab — see the DFlash paper and the original z-lab/Qwen3.5-4B-DFlash checkpoint.
• This repo is a strict weight-only quantization derivative; no retraining, no data, no architectural changes.

License

MIT, inherited from the upstream z-lab/Qwen3.5-4B-DFlash checkpoint.