[jesusjodar]

Summary

Related to #786. Add Recursive Flint support for BF16-layout safetensors, the ModelOpt NVFP4 export and the official NVFP4 GGUF.

• Expanded Q projections, per-head QK RMSNorm, learned sinks, and mixed sliding-window RoPE / full-attention NoPE layers.
• Sigmoid top-6 routing, selection-only expert bias, an ungated shared expert and two leading dense layers.
• GGUF metadata/tensor loading, the original three tokenizer splits and explicit EOG arrays.
• Automatic glm47 tool parsing and optional flint reasoning parsing.

Final review cleanup

Addressed the second review in 6ee048e3cca92da2a85cdc9ca6fd21cceed527ad:

• Remove docs/flint.md; the Flint name in docs/models.md is now plain text and the checkpoint links remain in the table.
• Restore tests/e2e/test_aime.py exactly to the base version. The generic UUID-aware free-memory probe is now isolated in fix(test): respect cuda device visibility in aime memory probe #821, based on current main, with no Flint implementation changes.

No tests were rerun for this final scope cleanup. The results below were recorded on 711a86ef11c67bbe6c69ff59370594f64c6ef1b8, before extraction. In particular, the historical UUID-selected AIME command used the probe now in #821; it is not a fresh run of the latest #797 tree. The reviewer also independently rechecked that revision on sm_120.

Changes following review

Rebased onto main at 507a8bfa81d5c305e047661576165a74068e5ab0, preserving the Qwen GGUF NFC normalization added upstream.

1. Use fused_topk_sigmoid from perf(moe): fuse sigmoid + selection-bias routing into a triton kernel #804. Keep self.gate.forward(x).float() so the router projection still produces BF16 logits before FP32 sigmoid/top-k. Keep the zero selection-bias tensor.
2. FlintForCausalLM now directly subclasses BaseLLMModel and owns its model/lm_head; no import of another family's model class.
3. Remove the extra flint tool-parser choice and use glm47. Register flint reasoning with ThinkReasoningParser; _make_reasoning_parser owns force_reasoning=False for this family. Streaming/non-streaming tests cover tools and generic thinking flags.
4. Remove the per-model e2e file, its README row and its documentation section. Remove config-guard/fixture-echo, shared-expert-call-order and mocked-attention tests. Keep independent fused-routing math, weight-reader, EOG and official GGUF/HF tokenizer/config comparisons.
5. Append Flint without reordering existing catalog/AOT entries; keep the llama.cpp llama_vocab::impl::load pointer and use a pre-tokenizer-to-splits dictionary.

The cache-rebuild gate needed small harness adjustments for this environment:

• Cache rebuild accepts a GGUF file and FREETOKEN_REBUILD_EXTRA_ARGS for offloaded models. Its short generation probe accepts text in either content or reasoning: Flint can spend the entire 32-token probe in its thought channel. Empty output in both channels still fails.

Local validation at 711a86e before final scope cleanup

Environment and checkpoints

• RTX 3060 12 GiB only; TP=1; Ryzen 5 5600X, about 32 GiB RAM.
• Linux 7.2.8-9-omarchy, NVIDIA 615.78.08, Python 3.12.14.
• PyTorch 2.14.1+cu130, Triton 3.8.0, FlashInfer 0.6.18.post1.
• Transformers 5.18.0, tokenizers 0.23.3, safetensors 0.8.0, apache-tvm-ffi 0.1.13.post3, pytest 9.1.1.
• Source through PYTHONPATH=python in an isolated uv-created venv. Native extensions (including ggml-cpu variants) and prebuilt kernels were reused from freetoken 0.2.0+gfe5a5ec9c / freetoken-kernel-cache 0.2.0+cu130.gfe5a5ec9c. No clean native-extension/CUDA build was performed for this revision; this PR changes no C++/CUDA source.

