From 711a86ef11c67bbe6c69ff59370594f64c6ef1b8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jes=C3=BAs=20J=C3=B3dar=20Piernas?= Date: Fri, 23 Oct 2026 11:46:35 +0200 Subject: [PATCH 1/2] feat(flint): add nvfp4 modelopt and gguf support --- docs/flint.md | 81 +++++++++++ docs/models.md | 1 + python/freetoken/kernel/aot_models.py | 11 ++ python/freetoken/models/flint/__init__.py | 9 ++ python/freetoken/models/flint/attention.py | 57 ++++++++ python/freetoken/models/flint/config.py | 66 +++++++++ python/freetoken/models/flint/gguf.py | 118 ++++++++++++++++ python/freetoken/models/flint/model.py | 65 +++++++++ python/freetoken/models/flint/moe.py | 37 +++++ python/freetoken/models/flint/weight.py | 77 +++++++++++ python/freetoken/models/gguf/config.py | 1 + python/freetoken/models/gguf/tokenizer.py | 17 ++- python/freetoken/models/register.py | 11 ++ python/freetoken/server/args.py | 6 +- python/freetoken/server/generation.py | 3 + python/freetoken/server/reasoning_parser.py | 1 + tests/README.md | 3 +- tests/e2e/test_aime.py | 37 +---- tests/e2e/test_cache_rebuild.py | 10 +- tests/models/test_flint.py | 130 ++++++++++++++++++ tests/server/test_parser_auto_selection.py | 5 + .../test_reasoning_parser_all_models.py | 20 +++ tests/server/test_streaming_model_matrix.py | 1 + 23 files changed, 727 insertions(+), 40 deletions(-) create mode 100644 docs/flint.md create mode 100644 python/freetoken/models/flint/__init__.py create mode 100644 python/freetoken/models/flint/attention.py create mode 100644 python/freetoken/models/flint/config.py create mode 100644 python/freetoken/models/flint/gguf.py create mode 100644 python/freetoken/models/flint/model.py create mode 100644 python/freetoken/models/flint/moe.py create mode 100644 python/freetoken/models/flint/weight.py create mode 100644 tests/models/test_flint.py diff --git a/docs/flint.md b/docs/flint.md new file mode 100644 index 0000000000..66c0935dbc --- /dev/null +++ b/docs/flint.md @@ -0,0 +1,81 @@ +# Recursive Flint + +FreeToken serves the Flint decoder with ModelOpt NVFP4 safetensors and the official +NVFP4 GGUF. Flint uses 49 layers (two dense, then 47 MoE), sigmoid top-6 routing +over 256 experts, an ungated shared expert, QK RMSNorm, learned attention sinks, +and hybrid full/sliding attention. Full-attention layers use NoPE; local layers +use RoPE and a 513-token window including the current token. + +Checkpoints: + +- [FutureMLS-Lab/flint-nvfp4-coldgptq](https://huggingface.co/FutureMLS-Lab/flint-nvfp4-coldgptq): + community ModelOpt export, approximately 25.16 GB (23.43 GiB) of tensors. Routed + experts are NVFP4; other weights remain BF16. +- [RecursiveSI/flint-nvfp4-gguf](https://huggingface.co/RecursiveSI/flint-nvfp4-gguf): + official approximately 22.9 GB GGUF, with more resident matrices quantized. + +`dgpp_mixed` FP8/NVFP4 exports and MLX-packed checkpoints use different layouts +and are not supported by this loader. The HF reader also accepts the original +BF16 layout, but full BF16 inference requires substantially more host memory and +was not validated on the 32 GB host used for the NVFP4 checks. + +## Run + +Start with one request and a short context on a 12 GB GPU. The routed expert +banks need about 20 GiB of host RAM, in addition to the process and resident +weights. The ModelOpt path uses `offload`, `cpu` or `hybrid` experts. These recipes +target TP=1. + +```bash +FLASHINFER_USE_CUDA_NORM=1 ft serve \ + --model FutureMLS-Lab/flint-nvfp4-coldgptq \ + --gpu 0 --moe-strategy hybrid --moe-hybrid-max-fetch 1 \ + --moe-cache-size 1024 --num-tokens 8192 \ + --max-seq-len-override 8192 --max-prefill-length 2048 \ + --max-running-requests 1 --host 127.0.0.1 --port 1919 +``` + +For the GGUF, download the file and use its path as `--model` with the same flags: + +```bash +hf download RecursiveSI/flint-nvfp4-gguf flint-nvfp4.gguf --local-dir models/flint-gguf +``` + +The reference SGLang kit explicitly selects FlashInfer's CUDA normalization +implementation. The environment variable above preserves that choice. Ampere +GPUs use FreeToken's compatible NVFP4 paths; they do not have Blackwell's native +W4A4 instructions. Memory fit and throughput therefore need local measurement. + +Recommended sampling defaults come from the checkpoint: temperature 1.0, +top-p 0.95 and top-k 20. The documented 262144-token context is a model limit; +the memory required to reserve it is separate from loading the weights. + +## Conversation protocol + +Use the checkpoint's chat template. The server selects the `glm47` tool parser and `flint` reasoning +parser automatically for HF and GGUF checkpoints. Flint chooses whether to emit +`...` on each turn; the parser never assumes the turn starts inside +reasoning, even when generic thinking flags or tools are present. Replay returned +`reasoning_content` and tool calls in the conversation history. + +Tool calls use the GLM-style `` / `` / `` envelope. +This integration uses FreeToken's existing parser and API behavior. It does not +port the separate Flint-specific XGrammar backend distributed in the SGLang kit. + +The GGUF pre-tokenizer is `flint`, with the original three regex splits and byte +level encoding. Its explicit EOG array contains four IDs: 150001, 150003, 150005, +and 150018. All four must stop generation, including the role boundaries. + +## Validate a local change + +The model tests compare fused sigmoid routing against independent tensor math and +check weight fusion across shards. With local checkpoints they also compare GGUF +tokenization, chat rendering and configuration against the original HF files. + +```bash +FREETOKEN_FLINT_MODEL=/path/to/flint-modelopt \ +FREETOKEN_FLINT_GGUF=/path/to/flint-nvfp4.gguf \ +uv run pytest tests/models/test_flint.py tests/server/test_parser_auto_selection.py \ + tests/server/test_reasoning_parser_all_models.py \ + tests/server/test_streaming_model_matrix.py +``` diff --git a/docs/models.md b/docs/models.md index 914e813bca..223f6642b0 100644 --- a/docs/models.md +++ b/docs/models.md @@ -20,6 +20,7 @@ for them; other checkpoints of the same architectures work too. | MiniMax-M2.5 | [nvidia/MiniMax-M2.5-NVFP4](https://huggingface.co/nvidia/MiniMax-M2.5-NVFP4) | | MiniMax-M3 | [nvidia/MiniMax-M3-NVFP4](https://huggingface.co/nvidia/MiniMax-M3-NVFP4) | | Muse-Glimmer | [meta-models/Muse-Glimmer-30B](https://huggingface.co/meta-models/Muse-Glimmer-30B), [RedHatAI/Muse-Glimmer-30B-NVFP4](https://huggingface.co/RedHatAI/Muse-Glimmer-30B-NVFP4) | +| [Flint](flint.md) | [FutureMLS-Lab/flint-nvfp4-coldgptq](https://huggingface.co/FutureMLS-Lab/flint-nvfp4-coldgptq), [RecursiveSI/flint-nvfp4-gguf](https://huggingface.co/RecursiveSI/flint-nvfp4-gguf) | ### Image input diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py index bf7e75b6e8..bc7b43707a 100644 --- a/python/freetoken/kernel/aot_models.py +++ b/python/freetoken/kernel/aot_models.py @@ -347,6 +347,17 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int moe_intermediate_size=2304, expert_formats=("ds_fp4",), ), + AotModel( + name="RecursiveSI/flint", + architecture="FlintForCausalLM", + hidden_size=2048, + kv_groups=((4, 128),), + top_k=6, + moe_intermediate_size=512, + expert_formats=("bf16", *_NVFP4_FORMATS), + aliases=("FutureMLS-Lab/flint-nvfp4-coldgptq",), + arch_aliases=("FlintGGUFForCausalLM",), + ), # ---- dense checkpoints (store/index only, no expert banks) ---- AotModel( name="Qwen/Qwen3.6-27B", diff --git a/python/freetoken/models/flint/__init__.py b/python/freetoken/models/flint/__init__.py new file mode 100644 index 0000000000..7fa0503006 --- /dev/null +++ b/python/freetoken/models/flint/__init__.py @@ -0,0 +1,9 @@ +from .config import parse_config +from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf +from .model import FlintForCausalLM +from .weight import iter_weights, nvfp4_expert_spec + +__all__ = [ + "FlintForCausalLM", "parse_config", "iter_weights", "nvfp4_expert_spec", + "iter_gguf_weights", "parse_gguf_config", "to_gguf", +] diff --git a/python/freetoken/models/flint/attention.py b/python/freetoken/models/flint/attention.py new file mode 100644 index 0000000000..6a720cc777 --- /dev/null +++ b/python/freetoken/models/flint/attention.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +import torch + +from freetoken.attention import AttentionSpec +from freetoken.core import get_global_ctx +from freetoken.distributed import get_tp_info +from freetoken.layers import BaseOP, LinearOProj, LinearQKVMerged, RMSNorm +from freetoken.layers.rotary import get_rope +from freetoken.models.config import ModelConfig, SWAAttentionGroupConfig +from freetoken.utils import div_even, nvtx_annotate + + +class FlintAttention(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): + self.layer_id = layer_id + group = config.attention_group_for_layer(layer_id) + self.num_qo_heads = div_even(config.num_qo_heads, get_tp_info().size) + self.num_kv_heads = div_even(config.num_kv_heads, get_tp_info().size, allow_replicate=True) + self.head_dim = config.head_dim + self.q_dim = self.num_qo_heads * self.head_dim + self.kv_dim = self.num_kv_heads * self.head_dim + self.qkv_proj = LinearQKVMerged( + config.hidden_size, self.head_dim, config.num_qo_heads, config.num_kv_heads, + has_bias=False, quant_config=config.quant, prefix=f"{prefix}.qkv_proj", + ) + self.o_proj = LinearOProj( + config.num_qo_heads * self.head_dim, config.hidden_size, + has_bias=False, quant_config=config.quant, prefix=f"{prefix}.o_proj", + ) + self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps) + self.softmax_offset = torch.empty(self.num_qo_heads, dtype=torch.float32) + self.sliding_window = group.sliding_window if isinstance(group, SWAAttentionGroupConfig) else None + rope = group.rotary_config + self.rotary = get_rope( + head_dim=self.head_dim, rotary_dim=rope.rotary_dim, + max_position=rope.max_position, base=rope.base, rope_scaling=None, + ) if rope.rotary_dim else None + + @nvtx_annotate("MHA") + def forward(self, x: torch.Tensor) -> torch.Tensor: + ctx = get_global_ctx() + qkv = self.qkv_proj.forward(x) + q, k, v = qkv.split((self.q_dim, self.kv_dim, self.kv_dim), dim=-1) + # Flatten heads for FlashInfer's 2D CUDA norm; normalizing the concatenated + # heads would silently change both the attention and the MoE routes. + q = self.q_norm.forward(q.reshape(-1, self.head_dim)).reshape(-1, self.q_dim) + k = self.k_norm.forward(k.reshape(-1, self.head_dim)).reshape(-1, self.kv_dim) + if self.rotary is not None: + q, k = self.rotary.forward(ctx.batch.get_attn_positions(), q, k) + out = ctx.attn_backend.forward( + q.view(-1, self.num_qo_heads, self.head_dim), k.contiguous(), v.contiguous(), + self.layer_id, ctx.batch, + attn_spec=AttentionSpec(sliding_window=self.sliding_window, sinks=self.softmax_offset), + ) + return self.o_proj.forward(out.reshape(-1, self.q_dim)) diff --git a/python/freetoken/models/flint/config.py b/python/freetoken/models/flint/config.py new file mode 100644 index 0000000000..916e8829cb --- /dev/null +++ b/python/freetoken/models/flint/config.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +from dataclasses import replace +from typing import Any + +from freetoken.models.config import ( + FullAttentionGroupConfig, + ModelConfig, + RotaryConfig, + SWAAttentionGroupConfig, + detect_expert_quant, +) + + +def parse_config(hf: Any) -> ModelConfig: + layers = int(hf.num_hidden_layers) + types = tuple(hf.layer_types) + if len(types) != layers or set(types) - {"full_attention", "sliding_attention"}: + raise ValueError("Flint requires one full_attention/sliding_attention type per layer") + full = tuple(i for i, kind in enumerate(types) if kind == "full_attention") + swa = tuple(i for i, kind in enumerate(types) if kind == "sliding_attention") + if set(hf.no_rope_layer_indices) != set(full): + raise ValueError("Flint supports RoPE on sliding layers and NoPE on full layers only") + dense = tuple(hf.mlp_only_layers) + if dense != tuple(range(len(dense))) or len(dense) >= layers: + raise ValueError("Flint dense MLP layers must form a leading contiguous prefix") + if int(getattr(hf, "decoder_sparse_step", 1)) != 1: + raise ValueError("Flint requires an MoE block in every layer after the dense prefix") + if getattr(hf, "rope_scaling", None): + raise NotImplementedError("Flint RoPE scaling is not supported") + if hf.hidden_act != "silu" or hf.router_score_function != "sigmoid": + raise ValueError("Flint requires SwiGLU experts and sigmoid routing") + if hf.softmax_type != "learnable" or bool(getattr(hf, "attention_bias", False)): + raise ValueError("Flint requires learned attention sinks and bias-free projections") + window = int(hf.sliding_window) + if window <= 0 or not full or not swa: + raise ValueError("Flint requires full and sliding attention with a positive window") + rope = RotaryConfig( + head_dim=int(hf.head_dim), rotary_dim=int(hf.head_dim), + max_position=int(hf.max_position_embeddings), base=float(hf.rope_theta), scaling=None, + ) + return ModelConfig( + num_layers=layers, num_qo_heads=int(hf.num_attention_heads), + num_kv_heads=int(hf.num_key_value_heads), head_dim=int(hf.head_dim), + hidden_size=int(hf.hidden_size), vocab_size=int(hf.vocab_size), + intermediate_size=int(hf.intermediate_size), hidden_act="silu", + rms_norm_eps=float(hf.rms_norm_eps), + tie_word_embeddings=bool(hf.tie_word_embeddings), rotary_config=rope, + num_experts=int(hf.num_experts), num_experts_per_tok=int(hf.num_experts_per_tok), + moe_intermediate_size=int(hf.moe_intermediate_size), + norm_topk_prob=bool(hf.norm_topk_prob), + model_type="flint", architectures=list(hf.architectures), moe_enabled=True, + expert_quant=detect_expert_quant(hf), first_k_dense_replace=len(dense), + shared_expert_intermediate_size=int(hf.shared_expert_intermediate_size), + use_qk_norm=True, + attention_groups=( + SWAAttentionGroupConfig( + name="swa", layer_ids=swa, num_kv_heads=int(hf.num_key_value_heads), + head_dim=int(hf.head_dim), rotary_config=rope, sliding_window=window, + ), + FullAttentionGroupConfig( + name="full", layer_ids=full, num_kv_heads=int(hf.num_key_value_heads), + head_dim=int(hf.head_dim), rotary_config=replace(rope, rotary_dim=0), + ), + ), + ) diff --git a/python/freetoken/models/flint/gguf.py b/python/freetoken/models/flint/gguf.py new file mode 100644 index 0000000000..3cc38a0ed1 --- /dev/null +++ b/python/freetoken/models/flint/gguf.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +import re +from dataclasses import replace +from types import SimpleNamespace +from typing import Iterator + +import torch + +from freetoken.models.gguf.reader import gguf_tensor_index +from freetoken.models.gguf.config import GgufConfigShim +from freetoken.models.config import ModelConfig +from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix +from freetoken.utils import cached_load_hf_config + +from .config import parse_config + +_LAYER = re.compile(r"^model\.layers\.(\d+)\.(.+)$") +_LEAVES = { + "self_attn.q_proj": "attn_q", "self_attn.k_proj": "attn_k", + "self_attn.v_proj": "attn_v", "self_attn.o_proj": "attn_output", + "mlp.gate": "ffn_gate_inp", + **{f"mlp.{p}_proj": f"ffn_{p}" for p in ("gate", "up", "down")}, + **{f"mlp.shared_expert.{p}_proj": f"ffn_{p}_shexp" for p in ("gate", "up", "down")}, + **{f"mlp.experts.0.{p}_proj": f"ffn_{p}_exps" for p in ("gate", "up", "down")}, +} + + +def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: + top = {"model.embed_tokens": "token_embd", "lm_head": "output"} + if name in top: + return (top[name],) + match = _LAYER.match(name) + if match is None: + return () + leaf = _LEAVES.get(match[2]) + return (f"blk.{match[1]}.{leaf}",) if leaf else () + + +def _hf_config(shim: GgufConfigShim) -> SimpleNamespace: + md = shim.metadata + def kv(name): + return md[f"flint.{name}"] + if md.get("flint.rope.scaling.type", "none") != "none": + raise NotImplementedError("Flint GGUF RoPE scaling is not supported") + if int(kv("expert_gating_func")) != 2 or float(kv("expert_weights_scale")) != 1.0: + raise NotImplementedError("Flint GGUF requires unscaled sigmoid routing") + check_scales(shim.model_path) + swa = tuple(kv("attention.sliding_window_pattern")) + rope = tuple(kv("attention.rope_pattern")) + if len(rope) != int(kv("block_count")): + raise ValueError("Flint GGUF requires one RoPE flag per layer") + return SimpleNamespace( + architectures=list(shim.architectures), model_type="flint", + num_hidden_layers=int(kv("block_count")), hidden_size=int(kv("embedding_length")), + intermediate_size=int(kv("feed_forward_length")), + num_attention_heads=int(kv("attention.head_count")), + num_key_value_heads=int(kv("attention.head_count_kv")), + head_dim=int(kv("attention.key_length")), vocab_size=shim.vocab_size, + max_position_embeddings=int(kv("context_length")), + rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")), + rope_theta=float(kv("rope.freq_base")), + layer_types=["sliding_attention" if v else "full_attention" for v in swa], + no_rope_layer_indices=[i for i, v in enumerate(rope) if not v], + mlp_only_layers=list(range(int(kv("leading_dense_block_count")))), + num_experts=int(kv("expert_count")), num_experts_per_tok=int(kv("expert_used_count")), + moe_intermediate_size=int(kv("expert_feed_forward_length")), + shared_expert_intermediate_size=int(kv("expert_shared_feed_forward_length")), + norm_topk_prob=bool(kv("expert_weights_norm")), + sliding_window=int(kv("attention.sliding_window")), + hidden_act="silu", router_score_function="sigmoid", softmax_type="learnable", + tie_word_embeddings=shim.tie_word_embeddings, + ) + + +def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig: + return replace(parse_config(_hf_config(shim)), expert_quant="gguf") + + +def iter_gguf_weights( + model_path: str, device: torch.device, *, include_moe_experts: bool, include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + if not include_non_moe: + return + hf = _hf_config(cached_load_hf_config(model_path)) + index = gguf_tensor_index(model_path) + def t(name): + return index[name if name in index else f"{name}.weight"] + def linear(module, names): + infos = [t(n) for n in names] + yield from fused(module, infos, [matrix(v) for v in infos], int(infos[0].ne[0])) + yield from linear("model.embed_tokens", ["token_embd"]) + if not hf.tie_word_embeddings: + yield from linear("lm_head", ["output"]) + yield "model.norm.weight", float_tensor(t("output_norm")) + for layer in range(hf.num_hidden_layers): + blk, pre = f"blk.{layer}", f"model.layers.{layer}" + yield f"{pre}.input_layernorm.weight", float_tensor(t(f"{blk}.attn_norm")) + yield f"{pre}.post_attention_layernorm.weight", float_tensor(t(f"{blk}.ffn_norm")) + sa = f"{pre}.self_attn" + yield from linear(f"{sa}.qkv_proj", [f"{blk}.attn_{p}" for p in ("q", "k", "v")]) + yield from linear(f"{sa}.o_proj", [f"{blk}.attn_output"]) + for p in ("q", "k"): + yield f"{sa}.{p}_norm.weight", float_tensor(t(f"{blk}.attn_{p}_norm")) + yield f"{sa}.softmax_offset", float_tensor(t(f"{blk}.attn_sinks")).float() + mlp = f"{pre}.mlp" + if layer in hf.mlp_only_layers: + yield from linear(f"{mlp}.gate_up_proj", [f"{blk}.ffn_{p}" for p in ("gate", "up")]) + yield from linear(f"{mlp}.down_proj", [f"{blk}.ffn_down"]) + else: + yield f"{mlp}.gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp")) + bias = index.get(f"{blk}.exp_probs_b.bias") + yield f"{mlp}.gate.expert_bias", ( + float_tensor(bias).float() if bias is not None + else torch.zeros(hf.num_experts, device=device, dtype=torch.float32) + ) + yield from linear(f"{mlp}.shared_expert.gate_up_proj", [f"{blk}.ffn_{p}_shexp" for p in ("gate", "up")]) + yield from linear(f"{mlp}.shared_expert.down_proj", [f"{blk}.ffn_down_shexp"]) diff --git a/python/freetoken/models/flint/model.py b/python/freetoken/models/flint/model.py new file mode 100644 index 0000000000..3f0b4208b5 --- /dev/null +++ b/python/freetoken/models/flint/model.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +import torch + +from freetoken.core import get_global_ctx +from freetoken.layers import BaseOP, OPList, ParallelLMHead, RMSNormFused, VocabParallelEmbedding +from freetoken.models.blocks import BaseLLMModel, GatedMLP +from freetoken.models.config import ModelConfig +from freetoken.utils import nvtx_annotate + +from .attention import FlintAttention +from .moe import FlintMoE + + +class FlintDecoderLayer(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): + self._layer_id = layer_id + self.self_attn = FlintAttention(config, layer_id, prefix=f"{prefix}.self_attn") + self.mlp = ( + GatedMLP(config, quant_config=config.quant, prefix=f"{prefix}.mlp") + if layer_id < config.first_k_dense_replace + else FlintMoE(config, layer_id, prefix=f"{prefix}.mlp") + ) + self.input_layernorm = RMSNormFused(config.hidden_size, config.rms_norm_eps) + self.post_attention_layernorm = RMSNormFused(config.hidden_size, config.rms_norm_eps) + + @nvtx_annotate("Layer_{}", layer_id_field="_layer_id") + def forward(self, x: torch.Tensor, residual: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor]: + x, residual = self.input_layernorm.forward(x, residual) + x = self.self_attn.forward(x) + x, residual = self.post_attention_layernorm.forward(x, residual) + return self.mlp.forward(x), residual + + +class FlintModel(BaseOP): + def __init__(self, config: ModelConfig, *, prefix: str = "model"): + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, config.hidden_size, quant_config=config.quant, prefix=f"{prefix}.embed_tokens", + ) + self.layers = OPList([ + FlintDecoderLayer(config, i, prefix=f"{prefix}.layers.{i}") for i in range(config.num_layers) + ]) + self.norm = RMSNormFused(config.hidden_size, config.rms_norm_eps) + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + x, residual = self.embed_tokens.forward(input_ids), None + for layer in self.layers.op_list: + x, residual = layer.forward(x, residual) + return self.norm.forward(x, residual)[0] + + +class FlintForCausalLM(BaseLLMModel): + def __init__(self, config: ModelConfig): + self.model = FlintModel(config) + self.lm_head = ParallelLMHead( + config.vocab_size, config.hidden_size, + tie_word_embeddings=config.tie_word_embeddings, + tied_embedding=self.model.embed_tokens if config.tie_word_embeddings else None, + quant_config=config.quant, prefix="lm_head", + ) + super().__init__() + + def forward(self) -> torch.Tensor: + output = self.model.forward(get_global_ctx().batch.input_ids) + return self.lm_head.forward(output) diff --git a/python/freetoken/models/flint/moe.py b/python/freetoken/models/flint/moe.py new file mode 100644 index 0000000000..4f7bf7bfdf --- /dev/null +++ b/python/freetoken/models/flint/moe.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from dataclasses import replace + +import torch + +from freetoken.kernel.triton.moe_router import fused_topk_sigmoid +from freetoken.layers import BaseOP, LinearReplicated, make_moe_layer +from freetoken.models.blocks import GatedMLP +from freetoken.models.config import ModelConfig + + +class FlintMoE(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): + self.top_k = config.num_experts_per_tok + self.renormalize = config.norm_topk_prob + self.gate = LinearReplicated(config.hidden_size, config.num_experts, has_bias=False) + self.gate.expert_bias = torch.zeros(config.num_experts, dtype=torch.float32) + self.experts = make_moe_layer( + config, layer_id=layer_id - config.first_k_dense_replace, + quant_config=config.quant, prefix=f"{prefix}.experts", + ) + shared_config = replace(config, intermediate_size=config.shared_expert_intermediate_size) + self.shared_expert = GatedMLP(shared_config, quant_config=config.quant, prefix=f"{prefix}.shared_expert") + + def _route(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + # The reference projects in model dtype, then computes sigmoid in fp32. + return fused_topk_sigmoid( + self.gate.forward(x).float(), self.gate.expert_bias, + self.top_k, self.renormalize, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + weights, ids = self._route(x) + # Some routed kernels overwrite their input; the shared branch needs it first. + shared = self.shared_expert.forward(x) + return self.experts.routed_forward(x, weights, ids) + shared diff --git a/python/freetoken/models/flint/weight.py b/python/freetoken/models/flint/weight.py new file mode 100644 index 0000000000..d867274ed9 --- /dev/null +++ b/python/freetoken/models/flint/weight.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +import re +from typing import Iterator + +import safetensors +import torch + +from freetoken.distributed import get_tp_info +from freetoken.models.loader import ( + MergeRule, drop_page_cache, iter_merged_tensors, iter_stacked_experts, + iter_weight_files, safetensors_weight_map, shard_tensor, +) +from freetoken.models.nvfp4_banks import Nvfp4ExpertSourceSpec +from freetoken.utils import cached_load_hf_config, download_hf_weight + +from .config import parse_config + +_EXPERT = re.compile(r"^(?Pmodel\.layers\.\d+\.mlp\.experts)\.(?P\d+)\.(?P.+)$") +_MERGE = { + ".q_proj": MergeRule(".qkv_proj", "q", ("q", "k", "v")), + ".k_proj": MergeRule(".qkv_proj", "k", ("q", "k", "v")), + ".v_proj": MergeRule(".qkv_proj", "v", ("q", "k", "v")), + ".gate_proj": MergeRule(".gate_up_proj", "gate", ("gate", "up")), + ".up_proj": MergeRule(".gate_up_proj", "up", ("gate", "up")), +} +_NVFP4 = Nvfp4ExpertSourceSpec( + key_pattern=re.compile( + r"^model\.layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." + r"(?Pgate_proj|up_proj|down_proj)\.(?Pweight|weight_scale|weight_scale_2)$" + ), + proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, + layer_to_bank=lambda layer, config: ( + layer - config.first_k_dense_replace if config.first_k_dense_replace <= layer < config.num_layers else None + ), + desc="Flint NVFP4 experts", +) + + +def nvfp4_expert_spec(model_path: str, config) -> Nvfp4ExpertSourceSpec: + return _NVFP4 + + +def iter_weights( + model_path: str, device: torch.device, *, include_moe_experts: bool, include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + model_path = download_hf_weight(model_path) + config = parse_config(cached_load_hf_config(model_path)) + if include_moe_experts and config.expert_quant != "none": + raise ValueError("Flint NVFP4 uses the offload, cpu or hybrid MoE strategy") + tp = get_tp_info() + names = safetensors_weight_map(model_path) + + def tensors(): + for file in iter_weight_files(model_path): + with safetensors.safe_open(file, framework="pt", device=str(device)) as reader: + for name in reader.keys(): + expert = _EXPERT.match(name) is not None + if (expert and not include_moe_experts) or (not expert and not include_non_moe): + continue + raw = reader.get_tensor(name) + if name.endswith(".softmax_offset"): + raw = raw.chunk(tp.size)[tp.rank].float() + elif name.endswith(".expert_bias"): + raw = raw.float() + else: + raw = shard_tensor(name, raw, rank=tp.rank, world_size=tp.size, num_kv_heads=config.num_kv_heads) + yield name, raw + drop_page_cache(file) + if include_non_moe: + for layer in range(config.first_k_dense_replace, config.num_layers): + name = f"model.layers.{layer}.mlp.gate.expert_bias" + if name not in names: + yield name, torch.zeros(config.num_experts, dtype=torch.float32, device=device) + + merged = iter_merged_tensors(tensors(), _MERGE, model_name="flint") + yield from iter_stacked_experts(merged, num_experts=config.num_experts, model_name="flint", expert_pattern=_EXPERT) diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py index 61198ead29..e718809f5e 100644 --- a/python/freetoken/models/gguf/config.py +++ b/python/freetoken/models/gguf/config.py @@ -20,6 +20,7 @@ # GGUF ``general.architecture`` -> FreeToken registry key (a GGUF-specific spec that # reuses the model classes but a GGUF parse_config / iter_weights). GGUF_ARCH_TO_REGISTRY: dict[str, str] = { + "flint": "FlintGGUFForCausalLM", "deepseek4": "DeepseekV4GGUFForCausalLM", "gemma4": "Gemma4GGUFForCausalLM", "glm4moe": "Glm4MoeGGUFForCausalLM", diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py index a9c77d7999..c78428dc9b 100644 --- a/python/freetoken/models/gguf/tokenizer.py +++ b/python/freetoken/models/gguf/tokenizer.py @@ -10,7 +10,8 @@ from .reader import load_gguf_metadata -# end-of-generation token names from llama.cpp (llama_vocab::impl::load in src/llama-vocab.cpp): a GGUF stores one eos id and no stop list +# End-of-generation names for GGUFs without an explicit eog_token_ids array. +# See llama.cpp llama_vocab::impl::load in src/llama-vocab.cpp. _EOG_NAMES = frozenset(( "<|eot_id|>", "<|im_end|>", "<|end|>", "<|return|>", "<|call|>", "<|flush|>", "<|calls|>", "", "<|endoftext|>", "", "<|eom_id|>", "", "_", "[EOT]", "[EOS]", "<|end_of_text|>", "", @@ -33,6 +34,15 @@ r'''[!"#$%&'()*+,\-./:;<=>?@\[\\\]^_`{|}~][A-Za-z]+|[^\r\n\p{L}\p{P}\p{S}]?[\p{L}\p{M}]+| ?[\p{P}\p{S}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+''', ) +# RecursiveSI/flint's tokenizer.json applies these splits in sequence. +_FLINT_SPLITS = ( + r"\p{N}{1,3}", + "[一-龥぀-ゟ゠-ヿ]+", + r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?| ?[\p{P}\p{S}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+", +) + +_CUSTOM_SPLITS = {"joyai-llm": _JOYAI_LLM_SPLITS, "flint": _FLINT_SPLITS} + def _add_bos(meta: dict) -> bool: if "tokenizer.ggml.add_bos_token" in meta: @@ -105,11 +115,11 @@ def load_gguf_tokenizer(model_path: str): from tokenizers import normalizers fast.normalizer = normalizers.NFC() - elif pre == "joyai-llm": + elif pre in _CUSTOM_SPLITS: from tokenizers import Regex, pre_tokenizers fast.pre_tokenizer = pre_tokenizers.Sequence( - [pre_tokenizers.Split(Regex(r), behavior="isolated") for r in _JOYAI_LLM_SPLITS] + [pre_tokenizers.Split(Regex(r), behavior="isolated") for r in _CUSTOM_SPLITS[pre]] + [pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False)] ) if pre in _IGNORE_MERGES_PRE: @@ -129,6 +139,7 @@ def gguf_eos_token_ids(model_path: str) -> set[int]: tokens = meta["tokenizer.ggml.tokens"] ids = {i for i, t in enumerate(tokens) if t in _EOG_NAMES} ids |= {int(meta[f"tokenizer.ggml.{k}_token_id"]) for k in ("eos", "eot", "eom") if f"tokenizer.ggml.{k}_token_id" in meta} + ids |= {int(i) for i in meta.get("tokenizer.ggml.eog_token_ids", ())} names = {tokens[i]: i for i in ids if i < len(tokens)} # llama.cpp's exceptions: <|end|> only ends a message when <|call|> and <|return|> end the turn (gpt-oss), and gemma4's is a normal token if "<|end|>" in names and names.keys() & {"<|call|>", "<|calls|>"} and names.keys() & {"<|return|>", "<|flush|>"}: diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index ad8ade862c..077b7e2dad 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -95,6 +95,17 @@ class ModelSpec: _DSV4_ENCODERS = (EncoderSpec("vision", "vision_n_layers", ("image",)),) _MODEL_REGISTRY: dict[str, ModelSpec] = { + "FlintGGUFForCausalLM": ModelSpec( + "freetoken.models.flint", "FlintForCausalLM", + parse_config="parse_gguf_config", iter_weights="iter_gguf_weights", + packed_modules_mapping=_DENSE_PACKED + _EXPERTS_PACKED, + unquantized_modules=("*.mlp.gate",), + ), + "FlintForCausalLM": ModelSpec( + "freetoken.models.flint", "FlintForCausalLM", + packed_modules_mapping=_DENSE_PACKED + _EXPERTS_PACKED, + unquantized_modules=("*.mlp.gate",), + ), "LlamaForCausalLM": ModelSpec( "freetoken.models.llama", "LlamaForCausalLM", diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index efe3b8c0bd..23144d6bd0 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -206,6 +206,8 @@ def _infer_tool_call_parser(model_path: str) -> str: " ".join(str(v) for v in text_cfg.get("architectures", []) or []), ] marker = " ".join(candidates).lower() + if "flint" in marker: + return "glm47" if "gpt_oss" in marker or "gpt-oss" in marker or "gptoss" in marker: return "gpt_oss" # M3 first: its marker also contains the bare "minimax" substring, but the @@ -257,6 +259,8 @@ def _infer_reasoning_parser(model_path: str) -> str | None: " ".join(str(v) for v in text_cfg.get("architectures", []) or []), ] marker = " ".join(candidates).lower() + if "flint" in marker: + return "flint" if "gpt_oss" in marker or "gpt-oss" in marker or "gptoss" in marker: return "gpt_oss" if "deepseek" in marker and any( @@ -638,7 +642,7 @@ def _infer_reasoning_parser(model_path: str) -> str | None: default="auto", choices=[ "auto", "off", "deepseekv32", "gpt_oss", "qwen3", "glm", - "minimax", "minimax_m3", "muse_glimmer", "gemma4", + "minimax", "minimax_m3", "muse_glimmer", "gemma4", "flint", ], help=( "Reasoning parser that splits chain-of-thought into reasoning_content " diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index 2e342ff899..ef90402aab 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -410,6 +410,9 @@ def _make_reasoning_parser(spec: GenSpec, state: Any) -> ReasoningParser | None: or bool(ctk.get("enable_thinking")) or bool(ctk.get("thinking")) ) + elif parser_name == "flint": + # Flint opens its own optional think block, including when tools are present. + force_reasoning = False elif parser_name == "minimax_m3": # M3's template pre-opens only in thinking_mode "enabled" (the # model then emits just the closing tag); "adaptive" (default) leaves the diff --git a/python/freetoken/server/reasoning_parser.py b/python/freetoken/server/reasoning_parser.py index 6788675ca9..411f9bfeed 100644 --- a/python/freetoken/server/reasoning_parser.py +++ b/python/freetoken/server/reasoning_parser.py @@ -874,6 +874,7 @@ class ReasoningParser: """Wraps a reasoning detector for streaming and non-streaming use.""" ReasoningParserEnum: Dict[str, Type[BaseReasoningParser]] = { + "flint": ThinkReasoningParser, "deepseekv32": DeepSeekV32ReasoningParser, "gpt_oss": GptOssHarmonyReasoningParser, "qwen3": ThinkReasoningParser, diff --git a/tests/README.md b/tests/README.md index 8b2ec8854b..3cd4cc9045 100644 --- a/tests/README.md +++ b/tests/README.md @@ -68,7 +68,8 @@ checkpoint is set: | `FREETOKEN_MINIMAX_M3_MODEL` | `tokenizer/test_mm_tokenize.py` — local MiniMax-M3 checkpoint (start/end expansion, token budget) | | `FREETOKEN_TEST_MOE_CACHE_SIZE` | `e2e/test_aime.py` — >0 switches to the offload MoE backend with this cache size | | `FREETOKEN_TEST_MEM_RATIO` | `e2e/test_aime.py` — offload-mode memory_ratio (default `0.9`) | -| `FREETOKEN_REBUILD_TEST_MODEL` | `e2e/test_cache_rebuild.py` — a SMALL local model dir; boots a real server (falls back to `FREETOKEN_TEST_MODEL`) | +| `FREETOKEN_REBUILD_TEST_MODEL` | `e2e/test_cache_rebuild.py` — local checkpoint directory or GGUF file; boots a real server (falls back to `FREETOKEN_TEST_MODEL`) | +| `FREETOKEN_REBUILD_EXTRA_ARGS` | `e2e/test_cache_rebuild.py` — extra server flags, e.g. `--moe-strategy hybrid --moe-cache-size 1024 --max-running-requests 1` for offloaded MoE | | `FREETOKEN_GEMMA4_GGUF_GLOB` | `models/test_gemma4_gguf_rope.py` — glob matching a local gemma-4 GGUF file | | `FREETOKEN_TEST_MODELS_ROOT` | `models/test_quant_config.py` — directory whose checkpoints the stored-tensor scheme scan checks (default `/mnt/nvme/models`) | | `FREETOKEN_DSV4_VL_MODEL` | `tokenizer/test_mm_tokenize.py` - local DeepSeek-V4 vision checkpoint | diff --git a/tests/e2e/test_aime.py b/tests/e2e/test_aime.py index 7507039668..8d0b41547f 100644 --- a/tests/e2e/test_aime.py +++ b/tests/e2e/test_aime.py @@ -198,41 +198,14 @@ def build_llm(model_path: Path) -> LLM: return LLM(**kwargs) -def visible_gpu_index() -> int: - visible = os.environ.get("CUDA_VISIBLE_DEVICES") - if not visible: - return 0 - first = visible.split(",", 1)[0].strip() - return int(first) - - def free_gpu_memory_gib() -> float: - if torch.version.hip is not None: - # Query logical device zero without initializing the engine's parent process. - result = subprocess.run( - [sys.executable, "-c", "import torch; print(torch.cuda.mem_get_info(0)[0])"], - check=True, - text=True, - capture_output=True, - ) - return int(result.stdout.strip()) / (1 << 30) - + # A child respects UUID visibility and CUDA_DEVICE_ORDER without initializing + # the engine's parent process, on both CUDA and ROCm. result = subprocess.run( - [ - "nvidia-smi", - "--query-gpu=index,memory.free", - "--format=csv,noheader,nounits", - ], - check=True, - text=True, - capture_output=True, + [sys.executable, "-c", "import torch; print(torch.cuda.mem_get_info(0)[0])"], + check=True, text=True, capture_output=True, ) - target = visible_gpu_index() - for line in result.stdout.splitlines(): - index, free_mib = [part.strip() for part in line.split(",", 1)] - if int(index) == target: - return int(free_mib) / 1024 - raise RuntimeError(f"GPU {target} is not visible in nvidia-smi") + return int(result.stdout.strip()) / (1 << 30) @pytest.mark.skipif(not torch.cuda.is_available(), reason="AIME e2e needs CUDA") diff --git a/tests/e2e/test_cache_rebuild.py b/tests/e2e/test_cache_rebuild.py index b17fc28ba9..8e6632a179 100644 --- a/tests/e2e/test_cache_rebuild.py +++ b/tests/e2e/test_cache_rebuild.py @@ -16,12 +16,14 @@ FREETOKEN_REBUILD_TEST_MODEL small local model dir (falls back to FREETOKEN_TEST_MODEL) FREETOKEN_REBUILD_MIN_FREE_GIB free-GPU-memory gate (default 20) FREETOKEN_REBUILD_BOOT_TIMEOUT seconds to wait for "serving" (default 300) + FREETOKEN_REBUILD_EXTRA_ARGS additional ft serve flags (e.g. offloaded MoE) """ from __future__ import annotations import json import os +import shlex import socket import subprocess import sys @@ -92,7 +94,7 @@ def _wait_until_serving(base: str, proc: subprocess.Popen, deadline: float) -> N def _generate(base: str) -> str: - """A short non-thinking completion — the point is that the engine still runs, not what it says.""" + """A short completion; reasoning models may spend the whole budget in their thought channel.""" code, body = _post( base, "/v1/chat/completions", @@ -105,7 +107,8 @@ def _generate(base: str) -> str: }, ) assert code == 200, body - return str(body["choices"][0]["message"]["content"]) + message = body["choices"][0]["message"] + return message.get("content") or message.get("reasoning_content") or "" @pytest.mark.skipif(not torch.cuda.is_available(), reason="cache rebuild e2e needs CUDA") @@ -113,7 +116,7 @@ def test_cache_rebuild_resizes_a_live_engine_and_keeps_serving(tmp_path): model_dir = _model_dir() if model_dir is None: pytest.skip("set FREETOKEN_REBUILD_TEST_MODEL to a small local model directory") - if not model_dir.is_dir(): + if not model_dir.is_dir() and not model_dir.is_file(): pytest.skip(f"model is not downloaded: {model_dir}") min_free = float(os.environ.get("FREETOKEN_REBUILD_MIN_FREE_GIB", "20")) @@ -133,6 +136,7 @@ def test_cache_rebuild_resizes_a_live_engine_and_keeps_serving(tmp_path): "--port", str(port), "--num-pages", str(BOOT_PAGES), "--enable-cache-report", + *shlex.split(os.environ.get("FREETOKEN_REBUILD_EXTRA_ARGS", "")), ], stdout=log, stderr=subprocess.STDOUT, diff --git a/tests/models/test_flint.py b/tests/models/test_flint.py new file mode 100644 index 0000000000..73d9131235 --- /dev/null +++ b/tests/models/test_flint.py @@ -0,0 +1,130 @@ +"""Flint contracts checked against independent tensor math and checkpoint layouts.""" +from __future__ import annotations + +import json +import os +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch +import torch.nn.functional as F + +from freetoken.models.flint.config import parse_config +from freetoken.models.flint.moe import FlintMoE + + +@pytest.fixture(autouse=True) +def _single_rank(monkeypatch): + from freetoken.distributed import info + monkeypatch.setattr(info, "_TP_INFO", info.DistributedInfo(rank=0, size=1)) + + +def _hf(**overrides): + return SimpleNamespace(**(dict( + architectures=["FlintForCausalLM"], model_type="flint", num_hidden_layers=3, + hidden_size=64, head_dim=128, num_attention_heads=4, num_key_value_heads=2, + intermediate_size=128, moe_intermediate_size=64, + shared_expert_intermediate_size=128, vocab_size=128, max_position_embeddings=2048, + layer_types=["full_attention", "sliding_attention", "full_attention"], + no_rope_layer_indices=[0, 2], mlp_only_layers=[0], + hidden_act="silu", router_score_function="sigmoid", softmax_type="learnable", + norm_topk_prob=True, num_experts=8, num_experts_per_tok=2, + rms_norm_eps=1e-6, sliding_window=513, rope_theta=10000., tie_word_embeddings=False, + ) | overrides)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="fused CUDA router") +@pytest.mark.parametrize("renormalize", [False, True]) +def test_sigmoid_routing_bias_affects_selection_only(renormalize): + from freetoken.layers import LinearReplicated + torch.manual_seed(42) + block = FlintMoE.__new__(FlintMoE) + block.top_k, block.renormalize = 6, renormalize + block.gate = LinearReplicated(128, 256, has_bias=False) + block.gate.weight = torch.randn(256, 128, device="cuda", dtype=torch.bfloat16) * .1 + block.gate.expert_bias = torch.randn(256, device="cuda") * .05 + x = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + weights, ids = block._route(x) + logits = F.linear(x, block.gate.weight).float() + scores = 1. / (1. + torch.exp(-logits)) + expected_ids = torch.argsort(scores + block.gate.expert_bias, dim=-1, descending=True)[:, :6] + assert [set(row) for row in ids.tolist()] == [set(row) for row in expected_ids.tolist()] + expected = scores.gather(1, ids.long()) + if renormalize: + expected = expected / expected.sum(1, keepdim=True) + torch.testing.assert_close(weights, expected, rtol=1e-6, atol=1e-7) + assert ids.dtype == torch.int32 + + +def test_weight_reader_merges_cross_shard_projections_and_preserves_sinks(tmp_path): + from safetensors.torch import save_file + from freetoken.models.flint.weight import iter_weights + hf = _hf() + (tmp_path / "config.json").write_text(json.dumps(vars(hf))) + q = torch.randn(512, 64, dtype=torch.bfloat16) + k = torch.randn(256, 64, dtype=torch.bfloat16) + v = torch.randn(256, 64, dtype=torch.bfloat16) + pre = "model.layers.0.self_attn" + sink = torch.tensor([.125, .25, .5, 1.]) + save_file({f"{pre}.q_proj.weight": q, f"{pre}.softmax_offset": sink}, tmp_path / "a.safetensors") + save_file({f"{pre}.k_proj.weight": k, f"{pre}.v_proj.weight": v}, tmp_path / "b.safetensors") + weights = dict(iter_weights(str(tmp_path), torch.device("cpu"), include_moe_experts=False, include_non_moe=True)) + torch.testing.assert_close(weights[f"{pre}.qkv_proj.weight"], torch.cat([q, k, v])) + torch.testing.assert_close(weights[f"{pre}.softmax_offset"], sink) + for i in (1, 2): + assert not weights[f"model.layers.{i}.mlp.gate.expert_bias"].any() + + +def test_gguf_explicit_end_of_generation_ids_are_not_lost(): + from freetoken.models.gguf.tokenizer import gguf_eos_token_ids + tokens = ["x"] * 19 + for i, token in [(1, "<|end_of_text|>"), (3, "<|user|>"), (5, "<|observation|>"), (18, "<|developer|>")]: + tokens[i] = token + meta = {"tokenizer.ggml.tokens": tokens, "tokenizer.ggml.eos_token_id": 1, + "tokenizer.ggml.eog_token_ids": [1, 3, 5, 18]} + with patch("freetoken.models.gguf.tokenizer.load_gguf_metadata", return_value=meta): + assert gguf_eos_token_ids("test.gguf") == {1, 3, 5, 18} + + +@pytest.mark.needs_weights +def test_official_gguf_matches_hf_tokenizer_and_config(): + from transformers import AutoTokenizer + from freetoken.models.gguf.config import build_gguf_shim + from freetoken.models.gguf.tokenizer import load_gguf_tokenizer, gguf_eos_token_ids + from freetoken.models.flint.gguf import parse_gguf_config + from freetoken.utils import cached_load_hf_config + gguf_path = os.environ.get("FREETOKEN_FLINT_GGUF") + hf_path = os.environ.get("FREETOKEN_FLINT_MODEL") + if not gguf_path or not hf_path: + pytest.skip("set FREETOKEN_FLINT_GGUF and FREETOKEN_FLINT_MODEL") + hf = AutoTokenizer.from_pretrained(hf_path, local_files_only=True) + gg = load_gguf_tokenizer(gguf_path) + texts = [ + "HelloWorld camelCase XMLHTTPRequest don't we're I'LL", + "1234567890 3.1415926535 -42 1e-10", + "漢字かなカナ 中文测试 café e\u0301 👩‍💻", + "def f(x):\n\treturn {'value': x + 1}\n\n# EOF ", + "<|user|>hello<|assistant|>check42<|end_of_text|>", + ] + for text in texts: + assert gg.encode(text, add_special_tokens=False) == hf.encode(text, add_special_tokens=False), text + tools = [{"type": "function", "function": {"name": "read", "parameters": { + "type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"], + }}}] + messages = [ + {"role": "user", "content": "Read the file."}, + {"role": "assistant", "content": "", "reasoning_content": "Inspect it first.", + "tool_calls": [{"type": "function", "function": {"name": "read", "arguments": {"path": "/tmp/a.py"}}}]}, + {"role": "tool", "content": "print(42)"}, + ] + assert gg.apply_chat_template(messages, tools=tools, add_generation_prompt=True) == hf.apply_chat_template( + messages, tools=tools, add_generation_prompt=True, + ) + assert gguf_eos_token_ids(gguf_path) == {150001, 150003, 150005, 150018} + cfg_hf = parse_config(cached_load_hf_config(hf_path)) + cfg_gg = parse_gguf_config(build_gguf_shim(gguf_path)) + for name in ("head_dim", "num_qo_heads", "num_kv_heads", "num_layers", "hidden_size", + "num_experts", "num_experts_per_tok", "moe_intermediate_size", + "shared_expert_intermediate_size", "first_k_dense_replace", "attention_groups"): + assert getattr(cfg_hf, name) == getattr(cfg_gg, name), name diff --git a/tests/server/test_parser_auto_selection.py b/tests/server/test_parser_auto_selection.py index 6bdc6579fd..49b0f1a6ce 100644 --- a/tests/server/test_parser_auto_selection.py +++ b/tests/server/test_parser_auto_selection.py @@ -103,3 +103,8 @@ def test_an_explicit_choice_beats_inference(): pinned, _ = parse_args(["--model", ANON_PATH, "--reasoning-parser", "qwen3"]) assert off.reasoning_parser is None assert pinned.reasoning_parser == "qwen3" + + +def test_flint_uses_optional_reasoning_and_xml_tool_calls(): + assert _inferred("FlintForCausalLM") == ("glm47", "flint") + assert _inferred("FlintGGUFForCausalLM") == ("glm47", "flint") diff --git a/tests/server/test_reasoning_parser_all_models.py b/tests/server/test_reasoning_parser_all_models.py index 5d115dc1a8..47e0f094c4 100644 --- a/tests/server/test_reasoning_parser_all_models.py +++ b/tests/server/test_reasoning_parser_all_models.py @@ -1,6 +1,11 @@ from __future__ import annotations +from types import SimpleNamespace + +import pytest + from freetoken.server.reasoning_parser import ReasoningParser +from freetoken.server.generation import GenSpec, _make_reasoning_parser def _stream(parser: ReasoningParser, chunks): @@ -15,6 +20,21 @@ def _stream(parser: ReasoningParser, chunks): return reasoning, content +@pytest.mark.parametrize("kwargs", [{}, {"enable_thinking": True}, {"enable_thinking": False}]) +@pytest.mark.parametrize("tools", [None, [{"type": "function", "function": {"name": "read"}}]]) +def test_flint_streaming_opens_reasoning_only_on_a_think_tag(kwargs, tools): + spec = GenSpec(messages=[], sampling_params=SimpleNamespace(), chat_template_kwargs=kwargs, template_tools=tools) + state = SimpleNamespace(config=SimpleNamespace(reasoning_parser="flint")) + thought = "The literal here is a quotation." + text = f"{thought}42" + for size in (1, 3, 7, len(text)): + parser = _make_reasoning_parser(spec, state) + assert _stream(parser, [text[i:i+size] for i in range(0, len(text), size)]) == (thought, "42") + parser = _make_reasoning_parser(spec, state) + assert _stream(parser, ["A direct ", "answer."]) == ("", "A direct answer.") + assert _make_reasoning_parser(spec, state).parse_non_stream("A direct answer.") == ("", "A direct answer.") + + # --------------------------------------------------------------- gpt-oss harmony ANALYSIS_FINAL = ( "<|channel|>analysis<|message|>The user says hi. Greet back.<|end|>" diff --git a/tests/server/test_streaming_model_matrix.py b/tests/server/test_streaming_model_matrix.py index eb6e3a2ed2..1ee7937951 100644 --- a/tests/server/test_streaming_model_matrix.py +++ b/tests/server/test_streaming_model_matrix.py @@ -126,6 +126,7 @@ # family, mirroring args.py auto-inference. think_open == "" when the chat template # opens the block implicitly (model emits only the closing marker). REASONING_FAMILIES = { + "flint": ("glm47", "flint", "", ""), "dsv4": ("deepseekv32", "deepseekv32", "", ""), "dsv4.1": ("deepseekv41", "deepseekv32", "", ""), "qwen3.5": ("qwen3_coder", "qwen3", "", ""), From 6ee048e3cca92da2a85cdc9ca6fd21cceed527ad Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jes=C3=BAs=20J=C3=B3dar=20Piernas?= Date: Sat, 24 Oct 2026 09:35:25 +0200 Subject: [PATCH 2/2] refactor(flint): address final review scope requests --- docs/flint.md | 81 ------------------------------------------ docs/models.md | 2 +- tests/e2e/test_aime.py | 37 ++++++++++++++++--- 3 files changed, 33 insertions(+), 87 deletions(-) delete mode 100644 docs/flint.md diff --git a/docs/flint.md b/docs/flint.md deleted file mode 100644 index 66c0935dbc..0000000000 --- a/docs/flint.md +++ /dev/null @@ -1,81 +0,0 @@ -# Recursive Flint - -FreeToken serves the Flint decoder with ModelOpt NVFP4 safetensors and the official -NVFP4 GGUF. Flint uses 49 layers (two dense, then 47 MoE), sigmoid top-6 routing -over 256 experts, an ungated shared expert, QK RMSNorm, learned attention sinks, -and hybrid full/sliding attention. Full-attention layers use NoPE; local layers -use RoPE and a 513-token window including the current token. - -Checkpoints: - -- [FutureMLS-Lab/flint-nvfp4-coldgptq](https://huggingface.co/FutureMLS-Lab/flint-nvfp4-coldgptq): - community ModelOpt export, approximately 25.16 GB (23.43 GiB) of tensors. Routed - experts are NVFP4; other weights remain BF16. -- [RecursiveSI/flint-nvfp4-gguf](https://huggingface.co/RecursiveSI/flint-nvfp4-gguf): - official approximately 22.9 GB GGUF, with more resident matrices quantized. - -`dgpp_mixed` FP8/NVFP4 exports and MLX-packed checkpoints use different layouts -and are not supported by this loader. The HF reader also accepts the original -BF16 layout, but full BF16 inference requires substantially more host memory and -was not validated on the 32 GB host used for the NVFP4 checks. - -## Run - -Start with one request and a short context on a 12 GB GPU. The routed expert -banks need about 20 GiB of host RAM, in addition to the process and resident -weights. The ModelOpt path uses `offload`, `cpu` or `hybrid` experts. These recipes -target TP=1. - -```bash -FLASHINFER_USE_CUDA_NORM=1 ft serve \ - --model FutureMLS-Lab/flint-nvfp4-coldgptq \ - --gpu 0 --moe-strategy hybrid --moe-hybrid-max-fetch 1 \ - --moe-cache-size 1024 --num-tokens 8192 \ - --max-seq-len-override 8192 --max-prefill-length 2048 \ - --max-running-requests 1 --host 127.0.0.1 --port 1919 -``` - -For the GGUF, download the file and use its path as `--model` with the same flags: - -```bash -hf download RecursiveSI/flint-nvfp4-gguf flint-nvfp4.gguf --local-dir models/flint-gguf -``` - -The reference SGLang kit explicitly selects FlashInfer's CUDA normalization -implementation. The environment variable above preserves that choice. Ampere -GPUs use FreeToken's compatible NVFP4 paths; they do not have Blackwell's native -W4A4 instructions. Memory fit and throughput therefore need local measurement. - -Recommended sampling defaults come from the checkpoint: temperature 1.0, -top-p 0.95 and top-k 20. The documented 262144-token context is a model limit; -the memory required to reserve it is separate from loading the weights. - -## Conversation protocol - -Use the checkpoint's chat template. The server selects the `glm47` tool parser and `flint` reasoning -parser automatically for HF and GGUF checkpoints. Flint chooses whether to emit -`...` on each turn; the parser never assumes the turn starts inside -reasoning, even when generic thinking flags or tools are present. Replay returned -`reasoning_content` and tool calls in the conversation history. - -Tool calls use the GLM-style `` / `` / `` envelope. -This integration uses FreeToken's existing parser and API behavior. It does not -port the separate Flint-specific XGrammar backend distributed in the SGLang kit. - -The GGUF pre-tokenizer is `flint`, with the original three regex splits and byte -level encoding. Its explicit EOG array contains four IDs: 150001, 150003, 150005, -and 150018. All four must stop generation, including the role boundaries. - -## Validate a local change - -The model tests compare fused sigmoid routing against independent tensor math and -check weight fusion across shards. With local checkpoints they also compare GGUF -tokenization, chat rendering and configuration against the original HF files. - -```bash -FREETOKEN_FLINT_MODEL=/path/to/flint-modelopt \ -FREETOKEN_FLINT_GGUF=/path/to/flint-nvfp4.gguf \ -uv run pytest tests/models/test_flint.py tests/server/test_parser_auto_selection.py \ - tests/server/test_reasoning_parser_all_models.py \ - tests/server/test_streaming_model_matrix.py -``` diff --git a/docs/models.md b/docs/models.md index 223f6642b0..9807df80db 100644 --- a/docs/models.md +++ b/docs/models.md @@ -20,7 +20,7 @@ for them; other checkpoints of the same architectures work too. | MiniMax-M2.5 | [nvidia/MiniMax-M2.5-NVFP4](https://huggingface.co/nvidia/MiniMax-M2.5-NVFP4) | | MiniMax-M3 | [nvidia/MiniMax-M3-NVFP4](https://huggingface.co/nvidia/MiniMax-M3-NVFP4) | | Muse-Glimmer | [meta-models/Muse-Glimmer-30B](https://huggingface.co/meta-models/Muse-Glimmer-30B), [RedHatAI/Muse-Glimmer-30B-NVFP4](https://huggingface.co/RedHatAI/Muse-Glimmer-30B-NVFP4) | -| [Flint](flint.md) | [FutureMLS-Lab/flint-nvfp4-coldgptq](https://huggingface.co/FutureMLS-Lab/flint-nvfp4-coldgptq), [RecursiveSI/flint-nvfp4-gguf](https://huggingface.co/RecursiveSI/flint-nvfp4-gguf) | +| Flint | [FutureMLS-Lab/flint-nvfp4-coldgptq](https://huggingface.co/FutureMLS-Lab/flint-nvfp4-coldgptq), [RecursiveSI/flint-nvfp4-gguf](https://huggingface.co/RecursiveSI/flint-nvfp4-gguf) | ### Image input diff --git a/tests/e2e/test_aime.py b/tests/e2e/test_aime.py index 8d0b41547f..7507039668 100644 --- a/tests/e2e/test_aime.py +++ b/tests/e2e/test_aime.py @@ -198,14 +198,41 @@ def build_llm(model_path: Path) -> LLM: return LLM(**kwargs) +def visible_gpu_index() -> int: + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if not visible: + return 0 + first = visible.split(",", 1)[0].strip() + return int(first) + + def free_gpu_memory_gib() -> float: - # A child respects UUID visibility and CUDA_DEVICE_ORDER without initializing - # the engine's parent process, on both CUDA and ROCm. + if torch.version.hip is not None: + # Query logical device zero without initializing the engine's parent process. + result = subprocess.run( + [sys.executable, "-c", "import torch; print(torch.cuda.mem_get_info(0)[0])"], + check=True, + text=True, + capture_output=True, + ) + return int(result.stdout.strip()) / (1 << 30) + result = subprocess.run( - [sys.executable, "-c", "import torch; print(torch.cuda.mem_get_info(0)[0])"], - check=True, text=True, capture_output=True, + [ + "nvidia-smi", + "--query-gpu=index,memory.free", + "--format=csv,noheader,nounits", + ], + check=True, + text=True, + capture_output=True, ) - return int(result.stdout.strip()) / (1 << 30) + target = visible_gpu_index() + for line in result.stdout.splitlines(): + index, free_mib = [part.strip() for part in line.split(",", 1)] + if int(index) == target: + return int(free_mib) / 1024 + raise RuntimeError(f"GPU {target} is not visible in nvidia-smi") @pytest.mark.skipif(not torch.cuda.is_available(), reason="AIME e2e needs CUDA")