Repository navigation
[Perf] [MI355X] deepseek_v41: DSpark step glue in fewer launches, Markov draft sampling, faster K1 and K2 - #95
Conversation
dfaa33e to
a4a5f13
Compare
|
Hi @limou102 Thanks! Awesome work really great improvements. Two small points before merge: 1. Simulated acceptance takes a path real serving never takes
So the simulated runs in the PR body (AL 3.51, temperature 1) measure work that real serving still does. At TP4 the accept piece saves 0.25 ms/round under simulation (4.505 -> 4.252 ms), and about half of that comes from the fold and the skipped argmax. Please keep the argmax under simulation and make no path depend on 2. Please add a precision point and a speed point (main vs PR, simulation off)
Other points
FYI, with #82 merged on top, every TP8 kernel check passes. TP8 greedy gives 4.027 ms/round at AL 4.80, 7% faster than TP4. |
…kov draft sampling, faster K1 and K2 - K1 / K2: earlier stage hand-offs; a layer's owed mHC post folded into the launch that reads it next (the next layer's K1, an Engram gate, the head's combine); DSpark's captures written into one packed tensor; a partial MXFP8 round stored to LDS without a branch; K1's wq_b weights issued off the wqkv path, and up to 7 rows a rank each wqkv row tile's SGLang-order chains on two CTAs (bitwise as before). - glue/: a DSpark verify step's and draft block's attention metadata in one launch each; the static verify's accept in SGLang's one-launch tail without the accept-site TP broadcasts, and under simulated acceptance a temperature-only sampled step accepted and committed in the graph; the verify's draft-KV inject in one HIP launch; the draft head in FP8; the prefill indexer kernels compiled once. Each piece keeps SGLang's path unless its first call agrees with SGLang; SGLANG_AMD_MONO_DSV41_GLUE=0 turns them all off. - A sampled DSpark draft's Markov steps in two launches each (HIP extension). - Fused posts decided at startup on random rows, not on a capture's warmup streams. - Tests: check_glue.py, check_markov.py and time_layer.py; the kernel test runs them. Parts are adapted from AKLR (https://github.com/AMD-AGI/AKLR, MIT), as each adapted file's header says. InferenceX AgentX (dsv41flash-fp4-mi355x-sglang-agentic-dspark, golden AL 3.51, 3600 s, MI355X, lmsysorg/sglang-rocm:v0.5.22-rocm10-mi35x-20261029), main -> this change: TP4 C1 p90 interactivity 675.8 -> 805.4 tok/s/user, throughput 7700 -> 8076 tok/s/chip; TP2 C4 339.5 -> 474.9, 18022 -> 19174; on the monokernel points p90 +14% to +58%, throughput +3% to +21%. Real acceptance (simulation off): gsm8k, humaneval and needles match main. Signed-off-by: limou <limou102@amd.com>
- accept: drop the in-graph fold of temperature-only sampled steps and the argmax skip, which ran only under SGLANG_SIMULATE_ACC_LEN; no path depends on it now. An explicit SGLANG_SPEC_TP_SYNC keeps SGLang's accept-site broadcast. - agreed(): one verdict on every rank of the DP rank's TP group (all-reduce MIN). - FP8 draft head off by default (SGLANG_AMD_MONO_DSV41_DRAFT_HEAD_FP8=1 turns it on). - K1: the qkv CTAs past the index head weights' (IW0), not on them. - inject.hip: rows whose position is past the RoPE table are not stored; the SWA slots must fit int32, as SGLang's committed_swa_locations makes them. - e2e: the verify server asserts the glue pieces run (real acceptance). - README: what hip_ext.py needs at runtime. Signed-off-by: limou <limou102@amd.com>
a4a5f13 to
881ec65
Compare
|
Thanks for the review. Addressed in a second commit ( 1. Simulated acceptance
2. Precision and speed, simulation offTP4, the AgentX TP4 C1 recipe flags (
Other points
Tests on
|
Motivation
DeepSeek-V4.1-Flash with DSpark on MI355X spent much of each decode round outside the monokernels: ~250 launches a round outside them, most of them ~4 µs SGLang glue kernels (attention metadata, accept, draft-KV inject, draft head); an eager sampled accept at temperature 1 (the AgentX default); and K1 waiting on its own weight loads. This PR moves that work into fewer launches and shortens K1, keeping SGLang's results bitwise where the kernels did before.
Modifications
wq_bweights off thewqkvpath. Up to 7 rows a rank, eachwqkvrow tile's SGLang-order chains run on two CTAs; results stay bitwise identical.deepseek_v41/glue/(new)SGLANG_AMD_MONO_DSV41_GLUE=0turns them all off; per-piece env vars are listed in the README.common/hip_ext.py).fused.py: fused mHC post verdicts are made at startup on random rows, not on a graph capture's warmup buffers.common/sync.py: addsMailbox.poll_issue/poll_wait, which issue a poll's first loads ahead of later loads.check_glue.py,check_markov.pyandtime_layer.py;test_deepseek_v41_kernels.pyruns the first two.Parts are adapted from AMD-AGI/AKLR, as each adapted file's header says.
Tested with
lmsysorg/sglang-rocm:v0.5.22-rocm10-mi35x-20261029(SGLang686cf984) on MI355X. Rebased onmain@8039635; the_engramconflict with #65 keeps this PR's structure, which already passesinput_idsand the image ids, and takes #65'sEngram.forward(input_ids)contract pin.Speed Tests and Profiling
InferenceX AgentX:
dsv41flash-fp4-mi355x-sglang-agentic-dspark, recipedsv41flash/sglang/mi355x-fp4-mtp/agentic.yaml.main=6777a89; this PR = the same changes on thatmain.Each point is one (TP, DPA, CONC) scenario, and each line connects one run's points; up and to the right is better.
Fixed-length decode, single user: the AgentX TP4 C1 server, random prompts, 2048 output tokens, simulated AL 3.51 (measured 3.50–3.53 tokens a round on both). Values are medians over the requests; the two runs are on different nodes.
Kernel-only layer time:
test/kernels/deepseek_v41/time_layer.py --widths 6,12, cold weights, median µs per layer (K1 + K2), main → PR.Accuracy Tests
Repo test suite, full, MI355X ×8, main vs this PR (before the rebase:
main@a0c3997)pytest test/unitpytest test/kernelspytest test/e2e -k deepseek_v41(TP4: GSM8K, KL vs prefill, long prompts, passkey, fault stop)mainand is outsidedeepseek_v41: theminimax/minimax_m31/glm53_flashspec and contract tests, andtest_minimax_m31_ffn::test_o_proj_is_sglangs_to_the_bit.main's extra kernel failure (TestMonoGlm53Kernels::test_indexer_selects_the_exact_top_k, retry limit exceeded) passes on this PR.check_k1passes at TP2, TP4 and preshuffled: the seam outputs and KV records are bitwise SGLang's, and q is within the spec's bound.8039635:pytest test/unit: 538 passed / 8 failed vsmain's 524 / 8, with the same 8 failures.Real DSpark acceptance (simulation off), AgentX server flags, main vs this PR
one launchon every rank.The
_engrammerge conflict: this PR at TP4 with the monokernel and DSpark, Engram host table in both layouts.shared: SGLang's lookup,engram(..., input_ids=...).per_rank: staged rows through the fused gate.Checklist
pre-commit run --all-filespasses.test/cover the change (check_glue.py,check_markov.py,time_layer.py)..github/CODEOWNERSthat names its owner. (n/a: no new model)