Checkpoint	Revision
FutureMLS-Lab/flint-nvfp4-coldgptq	56f37fcfc7bd17834512ec2e53696b8da15f3658
RecursiveSI/flint-nvfp4-gguf, flint-nvfp4.gguf	6ec43e19b565fc8d560e2d04985d7724f20673bf
RecursiveSI/flint reference config/tokenizer	2a247be0c0aec895ccb0ac7ca92229ec6df37288

Commands below ran from the source checkout; the sibling venv/checkpoint paths reflect the local validation layout:

export PYTHONPATH=python
export CUDA_VISIBLE_DEVICES=GPU-a0591354-3462-ad0e-6e65-e772fd9a045e
export FLASHINFER_USE_CUDA_NORM=1
PYTHON=../../flint-research/dev-venv/bin/python
MODEL="$PWD/../../flint-research/models/flint-modelopt"
GGUF=/home/jesusjodar/.freetoken/models/Flint-NVFP4-GGUF/flint-nvfp4.gguf

Targeted regressions: 709 passed, 24 skipped, 18 deselected, 1 warning

FREETOKEN_FLINT_MODEL="$MODEL" FREETOKEN_FLINT_GGUF="$GGUF" \
"$PYTHON" -m pytest \
  tests/models/test_flint.py tests/server tests/tokenizer \
  tests/models/test_models_registry.py \
  tests/models/qwen4_exp/test_weight.py::test_every_registry_architecture_is_claimed_by_an_aot_entry \
  tests/moe/test_fused_moe.py::test_fused_topk_sigmoid_matches_torch_reference \
  -m 'not slow' -q

Result: 709 passed, 24 skipped, 18 deselected, 1 warning in 14.43s.
The warning is the existing Starlette TestClient/httpx deprecation. The 18 slow router cases were then run explicitly:

"$PYTHON" -m pytest tests/moe/test_fused_moe.py::test_fused_topk_sigmoid_matches_torch_reference -q

Result: 18 passed in 12.35s.

Generic AIME gate: 1 passed

Dataset: math-ai/aime25, revision 563bb8404243c5f09de6ec262f2db674fe5bce9b, test.jsonl; convert problem to the gate's prompt key and preserve answer and row order. AIME 2025 problem 1, pass@3, 8192-token maximum per sample. This is one problem, not an AIME-suite accuracy claim. Sampling comes from the checkpoint: temperature 1.0, top-p 0.95, top-k 20.

FREETOKEN_TEST_MODEL="$MODEL" \
FREETOKEN_AIME25_JSONL="$PWD/../aime25.jsonl" \
FREETOKEN_AIME_REQ=1 FREETOKEN_AIME_MAX_TOKENS=8192 FREETOKEN_AIME_SAMPLES=3 \
FREETOKEN_AIME_MIN_FREE_GIB=6 \
FREETOKEN_TEST_MOE_CACHE_SIZE=1024 FREETOKEN_TEST_MOE_STRATEGY=hybrid \
FREETOKEN_TEST_MOE_CPU_THREADS=4 FREETOKEN_TEST_MEM_RATIO=0.9 \
"$PYTHON" -m pytest tests/e2e/test_aime.py -q -s

Result: 1 passed, 1 warning in 61.45s; existing CUDA allocator deprecation warning.

Generic cache-rebuild gate: 1 passed per format

Each format ran in a separate process. The gate exercises real generation, cold/warm prefix reuse and flush, KV growth from 4000 to 6000 pages with graph re-capture, and rejection of an unfittable target while preserving the serving engine.

for checkpoint in "$MODEL" "$GGUF"; do
  FREETOKEN_TEST_MODEL="$checkpoint" \
  FREETOKEN_REBUILD_MIN_FREE_GIB=6 FREETOKEN_REBUILD_BOOT_TIMEOUT=300 \
  FREETOKEN_REBUILD_EXTRA_ARGS='--moe-strategy hybrid --moe-cache-size 1024 --moe-hybrid-max-fetch 1 --moe-cpu-threads 4 --max-running-requests 1 --cuda-graph-max-bs 1 --max-prefill-length 2048 --max-seq-len-override 8192 --kv-reserve-tokens 6000' \
  "$PYTHON" -m pytest tests/e2e/test_cache_rebuild.py -q -s
