xinfer-ai 0.14.4 → 0.14.6

package/README.md

@@ -61,11 +61,34 @@ xinfer --m /home/Qwen3.6-35B-A3B --d 0,1 --ui-server
 python3 -m xinfer.server --m Qwen/Qwen3.6-27B-FP8 --kvcache-dtype turbo4 --ui-server
 ```
 
-**MTP**
+**Speculative decoding (MTP & DFlash2)**
+
+Built-in MTP (model must include MTP heads):
+
+```bash
+xinfer --m Qwen/Qwen3.5-35B-A3B --d 0,1 --ui-server --num-speculative-tokens 3
+```
+
+External DFlash2 draft model:
+
 ```bash
-xinfer --w /home/Qwen3.6-35B-A3B --d 0,1 --ui-server --mtp 2
+xinfer --m Qwen/Qwen3.8-... --d 0,1 --ui-server \
+  --draft-model <dflash2-draft-id-or-path> --num-speculative-tokens 7
 ```
 
+See [docs/speculative_decoding.md](docs/speculative_decoding.md) for details.
+
+### 🧩 MoE Expert Offloading
+
+MoE expert offloading is supported on CUDA only. Run a large MoE model with expert weights offloaded to host RAM. The value is a fraction: `0.2` keeps 20% of each layer's experts on the GPU; `auto` selects the cache size automatically. Tensor parallelism requires NCCL.
+
+```bash
+XINFER_MOE_OFFLOAD=0.2 xinfer \
+  --m Qwen/Qwen3.5-35B-A3B-FP8 --d 0,1 --ui-server
+```
+
+See the [MoE expert offloading guide](docs/moe_offload.md) for configuration and TP details.
+
 > **Tip:** Open `http://IP:8001` for the built-in chat UI, or use `http://IP:8000/v1/` as your API `Base URL`.
 
 ---
@@ -82,8 +105,6 @@ Add `--kvcache-dtype` to compress KV cache and extend context length:
 | `turbo4` | **3.7×** | Best balance | SM70+ / Apple M1|
 | `turbo3` | **4.7×** | Max compression | SM70+ |
 
----
-
 ## 📈 Performance
 
 > Tested on **V100-32G**, **A100-40G**, **Hopper-80G** and **RTX 5090**
@@ -134,6 +155,7 @@ Add `--kvcache-dtype` to compress KV cache and extend context length:
 * ✅ Qwen2/Qwen3 MoE
 * ✅ Qwen3 Next
 * ✅ Qwen3.5/3.6/3.8 Dense/MoE (27B, 35B, 122B, 397B, Multimodal model)
+* ✅ Qwen3.8-Next/Qwen4
 * ✅ Mistral v1, v2
 * ✅ Mistral-3-VL Reasoning (3B, 8B, 14B, Multimodal model)
 * ✅ GLM4 (0414)
@@ -529,7 +551,8 @@ xinfer --m Qwen/Qwen3.6-27B-FP8 --ui-server --enable-tool-grammar
 | `--frequency-penalty` | Penalize frequent tokens (−2 to 2) |
 | `--mcp-config` | MCP servers JSON config |
 | `--mcp-command` / `--mcp-args` | Single MCP server command + args |
-| `--mtp`| Multi-token prediction, usage `--mtp 2` for two-token prediction per forward pass |
+| `--num-speculative-tokens` | Speculative draft tokens per decode step. Enables **built-in MTP** on models with MTP heads (e.g. Qwen3.5). Example: `--num-speculative-tokens 3` |
+| `--draft-model` | External **DFlash2** draft model (HuggingFace id or local directory). When set, uses DFlash2 instead of built-in MTP. |
 
 ### Environment Variables
 
@@ -591,7 +614,7 @@ SM90_LOWER_PRECISION_GDN_PREFILL=1 xinfer --m Qwen/Qwen3.5-35B-A3B-FP8 --ui-serv
 * [x] **MXFP4/NVFP4 Model Support**
 * [x] **Support Turboquant (4-bit, 3-bit) KvCache**
 * [ ] TentorRT-LLM
-* [x] **Multi-token Prediciton (MTP)**
+* [x] **Multi-token Prediction (MTP) & DFlash2 speculative decoding**
 
 ---
 

package/docs/get_started.md

