feat: support GLM-5.3-Flash NVFP4

README.md

@@ -20,7 +20,7 @@ FreeToken is an edge-native Mixture-of-Experts (MoE) serving engine designed for
 - **Fast Edge-Native Runtime**: Provides efficient MoE serving with bandwidth-adaptive CPU–GPU co-execution ($q^\star$ policy), full-layer double-buffered prefill streaming, global LRU expert caching, graph-compatible execution, and the FTW fast weight format.  
 - **Semantic-Aware Caching**: Features semantic anchor checkpoints for recurrent state and KV caches, allowing agentic context edits (e.g., tool calls, thinking blocks) to avoid redundant context recomputation.  
 - **Elastic Memory Management**: Supports dynamic, runtime VRAM re-allocation between expert caches and KV memory without engine restarts or weight reloading.  
-- **Broad MoE & Ecosystem Support**: Supports frontier open-weight MoE models (e.g., DeepSeek-V4-Flash, Qwen3.6-35B-A3B, GLM-5.2) across various parameter scales and quantization formats (e.g., MXFP4, NVFP4, FP8, BF16), with Anthropic/OpenAI-compatible APIs for seamless integration with real-world coding and tool-calling agents (e.g., Codex, Claude Code, OpenCode, OpenClaw, DeepSeek Harness). 
+- **Broad MoE & Ecosystem Support**: Supports frontier open-weight MoE models (e.g., DeepSeek-V4-Flash, Qwen3.6-35B-A3B, GLM-5.3-Flash) across various parameter scales and quantization formats (e.g., MXFP4, NVFP4, FP8, BF16), with Anthropic/OpenAI-compatible APIs for seamless integration with real-world coding and tool-calling agents (e.g., Codex, Claude Code, OpenCode, OpenClaw, DeepSeek Harness).
 - **Diverse Consumer Hardware**: Scales across consumer laptops, gaming desktops, and workstation GPUs, with native support for NVIDIA RTX 30, RTX 40, and RTX 50 series GPUs.  
 
 ## Getting Started

benchmarks/bench_decode_moe.py

@@ -99,10 +99,27 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
         default=-1,
         help="hybrid: max PCIe fetches/layer; -1 = auto (benched pcie/cpu bandwidth fraction)",
     )
+    p.add_argument(
+        "--cpu-threads",
+        type=int,
+        default=0,
+        help="CPU MoE worker threads; 0 = runtime auto-selection",
+    )
+    p.add_argument(
+        "--kv-reserve-tokens",
+        type=int,
+        default=0,
+        help="KV token floor reserved before auto-sizing the expert cache; 0 = server default",
+    )
     p.add_argument("--mem-ratio", type=float, default=0.9, help="target VRAM utilization")
     p.add_argument("--gpu", default=None,
                    help="GPU for the serve: a UUID or nvidia-smi index (as ft serve --gpu)")
     p.add_argument("--no-graph", action="store_true", help="eager decode instead of CUDA graph")
+    p.add_argument(
+        "--disable-prefill-overlap",
+        action="store_true",
+        help="disable MoE prefill overlap (permits expert caches below 2 * num_experts)",
+    )
     p.add_argument(
         "--greedy",
         action="store_true",
@@ -187,6 +204,12 @@ def serve_cmd(args: argparse.Namespace, backend: str, port: int) -> list[str]:
     ]
     if args.gpu:
         cmd += ["--gpu", args.gpu]
+    if args.cpu_threads > 0:
+        cmd += ["--moe-cpu-threads", str(args.cpu_threads)]
+    if args.kv_reserve_tokens > 0:
+        cmd += ["--kv-reserve-tokens", str(args.kv_reserve_tokens)]
+    if args.disable_prefill_overlap:
+        cmd.append("--disable-moe-prefill-overlap")
     if args.cache > 0:
         cmd += ["--moe-cache-size", str(args.cache)]
     elif args.cache_rate is not None:

docs/models.md

