Repository navigation
[AMD] DeepSeek-V4.1: opt-in mono decode launches on gfx950 - #43497
amd-pedghazi wants to merge 14 commits into
Conversation
001865e to
c89b45d
Compare
With SGLANG_ROCM_MONO_DECODE=1, decode steps of 1-48 rows run each layer's attention all-reduce, FFN mHC seam, MoE and MoE all-reduce as one persistent FlyDSL launch. The kernels in dsv41_mono/ are copied from vllm-project/vllm#60397 (from ROCm/ATOM#2479, MIT) and keep their licence notices.
c89b45d to
fcc1a27
Compare
Add an opt-in attention launch (SGLANG_ROCM_MONO_K1=1, with SGLANG_ROCM_MONO_DECODE=1). On decode steps of 1-48 rows of the layers without compressor or indexer, one persistent FlyDSL launch runs the attention mHC seam and the attention front (wqkv, q and kv norms, RoPE, SWA KV insert, wq_b); its KV records are bitwise equal to SGLang's. Add its drift test and race stress test.
…ernel helpers, refuse speculative decoding - Import the mono decode module only on HIP; the weight hooks go through the HIP module. - Use the kernel helpers in common/ instead of local copies, and remove the timing trace code. - Fail at startup when speculative decoding is set; the launches do not run verify steps. - Raise errors instead of asserts on the mono attention paths. - Merge the two stress scripts into stress_dsv41_mono.py (--k1 runs K1 too) and add the speculative-decoding cases to the guard test.
|
A general comment: You have all the monokernel code in models/deepseek_common/amd/dsv41_mono/ - So in a vendor-specific directory under a model-specific directory. I expect we'll be seeing a lot of monokernels for AMD devices, and they will probably be based on the same flysDSL skeleton. Are there reusable parts of the code that could be placed in some shared AMD monokernel library directory? Or is the idea of monokernels that they are self-sustainable when it comes to code? Also, what is the maintainability when it comes to building optimally performning monokernels for gfx942, gfx950 and in the future also for gfx1250? |
Thanks @rkarhila-amd good questions! On the code location: AMD maintainers have been careful with PRs that touch common code, so I kept this self-contained on purpose. Apart from some minor changes in On other GPUs: this PR is gfx950 only and stops at startup with a clear error elsewhere, so gfx942 users see no change. gfx942 is possible but needs a port, and vLLM has a draft for it (vllm-project/vllm#60731). With 64 KB of LDS instead of 160 KB, it takes up to 42 rows instead of 48, runs at TP4 only and adds a few kernels. We can bring it over as a follow-up PR. For gfx1250, I'd wait: these kernels are built for fixed shapes and tuned per chip, so a small separate PR once the hardware and FlyDSL support are ready is cleaner. |
…el spans DeepSeek-V4.1-Flash's monokernels (sgl-project#43497 with K2 and the DSpark capture fix) moved onto the framework (RFC sgl-project#43949), without their env vars: - kernels/ops/mono/deepseek_v41_gfx950/: K1, K2 and the FFN launch, built on common_hip's guarded waits; FLYDSL KernelSpecs in the mono group. K1 takes the peer addresses too, so a timeout in any launch flags every rank; the error words the host reads are those flags, one copy a step. The runner checks every tensor, cache and key table it is given; weights.py holds the layout conversions (E8M0 codes, exact MXFP8, the RoPE table). - srt/layers/mono/models/deepseek_v41_hip.py: the spec and its glue. Standard layers take k1 (LAYER_IN -> CORE_IN) and k2 (CORE_IN -> LAYER_OUT); the compressor and indexer layers take ffn (ATTN_PARTIAL -> LAYER_OUT). LAYER_IN / LAYER_OUT carry the HIP loop's pending mHC post. Per-output TOLERANCE bounds, read by the kernel checks. The load hook keeps the FP8 shared expert, which --enforce-shared-experts-fusion converts to MXFP4; the other weights are read in place, and a layer whose FP8 linear method reorders an attention weight in place is refused. K1 binds its layer's caches at its first launch. - deepseek_v4.py / deepseek_v4_hip.py / deepseek_v4_fused_mhc.py: the seams (the load hook, attach, route once a forward, a kernel at each span start). In covered steps the DSpark capture reads the residual the fused boundary rebuilds, so the layers before DSpark's target layers keep their post pending for K1. With --enable-monokernel off nothing changes. - Tests: the spec on CPU; the K1, K2 and FFN checks under torchrun at TP2 and TP4; a TP4 DSpark server with GSM8K across the largest width, and a fault armed on one rank mid-request (every rank stops, no token is sent, the only faults are the deliberate ones at 0x0); the chained stress stays manual. MI350X TP4 DSpark, simulated acceptance 3.51, per-user decode tok/s at 1 / 4 / 8 concurrent: 447.3 / 338.2 / 274.0 without --enable-monokernel, 568.3 / 377.2 / 293.1 with it (+27% / +12% / +7%). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…el spans DeepSeek-V4.1-Flash's monokernels (sgl-project#43497 with K2 and the DSpark capture fix) moved onto the monokernel plugin (RFC sgl-project#43949), without their env vars: - kernels/ops/mono/deepseek_v41_gfx950/: K1, K2 and the FFN launch, built on common_hip's guarded waits; FLYDSL KernelSpecs in the mono group. K1 takes the peer addresses too, so a timeout in any launch flags every rank. The runner checks every tensor, cache and key table it is given; weights.py holds the layout conversions (E8M0 codes, exact MXFP8, the RoPE table). - deepseek_v41_gfx950/spec.py: the spec and its SGLang glue. Standard layers run whole (LAYER_IN -> LAYER_OUT) in two launches, K1 to the attention core and K2 from it; the compressor and indexer layers take ffn (ATTN_PARTIAL -> LAYER_OUT). Each span returns what the HIP fused mHC boundary returns, with the layer's post pending for the next layer. Per-launch TOLERANCE bounds, read by the kernel checks. The load hook keeps the FP8 shared expert, which --enforce-shared-experts-fusion converts to MXFP4; the other weights are read in place, and a layer whose FP8 linear method reorders an attention weight in place is refused. K1 binds its layer's caches at its first launch. - deepseek_v4.py / deepseek_v4_hip.py / deepseek_v4_fused_mhc.py: the seams (the load hook, attach, route once a forward, a span at each start point). Under a plan the covered layers take DSpark's captures (the residual each target layer's boundary rebuilds), so the layers before DSpark's target layers keep their post pending for K1. With --enable-monokernel off nothing changes. - Tests: the spec on CPU; the K1, K2 and FFN checks under torchrun at TP2 and TP4; a TP4 DSpark server with GSM8K across the largest width, and a fault armed on one rank mid-request (every rank stops, no token is sent, the only faults are the deliberate ones at 0x0); the chained stress stays manual. MI350X TP4 DSpark, simulated acceptance 3.51, per-user decode tok/s at 1 / 4 / 8 concurrent: 447.3 / 338.2 / 274.0 without --enable-monokernel, 568.3 / 377.2 / 293.1 with it (+27% / +12% / +7%). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…l launches DeepSeek-V4.1-Flash's monokernels (sgl-project#43497 with K2 and the DSpark capture fix) moved onto the monokernel plugin (RFC sgl-project#43949), without their env vars, and extended to all 40 backbone layers (vLLM #61082's layout), with ATOM sgl-project#2545's hand-off orderings: - kernels/ops/mono/deepseek_v41_hip/: K1 (the attention seam and front: the pending post, the mixes, wqkv_a, the q / kv norms, RoPE, the SWA KV store, wq_b) and K2 (the attention core, wo_a, wo_b, the FFN seam, the MoE and both all-reduces), on common_hip's guarded waits; the spec lists gfx950 as its only ISA. K1 takes the peer addresses too, so a timeout in any launch flags every rank. The runner (DSV41MonoRunner) checks every tensor, cache and key table it is given; weights.py holds the layout conversions (E8M0 codes, exact MXFP8, the RoPE tables). - Every backbone layer runs whole (LAYER_IN -> LAYER_OUT): - K1 folds the previous layer's pending post into the residual. Where none is owed (layer 0's broadcast embedding; the Engram layers 1 and 14, whose Engram write SGLang's loop runs on the settled residual) it takes the residual as it is, collapsed with the previous pre (copy 0 on layer 0). - The index sources (2, 8, 14, 20, 24, 28, 32, 36): K1 also writes the normed input and the bf16 qr, SGLang's forward_low_ratio_sources runs the layer's compressor and indexer on them between the launches, as SGLang runs them before the attention, and K2 builds the key table from the top-k they wrote (each split task waits on its token's flag). - The window-only layers 0 and 1 read no compressed cache. - deepseek_v41_hip/spec.py: the spec and its SGLang glue. Each span returns what the HIP fused mHC boundary returns, with the layer's post pending for the next layer. Per-launch TOLERANCE bounds, read by the kernel checks. The load hook keeps the FP8 shared expert, which --enforce-shared-experts-fusion converts to MXFP4; the other weights are read in place. Any --fp8-gemm-backend but aiter is refused at config time (the others reorder wqkv_a, wq_b and wo_b in place; K1 and K2 read them row-major). The kernels never quantize more than the normal path: K2 keeps the attention output in bf16 for wo_a (its FP8 weight widened to bf16 in registers), as SGLang's bf16 wo_a reads it. Each rank builds every width, starting at its own share (FlyDSL's disk cache hands it the rest). - From ATOM sgl-project#2545, bit for bit: the FFN seam's loads go out before the wait on the ranks' attention partials, a norm polls its rows once, after loading its weights, and the seams' Sinkhorn gates (read only by the next launch) run after K1's front and after K2's shared expert. - deepseek_v4.py / deepseek_v4_hip.py / deepseek_v4_fused_mhc.py: the seams (the load hook, attach, route once a forward, a span at each start point). Under a plan the covered layers take DSpark's captures (the residual each target layer's boundary rebuilds), so the layers before DSpark's target layers keep their post pending for K1. With --enable-monokernel off nothing changes. - Tests: the spec on CPU (test/registered/unit/mono); K1 (each input form: a pending post, a settled residual, layer 0, an index source's outputs), K2 (an index source's and a window-only layer's bit for bit against the K2 of the same keys) and the FFN launch under torchrun at TP2 and TP4, with graph replay across widths (test/registered/kernels/ops/mono, helpers in sglang.test.kernels.mono); serving at TP4 (test/registered/e2e/mono, also a nightly job): DSpark verify and plain decode steps past the largest width and back (GSM8K, and the covered steps' tokens scored again by the normal prefill path, KL), and a fault armed on one rank mid-request: every rank stops, no token is sent, the only faults are the deliberate ones at 0x0. The chained stress stays manual. - test/manual/kernels/mono/validate.py runs the CPU tests, the kernel checks, the serving tests and a paired decode A/B at 1, 4, 8 and 16 concurrent requests against baselines.json, and prints PASS or FAIL for each step. .agents/skills/monokernel gives agents the layout, the seams, the invariants and how to repair, tune or add a model. MI350X TP4 DSpark, simulated acceptance 3.51, per-user decode tok/s at 1 / 4 / 8 / 16 concurrent, both arms in one validate.py run: 453.5 / 328.9 / 283.7 / 219.7 without --enable-monokernel, 579.0 / 399.8 / 303.9 / 221.0 with it (+28% / +22% / +7% / +1%). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…l launches DeepSeek-V4.1-Flash's monokernels (sgl-project#43497 with K2 and the DSpark capture fix) moved onto the monokernel plugin (RFC sgl-project#43949), without their env vars, and extended to all 40 backbone layers (vLLM #61082's layout), with ATOM sgl-project#2545's hand-off orderings: - kernels/ops/mono/deepseek_v41_hip/: K1 (the attention seam and front: the pending post, the mixes, wqkv_a, the q / kv norms, RoPE, the SWA KV store, wq_b) and K2 (the attention core, wo_a, wo_b, the FFN seam, the MoE and both all-reduces), on common_hip's guarded waits; the spec lists gfx950 as its only ISA. K1 takes the peer addresses too, so a timeout in any launch flags every rank. The runner (DSV41MonoRunner) checks every tensor, cache and key table it is given; weights.py holds the layout conversions (E8M0 codes, exact MXFP8, the RoPE tables). - Every backbone layer runs whole (LAYER_IN -> LAYER_OUT): - K1 folds the previous layer's pending post into the residual. Where none is owed (layer 0's broadcast embedding; the Engram layers 1 and 14, whose Engram write SGLang's loop runs on the settled residual) it takes the residual as it is, collapsed with the previous pre (copy 0 on layer 0). - The index sources (2, 8, 14, 20, 24, 28, 32, 36): K1 also writes the normed input and the bf16 qr, SGLang's forward_low_ratio_sources runs the layer's compressor and indexer on them between the launches, as SGLang runs them before the attention, and K2 builds the key table from the top-k they wrote (each split task waits on its token's flag). - The window-only layers 0 and 1 read no compressed cache. - deepseek_v41_hip/spec.py: the spec and its SGLang glue. Each span returns what the HIP fused mHC boundary returns, with the layer's post pending for the next layer. Per-launch TOLERANCE bounds, read by the kernel checks. The load hook keeps the FP8 shared expert, which --enforce-shared-experts-fusion converts to MXFP4; the other weights are read in place. Any --fp8-gemm-backend but aiter is refused at config time (the others reorder wqkv_a, wq_b and wo_b in place; K1 and K2 read them row-major). The kernels never quantize more than the normal path: K2 keeps the attention output in bf16 for wo_a (its FP8 weight widened to bf16 in registers), as SGLang's bf16 wo_a reads it. Each rank builds every width, starting at its own share (FlyDSL's disk cache hands it the rest). - From ATOM sgl-project#2545, bit for bit: the FFN seam's loads go out before the wait on the ranks' attention partials, a norm polls its rows once, after loading its weights, and the seams' Sinkhorn gates (read only by the next launch) run after K1's front and after K2's shared expert. - deepseek_v4.py / deepseek_v4_hip.py / deepseek_v4_fused_mhc.py: the seams (the load hook, attach, route once a forward, a span at each start point). Under a plan the covered layers take DSpark's captures (the residual each target layer's boundary rebuilds), so the layers before DSpark's target layers keep their post pending for K1. With --enable-monokernel off nothing changes. - Tests: the spec on CPU (test/registered/unit/mono); K1 (each input form: a pending post, a settled residual, layer 0, an index source's outputs), K2 (an index source's and a window-only layer's bit for bit against the K2 of the same keys) and the FFN launch under torchrun at TP2 and TP4, with graph replay across widths (test/registered/kernels/ops/mono, helpers in sglang.test.kernels.mono); serving at TP4 (test/registered/e2e/mono, also a nightly job): DSpark verify and plain decode steps past the largest width and back (GSM8K, and the covered steps' tokens scored again by the normal prefill path, KL), and a fault armed on one rank mid-request: every rank stops, no token is sent, the only faults are the deliberate ones at 0x0. The chained stress stays manual. - test/manual/kernels/mono/validate.py runs the CPU tests, the kernel checks, the serving tests and a paired decode A/B at 1, 4, 8 and 16 concurrent requests against baselines.json, and prints PASS or FAIL for each step. .agents/skills/monokernel gives agents the layout, the seams, the invariants and how to repair, tune or add a model. MI350X TP4 DSpark, simulated acceptance 3.51, per-user decode tok/s at 1 / 4 / 8 / 16 concurrent, both arms in one validate.py run: 453.5 / 328.9 / 283.7 / 219.7 without --enable-monokernel, 579.0 / 399.8 / 303.9 / 221.0 with it (+28% / +22% / +7% / +1%). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…d routing (RFC sgl-project#43949) Adds python/sglang/experimental/monokernels/ without any kernels: --enable-monokernel (off by default), the handle_monokernel config hook (every refusal reason at once), the load_tensor / loaded / attach / route / kernel hooks, step-type plans compared across ranks, attach stages that every rank agrees on before any raises, widths limited to what a graph capture can reach, and the SGLANG_TEST_MONO_SPANS / SGLANG_TEST_MONO_FAULT_INJECT test hooks.
…rnels/hip/common Pure rename of dsv41_mono/common (epoch, mx, ops, peer_memory, plan, sync), adapted from ROCm/ATOM; the license notices move with the files. Imports follow in the next commit.
…gs and fail-stop Every mailbox poll and the epoch hand-off go through one guarded wait. The first poll reads no clock; after that, every 64 polls look at the 100 MHz clock and this rank's stop word. A wait past its deadline (half of --watchdog-timeout, at most 60 s) or one that sees a peer's flag records the reason and the site in this rank's stop block, flags every peer's stop block in peer memory and ends the launch with a store to address 0. Wait sites are named after the mailbox region they poll. A fault-injection build stalls one named site after a given launch. The guard block is found through two unused words of the epoch buffer, so the kernels' arguments do not change.
…experimental/monokernels/hip/deepseek_v41
Pure rename: dsv41_mono/{attention,stages,layer,runner,sources} to hip/deepseek_v41/gfx950/, and dsv41_mono_decode.py to hip/deepseek_v41/spec.py. The next commit rewrites the glue to the plugin API.
…nel plugin The DeepSeek-V4.1 spec: spans k1 (layer_in -> core_in, 30 decode / 27 DSpark verify layers; the three verify layers DSpark captures are logged with the reason) and ffn (attn_partial -> layer_out, 40 layers), with SGLang's attention core between them; refusal reasons at startup and after loading; load_tensor keeps the FP8 shared-expert shards before shared-expert fusion converts them; loaded() keeps the row-major wqkv_a / wq_b; every copy is logged with its size. The runner allocates every buffer at attach (two alternating FFN output sets, the step's positions and slots, scratch per reachable width) and compiles each reachable width without running it. srt: the plan through deepseek_v4_hip.py, start points at layer_in and attn_partial in deepseek_v4_fused_mhc.py, the four calls plus loaded() in deepseek_v4.py; the mono_q argument stays (its default path is unchanged). SGLANG_ROCM_MONO_DECODE and SGLANG_ROCM_MONO_K1 are removed: --enable-monokernel replaces them.
- unit (CPU): refusal with every reason, the deadline rule, span selection, the numerics metrics and keys, routing per step type, the load-hook order, staged rank agreement over two gloo ranks (a shape failure, a build failure, differing plans); DeepSeek-V4.1's plans from the checkpoint layout (decode: k1 on 30 layers, verify: 27, ffn on all 40), config refusals, the hook placement in the model and loader; the SGLANG_READS contract of every plugin file. - kernels (gfx950, TP2 and TP4): the FFN and K1 launches against the normal path at every width with the RFC metrics and the spec's bounds, graph replay bitwise against eager, the allocator flat over a capture and the replays; a stalled wait at every named site stops every rank with the reason and site readable. Moved from test/srt/amd/dsv41_mono; the old guard-flag test is replaced by test_guard.py. - e2e (TP4, DSpark): steps that cross the largest width and come back; a wait stalled mid-request stops the server and the request gets no answer. - manual: validate.py runs the steps in order and prints PASS or FAIL for each; the stress script. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Update: SGLang now runs DeepSeek-V4.1 monokernels through the sgl-amd-mono plugin (--monokernel-backend amd, #44353). We compared this PR with the plugin's deepseek_v41, and sent the parts the plugin does not have yet to sgl-amd-mono: sgl-project/sgl-amd-mono#36: the FFN seam builds its residual with SGLang's arithmetic, as the K1 seam already does. The residual now matches SGLang's bit for bit at all 48 widths, at TP2 and TP4, and K1 + K2 time per layer changes by less than 1 %. This draft stays open as the record of our in-tree work. It is not for merge. |
Motivation
DeepSeek-V4.1-Flash decode on MI355X runs many small kernels and two TP all-reduces in every layer. A monokernel runs a large part of a decoder layer, including its all-reduce, in one persistent launch, so a decode step needs far fewer launches.
This PR moves the two DeepSeek-V4.1 monokernel launches (FFN and K1) onto the plugin layout and API of RFC #43949, behind one switch:
--enable-monokernel. The switch replaces the env varsSGLANG_ROCM_MONO_DECODEandSGLANG_ROCM_MONO_K1. The PR is based on maind4daf52369and does not depend on #44405-#44408. #44033 (K1 on all layers, and K2) stacks on this PR and is moving to the same plugin.Modifications
python/sglang/experimental/monokernels/(35 files, RFC layout), plus theexperimentalpackage marker:guards.py,runtime.py,spec.py: startup refusal, load hooks, staged attach, routing, spans, numerics and the spec registry.hip/common/: the shared HIP runtime. Every wait between workgroups or GPUs goes through one guarded wait.hip/deepseek_v41/: the DeepSeek-V4.1 spec and its gfx950 kernels, moved withgit mvfromsrt/models/deepseek_common/amd/dsv41_mono/.k1(layer_in->core_in): the pending mHC post, the mHC pre, the norm, and the attention front (wqkv_a, q and kv norms, wq_b, RoPE, SWA KV write). It runs on 30 layers in decode: the layers without a compressor or indexer, except layer 0 and the Engram layers. DSpark captures the input of three of them, so a verify step runsk1on 27 layers.ffn(attn_partial->layer_out): the attention all-reduce, the FFN seam, the MoE with the shared expert, and the MoE all-reduce. All 40 layers.--watchdog-timeout, at most 60 s; a deadline below 5 s is refused. A timeout stops every rank, and no token of the failed step is sent.srtfiles, +75/-4: the flag, the config hook, three test env vars, and the load, attach, route and kernel calls in the DeepSeek-V4 model. Plus one CODEOWNERS line. With the switch off, SGLang runs its normal path.test/registered/monokernels/unit), kernel and fail-stop tests (kernels, nightly MI35x), serving tests (e2e, nightly MI35x), andtest/manual/monokernels/validate.py.--enable-monokernelto the DeepSeek-V4.1 MI355X recipe (cookbook section 3.6). Details:python/sglang/experimental/monokernels/README.md.hip/SOURCE.mdlists the code adapted from ROCm/ATOM (MIT) and vLLM (Apache-2.0), with their licenses. No code is taken from [Feat] [AMD] monokernel: Framework and shared HIP runtime, RFC #43949 (2/4) #44406 or [Feat] [AMD] dsv41: DeepSeek-V4.1 decode layers on monokernels (3/4) #44407.32e567e667.863fc1caffand0b265fdac8are pure renames. The earlier commits are the first version of this PR and merges of main.Accuracy Tests
MI355X, TP4, image
rocm/sgl-dev:v0.5.21-rocm720-mi35x-20261009, the cookbook section 3.6 recipe (DSpark block 5, Engram host tableper_rank,--enforce-shared-experts-fusion,--enable-decoder-swa-bounded-replay), run withvalidate.py. Requests are sent 8 at a time, so every verify step runs on the monokernels.End to end, switch on against switch off:
GPU kernel and fail-stop tests:
python3 test/manual/monokernels/validate.py --model <path> --tp 4runs every check above;--stepspicks a subset.ffn.residualbound (2^-7, now 2^-6).ffnoutput bound against SGLang's normal path (RFC section 10). It is still set against the test's reference.Speed Tests and Profiling
AgentX agentic benchmark, MI355X, TP4, DSpark block 5, 60 minutes per point, on the 393-entry
semianalysisai/cc-traces-weka-062126trace set. Every setup uses imagerocm/sgl-dev:v0.5.21-rocm720-mi35x-20261009, AITERperf/v41-indexer-bf16-logitsatcc177350,--fp8-gemm-backend aiterand--disable-shared-experts-fusion. All points were measured on the same node.Each cell: total throughput per GPU (tokens/s), then 1 ÷ ITL p90 (tokens/s per user).
d4daf52369)e4f3665be4)Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ci