[gdevenyi]

Merges upstream FlashML-org#169 (fix(engine): load prefill triton kernels at startup, not mid-request, head a8326af9). Stacked on #14.

Changes

• Merge of fix(engine): load prefill triton kernels at startup, not mid-request FlashML-org/FreeToken#169. One conflict in python/freetoken/engine/config.py: fix(engine): load prefill triton kernels at startup, not mid-request FlashML-org/FreeToken#169 adds prefill_warmup: bool = True next to a comment that main had reworded from --moe-backend offload to --moe-strategy offload. The resolution keeps the new field and main's wording.
• Behaviour:
  • The prefill warmup now runs on every attention backend, over a ladder of lengths derived from the config.
  • fla/l2norm stops specializing on the token count.
  • New flag --skip-prefill-warmup.

Testing

Unit tests:

• pytest tests/engine tests/server -m "not slow": 971 passed, 2 deselected.
• pytest tests/kernels/test_kda.py tests/models/test_glm5_next_kda_op.py tests/models/qwen4_exp/test_gdn.py (the l2norm users): 10 passed.

Before the merge, on a long-running server at 150c95d (HF checkpoint dir, 4096-token prefill chunks), each prompt length never seen before was sent twice. The first send was never slower, so there was no mid-request JIT stall to remove:

prompt tokens	first send	repeat
1,592	2.61 s	2.59 s
3,248	3.82 s	3.85 s
5,408	6.80 s	6.81 s
12,672	14.70 s	14.69 s

(TTFT, max_tokens 1)

After the merge, the deployment serves a bf16 FTW conversion (ft checkpoint --model ~/models/Qwen3.8-Flash-Next-NVFP4 --out ~/models/Qwen3.8-Flash-Next-NVFP4-FTW) with otherwise the same flags. The Triton kernel cache on disk was warm in all runs:

	FTW before merge	warmup on (default)	--skip-prefill-warmup
Startup to ready	18.6 s	35.6 s	17.1 s
First 1.1K-token request / repeat	2.66 / 2.28 s	2.58 / 2.29 s	2.61 / 2.27 s
First ~12.1K-token request / repeat	not measured	12.81 / 12.69 s	12.78 / 12.89 s

• Warmup: Prefill warmup over 10 lengths: [9, 16, 17, 32, 33, 48, 65, 80, 103, 4096], completed in 18.680 s. The ladder stopped at 4096 (this run had --max-prefill-length 4096); no OOM.
• Result: on this box with a warm kernel cache, the warmup costs ~18 s of startup and buys nothing measurable. It would matter after the kernel cache is invalidated (an upgrade), which was not tested. The deployment script runs with --skip-prefill-warmup.

---

Environment:

• RTX 5080 16 GB (sm_120), AMD Ryzen 9 9950X (16 cores), 123 GB RAM
• Ubuntu 24.04.4, kernel 7.1.9-x64v3-xanmod1
• Driver 610.57.04, CUDA 13.3, torch 2.11.0+cu130, triton 3.6.0, FreeToken 0.1.2 (editable install)
• Checkpoint: nvidia/Qwen3.8-Flash-Next-NVFP4

Stack: the 5080deploy branch is main e0886cc plus upstream PRs merged one at a time for a single-user, full-context (262,144-token) deployment on this card. Each PR in the stack is based on the previous one, so its diff is exactly one merge step. Merge them in order, or merge the top of the stack alone.

How tests were run: with the server stopped, since a running server holds almost all VRAM and GPU tests fail spuriously. Benchmarks are streamed chat completions with 256 tokens generated with ignore_eos, one request at a time. Prompts are salted so the radix cache cannot hit.

• Prefill = prompt tokens / TTFT.
• Decode = (completion - 1) / time after the first token.
• "94K recall": a 94,526-token prompt of repository source with a codeword planted in the first line.
• "248K recall": a 248,029-token prompt with facts planted at token ~0, ~124K and ~240K.

---

Rebased 2026-09-17 onto upstream main cac247a (v0.1.3). The whole stack was rebuilt merge by merge on the new base (previous base e0886cc); the merge resolutions were replayed unchanged via git rerere, and the rebuilt tip differs from the old one by exactly the e0886cc..cac247a file set. This PR's head is now c3a8f01.

Rebased 2026-09-19 onto upstream main cc1f5c2 (4 commits past cac247a: FlashML-org#471 greedy sampling in mixed batches, FlashML-org#518 WeightLoadError, FlashML-org#521 tvm-ffi jit arch, FlashML-org#524 install index). Same replay as before via git rerere; the replayed stack differs from the previous tip 8adde91 by exactly the cac247a..cc1f5c2 file set. This PR's head is now 84acf6e.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt

[Copilot]

🟡 Changes recommended

Unresolved warmup coverage and startup error-handling issues could cause request-time compilation or mask runtime failures.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds configurable startup prefill warmup and removes token-count specialization from FLA l2norm.

Changes:

• Adds --skip-prefill-warmup and FREETOKEN_WARMUP_MAX_LEN.
• Implements startup warmup length selection and execution.
• Updates warmup configuration and server environment handling.

File summaries

File	Summary
python/freetoken/server/args.py	Adds the warmup opt-out flag; CLI documentation should be updated.
python/freetoken/env.py	Adds the warmup length cap.
python/freetoken/engine/engine.py	Implements warmup selection and execution; several backend, bucket-selection, error-handling, and test-coverage issues remain.
python/freetoken/engine/config.py	Adds warmup configuration; the default-value comment should be clarified.

Review details

Suppressed comments (4)

python/freetoken/engine/config.py:54

• With the default value set to True, this comment describes the opposite of the field's normal behavior and makes programmatic configuration ambiguous. Document that True runs warmup and False skips it.

    # skip the startup prefill warmup forwards; first requests then pay the kernel module loads (and the full JIT compile on a cold start)

python/freetoken/engine/engine.py:1065

• When cap is in [2, 8] (for example, --max-seq-len-override 8), every entry in lens is filtered out, so _warmup_prefill() performs no forward at all. This regresses the previous min(80, max_seq_len) warmup and leaves the only reachable prefill sizes to compile on the first request; add a representative length for this small-cap range.

        lens = {9, 16, 17, 32, 33, 48, 65, 80}

python/freetoken/engine/engine.py:1065

• With MXFP8 dense layers, mxfp8_linear selects M_TILE = next_power_of_2(M) for M <= 256 (python/freetoken/kernel/triton/mxfp8_linear.py:205-216). This ladder reaches the 128 bucket at length 80 but never the 256 bucket: the 4096 sample takes the cuBLAS path (M > 256), so a first 129-256-token prefill still JIT-loads mid-request. Add a 129-token rung when the cap permits.

        lens = {9, 16, 17, 32, 33, 48, 65, 80}

python/freetoken/server/args.py:751

• --skip-prefill-warmup is a public server flag, but the repository's runtime option reference in docs/cli.md:37-50 does not list it. Please add the default and the startup versus first-request latency tradeoff so users can discover why startup behavior changes.

    parser.add_argument(
        "--skip-prefill-warmup",
        action="store_false",
        dest="prefill_warmup",
        default=ServerArgs.prefill_warmup,
        help=(
            "Skip the startup prefill warmup forwards. The first request at each "
            "new size then pays the Triton compile/load cost mid-request."

• Files reviewed: 4/4 changed files
• Comments generated: 3
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.