done

Results: ModelOpt 1 passed in 40.96s; official GGUF 1 passed in 51.10s.

Scope and limitations

• This revision's local serving checks cover the two NVFP4 formats on Linux / RTX 3060 / TP=1.
• Quantized ModelOpt fused MoE remains rejected; these e2e runs use hybrid offload.
• Full BF16 serving, native Blackwell W4A4, TP>1, Windows, full 262144-token context and concurrent serving were not locally validated in this revision. The reviewer's independent sm_120 checks are documented in their review.
• Mixed dgpp_mixed and MLX layouts remain unsupported. The separate Flint SGLang XGrammar backend is not ported.
• No end-to-end throughput improvement or broad-suite pass is claimed. The full repository suite was not rerun after this rebase.

Implementation and local test execution were assisted by Codex.

[jason-fxz]

Thanks for the thorough PR. I checked it on an RTX PRO 6000 (sm_120) at f48684f with all three checkpoints: the BF16 RecursiveSI/flint, the ModelOpt coldgptq export and the official NVFP4 GGUF.

What I verified

• I wrote an independent fp32 PyTorch reference on the BF16 weights, following the SGLang kit's semantics, and teacher-forced FreeToken's greedy outputs through it. The prompt set includes a 1028-token prompt, so SWA, NoPE and the sinks are all exercised.

  Path	Greedy tokens agreeing with the reference	First-token logit cosine
  BF16	311/316	0.99997
  ModelOpt NVFP4 (native W4A4 on sm_120)	314/325	≥ 0.9997
  GGUF NVFP4	300/325	≥ 0.998

  The BF16 divergences are all near-ties. llama.cpp's flint.cpp (#30928) has the same semantics.

• Per-module quant methods at runtime match the checkpoints. ModelOpt has NVFP4 experts and BF16 for everything else. In the GGUF, all linears and lm_head are NVFP4, while token_embd and the router are BF16.

• Your e2e passes 2/2 here, after raising the startup deadline locally. My own serving checks with default flags pass on both formats. They cover reasoning split, tool call plus tool result, streamed tool calls, 4 concurrent requests, and a 23k-token passkey both cold and with prefix reuse.

So the model support itself looks correct. Requests:

1. Router: use the fused kernel. Please rebase onto main and route with fused_topk_sigmoid from #804, as glm4_moe and minimax_m2 now do.

   • On sm_120 it picks identical experts, with weights equal to within 1.5e-8.
   • It cuts routing over the 47 layers from 786 to 187 µs per decode step at bs=1, and from 1921 to 328 µs for a 2048-token prefill.
   • The zero expert_bias can stay as the kernel's bias argument.
   • Please keep the kit's bf16 router logits (self.gate.forward(x).float()). fp32 logits, as glm4_moe uses, change the top-6 set for 2.3% of token-layer routes on Flint.

2. No per-model e2e. Please remove tests/e2e/test_flint.py, its row in tests/README.md, and the e2e part of docs/flint.md. We don't keep a separate e2e file per model. Run the generic gates (tests/e2e/test_aime.py and tests/e2e/test_cache_rebuild.py with FREETOKEN_TEST_MODEL) and put the results in the PR description.

3. Own CausalLM class. FlintForCausalLM subclasses qwen3_moe.Qwen3MoeForCausalLM. Please give Flint its own BaseLLMModel subclass, a few lines for embed and lm_head, as minimax_m2, gpt_oss and glm4_moe do. Model families shouldn't import each other's model classes.

