Qwen3.6-35B-A3B edge runtime: white-boxed prefill, SM110 support, and token-exact speculative decode - #169
Open
LiangSu8899 wants to merge 86 commits into
Open
Qwen3.6-35B-A3B edge runtime: white-boxed prefill, SM110 support, and token-exact speculative decode#169LiangSu8899 wants to merge 86 commits into
LiangSu8899 wants to merge 86 commits into
Conversation
Development tooling for a memory-constrained Qwen3.6-35B-A3B runtime: - probe.py projects per-category checkpoint sizes and expert-cache quotas for a given memory budget, and scores sampled experts under INT8/INT4. - quantize_experts.py writes routed experts as fixed-size INT8 or INT4 blocks. INT4 follows the sign-magnitude / UE4M3-per-16 contract and optionally applies the orthonormal H16 transform. - route_trace.py records router selections and simulates a bounded per-layer LRU to size the expert cache. The router trace hook in the shared decode path is eager-only and stays disabled unless a caller opts in, so CUDA Graph capture is unaffected.
The family was behind a single FLASHRT_ENABLE_QWEN35MOE gate that also required a Blackwell NVFP4 build, so a target that can run most of these kernels could not select them. Only four of the fourteen actually need block-scaled MMA. Three gates now describe what each kernel needs: - _CORE: layout/split, bf16 matvec, router top-k, activation fusion, GDN recurrence, weighted-sum reducer, bf16 GEMM. bf16 intrinsics plus cp.async and mma.m16n8k16.bf16, so SM80 and newer. - _W4A16: weight-only 4-bit matvec, grouped matvec and GEMM. Operands go through __nv_cvt_fp4x2_to_halfraw2 and accumulate in bf16, needing no block-scaled MMA; SM89 and newer get the hardware conversion. - _W4A4: grouped GEMV and the M16/M64/block-tile MMA kernels, which need the sm_120a/sm_121a CUTLASS path. _W4A4 refuses to configure without that path. CUTLASS compiles those translation units on other architectures but substitutes CUTE_INVALID_CONTROL_PATH for the MMA, so the build would succeed and the kernels would fail when called; the gate turns that into a configure error. FLASHRT_ENABLE_QWEN35MOE stays as the all-tiers switch, so existing configure lines and the frontend's fail-fast message are unchanged. The bindings are regrouped to match the three macros; no binding signature or symbol name changed. Verified: sm_120 with the alias enables all three tiers and exposes the same 16 bindings; sm_87 configures _CORE and _W4A16 and compiles all ten translation units; sm_87 with _W4A4 fails at configure.
A single per-layer LRU is the wrong model for this runtime. Prefill selects far more experts per layer than the cache can hold, so it evicts whatever decode is about to need and the measured hit rate reflects prompt churn rather than routing locality. route_trace now scores each quota under three policies: the existing single LRU, a warm set pinned from prompt-phase selection counts plus an evictable ring, and the same split with an oracle warm set taken from the decode phase. The oracle is not implementable and is there to bound what a better warm-set heuristic could add. Misses per token are also converted to a read volume and to the token rate each storage bandwidth would allow, since that is what decides whether a given memory budget is viable at all.
The INT8 block payload is 3,151,872 bytes, which is 769.5 times the 4096-byte logical block size, so neither its offset nor its length could be used with O_DIRECT. That matters because the target device's memory holds only a fraction of the experts: streaming them through the page cache would make the cache compete with the resident weights for the same physical memory, so the reader has no choice but to bypass it. Blocks now carry a trailing pad to the next 4096-byte boundary. The INT4 group-16 payload was already a multiple and is unchanged at 1,769,472 bytes; INT8 takes 2048 bytes of pad and becomes 3,153,920. manifest.json records the alignment and the padding entry so a loader can compute component offsets without reproducing the arithmetic. No bundle had been generated against the unpadded layout.
The weight-only 4-bit kernels pulled <cuda_fp4.h> for one function, __nv_cvt_fp4x2_to_halfraw2. That header first appears in CUDA 12.8, so the kernels could not be built for a Jetson image pinned to 12.6 even though nothing else in them needs a newer toolkit. On an architecture without cvt.rn.f16x2.e2m1x2 the header decodes each nibble in software anyway, by way of an intermediate E2M3 conversion. E2M1 has sixteen representable values, so a table of raw half bit patterns gives the same answer on any target. The new header uses <cuda_fp4.h> where it exists, so the emitted code on current targets is unchanged, and falls back to the table otherwise. FLASHRT_FP4_FORCE_TABLE selects the table explicitly. Verified on sm_120a, which has the hardware conversion: both paths produce identical half bit patterns for all 256 packed byte values. With the fallback available, all ten core and weight-only-4-bit translation units compile for sm_87 under CUDA 12.6 and gcc 11.4, and for sm_110 under CUDA 13.0.
The two-tier docstring implied that pinning a prompt-derived warm set is the correction to a plain LRU. Measurement on Qwen3.6-35B-A3B says otherwise: with a per-layer quota already in place the plain LRU wins from 16 slots up, and its margin grows with prompt length rather than shrinking -- 0.745 against 0.731 at 43 slots for a 32-token prompt, 0.711 against 0.664 for a 128-token prompt. Recency predicts this router's next selections better than prompt-phase frequency, and a longer prompt spreads the frequency estimate over more experts instead of sharpening it. The oracle warm set stays ahead of both, so the weakness is the prompt-derived choice of what to pin, not pinning. Both files now state the measured result and say to measure per checkpoint instead of assuming a policy.
The existing quality probe samples experts with torch.randn activations and quantizes the activations as well. Neither matches how the edge runtime uses these weights. At M=1 the activation is 4 KiB against a 1.7 MiB weight block, so quantizing it buys no bandwidth and the expert path is weight-only. And a scale calibrated against Gaussian noise is not the scale real post-norm hidden states need, so the probe can pass while the deployed path does not. expert_quality.py captures the activations the router actually sent to each expert, replays the expert in BF16 for the reference, and scores W8A16, W4A16 and W4A16 with the block-16 transform against it. It can also write a small bundle of activation and reference pairs so a device can check its own dequantization and kernel without loading the source checkpoint. The MoE input trace is a second opt-in hook alongside the router trace, off by default and eager-only for the same reason. dequantize_int4 moves from the probe to quantize_experts, next to the packing it inverts, so the format has one definition. The transform's effect is now asserted where it is supposed to appear: on weights with one outlier per group of 16 it cuts relative L2 by 38 %, while on Gaussian weights it is neutral, which is what an orthonormal rotation of an already-Gaussian group should do.
Scoring against real routed activations showed the single-level scale was unusable. Real expert weights have a per-group amax near 0.02, so amax/7 is about 0.003 -- below e4m3's smallest normal value of 2**-6. Every per-group scale in this checkpoint landed in e4m3's subnormal range, where the format keeps roughly three bits: the stored scale carried 18 % mean relative error, and in one sampled layer a quarter of the groups rounded to zero outright. That error multiplies every weight in the group, so it swamped the 4-bit value grid completely. The per-group byte is now a fraction of a per-tensor global scale, which is what the shipped NVFP4 expert path already does with its GEMM alpha. Stored bytes land in e4m3's normal range, and the round-trip error on Gaussian weights drops to the uniform-quantization-noise floor of the value grid itself, 8.6 % measured against 8.6 % predicted. On the experts that carry signal, output error roughly halves: 0.130 to 0.068 and 0.124 to 0.081. The global scales go in a per-layer sidecar rather than inside the blocks, so a block stays exactly block_bytes and 4096-aligned. They are eight bytes per expert and belong with the resident weights, where the kernel reads them as its alpha. Format tag moves to v2; no bundle had been generated. Two measurement corrections came with it: - E2M1 is now scored as a control. It is the format the shipped runtime uses for these experts with exact greedy reproduction, so it is the bar. Without it, per-expert numbers have nothing to be judged against. - Results pool by error energy over reference energy. Per-expert relative L2 is misleading here because expert output norms span three orders of magnitude, and the alarming values all came from experts with near-zero output, where the control scores just as badly and the router weights the contribution down to nothing anyway. The transform's measured benefit fell from 38 % to 9 % once the scale was fixed, because most of it had been compensating for scale error.
Three additions for sizing the latency of a streaming expert runtime. global_frequency and simulate_warm_lru cover filling each layer's cache at startup from offline selection statistics. On this checkpoint that removes 11 % of decode misses, and unlike the prompt-derived warm set it costs no adaptivity because the entries stay evictable. A set derived from the trace it is scored on is an oracle, so held-out traces are needed before trusting the figure. cold_prefill_blocks counts what prefill must read before the first token. Prefill routes every prompt token independently, so a layer costs the union of its tokens' selections; a 32-token prompt touched a mean of 74 experts per layer out of 256, or 4.88 GiB across the model. A resident set of 64 per layer brings that to 1.49 GiB. simulate_warm_lru also takes a window, which groups decode steps as a multi-token verification step would. It is here to document a negative result: grouping does not reduce reads, because an LRU already captures the reuse that a window's union would. Measured on this checkpoint, misses per token are 83.6 at window 1 and 83.7 at window 4, and rise to 85.7 at window 8 as a window's union starts evicting itself. An earlier estimate of 1.57x came from comparing a window's union against eight times its length with no cache present, which double-counts what the cache was already doing.
Filling each layer's cache at startup from offline selection statistics looked strongly positive, but the statistics had come from the trace being scored, which makes it an oracle rather than a predictor. A deployment builds the set from other traffic and then meets an unseen prompt. This traces several prompts in one model load and reports leave-one-out results: each prompt is scored against a set built only from the others. Measured over eight unrelated prompts, 32 prompt tokens and 32 decode tokens, mean over the eight held-out runs: | slots | metric | cold | held-out | oracle | |---|---|---|---|---| | 43 | decode misses/token | 117.19 | 105.43 | 96.25 | | 43 | cold prefill GiB | 4.81 | 3.11 | 2.35 | | 57 | decode misses/token | 107.53 | 87.79 | 73.51 | | 57 | cold prefill GiB | 4.81 | 2.66 | 1.65 | | 64 | decode misses/token | 104.57 | 80.04 | 63.09 | | 64 | cold prefill GiB | 4.81 | 2.44 | 1.34 | The set transfers: on topics from mechanical engineering to tort law, one built from seven prompts captures roughly 59 % of the oracle's decode benefit and 68 % of its prefill benefit on the eighth. The benefit also grows with slot count, from 10 % of decode misses at 43 slots to 23.5 % at 64.
The tier split was justified by compiling for sm_87 and sm_110, which says nothing about what those kernels compute. This exercises each binding in the core and weight-only-4-bit tiers, records its output, and diffs a run on one target against a reference from another. No checkpoint: shapes come from the Qwen3.6 geometry and inputs from a fixed generator. Inputs are stored in the reference and replayed rather than regenerated. CUDA RNG is not bit-reproducible across architectures -- the Philox thread mapping follows occupancy -- so regenerating on the target compares kernels on different data. Divergence appears only past the first launch block, which presents as small tensors agreeing and large ones diverging, and reads as a kernel fault. Each case also checks its kernel against a Torch expression on the local device. That separates a real kernel fault from a harness problem: a broken kernel fails its local check before it disagrees with a remote reference. Measured sm_120a against sm_110: all twelve recorded output tensors are bitwise identical, and all eight cases match Torch locally on both. Two harness faults found while building this, both from assuming a calling convention instead of reading the call site: - The weighted-sum reducer writes float32 into a flat buffer. A bfloat16 destination yields NaN. - The linear-attention split broadcasts q and k from 16 stored key heads to all 32 value heads, so all three of its outputs are 32 * 128 wide. Sizing q and k for 16 heads makes the kernel write past them into the next allocation, which surfaces as a corrupted third output. A structural test pins the case-to-binding mapping so a rename cannot turn a case into a silent skip.
The target device holds a fraction of the experts, so the cache is the runtime: what it holds and how fast it refills set both the token rate and the time to first token. Three properties are enforced rather than left to convention. The budget is a hard limit. On unified memory the weights, the cache, the staging buffers and the operating system draw on the same physical memory, so a runtime that merely intends to stay small cannot be measured. Construction computes its footprint and refuses to allocate over budget, naming the quota that would fit. plan() and max_slots_per_layer() answer the sizing question before anything is allocated. Reads use O_DIRECT. Streaming tens of GiB through the page cache would make it compete with the resident weights for that same memory, which is why the bundle pads blocks to 4096 bytes. Construction rejects a bundle whose blocks are not aligned, and checks that the pinned staging buffers are too. Misses are fetched concurrently, because a single reader leaves most of an NVMe device idle. Measured against the real bundle: 3.17 GB/s with one staging buffer, 5.70 with two, 5.85 with four, and flat after that -- matching an independent fio sweep of the same access pattern. Requiring a per-layer quota of at least experts_per_token makes a class of bug structurally impossible: one token's experts cannot evict each other, so a caller may hold every pointer from a get_many at once. An earlier prototype on another model lost most of its hit rate to exactly that. warm() preloads each layer's most frequent experts and leaves them evictable. On held-out prompts an offline set removes about a quarter of decode misses and half of the cold prefill read at this quota, where pinning it instead loses more adaptivity than it gains. close() releases the slots. It previously dropped only descriptors and the thread pool, which would leak the largest allocation in the process on any reconfiguration.
Whether O_DIRECT actually keeps the block stream out of the page cache was an assertion. It matters because on a device holding a fraction of the experts the page cache competes with the resident weights for the same physical memory, and the failure mode is an out-of-memory on the target rather than a wrong answer. A config flag selects buffered reads, which gives the measurement a control reading exactly the same bytes. Against the real bundle on Thor, 2.11 GiB of blocks: the direct path grew the page cache 0.07 GiB, the buffered control grew it 2.14 GiB. The difference is 2.07 GiB of memory an 8 GiB device does not have. The flag also covers filesystems without O_DIRECT.
Replaying a real trace through the cache on the target exposed a fidelity defect in the simulator that produced every projection in this work. The cache reported 5107 misses over 64 tokens; the simulator predicted 5120. Within one request, insertion order decides which entry is oldest, so it changes what the next eviction picks. The cache dedupes with an order-preserving map and therefore follows the router's ordering; the simulator iterated a set, whose order is unrelated. Twenty-two of the forty layers differed, some above and some below, netting 0.25 %. With order preserved the simulator returns 5107, matching the measurement exactly. No conclusion moves by 0.25 %, but the tool the conclusions came from should describe the cache it is modelling. The test pins the behaviour on a case where the choice is observable: at quota 2, requesting [7, 3] then 9 must evict 7 rather than 3, so the following request for 3 hits.
The leave-one-out figures were simulated. Replaying them on real hardware needs the traces themselves, so a warm set can be built from some prompts and measured against another one it has never seen.
Wiring the cache into a decode path needs the four parts of a block and the two scales that go with it. Doing that at the call site would mean reproducing the writer's offset arithmetic, which is how a reader drifts from a writer; these read the order from the manifest instead, and return views over the slot rather than copies. The global scales come from the per-layer sidecar, which exists so a block stays exactly block_bytes and 4096-aligned. Their file is size-checked, because a truncated sidecar would otherwise reshape into plausible nonsense. With these, the whole chain was measured on Thor against the reference pairs recorded on the SM120 machine -- a real routed activation and that expert's BF16 output. Pulling each expert's block through the cache, decoding it, and running the expert lands at cosine mean 0.974500 and relative L2 0.18510, against 0.974504 and 0.18508 for the same experts quantized straight from the checkpoint. Agreement to the fifth decimal across the on-disk format, the two-level scales, the sidecar, the direct read, the slot layout, the component split, and the transform on both GEMMs. The low minimum cosine, 0.574, falls on the same three near-zero-output experts as the checkpoint-direct run and at the same values, so it remains an artefact of an unpooled per-expert metric rather than anything the chain introduced.
The streaming path needs to get a block from a cache slot into a GEMM. The block-scaled 4-bit GEMMs cannot read it: they decode E2M1, and the bundle stores sign-magnitude integers, whose sixteen values are a different set, so no relabelling bridges them. They also want the scale bytes in the SM1xx swizzled tile layout, where the bundle's are linear. Decoding to bf16 first and handing that to the existing bf16 GEMM sidesteps both. It reads the bundle's own linear scale layout, which removes the swizzle step rather than implementing it, and it needs no second codebook inside a GEMM and no architecture beyond SM80. The cost is bandwidth on a block that is already resident, which is the cheap end of this system: the misses are what cost time. One thread per packed byte. Its two values always share a scale because the group size is even, so there is one e4m3 conversion per byte rather than two. The two-level scale is applied as the per-tensor float times the group's byte, matching how the quantizer chose them. It goes in the core tier: bit manipulation and bf16 writes, so it compiles and runs wherever that tier does. Verified bit-identical to the Python reference for (1024, 2048) and (2048, 512) at group 16 and 32, and it rejects a null pointer, odd columns, an odd group and a non-dividing group with distinct codes. Added to the cross-architecture parity harness, which now covers nine cases and thirteen output tensors, all bitwise identical between sm_120a and sm_110. That also puts it in the acceptance package's kernel test.
Assembles the pieces into a path where the experts are not resident. Three additive changes, each off unless asked for. The loader gains stream_experts, which skips building the per-layer stacked expert tensors. That skip is the point rather than an optimisation: those tensors are 16.9 GiB of a 21.4 GiB footprint, and attaching a cache without removing them would add to the total instead of replacing part of it. Their shapes are still checked, since a bundle is generated against them. The decode path gains a branch, reached only when the loader took that skip. It fetches a token's whole top-k in one call so the reads overlap, then decodes each block to bf16 and multiplies with the shared bf16 GEMV. The per-layer quota is at least the top-k, so no returned pointer can be invalidated by the others and the eight can be held at once. Decode scratch is two buffers allocated once, rather than eight allocations per layer per token. Qwen36MoeStreamingFrontend wires the two together. It sizes the cache against the resident bytes it actually measured after loading, not an estimate, and stays eager: a miss issues host reads, which a captured graph cannot replay. Everything else is unchanged -- same attention, same recurrence, same router, same reducer -- so comparing its tokens against the ordinary frontend isolates where the expert weights came from.
The frontend loaded a tokenizer while building weights, which made transformers a hard requirement of the runtime. It is not one: a caller with token ids of its own never needs it, and requiring it on a deployment target means installing a large dependency into an environment that may not want it. set_prompt_ids takes ids directly, so a target can run with the prompt tokenized elsewhere. The tokenizer property still loads on first use, so set_prompt and decode are unchanged. Also records what the pipeline actually depends on. The three tiers cover 14 of the 32 kernels it calls; resolving each call to the guard active where it is defined shows seven gates, with twelve kernels under FLASHRT_HAVE_QWEN36_KERNELS carrying the whole linear-attention path. That gate keys on NOT FLASHRT_SLIM_BUILD rather than on architecture, so a slim build removes them and the frontend refuses to start. Selecting tiers from the source's grouping was not enough to know what a target needs -- the call sites are.
The fail-fast check used one hardcoded list for every configuration. Running the streamed-expert path on a target where the block-scaled 4-bit tier is correctly not built showed why that is wrong: it refused to start over moe_blocktile_mma, which that path never calls. Those kernels serve the batched prefill, and streaming runs prefill through the per-token loop because a miss issues host reads. A list demanding more than a path uses turns a working build into a refusal; one demanding less lets a missing symbol surface mid-forward. So the list belongs with whatever decides which kernels get called: _require_kernels takes it, the frontend carries the default, and a subclass narrows or extends it. The streaming frontend drops the two MMA kernels and adds the two it does call. Behaviour for the existing frontends is unchanged.
Two defects, both found by running the assembled path on a non-SM120 target. The loader never received stream_experts. An unchecked string replacement had matched nothing while reporting success, so the feature was off and the run that followed reported a resident footprint of 21.436 GiB -- indistinguishable from the ordinary frontend's documented 21.44, and plausible enough to accept without comparing the two. A test now pins the wiring. The lm_head decode called fp4_w4a4_mma_sm120_full_n_bf16out, which is built only for GPU_ARCH 120 and 121. On every other target, the Orin one included, that path had no implementation: the symbol simply was not there. The W4A16 matvec reads the same swizzled weight and the same scale factors, leaves the activation in bf16 -- so it also drops the activation quantisation and its error -- and is in a tier that builds wherever the core does. It is selected when the W4A4 kernel is absent, so SM120 behaviour is unchanged.
Running the assembled path found three faults, all in the code that splices the cache in rather than in any of the pieces it splices together. It returned from the middle of the MoE layer, which dropped the shared expert and its sigmoid gate from every layer and returned (1, HID) float32 where the resident path returns (1, 1, HID) bfloat16. Streaming now replaces only the routed experts' own GEMVs and falls through to the shared tail, so the weighted sum, the shared expert and the gate are the same code on both paths. It never rotated the activation. The bundle stores H*W, so both GEMMs need the activation rotated the same way; without it the products are wrong while staying finite and plausible, which is the failure mode that hides. The transform is built once per state and applied to the hidden state entering gate_up and to the gated result entering down. Staging buffers were indexed by task number. The pool bounds how many tasks run at once, not the order they finish, so task N and task N + len(staging) could hold the same buffer and read into each other's memory. A task now takes one from a queue and returns it, owning it for the duration. Two diagnostics that turned a guess into a measurement, and are worth keeping. The router top-k return code was unchecked, so a failure left torch.empty memory to be used as expert indices; it now raises. And a direct read reports a misaligned offset, length or buffer with one indistinguishable EINVAL, so the cache reports which of the three it was, and rejects an out-of-range expert with the whole request quoted -- which is how the -1 was found.
Importing the FA2 module and finding its symbols proves neither that its kernel runs nor that it is right. Measured on an SM110 part: the module imported, every symbol was present, and at run time the kernel printed a complaint and returned without writing its output. Downstream that reads as plausible-but-wrong attention, not as a failure -- the model still produced 15 of 16 reference tokens with ten of forty layers contributing nothing, which is a coincidence of the residual stream rather than evidence of anything. Construction now runs one small case through the same launch the hot path uses and compares it against scaled_dot_product_attention. If it does not agree, or produces non-finite values, or raises, attention falls back to the reference implementation. One launch at construction, and it is the same launch, so the probe cannot pass while the real call fails. This is the third time in this work that a kernel compiled, linked and loaded while being unable to run: the block-scaled 4-bit tier substitutes an invalid control path off its own architecture, and this. Symbol presence is not a capability check.
The arch list omits FA2 for Thor on purpose: that target uses FA4, whose SM100-class CuTe-DSL kernel wants Blackwell tensor memory. Orin's SM87 is Ampere and has none, so it takes FA2, which the arch list does enable. The two targets differ by design, and a frontend that hard-requires FA2 refuses to start on one of them for no reason -- the backend already computes the same thing without it. The backend now treats an absent module as a fallback, and a frontend can declare that its attention can fall back. Existing frontends keep the requirement.
FA2 is absent from the Thor build on purpose and FA4 cannot serve SM87, so the two targets take different attention paths. Recorded alongside the reason the backend probes its kernel rather than trusting that a symbol implies a working one: three kernels in this work compiled, linked and loaded while unable to run, and the one that failed silently still produced 15 of 16 reference tokens with a quarter of the layers contributing nothing.
The decode attention backend already treats a missing FA2 as a fallback and probes whether the kernel computes; prefill called it unconditionally, so a target that builds FA4 instead could load the model and then fail on the first full-attention layer. The reference path builds the causal mask explicitly. FA2 aligns causal bottom-right, so a chunked block's queries attend to keys [0, Sk-Sq+i]; torch's is_causal aligns top-left and the two agree only when Sq == Sk. Prefill's MoE tile is chosen the same way: the block-scaled 4-bit MMA tier is a build tier, so ask the module for it rather than assume it, and fall through to the weight-only grouped GEMV when it is not there.
Both kernels stage the activation in shared memory at a 32-byte stride across lanes, which puts the eight lanes of a 128-bit load phase on four banks. On a 20-SM part the profiler reports 432,685 conflicts over 98,816 shared loads -- 2.41x the wavefronts the traffic needs -- with the kernel at 77% compute throughput against 37% memory, while the BF16 GEMV of the same shape reaches the memory roofline. Padding a block's footprint to 48 bytes lands the phase on eight distinct banks. The UE4M3 scale is decoded arithmetically at the same time. It was a 256-entry __constant__ LUT indexed by a per-lane byte, and constant memory serves one address per cycle, so a divergent index serialises; the decode is four integer ops and needs no table. Neither touches an arithmetic result, so the variants are accepted on bitwise equality against the kernels they stand in for, and the choice between them cannot move a token. They are added alongside rather than replacing, and a resolver picks the variant where the build has it (FLASHRT_QWEN35MOE_W4A16_EDGE=0 forces the original). Measured on sm_110: decode 70.07 -> 74.42 tok/s, 16/16 token-exact against the BF16 fixture.
With the bank conflict gone the dominant stall in situ is the global-load dependency -- 5.8 to 13 cycles per issued instruction against an ALU pipe at 27 to 50%. A warp owning one output row keeps only kUnroll eight-byte loads outstanding, and at K=512 not even that: K_BLOCKS is exactly 32, so the unrolled body never runs and the tail leaves a single load in flight. That is why the down projections sat at 38% of measured bandwidth while gate_up at K=2048 reached 67%. A warp now takes R consecutive rows and keeps R*kUnroll loads outstanding, R chosen so the product lands at 8 either way. The rows are 32-aligned by construction, so their scale offsets differ by a constant and cost no extra registers, and the per-row arithmetic is untouched -- same lane-to-block mapping, same order, same reduction -- so results stay bit-identical. Standalone at the shapes the decode actually issues, cold: experts down 38.6 -> 23.5 us (50% -> 80% of measured bandwidth), lm_head 1296 -> 1205 us (94%). End to end 74.42 -> 78.10 tok/s, 16/16 token-exact, time to first token 122 -> 110 ms.
The block-per-row form handed 256 threads a row of 512 values -- two elements each -- behind three barriers and three passes over shared memory. A block read two kilobytes and then waited, which measured 2.9x off what that traffic implies. A lane now owns one 16-element scale-factor group: it reads its own gate and up values, gates them, takes its own maximum and packs its own eight bytes. Nothing is shared, so there is no shared memory and no barrier, and a lane has sixteen values in flight where it had two. Same arithmetic in the same order, and the output is identical byte for byte -- packed data and scale factors both, over a routing layout with real group boundaries, since the scale-factor offsets depend on where each expert starts. 2.72x at the shape prefill issues: 1.778 -> 0.653 ms at 65536 slots, against a traffic bound of 0.627. That is 1.04x of the bound, so there is nothing left here. TTFT 394.3 -> 388.0 ms at 2048, 753.0 -> 734.7 at 4096, 217.6 at 1024. The block-per-row entry stays; this is a second one.
A prefill sends about sixty-four rows to the average expert, spread over 256 groups of unequal size. With an N tile of 256 the scheduler has 1024 blocks to balance those groups across twenty SMs; with 128 it has twice that. Paired against the wider tile on the same machine state, three runs each: 377.9 / 378.2 / 377.9 ms at 2048 tokens against 429.6 / 386.3 / 387.6. The worst run of the narrow tile beats the best run of the wide one, and the spread falls from 43 ms to 0.3 -- which is the load balance showing up directly. The cluster shape was swept first and does not move this. It is a runtime argument, so (1,1) (2,1) (1,2) (2,2) (4,1) (1,4) were measured without recompiling and the best two alternated three times each; the within-pair differences (+1.4%, -0.4%, +0.8%) came out smaller than the drift between runs of one setting. Left as a knob with that result written next to it rather than as a knob to try again.
The cuBLASLt wrapper picks its algorithm by timing eight candidates at first use. Timing is noisy, so different processes pick different algorithms, and different algorithms reduce in different orders -- which makes the model non-deterministic across processes. One binary gave the golden prefix 16/16 three times and 14/16 three times, flipping between exactly two token streams. Asking the wrapper for one candidate takes the heuristic's own choice instead. Five of five processes then agree, and six of six pass the golden gate. This is not a speed trade in the direction it looks. At 1024 tokens the timed pick is worth 1.6% warm -- 213.6 against 217.1 ms over three runs each -- and costs 25% of the cold time, about 1020 against 770 ms, because the timing loop runs inside the first call. A faster first token and a deterministic model, for 1.6% of the warm path. Set from this frontend rather than in the kernel, whose default is shared. The flip needed two things and neither is a defect alone: FA2 moved token 14 close to a decision boundary (it is bit-reproducible within a process, and with the timed pick left on but FA2 off, five of five processes agree), and the per-process algorithm choice then pushed it across, sometimes. The non-determinism is the part worth removing, and it predates this round.
The head takes the previous hidden state and the next token's embedding
concatenated, and fc is square in the concatenated width, so nothing in the
checkpoint says which half goes first. It was measured both ways when the head
was first loaded, and the measurement went into the notes rather than into the
default.
Over 48 decoded tokens:
cat[embed, hidden] first draft 0.896, chained 0.646, 0.417
cat[hidden, embed] 0.000, 0.000, 0.000
The wrong half drafts noise. Nothing is ever accepted, so every window pays for
a verify that keeps the one token it was going to emit anyway, and speculative
decode runs at a third of plain greedy. Expected tokens kept per window goes
1.00 -> 1.94 at K=1, 2.64 at K=2, 3.05 at K=3, which is what the acceptance
rate predicts.
The reference implementation of this head concatenates the embedding first as
well.
A speculative verify has to be the same function as the decode step it verifies, or the tokens it keeps are not the ones plain greedy emits. Ours was neither the same function nor cheap: it ran the prefill forward, which reads the dense weights at BF16 while decode reads them at four bits. Four times the traffic, and a different answer -- measured, logit cosine 0.988 against the decode path, which is why the emitted text diverged. The general W4A16 GEMM does read four bits and is not the answer either. At these shapes it is 7.5 to 9.3 times off the decode GEMV -- 250 us against 33 for an 8192x2048 projection -- and flat in M, so it is not reading the weight at bandwidth at all. This is the decode GEMV with M rows of activation. The weight stream, the lane-to-block mapping, the unroll and the reduction order are untouched; only the rows staged in shared memory and the accumulators a warp carries change. Each output row therefore accumulates in exactly the order the GEMV uses, and the output is bit-identical to running the GEMV once per row -- checked at six shapes by five window widths. Cost at the shapes a verify issues, against one decode row: 1.6x at two rows, 2.0x at three, 2.5x at four. Not the ratio the weight traffic alone implies, because the arithmetic grows with M while the weight read does not -- measured 20 us of weight plus 16 us a row on the largest projection. Still well under the 4.0x of running the GEMV per row, and a third of the general GEMM. Wired into the verify only, over the tensor the decode path caches rather than a second copy: same keys, same bytes. Logit cosine against decode goes 0.988 -> 0.994. Speculative decode at K=2 goes 0.29x of plain greedy to 0.73x. It does not pay yet, and the reason is now specific rather than general: the window still runs the MoE and the linear attention through prefill kernels, so a verify costs about 2.8 decode steps where it needs to cost one. Plain greedy, the golden fixture and the kernel preflight are unchanged.
A DeepSeek-V3-style draft head reads the pre-final-norm hidden state of the position before the token it is given. The per-token seeding path writes that buffer every step, so it was right there; the batched and chunked paths never wrote it at all. So the first speculative window of a generation drafted off whatever the previous generation had left in the buffer. Nothing crashes and nothing looks wrong -- the draft is simply predicted from an unrelated position, so the window keeps one token instead of two, the emitted sequence shifts, and every position after it is a different position. Which means the same prompt does not decode to the same text twice, and the tokens kept per window depend on what ran before. Both batched paths now take the hidden state the forward already computes and copy its last row. It is a 4 KB device copy on a path that runs once per prompt.
The verify was still the prefill forward. Only the dense projections had been
moved onto the decode path's kernels; the routed experts, the linear attention
and the full attention were all somebody else's arithmetic. So a window cost
about 2.8 decode steps where it should cost near one, and -- the part that
actually matters -- it was not the same function as the steps it verified, so
the tokens it kept were not the ones plain greedy emits.
The 27B solves this with a third forward and a parallel set of layer methods.
That is more than this needs, because the two kernels that looked like the hard
part already take the parameter: the grouped W4A16 GEMV takes a slot count, and
the MoE slots of w tokens are w*TOPK independent GEMVs. So the window is
decode_step with w rows, calling the kernels decode calls, over the weights
decode caches.
Three stages stay per token on purpose. The causal conv and the recurrence
carry state, and a window is accepted up to a prefix, so each token's state has
to be the state decode would have been in -- they run through the decode
kernels a token at a time, snapshotting as they go. The sequential-scan variant
would collapse the recurrence into one launch but it is cos 0.99999 against the
per-token kernel rather than equal to it, and this layer is 6% of a step.
Attention likewise runs at q_seq=1 per token: a batched q_seq=w call would need
a bottom-right causal mask and would reduce over a different tiling.
Measured against the decode step it stands in for, over three windows of four
at real decoded tokens: every logits row bit-identical, every per-token
recurrent and conv snapshot bit-identical, every KV row bit-identical.
Cost, against a captured decode step of 11.18 ms: the window costs 1.02 steps
at one row and about 0.20 of a step per row after that. The fixed part is the
dense weights and the lm_head, read once however wide the window is; the
per-row part is the routed experts, which do not amortise at all -- w tokens
pick up to w*8 distinct experts out of 256, so that traffic is the floor on
what a wider window can be worth.
End to end against plain greedy, one process per K, 64 tokens:
K=1 1.09x kept 1.88
K=2 1.13x kept 2.71 91.1 tok/s against 80.5
K=3 1.00x kept 3.20
K=4 0.82x kept 3.10
Off the default path: plain decode, the golden fixture and the kernel preflight
are untouched, and FLASHRT_QWEN35MOE_VERIFY_K_ROWS=0 puts the prefill forward
back. The lm_head quantisation moves to its own helper so the single-row head
and the window read one copy rather than two of the same bytes.
The same prompt did not decode to the same text twice. Four seedings of one
20-token prompt in one process: the logits differ in 232226 of 248320 elements,
the GDN recurrent state differs from linear rank 12 on, the KV differs. The
decode loop is fine -- eager and captured agree with each other and with
themselves over four runs -- so all of it comes from the seed.
Three of the four MoE paths finished with
out.index_add_(0, stok, d_dn.float() * sw.unsqueeze(-1))
which reduces through atomics: a token's eight expert outputs are added in
whatever order the blocks retire, and fp32 addition is not associative. Two of
them also sorted the routing with a plain argsort, so equal-expert ties changed
which rows were packed into a quantisation tile.
Neither is a new discovery in this file. The grouped-GEMM path already inverts
the permutation and sums with a kernel, and says why in a comment; the
block-tile path already sorts stably and says why. The other three had simply
never been brought along. They are now: invert the routing permutation, hand
the inverse to moe_weighted_sum_sm120_bf16 so each token's slots are summed in
k order, and sort stably everywhere.
It is not slower -- the comment on the path that was already converted records
index_add_ at 37.8 ms of a 1024-token prefill.
Worth knowing which prompts were affected: above 64 tokens the prefill already
took a converted path, so this bit short prompts, which is where a seeded
generation starts and where the fixture lives.
After: the seed is reproducible, plain greedy is reproducible over four runs,
and speculative decode emits exactly what plain greedy emits. Golden fixture
16/16 throughout.
The window's whole claim is that it is the same function as the step it verifies, so the conditions under which that holds should be asked about rather than assumed. Two were not. The GDN in_proj is gated by its own flag, separate from the rest of the dense path. With it off, decode reads that projection at BF16 and the window reads it at four bits -- in thirty of the forty layers. The window also fuses the router with the shared gate/up and reads every projection at four bits, which is what decode does only when the loader kept those weights BF16. One NVFP4 site among them and decode takes the W4A4 mma instead. Only two of the three were checked. Both fall back to the prefill forward, which is what the flag is for.
A_log and dt_bias are weights. -exp(A_log) and the fp32 bias are therefore the same on every decode step, and they were being rebuilt on every one of them, in each of the thirty linear-attention layers: a cast, an exp, a negate and a contiguous, all inside the captured region, all producing the bytes the previous replay had already produced. The profiler counts them. The elementwise/copy bucket of a step goes from 522 launches to 403 and from 791.0 to 623.3 microseconds, and the step from 11.238 to 11.059 ms. Small, but it is 119 launches of a step that is 99% kernel time, where each launch costs its dispatch quantum whether or not it computes anything. Same expressions in the same order, so the kernel is handed the same bytes: golden fixture 16/16, and the speculative window still emits plain greedy's sequence token for token.
The MoE tail -- sigmoid the shared expert's gate, scale the shared output by
it, add the routed sum, round to bf16 -- was five tensor ops a layer in decode
and in the speculative window, forty layers of them. A kernel that does all of
it already existed and prefill was already calling it; decode was not, and
could not, because it was not the same arithmetic.
The reason is worth recording. Written as
routed[i] + float(shared[i]) * g
the compiler contracts the multiply and the add into one fma. That is one
rounding where the tensor-op chain has two, so the kernel and the chain
disagree -- measured, one element in 16384 by one ulp, at scales where the
routed sum is small against the gated shared term. More accurate, and still
wrong for this purpose: decode's output is compared token for token against a
fixture, and the speculative verify may only keep a token because it computed
what the decode step would have. So the kernel now multiplies and adds as two
rounded operations, and is bit-identical to the chain at every shape and scale
the two paths issue.
That also removes a disagreement nobody had noticed, since prefill had been
using the contracted form all along while decode used the chain.
Per decode step, measured: the elementwise bucket goes from 403 launches and
623.3 us to 203 and 301.5, and the step from 11.059 to 10.827 ms -- 92.4 tok/s
against 90.4. Speculative decode at K=2 reaches 97.34 tok/s against 87.98 plain.
Golden fixture 16/16, the window still emits plain greedy's sequence token for
token, and the window's rows are still bit-identical to the decode steps.
ncu reports the single-token recurrence running at 39 registers per thread.
The kernel declares
float col[HD]; // HD = 128
and walks it five times. A 128-iteration loop unrolled by sixteen leaves the
index non-constant, so that array cannot be held in registers and is not: it is
in local memory, 512 bytes a thread, read and written across every pass. Against
2 MB of actual state traffic that is several megabytes of spill, and the kernel
lands at 51% of what the part can move while the profile shows 12.8% occupancy
and 0.13 waves per SM.
The column does not need to be held. The recurrence touches the state twice --
once to form the k-weighted sum, once to apply the rank-one update and emit the
output -- and the whole state is 1 MB, so the second read is an L2 hit. This is
the same kernel with the intermediate left where it belongs.
Every accumulation runs in the same order over the same rounded fp32 values, so
the two agree exactly: six independent trials of eight chained steps each,
outputs and final state compared with torch.equal, not a tolerance. Chained,
because a recurrence that drifts by an ulp would hide in a single step.
shipped 16.77 us 125 GB/s 51% of peak
edge 9.89 us 212 GB/s 87% 1.70x
Added alongside the existing entry rather than replacing it, and selected
through a dispatch helper, so a build without it keeps working.
End to end: the decode step goes 10.827 -> 10.379 ms, 92.4 -> 96.4 tok/s, with
the GDN bucket 746.5 -> 557.8 us. Plain greedy over 128 tokens reads 91.18
tok/s against 87.98. Golden fixture 16/16, the speculative window still emits
plain greedy's sequence token for token, and its rows are still bit-identical
to the decode steps they stand in for.
Selecting eight of 256 logits was costing 6.5 us to read 512 bytes. The kernel spreads the logits over 256 threads and runs k rounds, each with a warp reduction and three block-wide barriers to publish the winner and mask it -- 24 barriers for 512 bytes of input. One warp can hold all of it: eight logits a lane in registers, the same k rounds entirely in shuffles, nothing synchronised and nothing in shared memory. The output is identical, not merely equivalent, and the reason is worth stating because it does not hold for most reductions: argmax under a total order -- greater value wins, lower index breaks ties -- selects one specific element, so the answer cannot depend on the shape of the reduction tree the way a floating-point sum does. Which lane holds which logit is therefore free to change. Checked anyway, on inputs built to produce the case that would break it: 800 trials across four widths, 397 of them with a tie inside the top-8, indices and values compared with torch.equal. No mismatches. Worth less than the barrier count suggested. Standalone it is 8.30 us against 6.60; in the captured step the routing bucket goes 594.8 -> 499.2 us, about 2.4 us a layer. The rest is the per-kernel dispatch floor, which no amount of making this kernel faster will remove -- only launching it fewer times would. Step 10.379 -> 10.297 ms, 96.4 -> 97.1 tok/s. Speculative K=2 reads 98.57 against 91.98 plain. Fixture 16/16, same emitted text.
ncu on the grouped instance:
Registers Per Thread 121-124
Block Limit Registers 8 <- the only binding limit
Block Limit Shared Mem 18-25
Block Limit Warps / SM 24 / 24
Achieved Occupancy 31%
Every limit except registers allows 24 blocks a SM. The registers are not an
accident -- the loop deliberately keeps R * kUnroll eight-byte weight loads
outstanding, because an earlier round measured the dominant stall as the
global-load dependency with the ALU pipe at 27-50%. But wv[R][kUnroll] is R *
kUnroll uint64 before anything else, 64 registers at R=8, and R had been swept
while kUnroll never was.
Halving it halves that array. Loads in flight per lane drop from R*4 to R*2,
and twice as many lanes issue them; on this part the second is worth more.
Swept in the captured decode step, one build each:
kUnroll 4 3 2
step 10.384 10.360 9.743 ms
tok/s 96.3 96.5 102.6
Bit-identical, and provably so rather than by luck: the main loop advances by
32*kUnroll and the tail takes the remainder, so a lane visits the same
k-blocks in the same order for any kUnroll, and accumulates them in that
order. Checked as well -- the speculative window's rows are still bit-identical
to the decode steps, the kernel preflight passes 36/36, and the golden fixture
is 16/16.
Plain greedy over 128 tokens: 96.64 tok/s against 91.98. Speculative K=2 reads
100.66. The step's GEMV bucket goes 8392.8 -> 7783.5 us.
The speculative verify's dense projections run through an M-row form of the
decode GEMV, written by copying the single-row kernel's shape constants. When
the sweep moved that kernel from four packed-weight loads in flight to two, the
copy kept four -- so decode had traded registers for warps and the verify was
still paying for depth. It showed up as the speculative ratio sliding while
plain greedy improved: the window was not getting the win the step got.
Same constant, same reasoning, same invariance: the main loop advances by
32*kUnroll and the tail takes the remainder, so a lane visits the same
k-blocks in the same order and every output row stays bit-identical to running
the single-row GEMV once per row.
Preflight 31/31 exact, the window's rows still bit-identical to the decode
steps, fixture 16/16, and the emitted text is still plain greedy's.
Measured with nothing else on the device, one process per point:
K=1 105.22 tok/s against 96.75 plain
K=2 106.74 against 100.35
…ible Three review-standard gaps in the edge decode work. The new gated-DeltaNet recurrence entry copied its neighbour's habit of returning quietly when the head dim is not 128. The caller cannot tell that from success, and the output buffer is left undefined -- the exact "all-zero fallthrough" the error-handling rule rejects. It now returns a status and the binding raises with the operation name and the shapes it was given. The SM120 latency table in the model doc had no reproduction command, so nobody could check or refresh it. `benchmarks/qwen36_moe_edge_decode.py` produces every row it quotes and refuses to print throughput unless the eager and captured paths emit identical tokens, because a rate for a path that emits different text is not a rate for the same work. Re-measured on the same card and checkpoint: warm CUDA-graph decode 195.49 -> 238.55 tok/s, eager 48.14 -> 107.84, resident and peak allocation unchanged to the megabyte. Weight-load time and the old "first prefill" row are dropped rather than compared -- the first is page-cache state, the second used a different notion of warmup. Two doc claims had gone stale. Speculative decode is implemented and token exact, so "the MTP tensors are validated but not loaded" is wrong; and the runtime is no longer SM120-only now that the Thor path exists. Added a speculative-decode section with the measured operating point and the reason K=2 is where it sits, and listed the two optional kernels with the symbols they fall back to, since neither is part of the required set.
The constant was lowered from four to two on the strength of a sweep run entirely on a 20-SM, 244 GB/s part, where ncu shows the kernel register-limited to 8 blocks per SM while shared memory, warps and the SM limit all allow 24. Trading per-thread depth for occupancy is worth 6% there. It was applied globally, which is wrong twice over. A device with an order of magnitude more SMs and bandwidth may prefer the deeper per-thread parallelism, and the branch has no measurement for one -- the only card available to check it is shared with other work and every attempt was either interleaved with another process or killed by it. Shipping an unvalidated tuning change to an architecture on the strength of a different architecture's profile is exactly what the hardware-additivity rule exists to prevent. The kernels now read the value from the build and default to four, so every target that is not explicitly tuned compiles the code it compiled before. CMake sets two for sm_110 only, next to the reason. Either value is bit-identical -- the main loop advances by 32*kUnroll and the tail takes the remainder, so a lane visits the same k-blocks in the same order regardless. Also restores the SM120 first-light table. Its numbers were replaced with figures measured while another process held a third of the card, and compared against a row from a different generation length, which turned a regression into an apparent 22% gain. The table is quoted unchanged, with a note on what has and has not been re-measured.
The model doc led with a 32-token first-light row, which is not the figure to quote for this path and is easy to compare against by mistake. The same-shape `P=64, N=64` warm-graph measurement is now stated next to it, with its median, its range across eight runs, and the note that the two generation lengths are not interchangeable.
The text-inference path landed on main while this branch was extending it, so the model frontend, decode module and usage doc conflicted as add/add. This branch's side is a superset in every case: the kernel-requirement check gained the parameters a target without a vendored FA2 build needs, the frontend gained the narrowed requirement set and the speculative entry point, the decode module gained the speculative window state, and the doc gained the kernel-tier, attention-routing and speculative sections. Main's side of each conflict was either empty or the code this branch replaced, so nothing is dropped. The only doc lines unique to main were the two claims this branch corrects: that the runtime is SM120-only, and that the MTP tensors are not loaded.
The model had no entry in the performance section on either target. Adds one per hardware in the same shape the Qwen3.6-27B block uses, and a Thor numbers section in the model doc for the README to point at. The SM120 row is the same-shape measurement already on record for this path, not the older 32-token first-light row. The Thor rows are one prompt length per process; the decode figure there moves a few percent with what else the board is running, which the doc says next to the table.
Member
Author
|
Thank you for the substantial work here. The white-boxed prefill path, tiered Qwen3.5/3.6 build options, defensive capability probing, deterministic-state work, and the amount of real-device validation are all valuable contributions. I also checked the CPU-side test surface locally: This is a Codex-assisted maintainability review against FlashRT's long-term contribution rules. Because this PR is large and touches shared build/runtime surfaces, I do not think it is ready to merge yet. The remaining issues are mostly boundary and integration issues rather than objections to the feature itself. 1. Keep the default build unchanged when Qwen3.5/3.6 is disabledThe new
Please move the grouped MoE GEMM and grouped quantization implementation into dedicated Qwen3.5/3.6 translation units, and gate their sources, declarations, and bindings with the corresponding model tier. SM110 FA2 should likewise require a Qwen model-specific option (or a dedicated default-OFF Thor-FA2 option). A build with every Qwen3.5/3.6 option disabled should retain the previous source list and public symbol surface. 2. Complete the Thor integration instead of exposing it as a partial pathThe PR documents Thor/SM110 as supported, but the public registry still only registers Please add an explicit Thor frontend/registration and a resolver test. If the implementation is genuinely shared, a hardware-neutral base is preferable to requiring users to directly construct an The build UX also needs to be exact:
Please document the exact SM110 tier combination and make errors architecture-aware. Do not recommend the total switch on a target where it is known to fail. 3. Provide a supported speculative-decode API and a safe cache policy
Please add a public, default-OFF option such as The speculative graph cache also reuses 4. Isolate Qwen3.6/Thor policy from existing Nex-N2 behaviorThis PR makes large changes in shared Nex-N2 modules ( Most importantly, importing os.environ.setdefault("FLASHRT_BF16_CUBLASLT_AUTOTUNE_ALGOS", "1")That mutates process-global configuration merely by importing a shared frontend module and can affect models loaded later in the same process. Please pass this as frontend/runner configuration rather than changing Please also scope Qwen3.6/Thor performance policy (edge W4A16, router/GDN selection, fused combine, and related defaults) through an explicit frontend/config object. Common deterministic correctness fixes can remain shared when justified, but they need Nex-N2 regression coverage. Existing model defaults should not change because a newly built symbol happens to be present. 5. Reconcile the documentation with the implementationThere are direct contradictions that need one authoritative policy:
Please update the root README, model guide, frontend docstrings, build table, and error messages together after deciding the final gated FA2 policy. 6. Commit reproducible integration evidenceThe added kernel tests are useful, but I could not find a committed test for the new speculative runtime contract. Please add a reproducible test/harness covering:
Checkpoint/GPU tests may be skipped when prerequisites are absent, but the harness and assertions should live in the repository. Please also attach a build/symbol matrix for:
Merge checklist
Once these boundaries are closed, the feature can remain substantial without imposing compile, symbol, runtime-policy, or memory costs on unrelated models and platforms. |
A build with every qwen3_5_moe option off should compile the same sources
and export the same symbols it did before this branch. Three additions
did not:
- the grouped NVFP4 MoE GEMM was a second source in the SM100 W4A16
object, which every Thor build compiles. It gets its own object
library and its own gate, and its bindings follow it.
- the grouped activation quantisers were appended to the shared
quantize.cu under the broad NVFP4 gates. They move to a translation
unit built with the weight-only 4-bit tier. The element and scale
converters they share with the general quantiser move to a header, so
there is still one copy of them and quantize.cu compiles the same
definitions it did.
- FA2 on sm_110 was unconditional, so every Thor build paid its compile
time and carried its symbols for one model's benefit. It is now
FLASHRT_ENABLE_THOR_FA2, off by default.
Also make the missing-kernel error name the tiers the device in front of
the reader can compile, rather than recommending the switch that turns on
the block-scaled MMA tier -- which sm_110 refuses at configure time.
Importing the qwen3_5_moe forward set FLASHRT_BF16_CUBLASLT_AUTOTUNE_ALGOS in the environment. The variable is process-global and shared with every frontend, so a model loaded later in the same process inherited an autotune setting it never asked for -- from an import. The GEMM entry takes the bound as an argument instead, defaulting to 0, which is the environment-driven behaviour every existing call site keeps. Plans are cached per (M, N, K, max_algos), so one caller asking for the heuristic's own pick does not decide the algorithm for another. Gather the rest of this path's kernel choices into a KernelPolicy the frontend owns, rather than deciding each at its call site by asking the module which symbols it happens to export. Every field selects between implementations checked against each other with torch.equal, so a field decides speed and cannot decide output; the environment variables that predate it are its defaults. The attention backend likewise takes its decode-FA2 choice as an argument, keeping the per-architecture default when the caller does not say.
The frontend was registered for RTX SM120 only and named for it, while the branch documents Thor as supported -- so a Thor user had to construct an Rtx-named class directly. Nothing in it is architecture-specific: it moves to flash_rt.frontends.torch.qwen36_moe as Qwen36MoeTextFrontend and both architectures resolve to that one class. The previous module and class name re-export it, so an existing import keeps working. Speculative decode was reachable only by setting a private attribute on a subclass, which the error message then recommended. It is a constructor argument now, load_mtp=False by default, because the draft head is a transformer layer's worth of weights that plain generate never reads. Its graph cache no longer borrows the decode cap of 256. A speculative window covers k+1 positions through the whole stack, so its memory pool is several times a decode step's, and 256 of them is more than a 32 GB board has at a 2048-token context -- which is where it was measured running out. The default is 16, and spec_graph_cache_max sets it.
The model guide said the interface does not execute the MTP head and then documented generate_spec(); it said FA2 is deliberately absent on Thor while the build enables and tests it there; the Thor build line was "-DGPU_ARCH=110 ..." rather than a command; and the usage examples named the module and class as they were before the rename. Each of those now states one policy: MTP loads on load_mtp=True and nothing else reads it, FA2 on sm_110 is opt-in with the flag that builds it and the model runs either way, and both targets have an exact configure line. The README's two claims that FA2 is RTX-only are updated to match. The speculative table gains the paired long-context measurements and corrects the K=1 ratio, which was quoted from a different run than its own baseline.
Speculative decode had no committed test for its runtime contract. It has one now: the constructor arguments and their rejections, the graph cache's eviction policy, and -- against a checkpoint, skipped without one -- K=1 and K=2 equality with plain greedy, the window's logits and recurrent, conv and KV state against the decode steps they stand in for, the rejected-tail rewind, and boundary token counts. They are equality tests because the claim is equality; a tolerance would not be checking it. Both graph caches now insert through one function, which is also what makes the eviction policy testable without capturing a graph. The tier gates get a matrix. A build proves its own configuration works; only reading the gates can say what the four configurations nobody built contain, so the script walks CMake's tier blocks and the bindings' guards and fails if a tier source or symbol is reachable with every tier off. The five configure lines it prints are in the model guide with what each produces, including the one that is supposed to fail.
Nex-N2 and Qwen3.6 share this forward, so a change made for one can move the other. Two ways that happened here are now regression-tested: importing the module must not touch the process environment, and turning a policy field off must select the fallback even when the faster symbol is present in the build. The rest asserts the shipped defaults field by field and that the environment variables that predate the policy object still set them, so a later field cannot quietly move a value Nex-N2 was validated with.
The Thor table carried a 2 K decode figure higher than its 20-token one, which is not what any recorded sweep says, and the README repeated it. It came from a run taken while another job held the board -- its plain baseline read 79.9-90.4 tok/s against a known 100.4, which is the signal that a shared board gives and which should have retired the run on sight. Replaced with the recorded sweeps, each labelled with the protocol it was taken under, because two of them are not interchangeable: the vLLM comparison times a whole 64-token generate with TTFT subtracted, and the decode round times the captured steady step. Both are quoted, with what separates them stated rather than left for a reader to reconcile. The speculative table now carries a baseline per row, from that row's own process, rather than one ratio against a baseline from elsewhere.
The decode figures were quoted as two protocols that were "not interchangeable". They are interchangeable: on the same tree the sweep reads 87.1 tok/s at 1024 and the captured step reads 89.0 at 20, a 2% spread. The distance from 87 to 102.6 is the decode round, not the way either was timed. So the model guide labels each row with the tree instead of the timing method, and says plainly that the vLLM comparison predates the round and is therefore a lower bound on the current ratio. No figure is quoted for the current tree at 1024-4096, because none was measured.
It asserted on state read straight after the graph was captured, and on
every rank of the KV cache. Both are wrong, and against a real checkpoint
both failed:
- capturing snapshots everything the block mutates and restores it
afterwards, so the state at that point is the state *before* the
window. Only the replay advances it. Without one the test compared a
prefill against k+1 decode steps.
- the draft head owns an eleventh full-attention rank that the window
writes and a plain decode step has no counterpart for. Comparing it
asked decode to fill a slot it does not have.
Now it replays, compares the model's own ranks, and compares the emitted
argmax rather than the pre-final-norm hidden state the block returns --
argmax is what the accept decision actually reads. Both reasons are in
the test, since neither is visible from the call site.
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
1 participant
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.
Brings the Qwen3.6-35B-A3B (
qwen3_5_moe) text path from a working SM120runtime to one that also runs on Jetson AGX Thor, white-boxes the prefill onto
its own kernels, and adds speculative decode that is token-exact rather than
approximately so.
The branch is one model path end to end. It is large; the sections below are
ordered so a reviewer can stop after any of them.
What changes
Prefill, white-boxed. The prefill was reaching for tensor ops at 202 call
sites. Those are now kernels: fused GDN WY pack/norm/cumsum, a row-blocked
causal convolution that takes its own history so the chunked path stops
concatenating, a five-kernel MoE routing producer, and a fused shared-gate
combine. What the profiler still attributes to tensor ops at a 2048-token
prompt is 0.7 ms of KV-cache writes, against 29% of the prefill before.
FA2 on SM110. The designated Thor attention path materialises an
(S*heads, S_kv)score buffer -- gigabytes per layer at long context. FA2 wasexcluded from the arch list on build size, not capability, and the causal
hd256 instantiation was already in the tree. Enabling it makes the non-square
window of a chunked prefill 20x faster and takes the chunking penalty from 64%
to 4%. It is opt-in (
-DFLASHRT_ENABLE_THOR_FA2=ON), so no other Thor buildpays for it.
Decode kernels. Five changes, each bit-identical to what it replaced:
gating constants derived once instead of every step; the MoE tail fused into
one kernel; the gated-DeltaNet recurrence stopped spilling its state column to
local memory; the router's top-8 picked in one warp instead of
krounds ofblock-wide barriers; and the W4A16 GEMV's loads-in-flight halved, which trades
per-thread depth for occupancy the profiler said was register-limited. That
last one is scoped to sm_110 in CMake: it was measured on a 20-SM part and
a larger device may prefer the deeper parallelism, so every other target
compiles the constant it compiled before.
Speculative decode. A DeepSeek-V3-style MTP head, verified through the
decode kernels at
K+1rows over the weights the decode step caches, so averified row is the decode step it stands in for. Enabled with
load_mtp=True;generate()never reads the head.A prefill reduction that was not deterministic. Three of the four MoE paths
summed each token's experts with
index_add_, which reduces through atomics.The same prompt did not decode to the same text twice, and the 16-token fixture
never reached the divergence. The fourth path already had a fixed-order kernel
and said why in a comment; the other three now match it.
Performance
Jetson AGX Thor (sm_110) vs vLLM 0.26.0
Same part, same BF16 checkpoint, same pre-tokenized prompts, same protocol on
both sides, one length per process. Vision tower off on both.
TTFT is ahead at every length measured. 128 K context reaches this board at
2470 tok/s of prefill, which it could not before. The decode column was taken
before the round below, which added 15.3% to this side and nothing to vLLM's,
so those ratios are lower bounds.
What moved on Thor
Every decode change is bit-identical to what it replaced.
Speculative decode (Thor)
Plain and speculative in the same process, emitting plain greedy's sequence
token for token with the 16-token golden fixture passing on the same build.
For scale: vLLM runs the same MTP head and gets a larger relative gain --
31.59 -> 55.00, 1.74x -- because its step is three times heavier, so a fixed
per-draft cost is proportionally three times cheaper. Plain greedy decoding
here still beats that speculative figure by 1.8x, with no speculation at all.
SM120
No new SM120 latency claim is made. The card available for it is shared
with other work. The reference stays the same-shape
P=64, N=64warm-graphmeasurement already on record for this path:
The tuned constant discussed above is scoped to sm_110, so the SM120 kernels
are byte-identical to the ones that measurement covered. Per-change breakdowns
for both architectures are in
docs/qwen36_moe_usage.md.Correctness
Every kernel change is checked with
torch.equalagainst the path it replaces,not a tolerance, because the emitted tokens are compared against a fixture and
the speculative window is only allowed to keep a token because it computed what
the decode step would have.
The combine was not identical at first: the compiler contracts the multiply
and add into one fma, one rounding where the tensor-op chain has two, and that
showed up as one element in 16384 differing by one ulp. Fixed with explicit
rounded operations, which also removed a prefill-vs-decode disagreement that
had gone unnoticed.
Build boundaries
A build with every
qwen3_5_moeoption off compiles the same sources andexports the same symbols it did before this branch. Three additions did not
respect that and were reworked:
GPU_ARCH=110 AND FLASHRT_ENABLE_QWEN35MOE_W4A16, bindings onFLASHRT_HAVE_QWEN35MOE_GROUPED_SM100quantize.cuunder the broad NVFP4 gatescsrc/kernels/qwen35moe_grouped_quant.cu, built with the weight-only tierFLASHRT_ENABLE_THOR_FA2, default OFFThe element and scale converters the grouped quantisers share with the general
quantiser moved to a header rather than being copied, so
quantize.cucompilesthe same definitions it did.
scripts/qwen35moe_build_matrix.pyprints the sources and symbols each tieradds and fails if any of them is reachable with every tier off;
tests/test_qwen35moe_build_matrix.pyruns the same checks. The fiveconfigurations behind it, all verified configure-only:
-DGPU_ARCH=120FA2 ENABLED; noqwen3_5_moesource or symbol-DGPU_ARCH=110FA2 DISABLED; noqwen3_5_moesource or symbol-DGPU_ARCH=110 -DFLASHRT_ENABLE_QWEN35MOE_CORE=ON -DFLASHRT_ENABLE_QWEN35MOE_W4A16=ON -DFLASHRT_ENABLE_THOR_FA2=ONhdim={256} x dtype={bf16}-DGPU_ARCH=120 -DFLASHRT_ENABLE_QWEN35MOE=ON-DGPU_ARCH=110 -DFLASHRT_ENABLE_QWEN35MOE_W4A4=ONRuntime boundaries
flash_rt.frontends.torch.qwen36_moeasQwen36MoeTextFrontend-- nothing init is architecture-specific -- and both
("qwen36_moe", "torch", "thor")and("qwen36_moe", "torch", "rtx_sm120")resolve to that one class, with aresolver test asserting they are the same object. The previous module and
class name re-export it.
load_mtp=Falseby default onthe constructor, validated, documented with a complete example, and covered
by tests. It was previously reachable only by setting a private attribute on
a subclass, which the error message then recommended.
window covers
k+1positions through the whole stack, so its memory pool isseveral times a decode step's; 256 of them is more than a 32 GB board has at
a 2048-token context, which is where it was measured running out. The default
is 16,
spec_graph_cache_maxsets it, and the eviction policy is tested.FLASHRT_BF16_CUBLASLT_AUTOTUNE_ALGOS, which is process-global and sharedwith every frontend, so a model loaded later inherited an autotune setting it
never asked for. The GEMM entry takes the bound as an argument instead
(default 0 = the existing behaviour), and plans are cached per
(M, N, K, max_algos).KernelPolicygathersthe choices this path makes -- edge W4A16, the routing producer, the WY scan,
the fused combine, the warp router, the edge recurrence, the verify form --
with the existing environment variables as its defaults. Every field selects
between implementations checked against each other with
torch.equal, so afield decides speed and cannot decide output. Nex-N2 regression tests assert
the shipped defaults field by field and that turning a field off selects the
fallback even when the faster symbol is present.
compile, rather than recommending the switch that turns on the block-scaled
tier -- which sm_110 refuses at configure time.
Validation
SM110, Jetson AGX Thor, built with the supported tier combination
(
-DGPU_ARCH=110 -DFLASHRT_ENABLE_QWEN35MOE_CORE=ON -DFLASHRT_ENABLE_QWEN35MOE_W4A16=ON -DFLASHRT_ENABLE_THOR_FA2=ON), against theofficial BF16 checkpoint:
The speculative file is the one that matters for this branch's headline
feature: K=1 and K=2 against plain greedy, the window's emitted argmax and its
recurrent, conv and KV state against the decode steps it stands in for, the
rejected-tail rewind, boundary token counts, and the graph cache staying within
its bound -- all against the checkpoint, all equality rather than tolerance.
SM120, RTX 5090, PyTorch 2.9.1 / CUDA 12.8, BF16 checkpoint:
The remaining SM120-visible changes -- the fused MoE tail, the router top-8,
the recurrence rewrite -- are each checked bit-for-bit against the path they
replace, and the golden-fixture gate above is the end-to-end confirmation.
CPU-side, no GPU required:
benchmarks/qwen36_moe_edge_decode.pygives the doc's table a reproductioncommand; it reports a median and range over every repetition and refuses to
print throughput unless the eager and captured paths emit identical tokens.
Notes for review
policy above and falls back to the kernel it replaces, so a build without
them is slower and never different. Neither is in the required-symbol set;
the model doc lists both with the symbol they fall back to.
extends, and the implementing translation units are compiled under exactly
those gates.
raises with the operation name and shapes for anything else, rather than
returning quietly and leaving the output undefined.
exec/is untouched.implemented and reachable, and FA2 on Thor is opt-in rather than absent. The
root README's two RTX-only FA2 statements are updated to match.
Known gaps
reason given above. Closing it means a batched verify attention, which is a
different kernel from the one the decode step uses -- and using the decode
step's kernel is what makes the window bit-exact.
by attention's O(S^2), which is the one component still reached through
torch; the sweep's shape says so directly.
SM. The fix is a real retiling, not a constant, and is not in this PR.
2.7-2.8x is a lower bound rather than a measurement of the current tree. Only
this side needs re-sweeping; vLLM's is a different binary.