README.md

---
tags:
- sglang
- speculative-decoding
- dflash
- mirai
library_name: sglang
license: apache-2.0
pipeline_tag: other
---

<div style="display:flex;align-items:center;justify-content:space-between;gap:20px;flex-wrap:wrap">
  <div style="flex:1;min-width:260px">
    <h1 style="margin:0 0 12px">Trees from Marginals: Weaver</h1>
    <p style="margin:0">Weaver is a lightweight autoregressive Transformer that turns DFlash marginal predictions into proposal trees. With rollback-free Gated Delta Net tree verification, DFlash-TfM reaches 392.8 tokens/s/sequence on Qwen3.6-27B on a single B200: 4.37x over autoregressive decoding and 24.7% over tuned DFlash.</p>
    <p style="margin:10px 0 0;display:flex;gap:8px;align-items:center;flex-wrap:wrap">
      <a href="https://arxiv.org/abs/2607.06763"><img src="https://img.shields.io/badge/arXiv-2607.06763-b31b1b?style=for-the-badge&logo=arxiv&logoColor=white" alt="arXiv 2607.06763"></a>
      <a href="https://github.com/trymirai/sglang"><img src="https://img.shields.io/badge/GitHub-code-181717?style=for-the-badge&logo=github&logoColor=white" alt="GitHub code"></a>
    </p>
  </div>
  <img src="https://assets.trymirai.com/images/logo/ml_small_logo.svg" alt="Mirai Labs" width="64" height="64" style="width:80px;height:80px;object-fit:contain;flex:0 0 auto">
</div>

<p align="center">
  <img src="assets/tfm-throughput.png" alt="Per-dataset Qwen3.6-27B throughput comparison for autoregressive decoding, DFlash, DDTree, and DFlash-TfM." width="760">
</p>

### Method

A factorized drafter emits marginal distributions for several future positions in a single forward pass. Because it ignores conditional dependencies between positions, the acceptance rate falls as draft length grows. DFlash-TfM addresses this by extracting the top-K marginal predictions from DFlash and using Weaver, a lightweight autoregressive draft model, to restore dependencies conditioned on the previously selected draft tokens.

Weaver has 56.7M trainable parameters and operates exclusively over the DFlash candidate vocabulary, a tiny subset of the full vocabulary. It shares the verifier embedding and output projection matrices, but never projects to the full vocabulary: at each tree-expansion step, Weaver predicts residual logits over DFlash's top-512 candidates, adds them to the DFlash marginal logits, and normalizes over the candidate set.

<p align="center">
  <img src="assets/tfm-architecture.png" alt="DFlash-TfM uses DFlash marginals in parallel, then conditions tree proposals autoregressively with Weaver." width="760">
</p>

This repository contains the Weaver checkpoint for DFlash-TfM, short for DFlash - Tree From Marginals. Weaver is not a standalone language model: it is used during speculative decoding together with a Qwen3.6-27B target model, a Qwen3.6-27B DFlash drafter, and the Mirai Labs SGLang fork with the fused Gated Delta Net tree-verification kernels.

### Results

We evaluate on Qwen3.6-27B over chat, math, and code workloads: MTBench, ShareChat, GSM8K, MATH500, AIME25, HumanEval, MBPP, and LiveCodeBench. All runs use BF16 precision on a single B200 with batch size 1, temperature 1.0, reasoning enabled, maximum output length 4096, and the server cache flushed between requests.

Throughput is computed as total generated tokens divided by wall-clock runtime, including prefill, scheduling, and decoding. Speedup is measured against autoregressive decoding under the same dataset, temperature, and reasoning setting. Macro Avg. is the unweighted average across datasets.

| Method | Setting | Throughput | Speedup |
| --- | --- | ---: | ---: |
| Autoregressive | BF16 target only | 89.9 tok/s/seq | 1.00x |
| DFlash | tuned chain baseline | 315.0 tok/s/seq | 3.50x |
| DFlash-TfM + Weaver | tree budget 64 | 392.8 tok/s/seq | 4.37x |

<p align="center">
  <img src="assets/tfm-results-table.png" alt="Full DFlash-TfM table with speedup and accepted-token statistics across sampling and reasoning settings." width="900">
</p>

DFlash-TfM with Weaver is the fastest configuration on every task in this sweep. The gap comes from the superior acceptance probability: Weaver's trees lengthen the mean accepted draft by 77% relative to the chain DFlash baseline and by 32% relative to DDTree at the same tree size.

### Using this checkpoint

This repository contains the Weaver checkpoint used for DFlash-TfM:
`weaver/qwen36_27b_weaver.pth`

SHA-256: `71f540b143fb6bab14ba724c20e97a72ce198de103cfd228d31c3ce339227833`

Quick reproduction scripts are available in the Mirai Labs SGLang fork:

```bash
git clone https://github.com/trymirai/sglang
cd sglang
./reproduction.sh serve-tfm
```

In a second terminal, once the server is ready:

```bash
./reproduction.sh bench
```

### Citation

If you find our work helpful, feel free to give us a cite.

```bibtex
@misc{dflash-tfm,
    title  = {{Trees from Marginals}: Autoregressive drafting with factorized priors},
    author = {Yuma Oda and Ryan Mathieu and Roman Knyazhitskiy and Artur Chakhvadze},
    note   = {In collaboration with others at Mirai Labs},
    month  = {July},
    year   = {2026},
    eprint = {2607.06763},
    archivePrefix = {arXiv},
    primaryClass  = {cs.LG},
    url    = {https://arxiv.org/abs/2607.06763}
}
```