4. Parsers.

   • flint as a tool-call parser is only an alias of Glm47Detector. Auto-detection can return glm47 instead of adding a new user-facing choice.
   • For reasoning, FlintReasoningParser silently ignores the force_reasoning it is given. Please instead register flint with ThinkReasoningParser and add a parser_name == "flint" branch (force_reasoning = False) in generation._make_reasoning_parser, next to minimax_m3 and gemma4. That keeps the per-family rule in one place.

5. Tests. Please drop the tests that restate the implementation (see the "What earns a place here" rules in tests/README.md):

   • test_rejects_unimplemented_config_semantics only checks the config guards.
   • test_shared_expert_sees_input_before_routed_kernel_overwrites_it mirrors the call order in FlintMoE.forward.
   • test_attention_projection_norm_rope_and_sinks_match_torch mocks the attention backend, so the sinks and SWA kernels are not covered; the reference comparison above covers them.

   The weight-reader, EOG-ids and official GGUF tokenizer/config tests are worth keeping.

Nits:

• docs/models.md and SUPPORTED_MODELS in aot_models.py put Flint first; please keep the existing order.
• The gguf/tokenizer.py comment on _EOG_NAMES lost its pointer to llama.cpp's llama_vocab::impl::load; please keep it.
• pre in ("joyai-llm", "flint") with an inline ternary would read better as a pre -> splits dict.

[jesusjodar]

@jason-fxz Thanks for the detailed review and independent validation. Updated in 711a86e, rebased onto main at 507a8bf.

• Switched routing to fused_topk_sigmoid, preserving self.gate.forward(x).float() and zero expert selection bias.
• Made FlintForCausalLM a direct BaseLLMModel subclass with its own model and LM head.
• Reused glm47 for tool calls and ThinkReasoningParser for Flint, with force_reasoning=False handled centrally in generation.
• Removed the per-model e2e test and redundant tests; retained the independent routing, weight-loading, EOG and official GGUF/tokenizer checks.
• Addressed the model ordering, pre-tokenizer dictionary and llama.cpp reference comments, preserving the Qwen GGUF NFC normalization from main.

Validation on RTX 3060 / TP=1:

• 709 targeted tests and 18 fused-router cases passed (24 skipped).
• The generic AIME gate passed for AIME 2025 problem 1 with pass@3.
• The generic cache-rebuild gate passed separately for ModelOpt NVFP4 and the official NVFP4 GGUF.

The generic gates now support UUID GPU selection and GGUF/offload arguments. The short generation probe accepts nonempty visible text or reasoning, since Flint can spend its entire 32-token budget in the reasoning channel.

Full commands and validation limits are in the updated PR description. These runs reused installed native extensions and prebuilt kernel caches; they do not establish a clean native build or full-suite, multi-GPU or Blackwell validation.

Ready for another look.

[jason-fxz]

Thanks, all five points are addressed. Re-checked 711a86e on sm_120: greedy output still matches the fp32 reference as in round one, and unit tests, serving checks on both formats and the generic gates pass.

Two last requests:

1. Please remove docs/flint.md and make the Flint row in docs/models.md plain text like the others. We don't keep per-model pages in docs/, and the checkpoint list is already in the table.
2. Please move the tests/e2e/test_aime.py change (the child-process free-memory probe) out of this PR. It fixes UUID GPU selection for every model, not Flint, so it belongs in its own PR.

After that this is ready to merge.

[jesusjodar]

@jason-fxz Thanks for rechecking. Both final requests are addressed in 6ee048e:

1. Removed docs/flint.md and made the Flint name in docs/models.md plain text, keeping the checkpoint links in the table.
2. Restored tests/e2e/test_aime.py exactly to the base version. The child-process free-memory probe is now in the separate PR fix(test): respect cuda device visibility in aime memory probe #821, based on current main and limited to that file.

Updated the description to link the extracted fix and identify the validation revision. Tests were not rerun after this scope cleanup; the previously reported results remain those from 711a86e, which you independently rechecked.

[jason-fxz]

LGTM, thanks. docs/flint.md is gone, test_aime.py matches main, and no code changed since 711a86e, which I re-checked earlier.