Repository navigation
[Feat] [MI355X] muse_spark: Whole-layer monokernels for Muse Spark - #26
Merged
Merged
Conversation
…verify steps to SGLang's own layers)
…launch with SGLang's layer on the live rows
…oPE, fp8 KV write) on common/, widths to 64 rows; layout + K1 tests
…6 cache) on common/, R up to 8; layout + test
…m and gate copies logged)
…ory_fence, global_ptr via llvm) installed by the family package
…bedded TP all-reduce, attention boundary, router/pooler, top-8, pre-expert norm) on common/; K1 test tolerates bf16 flips
…+ embedded EP all-reduce, sharded unpool, FFN boundary with the last-layer rule) on common/; K-attn test lengths
…er (bf16 layers on SGLang's module); spec builds the views, binds the caches, steps the epoch
…, not at package import (CPU tests import no aiter)
…K1 in one kernel) behind SGLANG_AMD_MONO_MUSE_LAYER=1; K1 waits on the previous boundary's x flags in a merged launch
…ith the reference layer's
… step-time script
…_AMD_MONO_MUSE_SKIP timing diagnostic
…P4 quantization (aiter's rule, Quark layout), SGLANG_AMD_MONO_MUSE_QUANT_BF16=0 keeps them on the modules
… the bf16-layer views
…sliding 2047 / full 95k) by graph replay
…weight-copy log reads tensors), for the router split, the residual gates and the online MXFP4 experts
…nto a *_trace scratch region) for K1 / K-attn / K2 / K3, and the timeline script
…e fan-ins, the router logits in one poll batch, the attention merge on the idle CTAs folding four partitions a spin
…into the B operands (the LDS DMA path cost ~4.5 us a pass)
…pool CTAs; the LDS DMA path kept (direct loads measured the same)
… on idle CTAs (four partial slots a row); K3 unpool: the next job's weights behind this job's rows
…s behind them with a counted wait; dead jobs publish nothing
… tiles), 128 partitions, plain-word partials with a flag a task, a two-level parallel merge (direct for <= 8 partitions); the LDS as two aliasing views
K2: the attention boundary writes the pooler / router operand as two halves (b1 = bf16(pre_w alpha st), b2 = bf16(pre_w beta v)) and its four norm sums right after the cross-rank poll; the pooler / router CTAs poll the sums, compute each token's (rstd, rstd2) and form x = bf16(rstd2 (b1 + rstd b2)), so x no longer waits for the parts' sum exchange; the stream update follows, flagged for K3's boundary (polled with the routing flags). K1 (SGLANG_AMD_MONO_MUSE_DEFER_K1=1, merged launches): the next K1 forms its input from K2's c1 = bf16(npw alpha2 st) and the unpooled rows of its K part (c2 = bf16(npw beta2 bf16(u pfw))), the c1 MFMAs before the rows land, and folds K3's boundary scalars into its partials: no wait on the boundary's x.
…e a CTA In a merged launch at S <= 16, K2's xn task also writes the token's MXFP4 blocks (the same quantization stage_xq_cols applied to the bf16 row: bit-exact), K3 stages every token's blocks into LDS once before the first weights, and the ug items' B operands are read by the lane's token: a tile change is no longer a barrier, x loads and a requantization (~2 us a CTA that crossed a tile). SGLANG_AMD_MONO_MUSE_XQ_PRE=0 restages per tile.
…ueue (SGLANG_AMD_MONO_MUSE_UG_DYN=1)
…dynamic ifs); K2 test tie tolerance for the halves' rounding
…ne XCD: the second reads hit L2)
… step; neutral at 8)
TP experts remove the EP skew: 600 s AgentX (m15-21, cc-traces-weka 256k, simulated accept 4.05) C1 P90 interactivity 328.9 -> 354.5, C2 279.8 -> 342.6 with the 8-ballot plan. The per-row FP8 pooler / unpooler copies cost AIME26 pass@1 (0.908 vs native 0.958 / 0.967; 0.950 without them), so SGLANG_AMD_MONO_MUSE_FP8_DENSE stays 0.
…G_AMD_MONO_MUSE_PSHARD, TP 8) The replicated bf16 pooler (67 MB a layer a rank, 262 KB a CTA) gated the x arrival: the pooler CTAs issue its weights before polling the boundary's norm sums, and the poll returns behind them (CU loads in order): x landed at ~33-35 us, ~10 us after the boundary published it. As Kimi-K3's latent projection, each rank now computes its 512 rows (8.4 MB, one 16-row tile a CTA a K part) and pushes the K parts' partials to every rank's poolx; xn sums the same 8 parts in the same order from there (the same math).
…'s pooler job The chain CTAs ran the whole boundary (the norm sums' exchange and the stream's update included) before their pooler tile, so their tiles came last (passes 34 us at the last CTA vs ~30 for the stream CTAs). The part now publishes its x halves and sums, the CTA runs its pooler job, and stage_a_part_res (the partials re-polled, the exchange, the update, the flag K3's boundary polls) follows; nothing reads the stream before K3.
With the boundary split, a 64-row merged launch with the sharded pooler hangs at TP 8 (every rank's guard fires during graph capture); each change alone launches at widths 8..64. Wider launches keep the replicated pooler until it is understood; the served widths (8, 16) keep the shard.
A diagnostic: the wait that passes its deadline stops at page (site + 1) instead of the null page, so the runtime's fault message names it, and the trace prints each site's callers. Off by default (the stop stays at 0x0).
DFlash2's fc (five target layers' hidden states to the hidden size, 671 MB bf16, replicated: one 100 us GEMM a step on every rank) runs as each rank's 1/8 of the output rows, a view of the loaded weight, and the group's all-gather, then hidden_norm. Step A/B (95k, A B A B): C1 7.76/7.75 -> 7.69/7.68 ms, C2 10.14/10.13 -> 10.05/10.06; greedy outputs identical, accept length 3.739 vs 3.748. SGLANG_AMD_MONO_DRAFT_FC_SHARD=0 keeps the replicated projection.
…eir first two outputs SGLANG_AMD_MONO_MUSE_QUANT_BF16 now defaults to 0: the Quark MXFP4 and NVFP4 checkpoints both keep the bf16 experts of layers 0 and 61 on purpose, so those layers run SGLang's decoder layer instead of an online MXFP4 copy (about 5% slower at C1, 600 s: P90 530.1 -> 506.0, TPOT 1.86 -> 1.93 ms). The runner takes the first two outputs of a module layer, since newer SGLang Muse layers also return an FP8 copy of the next attention input (the module path failed to unpack three values).
Muse Spark's spec, draft spec and unit test import SGLang's monokernel interface from its new location (sglang.kernels.plugin_kernels.amd.spec, the binding in sglang.srt.kernel_plugins.amd) first; the specs register by (architecture, None); the runner declares its get_parallel read. Drops the WIP guard site-fault diagnostic.
# Conflicts: # sgl_amd_mono/__init__.py
…ell) and its CODEOWNERS line
zijiecode
marked this pull request as ready for review
November 1, 2026 06:02
zijiecode
requested review from
HaiShaw,
alexnails and
kevin-mii
as code owners
November 1, 2026 06:02
zijiecode
changed the title
HaiShaw
approved these changes
Nov 1, 2026
zijiecode
changed the title
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.
Motivation
Muse Spark Open (816B-A42B, Quark MXFP4 experts) with DFlash2 on 8x MI355X spends most of a low-concurrency verify step between small kernels and waiting on cross-rank all-reduces: a C1/C2 verify step carries only 8 or 16 rows. This PR adds whole-layer monokernels for Muse Spark at TEP8 (TP8 attention, EP8 experts), for the target model and its DFlash2 draft.
Modifications
sgl_amd_mono/muse_spark/: whole-layer monokernels for Muse Spark at TEP8, the target and the DFlash2 draft, on the 8- and 16-row verify steps (C1 / C2).test/: unit and kernel tests..github/CODEOWNERS:**/*muse_spark* @zijiecode.Accuracy Tests
Speed Tests and Profiling
InferenceX AgentX, 3600 s, TP8, C1:
Checklist
pre-commit run --all-filespasses.test/cover the change..github/CODEOWNERSthat names its owner.Review and Merge Process
lintandCheck PR titlechecks must pass..github/CODEOWNERSto review.main.