Qwen3.6-35B-A3B edge runtime: white-boxed prefill, SM110 support, and token-exact speculative decode

README.md

@@ -165,6 +165,32 @@ DGX Spark / GB10:
 | NVFP4, 128 | **170.1 ms** | **40.42 tok/s** | [Qwen3.6 Spark](docs/qwen36_spark.md#performance) |
 | NVFP4, 16 K | **8.545 s** | **54.94 tok/s** | [Qwen3.6 Spark](docs/qwen36_spark.md#performance) |
 
+#### Qwen3.6-35B-A3B
+
+RTX 5090:
+
+| Mode | Prefill | Decode | Source |
+|---|---:|---:|---|
+| NVFP4, 64 | **40.42 ms** | **257.95 tok/s** | [Qwen3.6-MoE usage](docs/qwen36_moe_usage.md#validation) |
+
+Jetson AGX Thor, against vLLM 0.26.0 on the same part and protocol:
+
+| Mode | Prefill | Decode | Source |
+|---|---:|---:|---|
+| NVFP4, 20 | **89.5 ms** (vLLM 102.3) | | [Qwen3.6-MoE Thor](docs/qwen36_moe_usage.md#jetson-agx-thor-numbers) |
+| NVFP4, 1 K | **216.0 ms** (vLLM 319.4) | **87.1 tok/s** (vLLM 31.6) | [Qwen3.6-MoE Thor](docs/qwen36_moe_usage.md#jetson-agx-thor-numbers) |
+| NVFP4, 2 K | **379.6 ms** (vLLM 495.0) | **86.3 tok/s** (vLLM 31.5) | [Qwen3.6-MoE Thor](docs/qwen36_moe_usage.md#jetson-agx-thor-numbers) |
+| NVFP4, 32 K | **7207.5 ms** (vLLM 7231.8) | | [Qwen3.6-MoE Thor](docs/qwen36_moe_usage.md#jetson-agx-thor-numbers) |
+
+TTFT leads at every length from 20 to 32768 tokens; 128 K context reaches the
+board at 2470 tok/s of prefill. The decode column was taken before a later
+round that moved the steady step from 89.0 to 102.6 tok/s, and vLLM's side is
+unaffected by it, so the ratios shown are lower bounds.
+
+Speculative decode with the MTP head reaches **106.74 tok/s** against 100.35
+plain in the same process, emitting the same tokens as greedy decoding. See
+[speculative decode](docs/qwen36_moe_usage.md#speculative-decode).
+
 #### Qwen3-8B
 
 | Hardware | Mode | Prefill | Decode | Source |
@@ -732,7 +758,7 @@ extension modules:
 | Artifact | Size | What it contains |
 |---|---|---|
 | `flash_rt/flash_rt_kernels.so` | ~3 MB | Hand-written memory-bound kernels (norm, activation, fusion, FP8 quant, cuBLASLt wrappers, Thor FMHA). **Always built.** |
-| `flash_rt/flash_rt_fa2.so` | ~135 MB | Vendored Flash-Attention 2 v2.7.4.post1 fwd (fp16 + bf16, SM80/86/89/120). **Built only on RTX targets** — Thor skips it and uses `fvk.attention_qkv_fp16` (cuBLAS-decomposed) for attention instead. |
+| `flash_rt/flash_rt_fa2.so` | ~135 MB | Vendored Flash-Attention 2 v2.7.4.post1 fwd (fp16 + bf16, SM80/86/89/120). **Built automatically on RTX targets.** Thor skips it by default and uses `fvk.attention_qkv_fp16` (cuBLAS-decomposed) instead; `-DFLASHRT_ENABLE_THOR_FA2=ON` builds it there for the one model whose long prefill needs it (Qwen3.6, bf16 head_dim 256 — a single instantiation). |
 
 **Crucially — no `pip install flash-attn` required.** The FA2 kernel
 is vendored at source level and built into `flash_rt_fa2.so` during
@@ -862,18 +888,20 @@ CMake reads `nvidia-smi --query-gpu=compute_cap` to pick the target
 arch. Override for cross-compilation or when auto-detect fails:
 
 ```bash
-cmake -B build -S . -DGPU_ARCH=110   # Jetson AGX Thor   (FA2 skipped, CUTLASS SM100 path ON)
+cmake -B build -S . -DGPU_ARCH=110   # Jetson AGX Thor   (FA2 opt-in, CUTLASS SM100 path ON)
 cmake -B build -S . -DGPU_ARCH=121   # DGX Spark / GB10   (FA2 sm_121 AOT, NVFP4 ON)
 cmake -B build -S . -DGPU_ARCH=120   # RTX 5090           (FA2 sm_120 AOT, NVFP4 ON)
 cmake -B build -S . -DGPU_ARCH=89    # RTX 4090           (FA2 sm_80 AOT natively runs on Ada)
 cmake -B build -S . -DGPU_ARCH=86    # RTX 3090 / A10     (FA2 sm_80 AOT)
 cmake -B build -S . -DGPU_ARCH=80    # A100               (FA2 sm_80 AOT)
 ```
 
-FA2 is enabled by CMake when `GPU_ARCH ∈ {80, 86, 89, 120, 121}`. Other
-arches (notably Thor SM110 and SM90 Hopper) route attention through
-the cuBLAS-decomposed `fvk.attention_qkv_fp16` path instead of FA2 —
-`flash_rt_fa2.so` simply isn't built, and no runtime error results.
+FA2 is enabled by CMake when `GPU_ARCH ∈ {80, 86, 89, 120, 121}`, and on
+Thor SM110 when `-DFLASHRT_ENABLE_THOR_FA2=ON` is passed. Other arches
+(notably SM90 Hopper, and Thor without that flag) route attention
+through the cuBLAS-decomposed `fvk.attention_qkv_fp16` path instead of
+FA2 — `flash_rt_fa2.so` simply isn't built, and no runtime error
+results.
 
 ### Build timing (one-time)
 

benchmarks/qwen36_moe_edge_decode.py

@@ -0,0 +1,139 @@
+#!/usr/bin/env python3
+"""First-light latency probe for the Qwen3.6-MoE (qwen3_5_moe) edge path.
+
+Reports the numbers quoted in ``docs/qwen36_moe_usage.md``: weight load time,
+resident and peak allocation, prefill latency, and decode throughput on the
+eager and the captured-graph paths. The two decode paths are compared token for
+token, because a throughput number for a path that emits different text is not
+a throughput number for the same work.
+
+Usage:
+
+    PYTHONPATH=. python benchmarks/qwen36_moe_edge_decode.py \\
+        --checkpoint /path/to/Qwen3.6-35B-A3B \\
+        --prompt-tokens 64 --max-new-tokens 32
+"""
+
+from __future__ import annotations
+
+import argparse
+import statistics
+import time
+
+import torch
+
+GIB = 2 ** 30
+
+
+def _sync(device: str) -> None:
+    torch.cuda.synchronize(device)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint", required=True,
+                        help="path to the BF16 checkpoint directory")
+    parser.add_argument("--prompt-tokens", type=int, default=64)
+    parser.add_argument("--max-new-tokens", type=int, default=64)
+    parser.add_argument("--prefill-reps", type=int, default=5)
+    parser.add_argument("--decode-reps", type=int, default=8)
+    parser.add_argument("--max-seq", type=int, default=512)
+    parser.add_argument("--device", default="cuda:0")
+    args = parser.parse_args()
+
+    from flash_rt.frontends.torch.qwen36_moe import Qwen36MoeTextFrontend
+
+    # Select and initialise the device before touching the memory stats: they
+    # are per-device counters and are not addressable until then.
+    torch.cuda.set_device(args.device)
+    torch.cuda.init()
+    torch.cuda.reset_peak_memory_stats(args.device)
+    t0 = time.perf_counter()
+    frontend = Qwen36MoeTextFrontend(
+        args.checkpoint, device=args.device, max_seq=args.max_seq)
+    _sync(args.device)
+    load_s = time.perf_counter() - t0
+
+    print(f"runtime weight load                  {load_s:8.2f} s")
+    print(f"resident allocated after load        "
+          f"{torch.cuda.memory_allocated(args.device) / GIB:8.2f} GiB")
+    print(f"peak allocated during load           "
+          f"{torch.cuda.max_memory_allocated(args.device) / GIB:8.2f} GiB")
+
+    base = frontend.tokenizer.encode(
+        "The quick brown fox jumps over the lazy dog. ")
+    ids = (base * (args.prompt_tokens // len(base) + 2))[:args.prompt_tokens]
+
+    # Prefill: the first call carries warmup and lazy weight packing, so it is
+    # reported separately rather than averaged into the steady-state figure.
+    frontend.set_prompt_ids(ids)
+    _sync(args.device)
+    t0 = time.perf_counter()
+    frontend.generate(max_new_tokens=1)
+    _sync(args.device)
+    first_ms = (time.perf_counter() - t0) * 1e3
+
+    warm = []
+    for _ in range(args.prefill_reps):
+        frontend.set_prompt_ids(ids)
+        _sync(args.device)
+        t0 = time.perf_counter()
+        frontend.generate(max_new_tokens=1)
+        _sync(args.device)
+        warm.append((time.perf_counter() - t0) * 1e3)
+
+    print(f"first prefill, including warmup      {first_ms:8.2f} ms")
+    print(f"subsequent prefill                   "
+          f"{min(warm):8.2f}-{max(warm):.2f} ms")
+
+    def run(fn) -> tuple[list[float], list[int]]:
+        # Median and range over every repetition, not a best-of: a single best
+        # sample hides both contention and variance, and the baseline this is
+        # compared against reports the same shape.
+        rates, toks = [], None
+        for _ in range(args.decode_reps):
+            frontend.set_prompt_ids(ids)
+            _sync(args.device)
+            t0 = time.perf_counter()
+            out = fn()
+            _sync(args.device)
+            rates.append(args.max_new_tokens / (time.perf_counter() - t0))
+            toks = list(out)
+        return sorted(rates), toks
+
+    def report(label: str, rates: list[float]) -> None:
+        med = statistics.median(rates)
+        print(f"{label:<44}{med:8.2f} tok/s   "
+              f"(range {rates[0]:.2f}-{rates[-1]:.2f} over {len(rates)} runs)")
+
+    state = frontend._decode_state
+    from flash_rt.frontends.torch import _nexn2_rtx_decode as dec
+
+    def eager():
+        t = torch.tensor(ids, dtype=torch.long, device=args.device)
+        with torch.no_grad():
+            return dec.generate_greedy(
+                state, t, args.max_new_tokens, frontend._fvk, args.device)
+
+    eager_rate, eager_toks = run(eager)
+    graph_rate, graph_toks = run(
+        lambda: frontend.generate(max_new_tokens=args.max_new_tokens))
+
+    report(f"{args.prompt_tokens}/{args.max_new_tokens} eager decode",
+           eager_rate)
+    report(f"{args.prompt_tokens}/{args.max_new_tokens} warm graph decode",
+           graph_rate)
+    free, total = torch.cuda.mem_get_info(args.device)
+    print(f"{'device free memory at exit':<44}{free / GIB:8.2f} GiB "
+          f"of {total / GIB:.2f} -- a shared device invalidates the timings")
+    same = eager_toks == graph_toks
+    print(f"eager and graph emit the same tokens {str(same):>8}"
+          f"   ({sum(a == b for a, b in zip(eager_toks, graph_toks))}"
+          f"/{len(graph_toks)})")
+    if not same:
+        raise SystemExit("eager and captured decode disagree; "
+                         "the throughput numbers are not comparable")
+
+
+if __name__ == "__main__":
+    main()