Skip to content

5080deploy 9/9: merge upstream #389 (+#385), load-time FP8 dense projections - #17

Open
gdevenyi wants to merge 89 commits into
5080/08-moe-stats-231from
5080deploy
Open

gdevenyi wants to merge 89 commits into
5080/08-moe-stats-231from
5080deploy

Conversation

@gdevenyi

@gdevenyi gdevenyi commented Sep 15, 2026

Copy link
Copy Markdown
Owner

Merges upstream FlashML-org#389 (feat(qwen4_exp): load-time per-tensor FP8 dense projections (W8A8 via _scaled_mm), head bd106c9a). FlashML-org#389 is stacked on FlashML-org#385 (Qwen3.8-Flash-Next tensor parallelism), which comes in with it. Both were closed by their author on 2026-09-14. Stacked on #16; this is the top of the stack, 5080deploy.

Changes

Merge of FlashML-org#389 (16 commits including FlashML-org#385). Six conflict hunks in three files, all against main's changes since FlashML-org#389 was based:

Local fixes in the merge commit:

  • models/config.py (silent semantic conflict): git merged cleanly, but main had removed import os and _VISION_TRUE, and feat(qwen4_exp): load-time per-tensor FP8 dense projections (W8A8 via _scaled_mm) FlashML-org/FreeToken#389's fp8_dense_enabled / fp8_lmhead_enabled still use both. Every qwen4_exp test failed with NameError: name 'os' is not defined (70 failed, 11 errors). Restored import os and renamed the set to _ENV_TRUE.
  • tests/models/qwen4_exp/test_build_cpu.py (a bug in the feat(qwen4_exp): load-time per-tensor FP8 dense projections (W8A8 via _scaled_mm) FlashML-org/FreeToken#389 branch itself):
    • _build() sets the rope device to CPU and never restores it. get_rope is functools.cached, so the CPU cos/sin table leaks into later CUDA tests.
    • test_skeleton.py::test_qsa_layer_matches_hf_dense and ::test_decoder_stack_prefill_and_decode then fail with FlashInfer's cos_sin_cache must be a CUDA tensor, but only when run after test_build_cpu.py; alone, test_skeleton.py passes 10/10.
    • Fixed by restoring the previous rope device and calling rotary.get_rope.cache_clear() in a finally.

Testing

Unit tests (-m "not slow"): tests/models/qwen4_exp tests/models/test_nvfp4_banks_tp.py tests/layers tests/moe tests/engine tests/kvcache tests/kernels/test_qsa_nvfp4.py tests/kernels/test_kv_nvfp4.py: 794 passed, 53 skipped, 2 failed (the rope-cache leak). After the test fix, test_build_cpu.py + test_skeleton.py: 12 passed.

Full suite on the top of the stack (2d56708), pytest tests/ -m "not slow": 2191 passed, 191 skipped, 15 deselected, 2 failed. The failures are the same two test_glm_dsa.py sm_120 shared-memory failures present since FlashML-org#408 (#2); nothing new.

Serving with FP8 off (default). bf16 FTW, 512 slots, 4096 chunks, --skip-prefill-warmup:

  • Startup 16.6 s; 1.19 GiB free after CUDA graphs, identical to before the merge.
  • 1,106-token prompt 486 / 31.31 tok/s; 90,107 tokens 960 / 30.48.
  • 94K recall correct. No regression.

Serving with FREETOKEN_FP8_DENSE=1 FREETOKEN_FP8_LMHEAD=1. nvfp4 KV, 262,144 tokens in every run:

HF dir, 512 slots, 4096 chunks HF dir, 1280 slots, 8192 chunks FP8 FTW, 1280 slots, 8192 chunks
Startup to ready 45.1 s 63.2 s 17.1 s (16.1 s on the final restart)
Free after CUDA graphs 4.12 GiB 2.16 GiB 2.27 GiB
1,106 tokens: prefill / decode 495 / 34.15 480 / 38.93 474 / 36.85
30,107 tokens 959 / 32.21 1167 / 38.39 1164 / 39.31
90,107 tokens 978 / 30.37 1198 / 39.12 1193 / 36.56
7K-prompt prefill peak (of 16,303 MiB) not measured 15,228 MiB 15,124 MiB
248K recall 3/3, 265 s 3/3, 215 s 3/3, 218 s

(prefill / decode in tok/s)

  • VRAM: FP8 dense + FP8 lm_head frees about 3 GiB (4.12 vs 1.09 GiB free at the same settings). On its own it moves decode only a little (30.25-30.71 -> 30.37-34.15 tok/s). Spent on 1280 expert slots and 8192-token chunks, decode rises 20-30% across these runs (36.56-39.31 tok/s) and prefill at 90K by 25% (1193-1198 tok/s). The comparison point is the nvfp4-only configuration at the base of this stack (953 prefill / 30.52 decode at 90K).
  • FP8 FTW: converted with the env vars set: FREETOKEN_FP8_DENSE=1 FREETOKEN_FP8_LMHEAD=1 ft checkpoint --model ~/models/Qwen3.8-Flash-Next-NVFP4 --out ~/models/Qwen3.8-Flash-Next-NVFP4-FTW-fp8. 70.45 GiB across 9 shards, about 5 minutes. A first attempt was killed by the host's memory watchdog; the successful run never dropped below 43.9 GB available RAM.
  • bf16 FTW with the FP8 env vars: dies at load with KeyError: 'model.layers.0.linear_attn.in_proj_qkvz.weight' (FREETOKEN_FP8_DENSE=1 with a bf16 FTW checkpoint dies with KeyError in_proj_qkvz.weight #8). An FP8 FTW needs the env vars at conversion and at serve time.
  • Accuracy: only the planted-fact recall checks above plus coherent short answers and tool calls. No GSM8K or similar on this card.

Deployment command (top of the stack, verified on the final restart: /v1/models reports 262144, a developer-role request answers, and a $ref integer tool argument comes back as {"limit": 7}):

FREETOKEN_FP8_DENSE=1 FREETOKEN_FP8_LMHEAD=1 ft serve \
    --model ~/models/Qwen3.8-Flash-Next-NVFP4-FTW-fp8 \
    --served-model-name Qwen3.8-Flash-Next-NVFP4 \
    --kv-cache-dtype nvfp4 --num-tokens 262144 --max-seq-len-override 262144 \
    --max-running-requests 1 --cuda-graph-max-bs 1 \
    --moe-cache-size 1280 --disable-moe-prefill-overlap \
    --memory-ratio 0.95 --max-prefill-length 8192 --skip-prefill-warmup \
    --host 127.0.0.1 --port 1919

Refs #2, #8.


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.

Later merges on this branch (after 2d56708)

This PR's head is 5080deploy, so it also carries the merges below. Branch tip 5767c99.

Scope rule for this branch: only PRs that are generic, specific to this hardware (RTX 5080 / sm_120), or for Qwen3.8-Flash-Next. Model families other than Qwen3.8 arrive only with a newer upstream base.

commit merge notes
2a98fa4..61128bb #18, #19, #20, #21, #22 The fork fix PRs, all clean merges.
d125f63 #23 (FlashML-org#17), cherry-picked onto this branch HF repo ids; applies without conflicts. Cherry-picked because #23's branch sits on upstream main 3faef36.
5767c99 #47 (FlashML-org#393), cherry-picked the same way One conflict in tokenizer/tokenize.py, keeping both the developer-role mapping (FlashML-org#391) and continue_final_message.

A merge of upstream main 3faef36 (FlashML-org#467 Gemma-4 images, FlashML-org#479 GLM-5.3 images) was briefly on this branch as bb37974 and was removed by rebuilding the branch from 2d56708. The removed file set is exactly e0886cc..3faef36. Qwen3.8-Flash-Next image input (FlashML-org#454) is part of the e0886cc base and still serves images.

Testing at 5767c99 (RTX 5080, production flags from above):

  • pytest tests/ -m "not slow": 2240 passed, 191 skipped, 0 failed. The two test_glm_dsa.py failures are gone with fix(triton): keep the float glm dsa sparse launches two-stage to fit sm_120 shared memory #21.
  • Startup ready with 2.27 GiB free after CUDA graphs, /v1/models 262144.
  • Vision: the 512x384 test image's text FREETOKEN 5080 is read correctly. A 1920x1080 image's text box is read exactly (2,070 prompt tokens, 3.6 s).
  • API features: a developer system instruction is followed (German answer); $ref tool argument {"limit": 7}; n=2 returns two choices; logit_bias -100 on the first token of the unbiased answer changes the generated tokens; /tokenize and /detokenize round-trip.
  • Speed: prefill 477 / 1164 and decode 37.81 / 38.05 tok/s at 1.1K / 30K-token prompts, within run-to-run noise of 2d56708.
  • Earlier checks at bb37974 (the same Qwen code plus the removed Gemma/GLM merge): 248K three-fact recall 3/3 in 217 s; 90K-token prefill 1193 tok/s.

Considered and not merged: #39 (FlashML-org#85, e4m3 rounding before the native downcast).


Rebased 2026-09-17 onto upstream main cac247a (v0.1.3); head is now 8adde91. The stack was rebuilt merge by merge on the new base (resolutions replayed via git rerere; the replayed stack (before the new merges) differs from the old backup tip bb37974 by exactly the 3faef36..cac247a file set). Upstream's 6 new commits: FlashML-org#489 release 0.1.3, FlashML-org#486 FTW/vision-encoder loading (ftw_lacks_vision refusal, iter_vision_weights), FlashML-org#485 docs, FlashML-org#480 MiniMax-M3 images, FlashML-org#481 Muse-Glimmer images, FlashML-org#445 GB10.

On top of FlashML-org#389 the branch now carries (branch 5080/09-fork-fixes-on-cac247a): fork #18-#22 replayed, #23 and #47 merged directly (their rebased/* branches sit on 3faef36, now an ancestor), dd2ebdf (block_ends test fake gets k_scale/v_scale) and cab9be3 (fork #19's FTW layout check skips a bare FTW without an HF config, which upstream FlashML-org#486's test_load_weight_text_only_skips_the_tower writes).

New upstream merges, one stacked PR each: #48 (FlashML-org#505 hybrid checkpoint across prefill chunks), #49 (FlashML-org#466 coalesced msgpack frames), #50 (FlashML-org#414 batch-memcpy probe order), #51 (FlashML-org#118 admission clamp), #52 (FlashML-org#476 --default-thinking-mode), #53 (FlashML-org#340 dummy page in the auto KV floor), #54 (FlashML-org#198 explicit KV pages in auto-sizing), #55 (FlashML-org#459 eager FTW validation), #56 (FlashML-org#487 hoist system messages), #57 (FlashML-org#222 AbortMsg on disconnect, plus 8adde91 turning the request-recorder middleware into pure ASGI so a non-streaming handler can see the disconnect at all).

Considered and not merged this round: FlashML-org#500 (prefill buffer invalidation) fails 3 of upstream's own tests/moe/test_offload.py overlap tests at its own head (the new Triton invalidate kernel receives the tests' CPU tensors), and only matters with prefill overlap on; FlashML-org#499 (host-RAM KV tier) and FlashML-org#300 (KV/MoE ladder) are promising but large, fresh and conflict with FlashML-org#408 / the runtime-rebuild path -- left for a later round; FlashML-org#464 (incremental stop-string matching) conflicts with #47's detokenizer changes and ships no tests; FlashML-org#498, FlashML-org#484, FlashML-org#456, FlashML-org#471, FlashML-org#439, FlashML-org#405, FlashML-org#429, #31, FlashML-org#266, FlashML-org#223, FlashML-org#224, FlashML-org#477, FlashML-org#504, FlashML-org#293 and the ROCm/GLM/DSV4/Gemma PRs are out of scope or superseded.

Testing (RTX 5080, production flags from above; unit suite at 69a5efa, everything live at 8adde91):

Round of 2026-09-19 (head 4bb35fc, base cc1f5c2). Two new merges, one stacked PR each: #58 (FlashML-org#499 host-RAM KV tier, carried as the base of the next one; its KV-page knobs are refused on an nvfp4 pool by 4bb35fc because the mirror does not carry block scales) and #59 (FlashML-org#525 host-resident GDN snapshot tier). The snapshot tier is now on in production (FT_GDN_HOST_TIER=1 FT_GDN_HOST_SLOTS=32, 3.7 GB pinned host RAM): a 43K-token conversation that took 30.5 s to resume after 12 aux calls of ~3.5K tokens (the 8 GPU snapshot slots were evicted; the KV pool had room) resumes in 2.76 s, facts intact, no page/slot leak over 3 rounds; 250K recall 3/3 in 183 s at a 15,502 MiB peak (+134 MB); bench unchanged (865 / 1410 / 1407 prefill, 38.1 / 38.4 / 38.3 decode).

Skipped this round: FlashML-org#527 (presence/frequency penalties -- already carried through SamplingParams and the sampler plan by #47), FlashML-org#500 (still fails 3 of its own test_offload.py tests at 0f26fd7), FlashML-org#300 (needs --moe-cache-auto, base of Aug 30 conflicting with the #22 rebuild path, and every rung growth is a pool rebuild that flushes the cache -- the opposite of the full-context goal), FlashML-org#464 (not perf/feature, conflicts with #47, no tests). Upstream's FlashML-org#471 (greedy rows in mixed batches) arrived with the base and is unioned with #47's sampler in engine/sample.py.

Full pytest tests/ -m "not slow" at 4bb35fc: 2416 passed, 206 skipped, the same 4 environment-only failures.

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 4bb35fc.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt

ascorb12 and others added 25 commits August 23, 2026 20:52
A prompt that fits the model's advertised max_seq_len but exceeds the
allocated KV pool passes admission and is queued forever with no error,
no log line, and an idle engine (issue FlashML-org#111). Clamp the admission bound
to num_pages * page_size and reuse the existing too-long rejection path
so such prompts fail loudly with context_length_exceeded instead.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
stream_with_cancellation reacted to a client disconnect with an unowned
asyncio.create_task(abort_user(uid)): request teardown could outrun delivery and a
failure inside the task degraded to a never-retrieved-exception warning, so the
scheduler kept decoding for a client that was gone. Await the abort inline behind
asyncio.shield (a second cancellation cannot kill the delivery task), and make
abort_user claim the uid first — exactly one AbortMsg even if cancellation runs
twice, and none at all when the stream already finished normally.

Found via the freetoken-mlx downstream audit (docs/AUDIT.md, defect 2).
Follow-up to the review evidence on FlashML-org#222 (benwilson): the non-streaming path
kept generating after the client was gone -- with --max-running-requests 1 an
abandoned request is a full outage for its remaining max_tokens (measured
repro: ~70 s of dead decode, the next client's first token 61 s late).

stream_with_cancellation was the only place a disconnect was observed. Give
the plain handlers its non-streaming twin: _await_watching_disconnect() runs
the generation drain as a task and polls request.is_disconnected() once a
second; when the client goes away it delivers the same shielded abort_user
(claim + AbortMsg first, so the drain task's own cleanup cannot swallow the
claim), then winds the drain down and answers 499 (client closed request --
for the access log; the wire is dead). Handler cancellation (server shutdown)
delivers the abort too, mirroring the streaming path.

Covers /v1/chat/completions and each prompt of a non-streaming
/v1/completions batch. Requests with request=None (adapter-internal callers)
are unaffected.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… floor

plan_cache_budget works in USABLE pages, while create_kv_pool allocates
num_pages + 1: every pool family keeps page 0 as an unreachable
dummy/sentinel. The plan prices the KV half as num_pages * cache_per_page
and hands experts everything the reserve does not protect (the fill is
greedy and deliberately leaves no headroom), so it under-counted the KV
pool by exactly one page's worth of bytes, and that page was then
allocated out of memory the split had already given to expert slots. On
a model with a large cache_per_page that page is not small, and the
failure mode is the one --moe-cache-auto exists to prevent: an
arithmetic plan that says it fits, followed by a CUDA OOM in the KV
allocation at boot.

The same off-by-one made --kv-reserve-tokens quietly under-deliver: it
is documented as a KV-cache token floor, but ceil(N / page_size) usable
pages minus the dummy is less than N tokens of reachable capacity.
Reserve the sentinel page on top of the user-visible floor and say so in
the --kv-reserve-tokens help text.

Tested (Windows, RTX 5090):

  python -m pytest -q tests/engine tests/server

  before: 643 passed, 2 failed   after: 644 passed, 2 failed

The two failures are pre-existing on this box and unrelated: both
tests/engine/test_cache_budget.py::test_adjust_config_* need flashinfer,
which is not installed here.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RG8BXfsSZi1nh4wMZnhJQK
_probe allocates its destination with torch.zeros, which enqueues the fill on
the current stream, then enqueues the verification copy on a fresh probe
stream. Nothing joins the two. When the current stream has a backlog the copy
completes first on the independent probe stream and the fill lands on top of
it, so the probe reads back zeros and load_batch_memcpy raises
"cudaMemcpyBatchAsync probe copied wrong bytes" on a GPU that supports the API.

OffloadMoeCache catches that and falls back to full-layer copies, so
--moe-prefill-hit-d2d silently does nothing whenever the current stream is busy
as the probe runs -- in practice during prefill warmup, which is exactly when
the flag is first exercised.

A cold process hides the bug: the first torch.zeros pays a cudaMalloc and the
first torch.cuda.Stream() populates the per-device stream pool, and each of
those synchronizes the device, draining the backlog before the copy is
enqueued. That is why the probe passes when run standalone and fails inside a
warmed-up server.

Join the probe stream to the current stream before the copy.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…inally

On a disconnect the CancelledError unwinds from the response task's
innermost await -- wait_for_ack -- whose finally pops ack_map and
event_map on the way out. stream_with_cancellation then called
abort_user, whose claim was `event_map.pop(uid) is not None`: always
False by that point, so the AbortMsg was never sent and the scheduler
kept decoding to max_tokens, with no log line either (the cancellation
interrupts the async for, not the is_disconnected branch). A body
closed after the response started arrives as GeneratorExit, which the
`except CancelledError` never saw either. Diagnosed with a live repro
by gdevenyi on the PR thread.

Key the claim on an insertion-ordered aborted_uids dict, FIFO-bounded
like the scheduler's abort tombstones (a duplicate AbortMsg is safe:
the scheduler acks aborts for uids it no longer has), and move delivery
into the stream wrapper's finally, gated on the stream not finishing.
One path now covers server-side cancellation, late closes, and
generator errors; the accounting drain's abort barrier also stops
missing requests whose maps were already emptied.

Assisted-by: Claude
…ackend)

Shard the dense weights per rank at load (attention qkv by head, GDN in_proj as
its six parts with the matching conv1d channels and A_log/dt_bias, shared-expert
gate_up per part; o_proj/out_proj/down_proj row-parallel; embed/lm_head by vocab
rows) and the NVFP4 expert banks along the intermediate axis, so every rank holds
half the experts and each MoE layer needs one all-reduce (routed + gate * shared
are combined before the reduce). Router, QSA indexer, norms, hyper-connections
and PLE stay replicated so all ranks select the same blocks and n-gram rows.

Also: LinearColParallelMerged(local_output_sizes=) for the kv-replicated case and
distributed_timeout 60 -> 1800 s (ranks reach their first collective minutes
apart behind a 100+ GiB load).

Limits: offload backend with bf16 dense projections; fp8_block / nvfp4 dense
checkpoints raise under TP.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
tests/models/qwen4_exp/test_weight.py feeds iter_weights a synthetic checkpoint whose
config.json has no model_type; at TP=1 nothing is sharded, so do not touch the config.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
… _scaled_mm)

Opt-in with FREETOKEN_FP8_DENSE=1 on a bf16-dense checkpoint (e.g. the
RadixArk NVFP4 build): the weight reader quantizes qkv_proj / o_proj, GDN
in_proj (q|k|v|z; the b|a gate rows stay bf16 as in_proj_ba) and out_proj to
per-tensor e4m3 after TP sharding, and layers/fp8_dynamic.py runs them as
cuBLASLt W8A8 GEMMs with a dynamic per-tensor activation scale (one fused
Triton launch at decode sizes; no host sync, CUDA-graph safe). Column-merged
and row-parallel variants, so it works at TP>1.

Why: on an RTX 6000 Ada (sm_89, torch 2.11.0+cu130) these projections are
2.67 GB of the ~4 GB a TP=2 rank reads per token; bf16 cuBLAS takes
3.2-3.4 ms per step per rank, raw _scaled_mm 1.9 ms, while the existing
Triton FP8 kernels are slower than bf16 there (measured, weights rotated
past the L2).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
…fore the cache planner runs

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
t[qkvz:].contiguous() on a contiguous row slice returns a view, so every GDN
layer's bf16 gate rows kept the whole sharded bf16 in_proj resident next to
the fp8 copy: 36 x 42 MB = 1.5 GiB per TP=2 rank, which is why the expert
cache planner saw no saving (22,594 -> 22,458 slots) after the FP8 switch.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
quant_per_tensor ran a single Triton program over the whole tensor, so its cost
grew linearly with the input: 9.8 us for an [8, 10240] activation and 68.6 us at
[64, 10240], against ~1.2 us of useful work. Above 16384 elements, split it into
a block-parallel partial amax and a reduce+cast, which flattens the cost to
~2.4 us. The arithmetic is the one the single-program kernel already used, so a
given tensor quantizes bit for bit as before; the old three-launch torch
reduction path above 65536 elements goes away with it.

The partial count is a runtime argument and the pass-1 grid is strided and capped
at _MAX_PARTS, so both kernels have exactly ONE compiled Triton variant. Letting
the partial count reach the kernel as a constexpr instead costs a fresh
compilation for every distinct input length -- unbounded variant growth in a
server that sees arbitrary prompt lengths, and a compile stall mid-generation.

Measured on an RTX 6000 Ada (sm_89, torch 2.11) under CUDA-graph capture,
one-program -> split: 20480 elts 2.59 -> 2.22 us, 81920 elts 9.83 -> 2.44 us,
655360 elts 68.60 -> 3.06 us. Below the threshold the single program still wins
(2560 elts: 1.25 vs 2.07 us) and is kept.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
(cherry picked from commit d254729)
The vocab-parallel head is the last large bf16 read on the decode path: 0.64 GB
per step per rank at the Qwen3.8-Flash-Next geometry ([124160, 2560] after the
TP=2 vocab split), 0.80 ms of the 9.9 ms step in an nsys trace of production.

Its own flag, not FREETOKEN_FP8_DENSE, because this one moves the logits: every
other quantized module feeds a norm or a sigmoid downstream, while a per-tensor
e4m3 vocab matrix changes each sampled token's score directly, so it carries its
own quality gate rather than riding along with the pure-throughput changes.

ParallelLMHead.forward grows a _logits() seam (the local vocab-shard GEMM);
Fp8ParallelLMHead overrides only that, leaving the all_gather of the logits
above it untouched. Untied embeddings only -- a tied head shares the bf16
embedding table, which the lookup side still reads as bf16.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
(cherry picked from commit 35e80b3)
(cherry picked from commit 0693735)
(cherry picked from commit 1797ab5)
…build

Qwen4ExpDecoderLayer builds its MoE as `Qwen4ExpMoE(config, layer_id, prefix=...)`,
but this PR's override of __init__ (added to hold the TP communicator) took only
(config, layer_id), so a server boot died with

    TypeError: Qwen4ExpMoE.__init__() got an unexpected keyword argument 'prefix'

The whole CPU test suite was green with that bug in place, because every test that
builds a decoder layer is behind requires_cuda -- nothing without a GPU ever
constructed the model. tests/models/qwen4_exp/test_build_cpu.py closes that: it
builds the full model on the meta device (no GPU, no memory) and asserts the state
dict has both layer families, an lm_head, and MoE weights on more than one layer,
so a dropped or shared prefix fails too.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
…el shard

With the expert piece stream sliced per rank and the banks sized from
MoEConfig.local_intermediate, a rank holds exactly its half of every expert, so this
kernel can serve TP>1 -- the routed output is a partial sum and the MoE layer already
reduces it (_maybe_all_reduce, or the single combined all-reduce in qwen4_exp's block).

Without this the whole selection table is empty under TP=2 on sm_89 and the server
refuses to start:

    KernelSelectionError: no usable kernel in table;
      triton: TP > 1 is not supported for this expert format;
      marlin: vLLM is not installed;
      b12x: b12x requires sm_120+, got sm_89

marlin and b12x keep tp_ok=False deliberately: their pack() repacks the native rows and
neither has been verified against a per-rank bank.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
…d loader alike

The model builder and the weight reader tested this separately and drifted, so on the
shipping checkpoint the reader emitted lm_head.weight_scale while the builder made a
plain ParallelLMHead, and startup died with

    RuntimeError: Unexpected keys in state_dict: ['lm_head.weight_scale']

The builder's extra condition was `config.quant is None`, which is wrong twice over:
what matters is whether the checkpoint declares a scheme for **lm_head**, not whether it
has a QuantConfig at all. This model ships NVFP4 routed experts (so quant is not None)
with lm_head in the modelopt ignore list (so it has no scheme and the synthetic FP8 head
is exactly what is wanted).

Both sides now call config.use_fp8_lmhead(), which owns the flag, the tied-embedding
exclusion and the scheme test in one place.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
The predicate moved into config.py but its import did not come with it, so building the
model raised NameError on the very first line of use_fp8_lmhead -- under every flag
setting, including the default one that wants no FP8 head at all.

Found by the CPU key check (build the model on meta, diff its slots against the names the
loader emits) before it could cost a GPU window. No behaviour change beyond the module now
importing.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
The rebase onto the quantization refactor dropped

    nk, nv = self._local_num_k_heads, self._local_num_v_heads

from the top of GatedDeltaNet.forward while keeping all eleven uses of nk and nv below
it. Every decode died at the first CUDA-graph capture:

    File "python/freetoken/models/qwen4_exp/gdn.py", line 189, in forward
      b, a = torch.split(ba, [nv, nv], dim=-1)
    NameError: name 'nv' is not defined

Also drops a duplicate `from freetoken.distributed import get_tp_info` two lines under
the first, left by the same merge.

`ruff check --select F821` reports both, and would have reported the missing
fp8_lmhead_enabled import a commit earlier. Lint the tree before asking for the GPU, not
after: this file's eleven undefined names cost a window that a sub-second check would
have saved. The eight F821 hits outside qwen4_exp are pre-existing on main.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
ZmqPullQueue/ZmqSubQueue/ZmqAsyncPullQueue decoded each received frame
with msgpack.unpackb, which raises ExtraData when a frame carries more
than one packed object. On Windows with offload MoE this surfaced in
production as two mid-batch worker deaths in 81 requests: the tokenizer
and scheduler processes crashed on the same helper
(utils/mp.py), taking the whole API server down.

Route every received frame through one buffered msgpack.Unpacker per
queue: the first object is returned by get() and any remainder is held
for subsequent calls instead of crashing. empty() accounts for the
buffer, and get_raw()/decode() stay unbuffered (the multi-rank
broadcast path pairs them with a rank-wide count that a buffered
remainder would desynchronize) -- mixing them with buffered get() now
raises instead of silently skipping messages.

Regression tests cover the coalesced frame on sync pull, sub queue,
decode buffer, and empty() accounting; the async variant is skipped on
Windows (zmq.asyncio needs a Selector loop; prod runs on Linux CI).

Fixes FlashML-org#452
…ng default

OpenAI-compatible clients that never send chat_template_kwargs (Vercel AI
SDK openai-compatible provider, llama-swap, LiteLLM proxies) currently get
the model's template default, which for reasoning models like Qwen3.6 means
thinking is always on — the token budget goes to reasoning_content and
content comes back empty.

--default-thinking-mode {auto,chat,thinking} lets the operator pick the
server-wide default:

- auto (default): current behavior, template decides
- chat: fills enable_thinking=False into every request that does not set
  any explicit thinking control
- thinking: fills enable_thinking=True the same way

An explicit per-request value (enable_thinking / thinking / thinking_mode
in chat_template_kwargs, or the DeepSeek thinking.type wire toggle) always
wins over the server default.

Applied to all three frontends: OpenAI chat completions, Anthropic
messages, and Responses API.

Closes FlashML-org#472
…acked on FlashML-org#385 TP)

# Conflicts:
#	python/freetoken/models/qwen4_exp/attention.py
#	python/freetoken/models/qwen4_exp/config.py
#	python/freetoken/models/qwen4_exp/weight.py
… read

Assisted-by: Claude Opus 5
(cherry picked from commit 82a21ae)
(cherry picked from commit 0369ed9)
… cpu build test

Assisted-by: Claude Opus 5
(cherry picked from commit 27ad92e)
(cherry picked from commit e152f66)
…s the quantized-kv backend reads

Assisted-by: Claude Opus 5
(cherry picked from commit 600819f)
(cherry picked from commit dd2ebdf)
…e Hugging Face repo ids to a local checkpoint dir
…verage

# Conflicts:
#	python/freetoken/tokenizer/tokenize.py

# Conflicts:
#	python/freetoken/engine/sample.py
#	python/freetoken/tokenizer/tokenize.py
…onfig

Fork #19 asks the model family to check the stored fp8 split against the env flags; a
synthetic FTW written straight from tensors (upstream's test_ftw_weights) carries no
config.json, so there is no family to ask and the read proceeds as before.

Assisted-by: Claude Fable 5.1
(cherry picked from commit cab9be3)
…the admission clamp reads

Assisted-by: Claude Fable 5.1
(cherry picked from commit 4584957)
…the --moe-cache-auto KV floor

# Conflicts:
#	python/freetoken/engine/cache_budget.py
#	tests/engine/test_cache_budget.py

# Conflicts:
#	python/freetoken/engine/cache_budget.py
#	tests/engine/test_cache_budget.py
…MoE auto-sizing

# Conflicts:
#	tests/engine/test_cache_budget.py

# Conflicts:
#	tests/engine/test_cache_budget.py
…ream cancellation

# Conflicts:
#	python/freetoken/server/api_server.py
#	python/freetoken/server/openai_api.py

# Conflicts:
#	python/freetoken/server/api_server.py
#	python/freetoken/server/openai_api.py
…isconnect

Guards the FlashML-org#222 resolution against FlashML-org#393's fan-out: the disconnect watcher takes the
whole uid list, not the first sample's uid.

Assisted-by: Claude Fable 5.1
(cherry picked from commit 93ff96349235ea3ca1504733a2abbd6e6bcccf94)
(cherry picked from commit 69a5efa)
…can see a client disconnect

Starlette's @app.middleware("http") (BaseHTTPMiddleware) wraps the handler's receive channel
and never yields the client's http.disconnect while a non-streaming request is still
computing, so request.is_disconnected() stays False and the disconnect watcher from
FlashML-org#222 could not stop an abandoned request. Streaming requests only
worked because the socket write failed. The recorder is the one such middleware; as pure
ASGI it records the same row at response start.

Assisted-by: Claude Fable 5.1
(cherry picked from commit 8adde91)
…ention models (--kv-host-pages)

# Conflicts:
#	python/freetoken/attention/qsa_sparse.py
#	python/freetoken/kvcache/__init__.py
#	python/freetoken/kvcache/qsa_pool.py
…cales they cannot mirror

Upstream FlashML-org#499/FlashML-org#525 mirror KV codes plus the fp8 pool's per-token _scale_buffer; the nvfp4
pool (FlashML-org#408) keeps per-block scales in _block_scale_buffer, which neither the offloader nor
the prefix tier copies, so --kv-host-pages or FT_PREFIX_HOST would rehydrate codes without
scales. Both constructors now raise a clear error naming the knob; the GDN-snapshot-only
tier (FT_GDN_HOST_TIER) never reads the KV pool and stays usable.

Assisted-by: Claude Fable 5.1
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.