feat(flint): add nvfp4 modelopt and gguf support

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 | [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
 

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",

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",
+]

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))

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),
+            ),
+        ),
+    )

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"])

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)

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