Repository navigation
MiMo-V2.6-Flash on CUDA: EXL3, DFlash, two GPUs - #728
ajensenwaud wants to merge 13 commits into
Conversation
…3 experts Checked ash hart PRs open+closed (read-only). Direct hits on our problems, all OPEN, branch BHCC2025/TensorFold exl3-cuda-experts: - ashhart#741 EXL3 routed experts behind experts.Grouped: gate+up+SwiGLU ONE kernel, down on pair rows, each expert's rows rotated by its own suh INLINE (register-resident rot for 1-2 pair chunks - the problem we had deemed impossible tile-incrementally). Rides on grouped.zig's Plan (in our tree). 297cu+124zig lines, with an exactness test harness (tf-cuda-test exl3-experts) - 367,872 values within FP64 bounds, alone==grouped byte-equal. - ashhart#740 EXL3 projections behind qlinear (the stacked base). - ashhart#718 branch-free trellis window extraction (perf). - ashhart#729 share decoded weights across prefill rows (perf). - ashhart#728 MiMo-V2.6 CUDA EXL3 two-GPU (another full CUDA EXL3 family). - ashhart#765 GLM CUDA multi-rank doorbells; ashhart#770/777/768 nemotron CUDA single-launch shared rounds + batched MTP drafting (patterns for our later concurrency work). PLAN: vendored exl3_experts.cu/.zig + exl3.cuh into our tree (this commit, unwired). Next: build wiring (fatbin + mod), port moeBlock to Grouped.decode with the router->plan plumbing from grouped.zig, run their exl3-experts test on our checkpoint, parity, bench. Our own grouped kernels (stage-1 parity) stay as fallback.
|
Thank you, @ajensenwaud, the shared-lane DFlash path and explicit optional FP8 policy give the MiMo work useful boundaries. We checked 83fd3f2. native.zig:374 still requires captured kernels, :379 hides bad TP settings with catch null, and rankWeights:412-416 divides bytes despite replicated embeddings/drafter. We can repair refusal and placement accounting while moving execution onto owned kernels and the shared EXL3 module. We will extract the depth rule and infrastructure first, then the family with prompt reuse and final-source numerical/peer/memory gates; 16-bit cache remains the default. Please provide the original/converted checkpoint and vision inventory with the pinned workload so image obligations and the final reference gate are explicit. Earlier pre-rebase receipts remain attached to their source. Posted by Ash's AI Assistant Synthena |
|
Thank you, @ajensenwaud, for #728 and the 24 concurrent-versus-solo checks on two GPUs, with @jayleaton and the common EXL3 contributors. MiMo-V2.6-Flash is not included in 1.1.1. Its current CUDA path depends on a Python-captured Triton set that TensorFold does not ship; without TENSORFOLD_CUDA_KERNELS it refuses to start. Native kernels for its eight Triton launches come next, followed by full-model fidelity, drafted/plain, resumed/fresh, concurrent/solo, two-rank failure and compiler/resource gates on the integrated EXL3/DFlash path. This is held for a later 1.1.x; captured-set receipts do not qualify the new native path. Posted by Ash's AI Assistant Synthena |
|
Thanks @ajensenwaud for MiMo-V2.6-Flash on CUDA. Main has moved a long way since this branch (it is stacked on #703), so rather than ask you to rebase, we'll port its MiMo-specific commits onto current main ourselves, keeping you as their author and saying in each message what was adapted. Then it goes to the GB10 kernel contract and served-equality runs on one and two GB10s before it lands. If you are partway through a change here, please say so, so we don't collide. - Synthema, Ash's AI Assistant. |
|
Thank you. Very keen to contribute. I know it's moving fast. I can rebase and recommit if that helps? |
|
@ajensenwaud, a question from the port before we combine the MiMo paths with main's shared EXL3 operators. In the branch, the arithmetic a row gets depends on how many rows or tokens arrive together: the router picks its tile by row count (BM 16 up to 16 rows, then larger tiles with BK 256 to 64 and BLOCK_E 32 to 64, - Synthema, Ash's AI Assistant. |
|
Thanks Synthema, and thanks for porting rather than asking for a rebase. Answers to both open questions. 1. Reference arithmetic. Please let main's shared EXL3 operators define the reference and regenerate the oracle fixtures against them. One arithmetic per operation is the right contract: our 24/24 concurrent-versus-solo result was measured, not guaranteed by construction, and I'd rather not keep a parallel MiMo numeric stack. Nothing MiMo-specific needs to survive in the linears or experts. The router tile switch came from the Python engine's Two requests so the port keeps what the branch was for:
I can run both on two Sparks against your port branch and post the receipts; point me at it. 2. Checkpoint, vision inventory, workload (your 17 Oct note)
Nothing in progress on the branch, so no collision risk. |
ajensenwaud
changed the title
|
Thanks Synthema. To avoid duplicated effort: I've rebased #728 onto current main ( I also fixed two points from your earlier review. Admission now uses main's
|
ashhart
left a comment
ashhart
left a comment
There was a problem hiding this comment.
Thanks @ajensenwaud. Commit 026bd20 adds an fp8 KV cache (--kv-cache fp8) that the PR itself notes is not bit-equal; TensorFold takes no lossy formats, even opt-in, so that commit has to go. Main has also moved MiMo-only kernels and geometric depth to a separate MiMo track (1964746), so the MiMo EXL3 kernels here, copies beside the shared cuda/exl3.zig, need to fit that plan and be registered with the kernel contract. It's about 150 commits behind and conflicts in docs/recipes/README.md and native/cuda.zig. After a rebase we'd check on two GB10s that a resumed conversation equals a fresh run across the 256-token boundary, and TP=2 equals TP=1.
|
Review against main e9b69ed. Commit 026bd20 adds an 8-bit key/value cache (
|
|
Yes in progress. Working on it. |
…one GPU weightShare read TF_TP_* with catch null, so a malformed world or rank was admitted as one GPU's share. It now returns BadSetting as open does, and the admission reports it through the family's explain. weightShare may return an error union (cuda_family_memory's share, the same change as MiMo-V2's in ashhart#728). Host test: unset, six malformed settings, and a two-rank share.
|
Hi @ajensenwaud, a status update on #728. What we've checked: the full diff, rebased onto current main without conflicts. The shared EXL3 prompt kernels Two small fixes are needed before it can land, and we have both ready if you'd like us to carry them:
What's still owed is our device run. Every cell needs two GB10s at once (TP=2 over RoCE): the full kernel Thank you for your patience with this one. |
|
Thanks @ashhart. Yes please, carry both fixes (Linux-only registration of the MiMo host tests, and I've rebased the branch onto current main (the only conflicts were the kernel and contract image lists); nothing else changed. |
…one GPU weightShare read TF_TP_* with catch null, so a malformed world or rank was admitted as one GPU's share. It now returns BadSetting as open does, and the admission reports it through the family's explain. weightShare may return an error union (cuda_family_memory's share, the same change as MiMo-V2's in ashhart#728). Host test: unset, six malformed settings, and a two-rank share.
|
@ashhart there appear to be merge issues on this one now. CI fails on Happy to push that fix, rebased onto current main, unless you're already looking into it alongside the macOS registration fix. |
|
Thanks @ajensenwaud for tracking that down. Yes, please push the |
|
On it ... fixing this now. |
…one GPU weightShare read TF_TP_* with catch null, so a malformed world or rank was admitted as one GPU's share. It now returns BadSetting as open does, and the admission reports it through the family's explain. weightShare may return an error union (cuda_family_memory's share, the same change as MiMo-V2's in ashhart#728). Host test: unset, six malformed settings, and a two-rank share.
A family whose head drafts a whole block in one pass (DFlash) pays for a draft pass once, whatever its depth, so the per-depth rule prices deep drafts too high. With `geometric_depth` set, a lone stream's depth is the one with the most expected tokens a ms under a geometric chain of its own acceptance rate (a moving average of kept over tried, from 0.6), the forward of its rows plus one draft pass; a stream that stops drafting tries one draft every 8 rounds. A round's cost can grow with the stream's context: `geometric_ctx` adds, per 1,000 positions past `geometric_ctx_base`, a step and a per-draft term, so long streams draft shallower as full-attention reads grow. `geometric_costs` and `geometric_draft_cost` give the per-depth rule of shared rounds the family's own cost table. The rule follows the Python engine's MiMo-V2 decoder (joint_widths). Everything defaults off: families that do not set it keep the per-depth rule unchanged.
…rows The prompt kernels' tile decode, weight-stage sizes and loads take the half-bit width and codebook as template arguments; tf_exl3_prompt_gu / _down stay the 4-bit mcg instance and tf_exl3_prompt_gu_k6_mul1 / _down_k6_mul1 are the 3-bit mul1 one (MiMo-V2.6-Flash's routed experts). eligible() takes either verified uniform width, Kernels picks the instance by layer, and run() takes the rows' dtype for the per-pick rotation (bf16 or fp16; the shared-rotation kernel stays bf16). GLM passes .bf16. On sm_121 the 4-bit gate/up kernel's SASS is unchanged and the down kernel's differs in one mask constant (0xe0 -> 0x3e0, the same value for the 256-thread launch).
… DFlash, two GPUs, prompt cache The mimo_v2 family serves an EXL3 pack of MiMo-V2.6-Flash on the native CUDA server, on one GPU or two DGX Sparks as tensor-parallel ranks, with the pack's DFlash drafter. Every kernel is built in: no captured set, no copies of the Python engine's kernels. - EXL3: projections and decode experts on cuda.exl3.Linear and cuda.exl3_experts; prompt experts on the common prompt kernels (their 3-bit mul1 instance, the commit before), then the family's slot sum. Every prompt token runs on the prompt path, so the decode/prompt cut is gone. - kernels/cuda/mimo_v2: ops.cu (RMS norms, SwiGLU, the fp32 router and its top 8, the experts' slot sum), attention.cu (queries and caches, chunked decode attention with sinks and its merge, a tensor-core prompt kernel for the full layers whose 64-key tiles start at 0, so a row's bits never depend on its chunk), glue.cu (embed, residual adds, RoPE tables, argmax), dflash.cu (the drafter's kernels). - families/mimo_v2: config and rank cuts, loading, the decode forward (eager and graphed by rows and key bucket), the prompt path, the DFlash drafter, the engine both ranks run, and kept prompt states (kept.zig: the slot's full-layer rows pinned, the sliding and drafter windows copied; a resume binds to the state's slot or copies the rows into a free one). - native.zig: the lane backend. Rank 0 serves and sends every GPU step (and cache step) over the plan link before launching it; rank 1 runs `tf-mimo follow`. --parallel sets the streams' cache slots; the prompt cache's marks are 2,048 rows apart; a lone stream's first copy window is 6 rows; draft depth uses the lanes' geometric rule with a context term. - native/cuda.zig: MiMo registered. Admission takes main's weightShare hook: the split matrices over the ranks plus the embeddings and the drafter whole on each; a bad TF_TP_* setting refuses instead of falling back to one GPU. - Checks: tf-mimo m1 (teacher-forced KL and argmax against an answer key and an earlier build's dump, one prompt chunk against two) and host tests for config, cuts, plans, kept-state runs and slot choice. The family's build is zig/build/mimo.zig (as kolibri.zig), its kernels are listed in cuda_kernels.zig, and the TP mailbox fatbin is built once for tf-tp-test and tf-mimo. The ranks run over Jay Leaton's tensor-parallel session and RoCE mailbox from his two-Spark DeepSeek-V4.1 work.
… for gate 5 contract_mimo*.zig registers the family's 28 kernel forms against float64 or their stated rules (ops, attention with the tensor-core prompt kernel, glue, the drafter's), served shapes and launches as the family makes them. mimo_cuda_records.zig carries their sm_121 usage from nvcc 13.3 (cuobjdump), with the two 3-bit mul1 instances of the common EXL3 prompt kernels.
… shallower top-8 tree mimo_attn_rows runs a chunk's G <= 16 query heads as one m16 tile on the prompt kernel's loads and mma, keys in 64-key tiles at multiples of 64 split over four warps whose partials fold in warp order; p in fp16 for o (the contract's allowance as the prompt kernel's). A row's bits still depend only on its own position and window. mimo_rms keeps its values in registers (unrolled; same bits), mimo_topk runs its tree's last five levels in warp 0's shuffles (the same pairs; same bits). Sliding rows go from 72 to 8.5 us, norms from 11 to 5.7 us.
load set each layer as it filled it, so a refusal before then (an fp16 embedding, say) had deinit free whatever the uninitialised layers' expert tables held. Every layer now starts blank, without experts. The host test refuses an fp16 embedding under the testing allocator; without the blank start it crashes in deinit.
… GPU weightShare read TF_TP_* with catch null, so a malformed world or rank was admitted as one GPU's share. It now returns BadSetting as open does, and the admission reports it through the family's explain. weightShare may return an error union (cuda_family_memory's share); GLM's and Nemotron's plain byte counts pass through unchanged. TF_TP_WORLD unset is still one GPU. Host test: unset, six malformed settings, and a two-rank share.
With TF_TP_WORLD unset, follow read the cfg of a session that world 1 never opens. The world and rank check now runs on the parsed settings before the driver opens.
Since 1911e7b GLM's CUDA family takes fabric/tp over the engine's cuda module, and a second family over the same runtime made its own tp module, so the engines failed with `zig/src/fabric/tp/tp.zig: error: file exists in modules 'tp' and 'tp0'`. tp.module now keeps one module per cuda module and returns it on every later call.
…one GPU weightShare read TF_TP_* with catch null, so a malformed world or rank was admitted as one GPU's share. It now returns BadSetting as open does, and the admission reports it through the family's explain. weightShare may return an error union (cuda_family_memory's share, the same change as MiMo-V2's in ashhart#728). Host test: unset, six malformed settings, and a two-rank share.
|
Thanks @ashhart. Pushed, rebased onto current main (
Locally on x86_64 Linux: |
|
All ready for your review, @ashhart |
|
Thanks @ajensenwaud, and sorry for the wait. Everything we asked for is in, and revision |
MiMo-V2.6-Flash on the native CUDA engine, on two DGX Sparks as tensor-parallel ranks, with the pack's DFlash drafter. Proposal and discussion: #725.
Admission is what was validated: GB10 (sm_121) at two ranks.
openrefuses any other GPU by name (MimoV2RequiresSm121), as DeepSeek-V4.1's kernel load does, and refuses one rank or any other world (MimoV2RequiresTwoRanks), at the weight share as well as atopen.tf-mimorefuses both too. Host tests cover each refusal. One rank, and the 2.27 bpw pack that fits one GB10, wait for their own receipts.Based on current
main(2836e9b2). Twelve commits:dfcd1b50lanes: a geometric draft depth whose round cost grows with context. About 100 lines incore/lanes, plus a host test. It is off unless a family setsgeometric_depth, so other families keep the per-depth rule unchanged. It can become its own PR if you prefer.6ca03954exl3 cuda: the common prompt path takes 3-bit mul1 experts over fp16 rows. The prompt kernels inexl3_prompt.cuh/exl3_experts.cutake the half-bit width and codebook as template arguments:tf_exl3_prompt_gu/_downstay the 4-bit mcg instance,_k6_mul1is the 3-bit mul1 one.exl3_prompt.zigpicks the instance by layer and takes the rows' dtype; GLM's call passes.bf16. On sm_121 the 4-bit gate/up SASS is unchanged and the down kernel's differs in one mask constant (the same value at 256 threads).1a7f7c4fMiMo-V2.6-Flash on CUDA: native kernels on the shared EXL3 operators, DFlash, two GPUs, prompt cache.cuda/mimo_exl3.zigand the vendored kernels are gone. Projections and decode experts run oncuda.exl3.Linearandcuda.exl3_experts, prompt experts on the common prompt kernels (commit 2).zig/kernels/cuda/mimo_v2):ops.cu(RMS norms, SwiGLU, fp32 router and top 8, the prompt experts' slot sum),attention.cu(queries and caches, chunked decode attention with sinks, and a tensor-core prompt kernel for the 9 full layers),glue.cu,dflash.cu.kept.zig): the slot's full-layer rows stay pinned, the sliding and drafter windows are copied (33 MiB a rank); a resume binds to the state's slot or copies the rows into a free one. Rank 1 follows every cache step over the plan link.zig/build/mimo.zig(module,tf-mimo,test-mimo), laid out likekolibri.zig; one line each incuda.zig(tf-mimo),cuda_engines.zig(the engine) andcuda_family_host.zig(host tests, beside Dory's).native/cuda.zig: one line, the registry entry; admission uses main'sweightSharehook.docs/recipes/mimo-v2.6-flash.md.5ab18d7ckernel contract: every MiMo-V2 kernel through gates 1-4, its launches for gate 5.contract_mimo*.zig(28 kernel forms against float64 or their stated rules, at the served shapes, launched as the family launches them) andmimo_cuda_records.zig(sm_121, nvcc 13.3, including the two mul1 prompt instances).4f56747fMiMo-V2 decode: attention rows on tensor cores, norms in registers, a shallower top-8 tree. Brings one-stream decode back level with the captured-set build.14bc442dMiMo-V2: a refused load frees only the layers it filled.7cd49cbeMiMo-V2: a bad TF_TP_ refuses the weight share instead of sizing one GPU.*f867006btf-mimo follow: refuse one rank before reading a session.525bb1d6MiMo-V2.6-Flash recipe: licences and sources of the model, the pack and the drafter.1d777654MiMo-V2: refuse every GPU but GB10 (sm_121) by name. The check is the first line ofopen, and oftf-mimo. A host test runsopenon a fake driver at sm_80 to sm_120, all refused, and at sm_121, which passes the check.dc53407fMiMo-V2: refuse one rank by name; only two were validated.rankView, the weight share andopenrefuseTF_TP_WORLDunset, 1, or anything but 2. The recipe drops the one-GPU and 2.27 bpw serving claims.dc92b94cMiMo-V2.6-Flash recipe: full revisions of the model and both packs.XiaomiMiMo/MiMo-V2.6-Flash-RLat5711b268169967567844e1e560e8a3966da959b1, the revision both packs were converted from and the one validated. The repository's head since then,dcd3fd83, changes only the visionpreprocessor_config.json.ajensenwaud/MiMo-V2.6-Flash-RL-exl3-3.0bpwat139d9a4e1eb32193890036fdddd521fe37033a30.benthecarman/MiMo-V2.6-Flash-RL-exl3(2.27 bpw, the drafter's source) at5ae87b830eb950c705967aea6130601a8e1f74a5.The grouped experts load-ahead that was commit 5 is now its own PR, #1505, with a GLM-5.3-Flash A/B. The gate below was measured with it in: without it MiMo's teacher-forced logits are byte-equal, so the identity and quality results stand, but decode tok/s is lower by the grouped kernel's difference (about 109 against 96.6 µs a launch) until #1505 lands.
The 8-bit KV cache stays out (fp16 only), as you asked.
Gate (the previous head with the load-ahead, two DGX Sparks)
GB10 (sm_121), checkpoint ajensenwaud/MiMo-V2.6-Flash-RL-exl3-3.0bpw,
--parallel 4 --context 608256 --top-k 20 --top-p 0.95, RoCE between ranks.Kernel contract (full CUDA runner on GB10, this head):
SUMMARY 446 kernels: 446 pass gates 1-4, 0 fail, 5 known findings; 312 controls: 312 red as planted, 0 missed; 214 limit problems. None of the limit problems are MiMo's or the shared experts': every MiMo launch and the threetf_exl3_groupedlaunches match their nvcc 13.3 records live on sm_121.Lane contracts (greedy, by the server's
token_sha): drafted =="draft": false6/6 (prompts of 29 to 10,806 tokens); concurrent == solo 24/24 (three groups of four, drafted and plain).Resumed == fresh (prompt cache on, 3.8 GiB): 117/117 replies identical to the cold serve's, 63 resumed (102,602 of 250,302 prompt tokens from the cache), including resume points at 63/64/65, 127/128/129, 255/256/257, 1,023/1,024/1,025, 2,047/2,048/2,049 and 4,095/4,096/4,097 tokens (18/18 resumed at the mark). With
"draft": false, 117/117 identical (16 resumed).Teacher-forced (
tf-mimo m1): against exllamav3 on the same pack, see Fidelity against exllamav3 below. Prompt path: one chunk == two chunks on the last row 8/8; against the decode rows, KL mean 0.0073 (different kernels, so not bit-equal).Quality (cache on): GSM8K 477/500, MMLU 190/228, passkey 10/10 (32K and 120K, five depths each).
Prefill, tok/s (spark-bench depth, cold):
The recipe's release build reads 16K at 1,309 and 64K at 1,183.
Decode, tok/s (spark-bench
--preset full, 512 tokens, thinking off, temperature 0; chat / code / agent / report):The bench tags each prompt at random, so one-stream figures move with the text. On fixed prompts (512 tokens, one stream, temperature 0, both builds served the same day), tok/s:
"draft": falseFidelity against exllamav3
The reference is exllamav3 1.6.0, the build that wrote the 3.0 bpw pack, run on the same pack. It is teacher-forced over the same 1,905 positions: 8 texts, 254 prompt rows and 1,651 generated rows. It runs its no-cache path (
flash_attn_nc), streamed one module at a time on one GB10, and its logits are its default fp16 head output. Ours come fromtf-mimo m1at TP=2 on this PR's kernels.shorttext, at steps 48 and 59.edit, step 50). exllamav3's lead over its second token there runs from 0.016 to 0.31 logits, median 0.13. The full list is below.Every position where the top tokens differ, with exllamav3's margin
Margin = exllamav3's top-1 minus top-2 (logit, prob). "rank" = where this PR puts exllamav3's top-1. Token = id
text. Next = the teacher token at that position.messagemysterious(16.8438, 0.1630)strange(0.1622)comingcoming(18.8438, 0.2429)storm(0.2410)TheHe(18.0312, 0.3508)The(0.3144)gratitudewonder(22.4688, 0.4402)gratitude(0.4338)aa(17.0625, 0.1783)from(0.2353)lmessage(17.3125, 0.3852)l(0.3987)wouldwould(23.4062, 0.4174)climbed(0.3757)lodgedlodged(22.2188, 0.2552)half(0.2813),in(25.7188, 0.4495),(0.4278)无法不能(21.8906, 0.2774)无法(0.3350)粒子两个(22.5000, 0.4304)粒子(0.4073)。的状态(25.8438, 0.3952)。(0.4682),.(23.8594, 0.3009),(0.3421)\n\nLet(14.8438, 0.0822)\n\n(0.0781)::\n(17.7812, 0.2198):(0.2269)**Rename(16.1406, 0.1151)**(0.0988)AnswerExpected(16.0625, 0.0680)Answer(0.0761)countcount(24.5781, 0.5373)cnt(0.5366):):\n(30.0312, 0.5592):(0.5490)TheRot(22.9062, 0.4174)The(0.4054)!I(8.7500, 0.0129),(0.0173)-g-g(17.5156, 0.0956)-m(0.0996)-First(14.1562, 0.0263)-(0.0298)II(20.8594, 0.3520)The(0.3544)correctlycorrectly(24.6719, 0.4150)fine(0.4299)tryingtrying(17.5000, 0.1722)not(0.1535)The two exact ties, both in the
shorttext. exllamav3's top two are equal, and our top token is one of them:SpriteSpriteand 51548Collision(16.4219)Sprite(0.1737)spriteplayerand 17051sprite(18.4219)sprite(0.3132)Host (this head, on
main2836e9b2):test-mimo,test-kernel-contract,test-native-options-host,test-kernel-symbols-host,test-shared-lane-regressions-hostandtest-tp-clipass;tf-mimo,native,compile-config-serverandcompile-cuda-testsbuild.test-cuda-hostpasses except two core prompt-cache spill tests,rows_async_test("the job permit bounds metadata ...") androws_close_test("spill-only flushes more than eight RAM states ..."). They fail the same way onmain2836e9b2in the same build container.zig fmt --checkandgit diff --checkare clean;tools/zig/lean_check.pyfinds 426 problems, the same as main.Next
Notes for review
cuda.zig,cuda_engines.zig,cuda_family_host.zig), so whichever lands second needs a one-line rebase there.