@@ -7,6 +7,7 @@ for them; other checkpoints of the same architectures work too.
 | Model | HF checkpoints |
 |---|---|
 | DeepSeek-V4 | [deepseek-ai/DeepSeek-V4-Flash-0731](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-0731) |
+| GLM-5.3-Flash | [LibertAIDAI/GLM-5.3-Flash-NVFP4](https://huggingface.co/LibertAIDAI/GLM-5.3-Flash-NVFP4) |
 | GLM-5.2 | [nvidia/GLM-5.2-NVFP4](https://huggingface.co/nvidia/GLM-5.2-NVFP4) |
 | GLM-4.7 | [nvidia/GLM-4.7-NVFP4](https://huggingface.co/nvidia/GLM-4.7-NVFP4) |
 | Qwen3.8-Flash-Next | [Qwen/Qwen3.8-Flash-Next-FP8](https://huggingface.co/Qwen/Qwen3.8-Flash-Next-FP8), [RadixArk/Qwen3.8-Flash-Next-NVFP4](https://huggingface.co/RadixArk/Qwen3.8-Flash-Next-NVFP4) |
@@ -38,5 +39,7 @@ for them; other checkpoints of the same architectures work too.
   FreeToken's fast-load format, and `ft serve --model` auto-detects the result.
 - DeepSeek-V4 checkpoints must keep the `inference/config.json` subdir — the
   authoritative model args are read from there.
+- GLM-5.3-Flash currently supports TP=1 with its NVFP4 routed experts served
+  from host RAM by the offload backend.
 - Qwen3.8-Flash-Next keeps a 47.7 GiB PLE n-gram table pinned in host RAM.
 - Multimodal checkpoints are served text-only.

python/freetoken/attention/dsa.py

@@ -31,9 +31,10 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Dict, List, Tuple
+from typing import TYPE_CHECKING
 
 import torch
+
 from freetoken.core import Batch, get_global_ctx
 
 from .base import AttentionSpec, BaseAttnBackend, BaseAttnMetadata
@@ -73,8 +74,8 @@ class DSAAttnBackend(DSAIndexerMixin, BaseAttnBackend):
     def __init__(self, config: ModelConfig) -> None:
         from freetoken.kvcache.dsa_pool import DSAKVCache, MLAKVCache
 
-        args = config.glm_dsa_args
-        assert args is not None, "dsa backend needs ModelConfig.glm_dsa_args (MLA dims)"
+        args = config.glm_dsa_args or config.glm5_args
+        assert args is not None, "dsa backend needs GLM MLA dimensions"
         self.config = config
         self.num_heads = config.num_qo_heads
         self.kv_lora_rank = args.kv_lora_rank
@@ -96,13 +97,29 @@ def __init__(self, config: ModelConfig) -> None:
         # layer -> group leader (most recent "full" layer); leader -> pool slot.
         # Only built when DSA serves: the dense ablation never consults indexer_types,
         # so a checkpoint with a malformed list cannot crash the ablation.
-        self._leader: Dict[int, int] = {}
-        self._idx_slot: Dict[int, int] = {}
+        self._leader: dict[int, int] = {}
+        self._idx_slot: dict[int, int] = {}
         if self.dsa_enabled:
             lead = None
+            # Hybrid models (GLM-5.3-Flash) carry an ``indexer_types`` entry for
+            # every decoder layer, including KDA layers that never enter this
+            # backend.  The index slab is deliberately sized only for MLA/DSA
+            # layers, so enumerate that group's layer ids rather than treating a
+            # ``full`` marker on a linear-attention layer as an allocated slot.
+            if getattr(config, "glm5_args", None) is not None:
+                full_group = next(
+                    group for group in config.attention_groups if group.name == "full"
+                )
+                served_layers = set(full_group.layer_ids)
+            else:
+                # GLM-5.2 is all-MLA and several backend-level tests intentionally
+                # pass a minimal config namespace without generic group metadata.
+                served_layers = set(range(config.num_layers))
             # Capped to the SERVED layer count (dev num_layers overrides must not
             # index slots past the pool the factory sized from the same cap).
             for lid, kind in enumerate(args.indexer_types[: config.num_layers]):
+                if lid not in served_layers:
+                    continue
                 if kind == "full":
                     lead = lid
                     self._idx_slot[lid] = len(self._idx_slot)
@@ -113,7 +130,7 @@ def __init__(self, config: ModelConfig) -> None:
         self._rows_buf: torch.Tensor | None = None
         self._kvlen_buf: torch.Tensor | None = None
         self.max_seq_len = 0
-        self.capture_bs: List[int] = []
+        self.capture_bs: list[int] = []
 
     def forward(self, q, k, v, layer_id, batch, attn_spec: AttentionSpec | None = None):
         raise NotImplementedError("MLA models use mla_forward(), not forward().")
@@ -128,7 +145,9 @@ def prepare_metadata(self, batch: Batch) -> None:
         # active_table_idx (which the decode path's addressing requires) for
         # phase == "decode". The prefill path handles extend_len == 1 fine.
         is_decode = getattr(batch, "phase", None) == "decode"
-        qo_indptr = torch.tensor([0] + seqlens_q, **_CPU_PINNED).cumsum_(0).to(torch.int32)
+        qo_indptr = (
+            torch.tensor([0] + seqlens_q, **_CPU_PINNED).cumsum_(0).to(torch.int32)
+        )
         kv_len = torch.tensor(seqlens_k, **_CPU_PINNED)
         last = (qo_indptr[1:].to(torch.int32) - 1).to(self.device, non_blocking=True)
         md = DSAMetadata(
@@ -151,8 +170,12 @@ def _attend(
         from freetoken.kernel.triton.glm_dsa_sparse import glm_dsa_sparse_attn
 
         return glm_dsa_sparse_attn(
-            q_cat, self.kvcache.latent_rows(layer_id), sel, self.sm_scale,
-            counts=cnt, d_v=self.kv_lora_rank,
+            q_cat,
+            self.kvcache.latent_rows(layer_id),
+            sel,
+            self.sm_scale,
+            counts=cnt,
+            d_v=self.kv_lora_rank,
         )
 
     def mla_forward(
@@ -177,7 +200,9 @@ def mla_forward(
         if self.dsa_enabled and indexer_qkw is not None:
             # Scatter index keys unconditionally: short prefills serve through the
             # identity path TODAY, but their keys must exist once decode passes topk.
-            self.kvcache.store_index_k(indexer_qkw[1], batch.out_loc, self._idx_slot[layer_id])
+            self.kvcache.store_index_k(
+                indexer_qkw[1], batch.out_loc, self._idx_slot[layer_id]
+            )
 
         if md.is_decode:
             return self._decode(md, layer_id, q_nope, q_pe, indexer_qkw)
@@ -194,7 +219,9 @@ def _decode(self, md, layer_id, q_nope, q_pe, indexer_qkw) -> torch.Tensor:
         else:
             if indexer_qkw is not None:
                 q_idx, _, w = indexer_qkw
-                s = self.dsa_decode_scores(q_idx, w, self._idx_slot[layer_id], rows, kvlen)
+                s = self.dsa_decode_scores(
+                    q_idx, w, self._idx_slot[layer_id], rows, kvlen
+                )
                 k_sel = min(self.index_topk, s.shape[-1])
                 picks = self.indexer_select_decode(
                     s.view(bs, 1, -1), valid=kvlen, topk=k_sel, offset=0
@@ -205,15 +232,21 @@ def _decode(self, md, layer_id, q_nope, q_pe, indexer_qkw) -> torch.Tensor:
                 md.sel.clear()
                 md.sel[layer_id] = (sel, cnt)
             sel, cnt = md.sel[self._leader[layer_id]]
-        q_cat = torch.cat([q_nope, q_pe], dim=-1).view(bs, 1, self.num_heads, self.latent_dim)
+        q_cat = torch.cat([q_nope, q_pe], dim=-1).view(
+            bs, 1, self.num_heads, self.latent_dim
+        )
         o = self._attend(q_cat, layer_id, sel, cnt)
         return o.view(bs, self.num_heads, self.kv_lora_rank)
 
     # ----- prefill / extend (eager) ------------------------------------------------------
     def _select_prefill(
-        self, slot: int, q_idx: torch.Tensor, w: torch.Tensor,
-        rows: torch.Tensor, positions: torch.Tensor,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        self,
+        slot: int,
+        q_idx: torch.Tensor,
+        w: torch.Tensor,
+        rows: torch.Tensor,
+        positions: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """Per-request causal top-k: ([1, m, K] physical rows, [1, m] counts)."""
         kv_len = rows.numel()
         k_all = self.kvcache.index_k_cache(slot).index_select(0, rows.long())
@@ -223,14 +256,20 @@ def _select_prefill(
         start_pos = int(positions[0])
         # Bound the fp32 [chunk, kv_len] logits transient (worst case is capped by the
         # model's max_position: floor 16 x 1M x 4 B = 64 MB, see _PREFILL_SCORE_BYTES).
-        chunk = max(16, min(_PREFILL_SCORE_CHUNK, _PREFILL_SCORE_BYTES // max(kv_len * 4, 1)))
+        chunk = max(
+            16, min(_PREFILL_SCORE_CHUNK, _PREFILL_SCORE_BYTES // max(kv_len * 4, 1))
+        )
         for s0 in range(0, m, chunk):
             s1 = min(s0 + chunk, m)
             scores = self.dsa_prefill_logits(q_idx[s0:s1], k_all, w[s0:s1])
             # Shared selection semantics (dsv4_indexer): token-granular == ratio 1.
             picks = self.indexer_select_prefill(
-                scores.unsqueeze(0), start_pos=start_pos + s0, seqlen=s1 - s0,
-                ratio=1, topk=k_sel, offset=0,
+                scores.unsqueeze(0),
+                start_pos=start_pos + s0,
+                seqlen=s1 - s0,
+                ratio=1,
+                topk=k_sel,
+                offset=0,
             )[0]
             sel[s0:s1] = self.dsa_map_rows(picks, rows.view(1, -1).expand(s1 - s0, -1))
         cnt = torch.clamp(positions + 1, max=k_sel).to(torch.int32)
@@ -249,7 +288,8 @@ def _prefill(self, md, layer_id, q_nope, q_pe, batch, indexer_qkw) -> torch.Tens
             md.sel[layer_id] = [
                 self._select_prefill(
                     self._idx_slot[layer_id],
-                    q_idx[qo[i] : qo[i + 1]], w[qo[i] : qo[i + 1]],
+                    q_idx[qo[i] : qo[i + 1]],
+                    w[qo[i] : qo[i + 1]],
                     page_table[r.table_idx, : r.device_len],
                     batch.positions[qo[i] : qo[i + 1]],
                 )
@@ -267,34 +307,50 @@ def _prefill(self, md, layer_id, q_nope, q_pe, batch, indexer_qkw) -> torch.Tens
                 # token at kv <= index_topk, and the ablation attends everything).
                 # One shared row list broadcast across queries (stride 0), causality
                 # through per-query counts.
-                sel = page_table[r.table_idx, : r.device_len].view(1, 1, -1).to(torch.int32)
-                cnt = (batch.positions[qo[i] : qo[i + 1]] + 1).to(torch.int32).view(1, m)
+                sel = (
+                    page_table[r.table_idx, : r.device_len]
+                    .view(1, 1, -1)
+                    .to(torch.int32)
+                )
+                cnt = (
+                    (batch.positions[qo[i] : qo[i + 1]] + 1).to(torch.int32).view(1, m)
+                )
             o[qo[i] : qo[i + 1]] = self._attend(
                 q_cat[qo[i] : qo[i + 1]].view(1, m, self.num_heads, self.latent_dim),
-                layer_id, sel, cnt,
+                layer_id,
+                sel,
+                cnt,
             ).view(m, self.num_heads, self.kv_lora_rank)
         return o
 
     # ----- CUDA graph (decode) ----------------------------------------------------------
-    def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None:
+    def init_capture_graph(self, max_seq_len: int, bs_list: list[int]) -> None:
         self.max_seq_len = max_seq_len
         self.capture_bs = sorted(bs_list)
         max_bs = max(bs_list)
         width = get_global_ctx().page_table.shape[1]
-        self._rows_buf = torch.full((max_bs, width), -1, dtype=torch.int32, device=self.device)
+        self._rows_buf = torch.full(
+            (max_bs, width), -1, dtype=torch.int32, device=self.device
+        )
         self._kvlen_buf = torch.zeros(max_bs, dtype=torch.int32, device=self.device)
 
     def _decode_rows(self, batch: Batch) -> torch.Tensor:
         """This decode step's per-request page-table rows [bs, W], gathered off the
         scheduler-staged ``active_table_idx`` (a device tensor -- no host loop)."""
-        assert batch.active_table_idx is not None, "decode batch is missing its page-table rows"
-        return get_global_ctx().page_table.index_select(0, batch.active_table_idx.to(torch.int64))
+        assert batch.active_table_idx is not None, (
+            "decode batch is missing its page-table rows"
+        )
+        return get_global_ctx().page_table.index_select(
+            0, batch.active_table_idx.to(torch.int64)
+        )
 
     def _stage_decode(self, batch: Batch, bs: int, table_idx: torch.Tensor) -> None:
         """Copy this step's addressing into the static graph buffers and point the
         metadata at them (restage-per-replay, same shape as the generic backends)."""
         md = batch.attn_metadata
-        self._rows_buf[:bs].copy_(get_global_ctx().page_table.index_select(0, table_idx))
+        self._rows_buf[:bs].copy_(
+            get_global_ctx().page_table.index_select(0, table_idx)
+        )
         self._kvlen_buf[:bs].copy_(md.kv_len_cpu.to(self.device, non_blocking=True))
         md.rows = self._rows_buf[:bs]
         md.kvlen = self._kvlen_buf[:bs]
@@ -311,7 +367,9 @@ def prepare_for_capture(self, batch: Batch) -> None:
         self._stage_decode(batch, bs, dummy)
 
     def prepare_for_replay(self, batch: Batch) -> None:
-        assert batch.active_table_idx is not None, "decode batch is missing its page-table rows"
+        assert batch.active_table_idx is not None, (
+            "decode batch is missing its page-table rows"
+        )
         self._stage_decode(
             batch, batch.padded_size, batch.active_table_idx.to(torch.int64)
         )

python/freetoken/kernel/aot_models.py

@@ -269,6 +269,17 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
         moe_intermediate_size=2048,
         expert_formats=_NVFP4_FORMATS,
     ),
+    AotModel(
+        # Hybrid KDA + MLA sparse attention. KDA state and the MLA latent/index
+        # slabs bypass store_cache, so this family contributes no paged-KV row.
+        name="zai-org/GLM-5.3-Flash",
+        architecture="Glm5NextForConditionalGeneration",
+        hidden_size=4096,
+        kv_groups=(),
+        top_k=8,
+        moe_intermediate_size=2048,
+        expert_formats=("fp8_block",),
+    ),
     AotModel(
         # MiniMaxAI/MiniMax-M2.5 ships block-fp8, which has no expert-bank
         # provider for this arch on main -- the NVFP4 release is the servable
@@ -317,7 +328,13 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
         architecture="Qwen3_5ForConditionalGeneration",
         hidden_size=5120,
         kv_groups=((4, 256),),
-        aliases=("Qwen/Qwen3.6-27B-FP8", "nvidia/Qwen3.6-27B-NVFP4"),
+        aliases=(
+            "Qwen/Qwen3.6-27B-FP8",
+            "nvidia/Qwen3.6-27B-NVFP4",
+            "Qwen/Qwen3.8-27B",
+            "Qwen/Qwen3.8-27B-FP8",
+            "RadixArk/Qwen3.8-27B-NVFP4",
+        ),
     ),
     AotModel(
         name="google/gemma-4-12B-it",