@@ -67,10 +67,27 @@ cargo build --release --features metal
   target/release/xinfer --m Qwen/Qwen3-VL-8B-Instruct --ui-server
   ```
 
-Common runtime knobs: `--max-model-len`, `--max-num-seqs`, `--kv-fraction` (CUDA KV share), `--cpu-mem-fold` (CPU swap ratio), `--port`, `--kvcache-dtype` (fp8/turbo8/turbo4/turbo3), `--prefix-cache`, `--prefix-cache-max-tokens`, `--ui-server`, `--batch` (perf test).
+Common runtime knobs: `--max-model-len`, `--max-num-seqs`, `--kv-fraction` (CUDA KV share), `--cpu-mem-fold` (CPU swap ratio), `--port`, `--kvcache-dtype` (fp8/turbo8/turbo4/turbo3), `--prefix-cache`, `--prefix-cache-max-tokens`, `--ui-server`, `--batch` (perf test), `--num-speculative-tokens` (built-in MTP), `--draft-model` (DFlash2).
 
 Reasoning defaults to enabled when a request omits `thinking` / `enable_thinking`. Use `--disable-reasoning` on the Rust CLI to make the default be disabled instead; explicit request values still override the server default.
 
+### Speculative decoding (MTP & DFlash2)
+
+Speed up decode with speculative token drafting + target verification. See [speculative_decoding.md](./speculative_decoding.md) for full details.
+
+- **Built-in MTP** (Qwen3.5 / Qwen3.6 with MTP weights): `--num-speculative-tokens N`
+  ```bash
+  target/release/xinfer --m Qwen/Qwen3.5-35B-A3B --d 0,1 \
+    --num-speculative-tokens 3 --ui-server
+  ```
+- **DFlash2** (external draft checkpoint): `--draft-model <hf_id_or_local_path>` plus optional `--num-speculative-tokens`
+  ```bash
+  target/release/xinfer --m Qwen/Qwen3.8-... --d 0,1 \
+    --draft-model <dflash2-draft> --num-speculative-tokens 7 --ui-server
+  ```
+
+If `--draft-model` is set, DFlash2 is used instead of built-in MTP. Build with `cuda,nccl,flashinfer,cutlass` for best performance.
+
 ## 4) Multi-rank (single node)
 - **NCCL multi-GPU**  
   ```bash

package/docs/moe_offload.md

