MoS-DFlash-Qwen3-8B

MoS-DFlash-Qwen3-8B

This repository contains a routed Mixture of Speculators (MoS) checkpoint built on DFlash for speculative decoding with Qwen3-8B as the target model. MoS shares the DFlash attention parameters across domains while retaining a separate MLP path for each domain. A request is assigned to one path before drafting, and that path is used throughout the response.

Important: this is a draft/speculator model, not a standalone language model. It cannot be used with the standard text-generation pipeline by itself. The Qwen3-8B target model is required for generation.

What is included

• One routed DFlash checkpoint with five domain-specific MLP paths (num_domains=5)
• Shared DFlash attention parameters
• Custom Transformers model code in dflash.py
• Two safetensors weight shards

The learned request router, the Qwen3-8B target weights, exported standalone experts, and the optimized serving harness are not included in this repository.

Checkpoint summary

Only this selected checkpoint is published here. Other epochs, optimizer states, baselines, 250K experiments, 4B experiments, serving logs, and evaluation caches are excluded.

Requirements

The checkpoint was validated with:

• Python 3
• PyTorch 2.9.1 with CUDA 12.8
• Transformers 4.57.1
• Hugging Face Hub 0.36.2

Install PyTorch for your CUDA environment first, then install the remaining dependencies:

pip install "transformers==4.57.1" "huggingface_hub>=0.36" accelerate safetensors typing_extensions

Load the routed checkpoint

Because this repository contains custom model code, loading requires trust_remote_code=True:

import torch
from transformers import AutoModel

repo_id = "ryan-0608/MoS-DFlash-Qwen3-8B"

draft = AutoModel.from_pretrained(
    repo_id,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
)
draft.eval()

print(type(draft).__name__)  # DFlashDraftModel
print(draft.num_domains)     # 5
print(draft.block_size)      # 16

For reproducible deployments, pin revision to a specific repository commit when loading remote code.

Select a domain path

This routed checkpoint requires one domain ID per batch element. Set the IDs before every forward or generation call:

DOMAIN_TO_ID = {
    "code": 0,
    "math": 1,
    "factual_qa": 2,
    "creative_writing": 3,
    "general": 4,
}

device = next(draft.parameters()).device
domain_ids = torch.tensor([DOMAIN_TO_ID["code"]], device=device)
draft.set_domain_ids(domain_ids)

For a batch of size B, pass a tensor of shape [B] containing one ID for each sample. set_domain_ids() only sets transient routing state; the IDs are not stored in the checkpoint.

If num_domains > 1 and no domain IDs are set, the model intentionally raises an error instead of silently selecting an expert.

Minimal greedy decoding example

The following example uses the reference spec_generate() implementation included in dflash.py. It is intended to verify model behavior; it is not the optimized serving benchmark used for throughput measurements.

import torch
from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer

target_id = "Qwen/Qwen3-8B"
draft_id = "ryan-0608/MoS-DFlash-Qwen3-8B"
device = "cuda:0"
dtype = torch.bfloat16

tokenizer = AutoTokenizer.from_pretrained(target_id)
target = AutoModelForCausalLM.from_pretrained(
    target_id,
    torch_dtype=dtype,
).to(device).eval()
draft = AutoModel.from_pretrained(
    draft_id,
    trust_remote_code=True,
    torch_dtype=dtype,
).to(device).eval()

# Manual request-level routing for this example: select the code path.
draft.set_domain_ids(torch.tensor([0], device=device))

prompt = "Write a Python function that returns the longest common prefix."
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)

with torch.inference_mode():
    output_ids = draft.spec_generate(
        target=target,
        input_ids=input_ids,
        max_new_tokens=128,
        stop_token_ids=[tokenizer.eos_token_id],
        temperature=0.0,
    )

print(tokenizer.decode(output_ids[0], skip_special_tokens=True))

Both models must fit in accelerator memory. The target and draft should use the same tokenizer and compatible Qwen3-8B configuration.

Routing and production serving

MoS performs routing once per request, before the token-level drafting loop:

1. The router reads the request prompt and predicts a domain ID.
2. The corresponding MLP path is selected.
3. DFlash generates draft blocks using the shared attention parameters and the selected MLP path.
4. Qwen3-8B verifies the proposed tokens as in speculative decoding.

The route remains fixed for the response. The direct checkpoint interface exposes this decision through set_domain_ids() for manual or externally supplied routing.

For optimized serving, the five paths can instead be exported as five standalone DFlash drafters, each combining the shared attention parameters with one domain-specific MLP path. A separate request router then dispatches each request to the corresponding target-plus-drafter service. The router checkpoint and export/serving utilities are part of the MoS experimental stack and are not packaged in this model repository.

Repository files

• config.json: DFlash configuration, including num_domains=5, block size 16, and the target hidden-state layers
• dflash.py: custom DFlashDraftModel implementation and reference speculative-generation loop
• model-00001-of-00002.safetensors
• model-00002-of-00002.safetensors
• model.safetensors.index.json

Intended use

This checkpoint is released for research on multi-domain speculative decoding, request-level speculator routing, and MoS evaluation with Qwen3-8B. It is intended for researchers who already have a compatible Qwen3-8B target and DFlash inference environment.