qwen36-fp4

Single-stream inference engine for sakamakismile/Qwen3.6-27B-Text-NVFP4-MTP, targeting RTX 5090 / Blackwell SM120.

This repo implements the project structure from doc.md: Rust host orchestration, safetensors discovery, tokenizer wrapper, runtime state planning, MTP rollback contracts, TurboQuant policy, CUDA ABI, and baseline CUDA kernels for bring-up.

Status

Implemented:

• Rust workspace with core, loader, tokenizer, kernels, runtime, mtp, and cli crates.
• qwen36 discover path that mmaps .safetensors and writes model_layout.json.
• Zero-copy MappedModel tensor access for safetensors shards.
• HF config.json parser and Qwen3.6 topology validation.
• Runtime weight manifest validation for all layer, MTP, lm-head, embedding, and NVFP4 scale tensors.
• CUDA runtime memory ABI plus Rust RAII buffers for device allocation, copy, memset, and synchronization.
• Real-checkpoint GPU upload path for required manifest tensors and runtime buffers.
• Hybrid-aware TurboQuant attention skip policy: first and last full-attention layer.
• KV-cache and DeltaNet-state memory planning.
• CUDA shared-library ABI and baseline kernels for FP4 GEMM, attention decode, int8 KV quantization, quantized attention, RMSNorm, partial RoPE, SwiGLU, greedy sampling, and DeltaNet decode.
• MTP speculative controller with rollback/replay tests.

Not final yet:

• DeltaNet decode has an exact single-token recurrence path when gate/beta tensors are supplied, but conv/projection fusion and prefill are not final.
• TurboQuant is currently int8 per-vector KV quantization, not the full rotation/QJL implementation.
• End-to-end reference decode is wired through all layers, but it uses slow scalar CUDA matvecs and still needs numerical parity work against vLLM/Transformers.
• MTP speculative execution is not yet wired into the runtime scheduler.
• CUDA Graph capture and final hot-path tuning are pending.

Quick Start

git clone <repo-url> qwen36-fp4
cd qwen36-fp4

rustup show
./scripts/build_cuda.sh
./scripts/smoke_cuda.sh

export QWEN36_FP4_KERNEL_LIB_DIR="$PWD/target/cuda"
export LD_LIBRARY_PATH="$QWEN36_FP4_KERNEL_LIB_DIR:${LD_LIBRARY_PATH:-}"
cargo check --workspace --features qwen36-fp4-kernels/cuda

Download the model, then inspect it:

hf download sakamakismile/Qwen3.6-27B-Text-NVFP4-MTP \
  --local-dir /models/Qwen3.6-27B-Text-NVFP4-MTP

cargo run -p qwen36-fp4 -- discover \
  --model-dir /models/Qwen3.6-27B-Text-NVFP4-MTP \
  --output model_layout.json

cargo run -p qwen36-fp4 -- validate-weights \
  --model-dir /models/Qwen3.6-27B-Text-NVFP4-MTP

cargo run -p qwen36-fp4 --features cuda -- gpu-load \
  --model-dir /models/Qwen3.6-27B-Text-NVFP4-MTP \
  --max-context 2256

Documentation

• Installation: full host setup for Rust, CUDA, model download, and first commands.
• Model Setup: checkpoint layout, model_layout.json, and validation checks.
• Kernel Validation: CUDA build, smoke tests, and numerical validation plan.
• Repository Layout: crate responsibilities and source tree.
• Development: local workflow, ABI rules, and PR checklist.
• Troubleshooting: common build/runtime failures.
• Roadmap: path from baseline kernels to final optimized inference.
• Research Notes: verified external assumptions and links.

Core Commands

cargo run -p qwen36-fp4 -- inspect-config --model-dir /path/to/model
cargo run -p qwen36-fp4 -- budget --ctx 32768 --kv fp8
cargo run -p qwen36-fp4 -- tokenize --model-dir /path/to/model --text "Bonjour"
cargo run -p qwen36-fp4 -- validate-weights --model-dir /path/to/model
cargo run -p qwen36-fp4 --features cuda -- gpu-load --model-dir /path/to/model --max-context 2256
cargo run -p qwen36-fp4 -- validate-eagle3-drafter --drafter-dir /path/to/EAGLE3
cargo run -p qwen36-fp4 --features cuda -- chat --model-dir /path/to/model --prompt "Bonjour" --drafter eagle3 --drafter-dir /path/to/EAGLE3

CUDA Runtime Knobs

The list below is a partial selection. The complete reference (90 env vars, defaults verified in code, plus the component status map) is in docs/code-inventory.md.

• QWEN36_KV_CACHE_DTYPE=tq35 selects the TurboQuant 3.5-bit KV cache.
• QWEN36_LONG_CONTEXT_MODE overrides the automatic long-context policy. When unset, contexts at or above 8192 tokens disable the large fused weight stores to save VRAM; set QWEN36_LONG_CONTEXT_MODE=0 to force the fused stores back on, or =1 to force long-context mode on smaller runs.
• QWEN36_LONG_CONTEXT_AUTO_MIN_CONTEXT=<tokens> changes the automatic threshold.
• QWEN36_PREFILL_CAPACITY=<tokens> overrides the prefill chunk buffer capacity. By default, contexts up to 32K allocate an 8K prefill chunk to avoid splitting 8K+ prompts into many passes; larger reserved contexts keep the older 2K chunk to limit VRAM growth.
• QWEN36_PREFILL_SPLIT_MAX_TOKENS=<tokens> raises the opt-in split-KV prefill chunk limit for long-context experiments. The default remains tuned for short MTP verify chunks.
• QWEN36_PREFILL_SPLIT_MIN_SPLITS=<count> changes the split-KV prefill activation threshold when the max-token override is used.
• QWEN36_DECODE_ATTENTION_BUCKET_MIN_CONTEXT=<tokens> changes the minimum active-context bucket used to size decode split-KV launches. The default is 8192, so a run reserving 262K context no longer launches 262K-sized decode attention splits while decoding near 8K.
• QWEN36_DECODE_ATTENTION_BUCKET_DISABLE=1 restores the older behavior where decode split-KV launch shapes are sized from the configured max_context.
• QWEN36_MTP_ASSUME_ACCEPT=1 skips per-draft verification checks in the multi-token MTP graph path. It also skips recurrent snapshots for that path, so use it only when rejection is known to be impossible or acceptable for the run.
• QWEN36_MTP_FUSED_ARGMAX=1 enables the experimental fused BF16 lm_head + greedy argmax path for MTP prefill rows, avoiding the intermediate logits buffer on those samples.
• QWEN36_EAGLE3_DRAFT_TOKENS=<tokens> changes the number of chain draft tokens proposed per EAGLE3 speculative iteration. The default is 8.
• QWEN36_EAGLE3_P_MIN=<prob> changes the minimum top-1 drafter confidence required to keep extending an EAGLE3 chain. The default is 0.5; set it to 0 to restore the previous no-cutoff fast path.
• QWEN36_PREFILL_GQA_TILE2=1 enables the experimental two-token GQA prefill attention tile for BF16/FP8 KV layouts. QWEN36_PREFILL_GQA_TILE_TOKENS=2|4 selects the tile size directly. This is an opt-in first step toward a FlashAttention-style tiled prefill path.
• QWEN36_PROFILE_PREFILL_CHUNKS=1 prints per-prefill-chunk timing buckets for profiling. It synchronizes between buckets and should not be used for throughput numbers.

License

Apache-2.0. See LICENSE.