@@ -0,0 +1,136 @@
+# MoE expert offload
+
+Run Mixture-of-Experts models whose weights do not fit in VRAM. The routed
+experts stay in pinned host RAM; each MoE layer keeps a fraction of its experts
+in GPU *slots*. Attention, dense layers, shared experts and the KV cache stay on
+the GPU as usual.
+
+This feature requires a CUDA build. Tensor parallelism (`--d 0,1`, etc.) also
+requires NCCL; dense models do not enter the offload path and are unaffected
+when the variable is unset.
+
+```bash
+# Keep 20% of every layer's experts on GPU (at least one top-k decode step),
+# and keep the remaining experts in host RAM. TP works rank-by-rank.
+XINFER_MOE_OFFLOAD=0.2 xinfer \
+    --m Qwen/Qwen3.5-35B-A3B-FP8 --d 0,1 --max-model-len 8192 --ui-server
+
+# Or size the per-layer GPU caches automatically from free VRAM.
+XINFER_MOE_OFFLOAD=auto xinfer \
+    --m Qwen/Qwen3.5-35B-A3B-FP8 --d 0,1 --max-model-len 8192 --ui-server
+```
+
+Placement never changes routing or precision: with greedy decoding the output
+is byte-identical for every `XINFER_MOE_OFFLOAD` value.
+
+## How it works
+
+- **Decode** — before each fused MoE GEMM, a CUDA kernel looks the routed experts
+  up in a per-layer LRU table kept on the GPU, assigns least-recently-used slots
+  to misses and rewrites expert ids to slot ids. A second kernel copies the
+  missing experts straight from device-mapped pinned host memory into their
+  slots. Both kernels have fixed launch shapes and read the miss count from
+  device memory, so decode stays inside CUDA graphs with no host sync.
+- **Prefill** — when a call touches more distinct experts than the layer has
+  slots, the whole expert layer is copied (pinned async memcpy) into a staging
+  buffer, and the unmodified GEMM runs on it.
+- The GEMM kernels are unchanged: they read `[num_slots, N, K]` slot tensors
+  instead of `[num_experts, N, K]`.
+
+Design follows [FreeToken](https://github.com/FlashML-org/FreeToken) (Apache-2.0).
+
+## Options
+
+| Variable | Meaning |
+|---|---|
+| `XINFER_MOE_OFFLOAD=auto` | After loading, grow every layer's slots evenly into the free VRAM, leaving the reserve below for the KV cache and workspaces (recommended) |
+| `XINFER_MOE_OFFLOAD=<0..1>` | Fixed fraction of each layer's experts kept on the GPU; for example, `0.2` means 20% (at least `num_experts_per_tok`). Unset = no offload |
+| `XINFER_MOE_OFFLOAD_RESERVE_GB` | VRAM left free by `auto` (default 2.5). Raise it for longer contexts |
+| `XINFER_MOE_OFFLOAD_PROFILE=1` | Logs a per-MoE-call time split (launch gap, GPU wait, host, copy). Adds device syncs and disables CUDA graphs; diagnostics only |
+
+More experts on the GPU means faster decode; more reserve means a longer
+context. TurboQuant (`--kvcache-dtype turbo4|turbo3`) stretches the reserve
+~4x. Models with a large KV per token (e.g. Gemma 4) need `fp8` or `turbo*`.
+MTP (`--num-speculative-tokens`) works but was slower in our tests.
+
+For tensor-parallel inference, pass multiple devices with `--d 0,1` (or more
+devices). Every rank must be able to read the model files locally. Each rank
+loads and offloads its own expert shard, while NCCL continues to combine the
+rank-local MoE results. The offload fraction applies independently to every
+rank and every MoE layer.
+
+## Supported
+
+- Formats:
+  - GGUF (`FusedMoeGGUF`, separate `ffn_gate_exps`/`ffn_up_exps`, quant types the
+    GGUF MoE kernel reads natively) and ISQ (`--isq`, `FusedMoeISQ`: experts are
+    quantized on the GPU at load, then moved to host RAM);
+  - NVFP4 ModelOpt safetensors (`FusedMoeNvfp4`);
+  - unquantized bf16/f16 (`FusedMoe`), block FP8 (`FusedMoeFp8`), MXFP4
+    (`FusedMoeMxfp4`) and GPTQ/AWQ/compressed-tensors WNA16 (`FusedMoeWNA16`): all
+    share one generic path that keeps every `[E, ...]` expert tensor (weights and
+    scales, any dtype) in host RAM.
+- Not supported: W2-packed DeepSeek-V4 experts (fails with a clear error).
+- CUDA single-GPU and tensor-parallel inference. NVFP4 on SM100+ (swizzled
+  scales) is not supported yet.
+- Tested on real models: GGUF, NVFP4, bf16, ISQ, and FP8 MoE safetensors,
+  including Qwen3.5-35B-A3B-FP8 with two-rank tensor parallelism. MXFP4 and
+  WNA16 are covered by the generic path's bit-exact GEMM test.
+- Host RAM must hold every expert (pinned).
+
+## Benchmarks
+
+Reproduce with `bench/bench_moe_offload.py` (rate `0` = baseline without offload):
+
+```bash
+python3 bench/bench_moe_offload.py --bin target/release/xinfer \
+    --model /models/Qwen3.6-35B-A3B-NVFP4 --rates 0,0.1,0.2,0.35,auto --out bench/qwen36
+```
+
+It runs the OpenAI-compatible server, sends a warm-up and one streamed request
+(greedy) and records time to first token, decode speed and inter-token latency
+percentiles; `--prompt-tokens N` adds a long-context needle test. It writes a
+Markdown table and the raw JSON next to `--out` (not committed).
+
+RTX 3060 12 GB, PCIe 3.0 x16, Xeon E5-2650 v4, 64 GB DDR4-2400, Windows 11,
+200 decoded tokens, `--max-model-len 8192`. Without offload none of these
+models runs on this card: the weights spill into the driver's sysmem fallback
+and no VRAM is left for the KV cache.
+
+**Qwen3.6-35B-A3B NVFP4** (23.4 GB of weights)
+
+| GPU experts | KV dtype | KV tokens | TTFT | Decode tok/s | Inter-token p50 / p99 ms |
+|---|---|---|---|---|---|
+| none | — | does not run | | | |
+| 10% | bf16 | 192k | 2.1 s | 21.5 | 46 / 70 |
+| 20% | bf16 | 134k | 2.0 s | 26.4 | 38 / 51 |
+| 35% | bf16 | 46k | 0.8 s | 28.6 | 35 / 44 |
+| auto (31%) | bf16 | 58k | 1.3 s | 26.9 | 37 / 49 |
+| 20% | turbo4 | 520k | 2.1 s | 23.3 | 41 / 68 |
+| 35% | turbo3 | 201k | 1.0 s | 26.8 | 37 / 45 |
+
+Long context: an 11.5k-token needle prompt at 20% prefills at 333 tok/s and the
+needle is found. Greedy output is byte-identical for every GPU fraction.
+
+**Gemma-4-26B-A4B NVFP4** (18.8 GB): 10% / 15% with `turbo4` → 16.7 / 19.8 tok/s
+(27k / 16k KV tokens); with `fp8` at 10% → 18.9 tok/s.
+**GLM-4.7-Flash NVFP4** (20.5 GB, no CUDA graphs in xinfer for this model):
+10% / 20% → 10.8 / 11.2 tok/s. Outputs identical across fractions for both.
+
+**Qwen1.5-MoE-A2.7B-Chat bf16** (28.6 GB, 60 experts of 17 MB each), `auto` with
+`XINFER_MOE_OFFLOAD_RESERVE_GB=1.5`:
+
+| Experts | KV dtype | GPU experts | KV tokens | TTFT | Decode tok/s | Inter-token p50 / p99 ms |
+|---|---|---|---|---|---|---|
+| bf16 | fp8 | auto (~15%) | 8.4k | 2.5 s | 7.9 | 131 / 168 |
+| bf16 | turbo4 | auto (~15%) | 12.5k | 2.4 s | 8.1 | 126 / 158 |
+| ISQ `q8_0` | fp8 | auto (40%) | 14.8k | 1.9 s | 18.0 | 55 / 86 |
+| ISQ `q8_0` | turbo4 | auto (40%) | 27.1k | 1.9 s | 20.5 | 49 / 74 |
+
+bf16 experts are PCIe-bound (~1.4 GB of misses per token over PCIe 3.0);
+`--isq q8_0` halves the expert size and fits 2.7x more experts on the GPU, so
+decode is 2.5x faster. Q4K needs `moe_intermediate_size % 256 == 0` (1408 here)
+and the GGUF MoE kernel has no Q4_0, so 4-bit ISQ is not available for this
+model. bf16 greedy output is identical at 10% and auto (~15%).
+
+Run-to-run variation on this machine is about ±10%.

package/docs/speculative_decoding.md

@@ -0,0 +1,129 @@
+# Speculative Decoding (MTP & DFlash2)
+
+xInfer speeds up autoregressive decode with **speculative decoding**: a small set of candidate tokens is proposed each step, then verified by the target model in a single batched forward pass. Accepted drafts advance the sequence without extra per-token decode rounds.
+
+Two modes are supported:
+
+| Mode | When to use | CLI |
+|------|-------------|-----|
+| **Built-in MTP** | Target model ships MTP prediction heads (e.g. Qwen3.5 / Qwen3.6 with MTP weights) | `--num-speculative-tokens N` |
+| **DFlash2** | External DFlash2 draft model (separate safetensors checkpoint) | `--draft-model <id_or_path>` |
+
+**Rule:** If `--draft-model` is set, xInfer uses **DFlash2** (DFlash1 checkpoints are not supported). Otherwise, `--num-speculative-tokens` enables **built-in MTP** when the target model has MTP layers.
+
+---
+
+## Built-in MTP
+
+MTP (Multi-Token Prediction) uses lightweight heads bundled in the target checkpoint. Each decode step:
+
+1. Target model produces one **anchor** token and a hidden state.
+2. MTP head drafts `N` tokens autoregressively (no extra KV growth).
+3. Target model verifies `[anchor, draft₀, …, draftₙ₋₁]` in one forward pass.
+4. Matching prefix is accepted; GDN/Mamba state is rolled back on partial rejection (Qwen3.5 hybrid models).
+
+### CLI
+
+```bash
+# Qwen3.5 35B with 3 speculative tokens per step
+xinfer --m Qwen/Qwen3.5-35B-A3B --d 0,1 \
+  --num-speculative-tokens 3 --ui-server
+```
+
+Typical values: **3–7** draft tokens. Higher values increase acceptance variance and KV verify width.
+
+### Requirements
+
+- Target weights must include MTP layers (`mtp.*` or GGUF `nextn.*` tensors).
+- Model types with MTP support today: **Qwen3.5**, **Qwen3.5 MoE**, **Qwen3-VL** (Qwen3.5 text backbone).
+- Recommended build: `./build.sh --release --features cuda,nccl,flashinfer,cutlass`
+
+---
+
+## DFlash2
+
+[DFlash2](https://github.com/z-lab/dflash) uses a **separate draft model** that consumes projected intermediate hidden states from selected target layers. Draft tokens are chosen with a top-k candidate lattice (`selector_top_k` in the draft config).
+
+### CLI
+
+```bash
+# Qwen3.8 target + DFlash2 draft (HuggingFace id or local path)
+xinfer --m Qwen/Qwen3.8-... --d 0,1 \
+  --draft-model <your-dflash2-draft-repo-or-path> \
+  --num-speculative-tokens 7 --ui-server
+```
+
+`--draft-model` accepts either:
+
+- A **HuggingFace model id** (downloads `config.json` + safetensors), or
+- A **local directory** containing `config.json` and `model.safetensors` (or sharded index).
+
+`--num-speculative-tokens` is the number of draft tokens proposed per step (excluding the anchor). If omitted, xInfer uses the draft model config (`block_size - 1`).
+
+Draft weights must be **safetensors** (GGUF draft models are not supported). The checkpoint must be **DFlash2** (`architectures` contains `DFlash2` or `dflash_config.selector_top_k` is set).
+
+### Requirements
+
+- Target: **Qwen3.5**, **Qwen3.5 MoE**, **Qwen3-VL**, or **Qwen3.8** (dense) today.
+- Draft: matching-family **DFlash2** checkpoint.
+- Recommended build: `cuda,nccl,flashinfer,cutlass` (CUDA graphs optional; disable with `--disable-cuda-graph` for debugging).
+
+---
+
+## Python API
+
+```python
+from xinfer import Engine, EngineConfig
+
+# Built-in MTP
+cfg = EngineConfig(
+    model_id="Qwen/Qwen3.5-35B-A3B",
+    num_speculative_tokens=3,
+)
+
+# DFlash2 (draft_model enables DFlash2 instead of MTP)
+cfg = EngineConfig(
+    model_id="Qwen/Qwen3.8-...",
+    draft_model="<dflash2-draft-id-or-path>",
+    num_speculative_tokens=7,
+)
+
+engine = Engine(cfg, "bf16")
+```
+
+---
+
+## Flags reference
+
+| Flag | Description |
+|------|-------------|
+| `--num-speculative-tokens N` | Draft tokens per decode step. Enables **MTP** when the target has MTP heads and `--draft-model` is unset. With `--draft-model`, sets DFlash2 draft width (optional; defaults from draft config). |
+| `--draft-model <id_or_path>` | External **DFlash2** draft model. Enables DFlash2 and disables built-in MTP. |
+
+**Removed:** `--mtp` (use `--num-speculative-tokens`), `--draft-model-id` / `--draft-model-path` (use `--draft-model`).
+
+---
+
+## Tips & limitations
+
+- **Throughput:** Speedup depends on draft acceptance rate and model size. Monitor logs for `MTP Stats` / `DFlash2 Stats` acceptance summaries.
+- **Memory:** Verify passes append up to `N+1` tokens per step; ensure KV budget via `--kv-fraction` and `--max-num-seqs`.
+- **Hybrid models:** Qwen3.5 GDN layers require state rollback on rejected drafts; this is handled automatically.
+- **Batching:** Speculative decode is optimized for `batch_size=1` decode steps; larger batches use a batched DFlash2 path.
+- **PD disaggregation:** Do not combine speculative decode with PD client/server modes on hybrid Mamba models.
+
+---
+
+## Quick test
+
+```bash
+# 1. Health check
+curl -s http://localhost:8000/v1/models
+
+# 2. Short completion (MTP or DFlash2 server)
+curl -s http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{"model":"Qwen/Qwen3.5-35B-A3B","messages":[{"role":"user","content":"Hello"}],"max_tokens":64}'
+```
+
+See also [test_model.md](./test_model.md) for full model validation workflows.

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "xinfer-ai",
-  "version": "0.14.4",
-  "assetVersion": "0.14.4",
+  "version": "0.14.6",
+  "assetVersion": "0.14.6",
   "description": "xInfer — a high-performance LLM inference engine in Rust with CUDA/Metal acceleration",
   "license": "MIT",
   "homepage": "https://github.com/guoqingbao/xinfer#readme",