Repository navigation
Muse Spark Open on CUDA: EXL3, DFlash2, two GPUs - #1295
ajensenwaud wants to merge 13 commits into
Conversation
|
Thank you for this. Bringing Muse Spark Open up on two GB10s with drafts and exact repeats is serious work. To fit how we land things, it needs splitting before we can run it:
After that, the family PR will just be the Muse Spark family, and we'll run it on two GB10s. - Synthema, Ash's AI Assistant. |
ashhart
left a comment
There was a problem hiding this comment.
Thanks @ajensenwaud. We rebased it onto current main (conflicts in chat.zig next to #1298's cut tool calls, build lists and the CUDA registry; all kept). Host suites and both Linux CUDA compiles pass, and lean adds no findings. Before a device run:
TF_MUSE_KV8and the e4m3 attention mode are still here (forward.zig,engine.zig,ops.zig,kept.zig,attention.cu, the recipe). The engine takes no lossy KV, even opt-in, so please remove them. The served exactness results were taken with int8 KV on, so they need rerunning with fp16.- Please drop the recipient-replies and sampler-race commits now that they're #1336 and #1337, and stack on #1336.
- Register the
muse_*kernels with the CUDA kernel contract. - Where the shared EXL3 operators cover a shape, please use them instead of the vendored copies (
muse_spark/exl3_linear.cu,exl3_experts*.cu,experts_grouped.cuh,decode.cuh). memory_reserve_gib = 3lowers the OS reserve from about 12.8 GB to 3 GB by default on a shared GB10. We'd rather this be the operator's explicit choice, with measured per-rank peak footprint showing the headroom, than a family default.- Please give the full revisions of the pack and the DFlash2 drafter, and add host tests for the lane backend's verify/keep/draft.
Then we'd run, on our GB10 pair: shared fatbins identical, the contract SUMMARY, drafted vs plain (greedy and seeded), concurrent vs solo, resumed vs fresh, graphed vs eager, and a GLM control.
85420d6 to
9168996
Compare
|
Thanks. All four done:
Output is byte-identical to before, so the parity numbers stand. The branch is just the family now, rebased on main. Happy to run anything on our two GB10s. |
9168996 to
9b33491
Compare
9b33491 to
e2231f5
Compare
|
Re-review of e2231f5 against main 3f408b6. Thank you, every point from the earlier review is addressed: the fp8/int8 KV is gone, the vendored EXL3 copies are replaced by the shared operators (registered as
|
|
Thanks Ash. All six from your review are addressed on the branch (
Host suites pass on the new head ( |
e2231f5 to
7db0640
Compare
|
The |
|
Please let me know what more is pending. |
951d536 to
92e7cae
Compare
|
Thanks @ajensenwaud. The Muse Spark family files are byte-for-byte what we re-reviewed on 26 Oct. Only the registration lists and #1336's two newer commits changed, so that review stands. Rebased on main Three small things, then our GB10 run:
On our pair we'll |
92e7cae to
aa9588a
Compare
|
Thanks @ashhart. All three are done. The branch is rebased on main
Every miss on the four short texts, with both sides' top 5, margins and gapsEnglish prose (4 misses of 157)
Chinese prose (1 misses of 138)
Code (2 misses of 214)
Agent turn (9 misses of 239)
|
|
Thanks @ajensenwaud. The licence paragraph, the per-miss tables and the concurrent-alone runs answer our three asks, and the family's code is unchanged since 26 Oct apart from the recipe. We merged aa9588a onto main a1e6d34 (clean). fmt and lean are clean, and the server changes are byte-for-byte #1336's, which pass Three things before our GB10 run:
Then on our pair, with the 1.8 bpw pack and a 3 GiB reserve on both ranks, we'll run:
|
aa9588a to
170b0a9
Compare
|
Thanks, Ash. All three are addressed on the branch (rebased onto current main; #1336's server commits are unchanged). 1. The drafter's licence. You're right: RadixArk's 2. The misses. The long text's 64 are below, without the context column, in the same columns as the short texts'. Position 124 of the agent turn: I reran exllamav3 over that text and kept the 222-token prefix run's logits this time (its full run reproduced the stored reference byte for byte). exllamav3's own prefix run does move 124, but not across: its lead of 3. Admission. The long text's 64 misses (no context column; 11 token ids redacted)Same run and columns as the short texts' table: exllamav3's Summary: 64 misses of 2,600 rows (top-1 0.9754); exllamav3 against itself (prefix vs full forward) 0.9777 over the Redactions: the text is an internal note, so every token listed below (true next, both tops, both top-5s) was Long mixed text (64 misses of 2600)
|
|
@ashhart could you rerun |
…g, answer and tool calls A tokenizer with <|start|>, <|message|>, <|eom|> and <|eot|> as single tokens and a chat template that writes to=self messages (Muse Spark Open's) replies as messages to recipients. self is reasoning_content; user, commentary and no recipient are content; any other recipient is a tool call whose <atem:function_calls> block becomes the <tool_call><function=...> markup tool_parse and tool_stream already read, so its arguments take the tools' schema types. While a reply streams, an unfinished header, a call message that has not closed and a partial marker are held back, so every split extends the last and the finished split extends them all; the finished reply feeds the call streamer once more, so a last call ended by the end token streams too. Harmony's tokenizers have no <|eom|> or <|eot|>, so gpt-oss replies keep their channel split; tests cover both formats whole and streamed.
…ompt's open message read, thinking off kept off A call message max_tokens cuts stays open as far as it was written, named by its header until its invoke names the tool, so tool_stream reads it as the reply's cut call, whole and streamed, as ashhart#1299 made a GLM call; content never shows its markup. A stop string inside a call keeps the call's text, as the other formats keep theirs. The reply goes on from the prompt's last unclosed <|start|> header, so a generation prompt that already opened to=self<|message|> keeps the reasoning after it; Muse Spark Open's template ends at <|start|>assistant, with or without reasoning_effort. Only calls to offered tools become <tool_call> markup: with no tools, or a tool not offered, a call message stays as the model wrote it, as DSML, LFM2 and Qwen's XML calls do. Thinking off no longer turns the split's reasoning on: self messages stay out of the reply, and the template, which has no thinking switch, hears its lowest reasoning effort rather than its default (xhigh). Raw completion prompts keep the model's text. Tests cut a call before its invoke names the tool, inside a value, a value's closer, a parameter's name, between parameters, after the invoke closed, and in a call's header, at strides 1, 3, 7 and whole; and cover an opened prompt, no tools, an unoffered tool and thinking off.
…er that embeds it
…TP=2 over fabric/tp, DFlash2, prompt cache The family serves an EXL3 pack of Muse Spark Open's text model on the native CUDA server, one GPU or two DGX Sparks as tensor-parallel ranks (rank 1: tf-muse follow). - EXL3 projections and routed experts run on main's shared operators: cuda.exl3 Linear (tf_exl3_rot_in, tf_exl3_linear_*) and cuda.exl3_experts (group, rotate, tf_exl3_grouped_2, the fp32 SwiGLU epilogue, down). - Family kernels: attention.cu (head dim 64, GQA 8:1 decode attention, fp16 cache), glue.cu (norms, gated residuals, fp32 router, top-8, expert norms and sums, softcap, argmax, RoPE tables), linear_multi.cu (q/k/v/gate in one launch and 32-row passes over exl3_decode.cuh, the shared linear's bits), experts2.cu (32-member tiles of the shared grouped GEMM for prompt windows, its bits), dflash.cu and dflash2.cu (the drafter). - A row's bits do not depend on its window, the graph bucket or prompt chunking; the prompt cache keeps states through the snap_* hooks; the family declares its prompt-cache budget default (1 GiB a rank, --prompt-cache-gib overrides it). The memory reserve stays the operator's (TENSORFOLD_MEMORY_RESERVE_GIB). - tf-muse m1 / prompt / dflash / drafted checks and test-muse host tests; docs/recipes/muse-spark.md.
…hes for gate 5 Glue, attention and drafter kernels against float64 (rows and dense families where the formulas match) or their stated rules; the cluster norm forms and the EXL3 launch forms as bit-preserving replacements.
…), with host tests rounds.zig holds what verify, keep and draft do before and after their GPU calls, and native.zig's backend calls it: a round's rows (each window's pending token, then its held drafts or the host's, at position - 1 on the stream's slot), whether the forward may return ids (no window samples), the drafts a window verified, the kept-path check (chains from row 0; no trees), each draft request's entry (the last follow token or the fed first token, the depth capped by the drafter's block, held drafts cleared) and the pass output into each stream's held drafts. Behaviour is unchanged. Host tests cover each, with the refusals: trees, unknown streams, windows past the rows, drafts not held, empty or branching paths, no drafter, no first token.
…e the last call; a reply with no markup streams as it comes
…espace back, as the finished reply strips it
Meta's Muse Model License for the model, AMD's export and the EXL3 pack; none stated for RadixArk's DFlash2 drafter, which is built locally rather than published; MIT for exllamav3, the quantiser and parity reference.
With main's loop-guard stop, chat.zig plus the recipient-reply hook reached 603 lines, past lean's 600. flag (Python's bool(fields.get(key))) is a fields helper, so it moves there unchanged; chat.zig calls fields_mod.flag.
…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.
…ame otherwise Muse Spark was gated and run only on two GB10s at TP=2. open now refuses any other GPU (MuseSparkRequiresSm121) and an unset TF_TP_WORLD or any world but 2 (MuseSparkRequiresTwoRanks), as does tf-muse follow; the weight share and the target's and drafter's rank views refuse one rank by the same name, so admission names it before sizing memory. Host tests: the admit rule, rankView(1), and the weight share at one, two and four ranks.
…ates no licence RadixArk's Muse-Spark-Open-DFlash2 states no licence, so all rights are reserved by default: the recipe no longer names a drafter pack or a --drafter command, and serves two GB10s only, as the family now admits.
170b0a9 to
1d40125
Compare
What this changes
Adds Meta's Muse Spark Open (
model_type = "muse_spark", 816B total / 42B active MoE, text) to the native CUDAengine: an EXL3 pack (mul1, 1.8 bpw,
ajensenwaud/Muse-Spark-Open-816B-A42B-EXL3-1.8bpw) served at TP=2 across twoGB10s with plain decode, and the prompt cache through the family's
snap_*hooks. Engine notes:docs/recipes/muse-spark.md.The drafter is not part of the served setup. The family also reads a DFlash2 drafter (
--drafter), written forRadixArk's
Muse-Spark-Open-DFlash2, which states no licence (so all rights are reserved by default). The recipeserves plain decode, names no drafter pack and says the drafter is unlicensed; RadixArk will be asked to state one.
Stacked on #1336 (the server's recipient-reply parsing, which this family's chat format needs): #1336's commits
are the
server:ones, and their change tozig/src/server/is byte-for-byte #1336's; the rest are this PR's. The sampler change (#1337) is not needed. KV is fp16 only.cuda.exl3.Linear,exl3_experts, main'sprompt kernels); the vendored copies are gone. Two launch forms stay because no shared form covers them, each
registered as a byte-equal replacement of main's kernel: q/k/v/gate in one launch with 32-row passes
(
linear_multi) and 32-member expert tiles (experts2).with sliding and full layers, norms and gated residuals, the fp32 router and sigmoid top-8, the per-expert norm
and sum, the softcapped head, and the DFlash2 drafter's convolutions, top-k and candidate walk.
openrefuses any GPU but sm_121 (MuseSparkRequiresSm121) and one rank or any world but2 (
MuseSparkRequiresTwoRanks, also for TF_TP_WORLD unset), asdeepseek_v41/kernels.zigrefuses by name;tf-muse followand the target's and drafter'srankViewdo the same.weightSharerefuses a malformed TF_TP_*(
BadSetting) and one rank instead of sizing one GPU's share (MiMo-V2.6-Flash on CUDA: EXL3, DFlash, two GPUs #728's fix for MiMo-V2).zig/build/{cuda,cuda_kernels,muse,options,tp}.zig,zig/src/cuda/kernels.zig, the registry line inzig/src/native/cuda.zig, andzig/src/native/cuda_family_memory.zig: one optional family declaration,prompt_cache_gib(Muse: 1 GiB a rank when--prompt-cache-gibis not given; families without it are unchanged),and
weightSharemay return an error union, the same hunk as MiMo-V2.6-Flash on CUDA: EXL3, DFlash, two GPUs #728's (the two merge cleanly there).server's default reserve (a tenth of unified memory, 12.2 GiB) does not leave room for; the recipe passes
TENSORFOLD_MEMORY_RESERVE_GIB=3(per-rank figures under Receipt).rounds.zig: verify rows, kept paths, draft entries, held drafts) hasits own host tests, refusals included.
Receipt
receipts below are from this family on 2ff0191; the branch is now on server: replies written as messages to recipients split into reasoning, answer and tool calls #1336 over db3fd5a (the family takes one line each in
cuda.zig,cuda_engines.zigandcuda_family_host.zig; the sampling-fieldflagreader moves tofields.zigtokeep
chat.zigunder 600). Concurrent == solo: see the review reply. Receiptsmarked "drafter on" were taken before the drafter left the served setup; plain-decode speed on this head is not in
this receipt yet (it needs a GPU window; both GB10s are serving).
docs/recipes/muse-spark.md):ajensenwaud/Muse-Spark-Open-816B-A42B-EXL3-1.8bpwat38ad2925945840c3a677bd3f213867d30368636c(weightsas of
efcdf2ec4222c4597e291f1b07edd19c3c02c652; the later commit changes the card only), quantised fromamd/Muse-Spark-Open-816B-A42B-Quark-MXFP4at4885eb2a901c2638854ef35cf6833e6447622afb.RadixArk/Muse-Spark-Open-DFlash2at46c90f412549c472941a68382d91a99e5e6372a0, quantised locally to EXL3 4.0 bpw.reserve 3 GiB): 107.7 GiB a process idle,
108.8 GiB with a 24,160-token prompt running; MemAvailable then bottomed at 5.9 and 5.5 GiB on the two nodes.
/healthpeak on rank 0: 107.8 GiB. The server reports 2.7-3.9 GiB left for streams after a 3 GiB reserve, so thedefault 12.2 GiB reserve would leave none.
muse_sparkmodule (streamed full-model reference, 1.8 bpw pack, TP=2),per text KL 0.67-1.24x exllamav3's own prefix-vs-full noise floor, true-token NLL within +0.61 % / -0.08 %, top-1
0.96-0.99 (prose, Chinese, code, agent turn, 2.6K-token mixed text). Every top-1 miss is listed in the review
replies, the long text's without its context (it is an internal note; 11 token ids redacted). The logits this branch produces are
byte-equal to the engine that measured those numbers: 3,503 teacher-forced decode rows (10/10 files), drafted
windows (10/10, acceptance 953/2210 on both) and prompt chunks (0 plans differ from decode rows).
"draft": falseon 12decode prompts, serial and 4 at a time (12/12 equal token sha); resumed == fresh: 115/115 replies identical to a
cache-off serve over multi-turn, agent, document and seam-crossing conversations, 49 resumed, 130,544 of 285,822
prompt tokens from the cache.
zig build kernel-contract, GB10):SUMMARY 381 kernels: 379 pass gates 1-4, 2 fail, 3 known findings; 89 controls: 89 red as planted, 0 missed; 291 limit problems, exit 1. Every Muse registration (45)passes gates 1-4 and all 46 Muse launch records pass gate 5. The 2 fails (w4a8 nonfinite), 3 known findings and
291 limit problems are all in main's own kernels and records.
muse_topkstates and tests its NaN rule.code / agent / report); 4 streams aggregate 42.2 / 43.3 / 45.5 / 41.6. Context that fits: 1 stream x 471,040
tokens or 4 x 110,592. llama.cpp on Unsloth's UD-IQ2_XXS GGUF over RPC on the same two Sparks: 7.4 tok/s at 1
stream, 17.4 at 4.
15,375 tokens at 196.4 tok/s (TTFT 78.3 s).
193/200.
test-muse(26 tests: 2 new for admission and the weightshare, failing on the previous code),
tf-muse,native,test-server-cpu,test-golden,compile-config-server,compile-cuda-tests,test-kernel-contract,test-kernel-symbols-host,test-shared-lane-regressions-hostandtest-native-options-hostpass: 317/319 steps, 738/740 tests. The one failure istest-cuda-host'srows_close_test"spill-only flushes more than eight RAM states", which fails the same way on main db3fd5a on thishost.
zig fmt --checkclean;lean_check.py428 problems, the same as main. The CUDA build for sm_121 and thekernel-contract run above are from an earlier head (no kernel changed since).
Metal (CUDA only); vision (text only; image parts are refused).
Checklist
"draft": false, concurrent == solo, resumed == fresh (counts above).SUMMARYabove.