Repository navigation
CI run: qwen36-moe on windows-native - #16
Closed
BobClawblaw wants to merge 4 commits into
Closed
BobClawblaw wants to merge 4 commits into
BobClawblaw wants to merge 4 commits into
Conversation
…ng with a shared expert The grouped affine-4 expert kernel ran only Nemotron's relu^2 up; its SwiGLU epilogue over two accumulators is now instantiated for decode and prompt rows (M = 2, gate and up interleaved per group), and the pack kernel takes a table stride and index so gate and up pack into one interleaved table. Layer.swiglu picks that form; the shared expert sits last in the same tables. cuda/moe.zig and moe.cu are the routing the frozen tensorfold/cuda/moe.py did: fp32 logits over the experts and the shared expert's gate row, each row's top-k by logit with ties to the lower id, weights exp(l - l_top) over the k rounded to bf16, the shared expert as a slot weighted by bf16(sigmoid(bf16(gate))), and the slots' weighted sum rounded once to bf16. tf-cuda-test moe checks each against a host reference: the router and the combine byte-equal, the picks equal with weights within a bf16 step (a tie, a flat row and a NaN included), and the pair experts against fp64 over random affine-4 tables.
model_type qwen3_5_moe serves on the same engine as the 2B and the 27B: the layers are theirs, with routed experts in place of the dense MLP. The geometry reads num_experts, num_experts_per_tok and moe_intermediate_size (the shared expert's width must equal it, as it runs in the routed experts' tables). The loader dequantizes the 8-bit router and shared-expert gate rows to bf16 on the host, and packs each layer's 256 experts plus the shared one as a SwiGLU pair table and a down table. A layer's MLP is a union of the dense projections and the routed set; the forward routes, plans, runs the pair and down kernels and combines into the residual, inside the captured graphs as before. Prompt chunks past a window's rows take the prefill kernels, shorter ones the decode form, so a short chunk keeps one-row steps' bits. A drafter is refused (takes_drafter false): the checkpoint drafts with its own MTP layer, which is not on CUDA yet. tf-qwen35-cuda gains a probe mode that prints one layer's picks, weights and logits for a reference to compare. On an RTX 4090, TensorFold/Qwen3.6-35B-A3B-MLX-4bit-MTP: one-row steps decode at 141 tok/s; windows and trees of up to 16 rows equal the steps byte for byte over 300 ids; the 2B's and 27B's bits are unchanged. Against an fp32 PyTorch reference the steps agree within the model's own rounding sensitivity (mean KL 0.015 over 40 steps, where two references differing only by bf16 row rounding differ by 0.016; the dense 2B agrees at 0.0005).
--drafter <the checkpoint's directory> loads mtp-4bit.safetensors beside the weights: one attention-plus-experts layer over [norm(embedding of the next token) | norm(the row's final normed state)], the target's head for its draft, as the frozen engine's cuda/mtp.py ran it. The drafter taps the target's last layer and applies the final norm itself; each stream keeps the layer's keys and values over its whole window in a ring, absorbs the kept rows with the token after each (the newest row's output is the first draft), and holds a chain of up to TF_MTP_DRAFTS (default 3) more, each row the previous output with the draft before it. Drafts only propose: the target's windows verify them, so replies equal "draft": false ones. qwen36_mtp.cu holds the attention at the target's head width (256) with its partial rotary and output gate over the ring and the causal block; the ring kernels are the DFlash2 drafter's. The expert tables and router rows pack through the same loader as the target's, now writing into any arena. On an RTX 4090, TensorFold/Qwen3.6-35B-A3B-MLX-4bit-MTP with --parallel 2 --context 8192: one stream decodes 258 / 229 tok/s greedy (code / chat) and 252 / 181 sampled, against 101 / 98 plain; two streams 423 tok/s aggregate on code. bench_concurrent --alone --serial: every drafted reply equals its "draft": false reply and its lone run, 0 unequal. A chained draft costs 0.5 ms; a hold 1.2 ms, an absorb 0.7 ms.
Owner
Author
|
All workflows green at 178414b (no-GPU on ubuntu and windows, fmt, app); closing, the real pull request is on ashhart/TensorFold. |
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.
Fork-internal draft to run the workflows on the Qwen3.6 MoE commits; the real pull request goes to ashhart/TensorFold.