Conversation
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).
Adds the engine half of OpenAI logprobs support. SamplingParams gains logprobs/top_logprobs; the sampler computes log_softmax over the PRE-temperature logits (raw model distribution, so temperature/top-k/top-p do not change reported values), gathers the chosen token and batch-max top-k, and ships CPU copies covered by the existing copy_done_event. The scheduler attaches per-request values (cut to each request's own top_logprobs) to DetokenizeMsg; the detokenizer builds a neutral entry (token text via single-id decode + UTF-8 bytes, so clients can reassemble partial-UTF-8 pieces) onto UserReply.logprobs. Zero cost when off: no row asked -> no mask tensor, no log_softmax, ForwardOutput carries None. Message fields default to None, so old and new peers interoperate. Stop-string trimming can hide final visible text; entries still cover every sampled token.
API half of logprobs support. Chat: logprobs + top_logprobs (0..20) yield
choice.logprobs.content entries {token, logprob, bytes, top_logprobs[]},
streaming and non-streaming. Completions: the legacy integer field (0..5) yields
{tokens, token_logprobs, top_logprobs, text_offset} with absolute offsets across
the stream; echo+logprobs stays rejected (prompt logprobs need prefill logits
and are out of scope here).
The protocol-neutral event layer carries entries on ContentDelta (a list -- parser
buffering can release several tokens' text in one delta); entries always
accumulate on GenResult for the non-streaming path, and reasoning/tool-call
buffering carries pending entries onto the next content delta. Formatting lives
in server/logprobs.py; wire compatibility follows the OpenAI shapes.
Review finding on FlashML-org#224 (HaileyStorm): streaming chat carried logprob entries of hidden reasoning tokens onto the next visible content delta -- pending entries accumulated before the reasoning parser classified the text, and _content_delta() drained the whole list into whatever visible chunk came next. That broke token<->content alignment and leaked hidden reasoning token strings through logprobs entries; non-streaming meanwhile kept every sampled entry and end-of-stream dropped undrained ones -- three behaviors for one surface. Close it at both layers: - /v1/chat/completions rejects logprobs=true with a 400 when the server runs a reasoning parser or the request enables tool parsing: the semantic layer can hide or reclassify tokens, so entries cannot be aligned 1:1 with visible content tokens. tool_choice="none" keeps logprobs available; the raw /v1/completions path (no semantic layer) is unchanged. - _generate_events_impl now collects entries only on the passthrough path (no reasoning parser, no tool parsing), so a hidden-token entry can never ride a later visible delta even for a caller that skips API validation. Replaces test_reasoning_logprob_is_carried_to_next_content_delta (which enshrined the carry) with fail-close coverage for both conflict cases, the tool_choice="none" pass-through, and a no-carry guard test. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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
The logprobs change moved the drain from positional unpacking to named ForwardOutput attributes (next_tokens_cpu, copy_done_event, the *_cpu logprob columns), which the bare-tuple fake in the in-flight-abort scheduler tests predates: three of them died with 'tuple' object has no attribute 'next_tokens_cpu' on a GPU run. Caught by gdevenyi on the PR thread; a macOS run cannot see it (the file already fails on the missing kernel packages there, identically on both sides of the series). A namespace mirroring the drain-facing fields keeps the test off the engine import chain, whose kernel packages a scheduler unit test should not require. 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
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
…a pooled one torch.cuda.Stream() hands out 32 pooled streams per priority round robin, so the side stream could come back as the compute or capture stream (it did in the full test run): the shared expert then ran serially, outputs unchanged but the overlap gone. Create it with cudaStreamCreateWithFlags and wrap it in an ExternalStream. Assisted-by: Claude Opus 5.5
…ss copy runs on Assisted-by: Claude Opus 5.5
…expert beside the miss copy (FREETOKEN_MOE_COPY_OVERLAP), opt-in slim copy
…KEN_MOE_PREFETCH Stage A of cross-layer expert prefetch: in GPU decode, layer L-1's router input times layer L's router predicts layer L's experts on a dedicated stream forked after ensure(L-1); the first BUDGET non-resident candidates of the top-K are joined and counted against layer L's routing after its ensure. No copies and no cache writes; off by default and launch-identical when off. Assisted-by: Claude Opus 5.5
…d its measure mode Assisted-by: Claude Opus 5.5
Assisted-by: Claude Opus 5.5
…gh layer l's router The counter test scores the spied logits, so it cannot see a wrong router or input; this one feeds distinct per-layer inputs and recomputes each prediction's logits on the host. Assisted-by: Claude Opus 5.5
…FREETOKEN_MOE_PREFETCH=on The predictor stream installs layer L's kept candidates with a prefetch mode of the scored ensure kernel (no clock bump, usage == step stays pinned, usage 0 marks the slot low priority) and a third dedicated stream slim-copies them after copy(L-1). The compute stream joins the install before ensure(L) and the copy before copy(L); ensure(L) counts useful hits, calls, misses and late copies. LRU runs the vendored kernel while prefetch is on. Assisted-by: Claude Opus 5.5
…der delayed streams A delayed install that lands after ensure(L) changes no output, so the race test now also compares the delayed run's slot maps and counts with an undelayed run's. Assisted-by: Claude Opus 5.5
…a cache frees its banks The prefetcher's back-reference made a cycle, so a dropped cache kept its slot banks in VRAM until the next gc pass. Assisted-by: Claude Opus 5.5
The slim copy is an SM kernel. At default priority a copy that starts even 50 us late waits for SMs behind the next compute grid while the compute stream waits on it: +13.7% per step vs off at bs 1 in bench_moe_prefetch --on, against -16.3% at high priority. No measurable cost when nothing is late. Assisted-by: Claude Opus 5.5
…the prefetch bench Assisted-by: Claude Opus 5.5
…ls nothing Assisted-by: Claude Opus 5.5
…d graph interleaving Adds a stall before each decode GEMM (copy(L+1) lands before GEMM(L) reads; flood candidates reach the pinned slots even under lru), prefills between decodes on every prefill path (materialize, overlap, overlap + hit-D2D, small ensure), bs 1/2 graphs in one pool replayed between eager steps, a per-layer delay fuzz over all four streams, the per-bank copy fallback, and the kernel against the reference at 48 x 512 experts, 1480 slots and 8 rows. Assisted-by: Claude Opus 5.5
prefetch_ensure(L) scored slots as if layer L had just run, so L's own residents looked a whole pass away and were the first residents it evicted, right before L routed them. On the steady-state trace bench this cost +0.18 (bs 1) and +0.42 (bs 2) extra demand misses per layer; scored from layer L-1 the excess is gone and on-mode steps are 1.6-2.5 pts faster. Assisted-by: Claude Opus 5.5
--steady replays a recorded routing trace with no cache restore, a noisy copy of each layer's traced logits as the lookahead and per-layer-kind fillers, and interleaves the arms step by step (each on its own copy of the cache state) so drift hits them alike. It reports pollution and the late fraction before GDN and QSA layers next to ms/step. Assisted-by: Claude Opus 5.5
…and relax the eager late bound Every offload cache under test makes dedicated streams and torch keeps a cuBLAS workspace per stream for the session (~16 MiB a prefetch stack): the prefetch suites grew one pytest process past 8 GiB and hit OOM beside another GPU user; now the peak is ~1 GiB. The eager delayed-copy case asserted late >= 0.9 of copies, which a loaded host misses (3 of 10 loops); replay keeps the exact late == copied check. Assisted-by: Claude Opus 5.5
…teardown Assisted-by: Claude Opus 5.5
… the gpu A process-wide FREETOKEN_MOE_PREFETCH made every hybrid or CPU-decoded cache raise at build (two hybrid-fetch and two overlap-flag tests failed under the flag). The flag now logs and leaves those caches off, as FREETOKEN_MOE_COPY_OVERLAP does; an explicit prefetch_mode still raises. Assisted-by: Claude Opus 5.5
…behind FREETOKEN_MOE_PREFETCH_VERIFY Debug detector for the prefetch-on nondeterminism. For every GPU decode layer call, on the compute stream and without any event or wait: the routed slots map to the router's experts (before and after the GEMM) and, with =1, every routed slot's bank rows equal the host rows before and after the GEMM (a 53 MiB scratch filled by the fused copy). Per-layer device counts and the first 64 offending records (with the slot owner, usage and which copy plan held the slot) are logged every 256 decode steps and at shutdown. =meta checks only the maps (~1% decode cost); =1 costs 2.4-3.3x. Assisted-by: Claude Opus 5.5
…een decode steps Assisted-by: Claude Opus 5.5
…d FREETOKEN_MOE_SLOT_AUDIT Assisted-by: Claude Opus 5.5
Assisted-by: Claude Opus 5.5
… sizes and eager steps Assisted-by: Claude Opus 5.5
…write has one writer The predictor stream ran prefetch_ensure(L), a slot map writer, guarded from the demand ensure(L) only by a cross-stream event; a production verify record shows the two overlapped. The predictor now forks at the start of layer L-1's MoE block and only computes the merged lookahead candidates (it reads no cache state); right after ensure(L-1) the compute stream waits for them and runs prefetch_ensure(L), which applies residency and the budget itself, so the installs equal the old ones and do not depend on stream timing. join_ensure is gone. Assisted-by: Claude Opus 5.5
… right after a prefill Assisted-by: Claude Opus 5.5
…-writer test At d6a0a99 the select assertion came first and raised AttributeError on _none_resident, so the writer-stream checks never ran against the old engine. Now it fails there on prefetch_ensure off the compute stream (16/16). Assisted-by: Claude Opus 5.5
…its negative control prefetch_ensure(L) runs on the compute stream before join_copy(L-1), so it can take over a slot copy(L-1) is still writing; count that on the device under late copies and check no bit changes. Ordering copy(L) only after its own install, on a second stream, corrupts the slot. Assisted-by: Claude Opus 5.5
…writer slot maps (FREETOKEN_MOE_PREFETCH), plus decode verify and slot audit debug modes
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Merges upstream FlashML-org#389 (
feat(qwen4_exp): load-time per-tensor FP8 dense projections (W8A8 via _scaled_mm), headbd106c9a). 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:models/qwen4_exp/attention.py:mainpassesbatch.get_attn_positions()(mRoPE, feat(mm): serve image input on the Qwen families FlashML-org/FreeToken#454), while feat(qwen4_exp): tensor parallelism for Qwen3.8-Flash-Next (offload backend) FlashML-org/FreeToken#385 passesbatch.positionswith per-shard sizes. Resolved tobatch.get_attn_positions(), q.view(-1, self._local_qo_dim), k.view(-1, self._local_kv_dim).models/qwen4_exp/config.py: imports from both sides (fp8_dense_enabled,fp8_lmhead_enabled,mrope_layout_from_rope_params).models/qwen4_exp/weight.py, the only non-mechanical resolution:mainreplaced_try_fusewith the QuantConfig-driven_DenseFuser(block-FP8 dense support) and addedinclude_visionfiltering. feat(qwen4_exp): load-time per-tensor FP8 dense projections (W8A8 via _scaled_mm) FlashML-org/FreeToken#389's loop still called_try_fuse.main's loop (_DenseFuser.fuse/check_unfused, vision skip) and routes every tensor it yields through feat(qwen4_exp): load-time per-tensor FP8 dense projections (W8A8 via _scaled_mm) FlashML-org/FreeToken#389'semit(), which shards for TP and then applies_fp8_densewhen enabled.torch.cuda.empty_cache()after FP8 quantization is kept.Local fixes in the merge commit:
models/config.py(silent semantic conflict): git merged cleanly, butmainhad removedimport osand_VISION_TRUE, and feat(qwen4_exp): load-time per-tensor FP8 dense projections (W8A8 via _scaled_mm) FlashML-org/FreeToken#389'sfp8_dense_enabled/fp8_lmhead_enabledstill use both. Every qwen4_exp test failed withNameError: name 'os' is not defined(70 failed, 11 errors). Restoredimport osand 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_ropeisfunctools.cached, so the CPU cos/sin table leaks into later CUDA tests.test_skeleton.py::test_qsa_layer_matches_hf_denseand::test_decoder_stack_prefill_and_decodethen fail with FlashInfer'scos_sin_cache must be a CUDA tensor, but only when run aftertest_build_cpu.py; alone,test_skeleton.pypasses 10/10.rotary.get_rope.cache_clear()in afinally.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 twotest_glm_dsa.pysm_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:Serving with
FREETOKEN_FP8_DENSE=1 FREETOKEN_FP8_LMHEAD=1. nvfp4 KV, 262,144 tokens in every run:(prefill / decode in tok/s)
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.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.Deployment command (top of the stack, verified on the final restart:
/v1/modelsreports 262144, adeveloper-role request answers, and a$refinteger tool argument comes back as{"limit": 7}):Refs #2, #8.
Environment:
nvidia/Qwen3.8-Flash-Next-NVFP4Stack: the
5080deploybranch ismaine0886cc 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.Later merges on this branch (after 2d56708)
This PR's head is
5080deploy, so it also carries the merges below. Branch tip5767c99.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.
2a98fa4..61128bbd125f635767c99tokenizer/tokenize.py, keeping both the developer-role mapping (FlashML-org#391) andcontinue_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
bb37974and was removed by rebuilding the branch from 2d56708. The removed file set is exactlye0886cc..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 twotest_glm_dsa.pyfailures are gone with fix(triton): keep the float glm dsa sparse launches two-stage to fit sm_120 shared memory #21./v1/models262144.FREETOKEN 5080is read correctly. A 1920x1080 image's text box is read exactly (2,070 prompt tokens, 3.6 s).developersystem instruction is followed (German answer);$reftool argument{"limit": 7};n=2returns two choices;logit_bias-100 on the first token of the unbiased answer changes the generated tokens;/tokenizeand/detokenizeround-trip.Considered and not merged: #39 (FlashML-org#85, e4m3 rounding before the native downcast).
.to(tl.float8e4nv)already matches torch's RNE byte for byte: 0 mismatches over 4,194,304 uniform values in [-448, 448], including the clamp-then-cast shapelayers/fp8_dynamic.pyuses.round_e4m3maps tiny negative values to+0.0where torch stores-0.0(5 of the 4,194,304 values). Equal in value, but not byte-identical.Rebased 2026-09-17 onto upstream main
cac247a(v0.1.3); head is now8adde91. The stack was rebuilt merge by merge on the new base (resolutions replayed viagit rerere; the replayed stack (before the new merges) differs from the old backup tipbb37974by exactly the3faef36..cac247afile set). Upstream's 6 new commits: FlashML-org#489 release 0.1.3, FlashML-org#486 FTW/vision-encoder loading (ftw_lacks_visionrefusal,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 (theirrebased/*branches sit on 3faef36, now an ancestor),dd2ebdf(block_ends test fake getsk_scale/v_scale) andcab9be3(fork #19's FTW layout check skips a bare FTW without an HF config, which upstream FlashML-org#486'stest_load_weight_text_only_skips_the_towerwrites).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, plus8adde91turning 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.pyoverlap 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 at8adde91):pytest tests/ -m "not slow"at69a5efa(the tip before8adde91, which touches onlyserver/api_server.pyand adds one test;tests/serverrerun at8adde91: 871 passed): 2412 passed, 206 skipped, 4 failed -- all environment-only: 3xtests/models/test_muse_glimmer_vision.pyfail identically on pristinecac247a(transformers 5.15.1 has nopaddingkwarg onget_vision_interpolation_indices_and_weights), andtest_async_pull_coalesced(fix(utils): tolerate coalesced msgpack frames in zmq pull queues FlashML-org/FreeToken#466) is anasync deftest needingpytest-asyncio, which the project does not declare (passes 5/5 withuv run --with pytest-asyncio)./v1/models262144; the FTW-fp8 index passes fix(checkpoint): validate FTW v1 indexes eagerly FlashML-org/FreeToken#459's eager validation.$reftool arguments,n=2,/tokenize+/detokenize,logit_bias, image OCR (FREETOKEN 5080) and a shapes image all as before.--disable-moe-prefill-overlapwas dropped from the launch script after an A/B on this tip: prefill 866 / 1407 / 1411 tok/s vs 526 / 1175 / 1203 at 1.2K / 30K / 90K, decode unchanged (38.5 / 38.9 / 38.5), 248K recall 3/3 in 182 s vs 215 s, peak VRAM 15,368 MiB of 16,650 at both 90K and 248K (the peak is the 8,192-token chunk working set, not context-dependent).Correctness re-review, 2026-09-23 (head
8975217). A full re-review of the stack found 11 bugs, fixed in #64 with a failing-then-passing test each. The worst: the GDN snapshot tier (#59, FlashML-org#525) could revive a node from a shorter prefix's snapshot (silently wrong output after a conversation diverged), and FlashML-org#338's PLE index cache let the decode CUDA graph read freed memory after ~64 distinct prompt shapes. Full suite at8975217: 2462 passed, the same 4 environment-only failures.Round of 2026-09-23b (head
77d8d5c). Feature merges after a review of all open upstream PRs: #61 (FlashML-org#108GET /v1/models/{id}), #62 (FlashML-org#504 opt-in--enable-metrics, off in production), #63 (FlashML-org#224 logprobs -- on/v1/completionshere; chat logprobs are refused by design while theqwen3reasoning parser is active). Nothing else open upstream improves this deployment's speed; the evaluation is summarized in the session report. Full suite at77d8d5c: 2446 passed, the same 4 environment-only failures.Round of 2026-09-23 (head
2419ebc, base71eb1f0). Upstream moved 2 commits (FlashML-org#545 flashinfer version pin for installs, FlashML-org#469 GLM-4 router bias); the stack replayed with no conflicts and its diff to4bb35fcis exactly that delta. New merge: #60 (FlashML-org#500, now passing its own tests). Skipped: FlashML-org#531 (clamps/v1/modelscontext to the KV pool -- a no-op here, pool = 262,144), FlashML-org#491 (refactor), ROCm/GLM/GGUF updates. Full suite at2419ebc: 2417 passed, the same 4 environment-only failures. The stack branches' rebase notes on #9-#16 and #48-#59 refer to the 09-19 heads; every branch was force-pushed onto71eb1f0in the same order.Round of 2026-09-19 (head
4bb35fc, basecc1f5c2). 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 by4bb35fcbecause 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
SamplingParamsand the sampler plan by #47), FlashML-org#500 (still fails 3 of its owntest_offload.pytests at0f26fd7), 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 inengine/sample.py.Full
pytest tests/ -m "not slow"at4bb35fc: 2416 passed, 206 skipped, the same 4 environment-only failures.Rebased 2026-09-19 onto upstream main
cc1f5c2(4 commits pastcac247a: 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 viagit rerere; the replayed stack differs from the previous tip8adde91by exactly thecac247a..cc1f5c2file set. This PR's head is now4bb35fc.🤖 Generated with Claude Code
https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt