Repository navigation
[DSV4.1][gfx1250] Single-launch decode step (FlyDSL) - #3109
Conversation
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
valar-cop
left a comment
There was a problem hiding this comment.
some changes needed..
acceptance for the long-context single-launch step (agentic first turn, ~26.6k tokens)?
ci Check Code Style with Ruff failed — looks like the pr itself (high): reviewdog reports 5 ruff findings in the PR's new files (C408 in kernels/layout.py:57; RUF100 in the three plan tests; RUF059 in test_deepseek_v41_gfx1250_index_plan.py). The job log captured here is empty, but the inline comments are from this run.
⚠️ [ruff] Unnecessarydict()call (rewrite as a literal)
⚠️ [ruff] Unusednoqadirective (non-enabled:E402)
next: Apply the bot's suggestions (dict literal, drop the unused noqa,_k) and rerun. https://github.com/ROCm/ATOM/actions/runs/42573852051/job/129017056864
ci Check Code Style with Black failed — not sure yet (low): The captured log is 0 bytes, so the failing file/line cannot be determined; the sibling Ruff job failing suggests formatting, but that is unverified.
next: Fetch the job log (or run black --check on the changed files) and rerun. https://github.com/ROCm/ATOM/actions/runs/42573852051/job/129017056992
ci Run non-GPU unit tests failed — not sure yet (low): The captured log is 0 bytes; cannot tell whether a new gfx1250 plan test failed or an existing test broke. Note test_deepseek_v41_gfx1250_single.py skips under this job (no aiter) while the three plan tests do run.
next: Fetch the job log / JUnit artifact, or reproduce with .github/scripts/run_unit_tests.sh. https://github.com/ROCm/ATOM/actions/runs/42573852051/job/129017056997
valar-cop
left a comment
There was a problem hiding this comment.
some changes needed..
ci Check Code Style with Ruff failed — not sure yet (low): The captured log is 0 bytes, so the failing file/rule cannot be determined. The findings from the previous head (C408 in layout.py, RUF100/RUF059 in the plan tests) are fixed in this revision, so this run fails for a cause not visible in the artifacts.
next: Fetch the job log or the reviewdog annotations for this run and fix the reported lines before merge. https://github.com/ROCm/ATOM/actions/runs/42579308448/job/129034011861
ci Run non-GPU unit tests failed — not sure yet (low): The captured log is 0 bytes. The new gfx1250 tests are written to skip their aiter-backed cases in this job and the remaining plan/layout tests look consistent with the kernels, so no PR-caused failure can be identified; this job also failed on another recent run.
next: Fetch the JUnit artifact (unit-report.xml) for this run and rerun to see whether it reproduces. https://github.com/ROCm/ATOM/actions/runs/42579308448/job/129034011870
…erved per run_model, one timeout event From the review of #3109: - A served step must be 6 rows a request: step_override declines a step with fewer (DSpark's confidence schedule or ragged verify), and config_refusal refuses both switches, so such a step never reaches run_step and raises mid-forward. - ModelRunner clears _served at the top of run_model, which every forward path (RapidServeModelRunner's too) runs, not in ModelRunner.forward. - The poll-timeout check reuses one event and one pinned buffer instead of a new event every launch. - The draft temperature scale env read (ATOM_DSPARK_PROB_DRAFT_TSCALE, default 1, not a main switch) is gone: the drafts use the request's temperature. - recipes: the kernel's own copy of the routed experts (~81 GB a rank) and what it costs the KV pool (~275 -> ~199 GB a rank on MI455X TP4). Checked: ruff, black, CPU tests (aiter tests on the GPU host pending). Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
…erved per run_model, one timeout event From the review of #3109: - A served step must be 6 rows a request: step_override declines a step with fewer (DSpark's confidence schedule or ragged verify), and config_refusal refuses both switches, so such a step never reaches run_step and raises mid-forward. - ModelRunner clears _served at the top of run_model, which every forward path (RapidServeModelRunner's too) runs, not in ModelRunner.forward. - The poll-timeout check reuses one event and one pinned buffer instead of a new event every launch. - The draft temperature scale env read (ATOM_DSPARK_PROB_DRAFT_TSCALE, default 1, not a main switch) is gone: the drafts use the request's temperature. - recipes: the kernel's own copy of the routed experts (~81 GB a rank) and what it costs the KV pool (~275 -> ~199 GB a rank on MI455X TP4). Checked: ruff, black, CPU tests (aiter tests on the GPU host pending).
2613756 to
930889a
Compare
|
Thanks for the review, replies are inline. Long-context acceptance: in the description: 8 concurrent Agentic first turns (~26.6K tokens), TPOT 1.974 / 2.001 ms vs native 2.179 ms, acceptance 0.507 / 0.519 vs 0.519. CI: Ruff and Black have passed on e2c7bd0; non-GPU unit tests are still running (CI run). The earlier Ruff failure was in "Install reviewdog"; the non-GPU unit tests failed at collection because a FlyDSL stand-in was left in sys.modules, fixed in 57c6e01. |
…e refactor) The gfx1250 DeepSeek-V4.1 single-launch decode step (TP4, DSpark5, BF16 KV, FP8 index) as validated before the refactor, on origin/main: the gfx1250_single package (host runner, pointer tables, FlyDSL kernels as synced from the development tree) and its hooks (mono dispatch, the routed experts' pre-shuffle capture, fine-grained peer buffers, ModelRunner's served-step path, ATOM_MONO_SINGLE). Includes the three fixes of 2026-11-09: the MoE route weights' sum floor (as aiter topk_gating), the draft attention's last rows (ceil clusters) and the PV reads past a row's keys (zeros). Not for upstream as is: the refactor follows on this branch.
…ug.py (H2) The step timer, the shadow comparators, the checkpoint row check and the per-step dump move to debug.py (RunnerDebug, StepTimer), all behind their ATOM_MONO_SINGLE_* switches. Removed: the bring-up only paths (pointer log, address map, input checksums and syncs, one-hot verify, KV scan, plane check) and the switches for paths that are always on (GRAPH, SPLIT, SERVE, PRECOMPILE, FASTLAUNCH). runner.py 1277 -> 565 lines; kernels unchanged. Checked: CT16 serve, c8 / c1 perf (T=1, 3 rounds) on this branch: served on every step, 48-row kernel 7.48 ms, c8 TPOT 2.25-3.06 ms, c1 2.15-2.44 ms.
…py (H2)
ATOM_MONO_SINGLE_{SHADOW,SHADOW_NL,SHADOW_AFTER,SHADOW_NOENGRAM,TIMING,
TIMING_EVERY,CHECK_ROWS,PROWDUMP,PROWDUMP_N,PROWDUMP_QLOG} are read through
atom.utils.envs instead of os.environ.
Checked: CT16 serve smoke (c8 / c1, timing and dump on).
…kage (H1)
install.py splits into config.py (the deployment constants and
config_refusal, formerly single_refusal) and dispatch.py (SingleStepModel,
install_single_step); the shadow forward moves to debug.ShadowModel and
SingleStepModel.forward only routes to it. planes / layer_weights /
draft_weights move to weights/{planes,layer,draft}.py. The runner and
debug take NL / NDR / ROWS_PER_REQ / TP from config. Imports sorted.
Checked: CT16 serve smoke (c8 / c1) on this tree: every width served.
… not runner.split (H3) A served step (``_single_served``) is always the two-launch step, so the id broadcast skip and the output D2H event no longer read the runner's split flag; the runner's verify_ready is marked as its own event. Checked: CT16 serve smoke (c8 / c1).
hx -> device, attn_stages -> attention, ffn_stages -> front, moe_stages -> moe, ctx_stages -> draft_ctx, head_fly -> head, sample4r_fly -> sample_common, sample4p_fly -> sample, markov_prob_fly -> markov, step_tail_ct3 -> tail, layer4l_step_fly -> step, idx_stages -> index_query, score_stages -> index_score, topk_stages -> topk, comp_stages -> compressor, engram_stages -> engram; the unused markov_fly, verify_fly, sample_fly and sample4_fly removed. The host's module references follow. Checked: CT11 traced-MLIR hashes of every width (target / draft) equal to the baseline's; CT16 serve smoke (c8 / c1) served on every width.
…s in layout.py (K2, CT11) The device-wide handoff counters (each layer's block and the tail block), the entry counters and the fixed LDS regions the step and the tail used as literals are named in kernels/layout.py and read from there. Checked: CT11 traced-MLIR hashes of every width (target / draft) equal to the baseline's; CT16 serve smoke (c8 / c1) served on every width.
…e shared files Formatting (black) and lint (ruff: import order, dict literals, a file opened in a with block) of the gfx1250_single host modules and of the shared files this branch changes; the kernels are untouched. The noqa markers of the shared runtime's error handling are kept as on main.
… routing (T1) tests/test_deepseek_v41_gfx1250_single.py (31 cases): config_refusal accepts the supported deployment and names why every other is refused; SingleStepModel serves every whole-request width up to 8 requests and nothing else and forwards to the production model; off unless ATOM_MONO_SINGLE; a refused deployment keeps the model; the V4.1 mono dispatch routes to the single launch when enabled. Checked in the gfx1250 image (CT8): 31 pass; the V4.1 mono dispatch, index plan and mono plan / runtime tests pass as on main (two aiter-stub teardown failures of test_mono_runtime happen on main there too).
recipes/DeepSeek-V4.1-Flash.md: the gfx1250 single-launch decode step (what runs in the kernel, the deployment it serves, the steps it leaves to the production path, build, debug switches, MI455X GSM8K / step time / TPOT against the original path). docs/environment_variables.md: ATOM_MONO_SINGLE and its debug switches.
…me hook, notes trimmed (K3, CT11) The stages' phase stamps go through kernels/stamps.py (no code unless a trace sets its hook; the production build leaves it unset); the layer stamps the step timer reads stay. Development notes (machine names, dated measurements) removed from the kernel comments. Checked: CT11 traced-MLIR hashes of every width (target / draft) equal to the baseline's; CT8 serve smoke (c8 / c1) served on every width.
…ts.py (H1) precheck, the host / device input preparation, the sampling parameters and forget_drafts move to inputs.StepInputs (SingleStepRunner inherits it); the unused _step_inputs is removed. runner.py 565 -> 425 lines. Checked: CT8 serve smoke (c8 / c1) served on every width.
…erves (H3) ModelRunner._single_step only asks SingleStepModel.single_step_runner(batch, rows); the conditions (a DSpark verify step, no structured output / logprobs / required channel, a width it serves, request shapes precheck takes, not in shadow mode) move there, with CPU tests. Checked: tests 111 pass (gfx1250 single + V4.1 mono dispatch) in the image; CT8 serve smoke (c8 / c1) served on every width.
…ry front (K4, CT11) The counter waits share device.wait_count (bounded, timeout noted) and the two segment entries' attention fronts share step.entry_front. Checked: CT11 traced-MLIR hashes of every width (target / draft) equal to the baseline's.
…K5, CT11) The router shape (target 384 / 6, draft 128 / 3) and the exchange store go through the front / MoE contexts instead of module globals set while tracing (front.set_router, MS.TOPK, FS.STX). Checked: CT11 traced-MLIR hashes of every width (target / draft) equal to the baseline's.
…nctions, CPU tests (K6, CT11)
kernels/index_plan.py (selection tree levels / lists, score tasks, REINDEX
columns and the key-row clamp) and kernels/moe_plan.py (unit striping,
early shared expert) hold the step's plan arithmetic; the kernel calls them
on traced Int32, tests/test_deepseek_v41_gfx1250_{index_plan,moe_plan,layout}.py
on Python ints (with the REINDEX > 16384-row regression).
Checked: CT11 traced-MLIR hashes of every width (target / draft) equal to
the baseline's; 45 CPU tests.
Each kernel module says which stages it holds, the counters it hands off on (layout.py names) and its LDS regions. Checked: CT11 traced-MLIR hashes of every width (target / draft) equal to the baseline's.
…n the build key, no env reads (K7, CT11) step.build(..., sources, moe_ids): every build parameter reaches FlyDSL's JIT cache key, the runner passing the kernel directory's source digest (atom.mono.plan.build_key.source_digest), so an edit to any kernel module rebuilds; the fixture's fixed MoE routes are moe_ids="mroute" instead of the MROUTE_FIXED environment read. CT11's plan / layout tests formatted (black, ruff). Checked: CT11 traced-MLIR hashes of every width (target / draft) equal to the baseline's; CT8 full retest on this tree (GSM8K 512: 451, native 451; 48-row kernel 7.310 ms vs 7.393 pre-refactor) and serve smoke (c8 / c1).
…p override; drafts handled in mono; tests in one file ModelRunner (+40/-3, no getattr) knows only a generic base, atom.mono.runtime.step_override.StepOverrideModel, checked once at init: a model that runs a whole decode step itself (sampling, verification and the next drafts included) returns from `step_override(...)` an object with `verify_outputs(bs)` and `output_stream`, or None for the production path. Per step ModelRunner checks `self._served`: - run_model's decode branch: no logits on a served step; - postprocess: the verdict and ids from `verify_outputs`; - the TP id broadcast is skipped (the ids are identical on every rank, and the broadcast would wait for the draft launch queued after them); - forward_done_event is recorded on the output stream, so the output D2H waits for the verdict copies, not the draft launch. Everything else moved into gfx1250_single/dispatch.py: the step's batch is taken in single_step_runner; the drafter's compute_draft_kv (a no-op on a served step) and propose (the kernel's drafts on a served step; otherwise the production drafter's, and the kernel's kept q are dropped) are wrapped at install. tests: the plan / layout tests merged into test_deepseek_v41_gfx1250_single.py (CT11; integer backend inline, plan tests run without aiter); _gfx1250_plan_backend.py and the three plan / layout test files removed; two tests of the drafter wrap added.
…fined once in kernels/dims.py (K8, CT11) kernels/dims.py holds the model's sizes (hidden, q / kv latent, heads, experts, vocab, layers, mHC, indexer, draft), the per-rank sizes they give over TP4 and the deployment's capacities (8 requests, 6 rows each, 48 rows; 256 CTAs of 256 threads); each kernel module binds its local names from it instead of repeating the literals. layout.layer_slots() written as a dict literal (ruff C408). Checked: CT11 traced-MLIR hashes of every width (target / draft) equal to K7's; ruff, black, CPU tests (45 pass without aiter).
…gfx950 mono, no hook in the shared MoE - The package is atom/models/deepseek_v41/mono/gfx1250 (was gfx1250_single), its test tests/test_deepseek_v41_gfx1250.py. - Switches: the single launch takes the gfx950 mono decode's own -- on with ATOM_MONO_ENABLE (default 1) on a gfx1250 device for the deployment it is written for (gfx1250.serves: config_refusal); ATOM_MONO_CHECK runs it beside the production forward, compared (and layer 0's table row against the checkpoint); ATOM_MONO_TIMELINE logs where a step's time goes; ATOM_MONO_TRACE=<dir> dumps the verification per step. The ATOM_MONO_SINGLE* variables and their bring-up sub-switches are gone (docs and recipe updated). - The shared atom/model_ops/moe.py is back to main: the V4.1 MoE packs its routed experts into the step's planes in its own process_weights_after_loading, which the loader runs before the experts' (parent modules first), i.e. before their in-place shuffle -- and only when the step will serve this deployment. - Code comments no longer name machines.
…es (CT11) debug.py 839 -> 270 lines. ATOM_MONO_CHECK (eager): the launch beside the production forward, its head input / logits (relL2, top-1, top-5), per-layer index selection, taps and drafts compared. ATOM_MONO_TIMELINE=<prefix>: the step timer, averaged over 64 steps, logged and appended to <prefix>_r<rank>.jsonl. Removed: the table-row check, the per-step verification dump (ATOM_MONO_TRACE is not used by this step), partial-layer shadow runs, the draft context / hidden comparisons and the bring-up switches. runner.py always builds every layer and draft stage. Checked: ruff, black, CPU tests (45 pass without aiter).
… stand-in behind
Without FlyDSL installed (the non-GPU CI) the plan / layout tests put a
stand-in `flydsl` module, which has no __spec__, into sys.modules on the first
load; the second load's find_spec("flydsl") then raised ValueError at
collection, and later test files that probe for FlyDSL (the Kimi-K3 mono plan
tests) saw the stand-in instead of nothing. The stand-in now lives only while
a kernel file loads (mock.patch.dict).
Checked: the CI's non-GPU suite on Python 3.12 / CPU torch without aiter or
FlyDSL: 11878 passed, 0 failed.
…erved per run_model, one timeout event From the review of #3109: - A served step must be 6 rows a request: step_override declines a step with fewer (DSpark's confidence schedule or ragged verify), and config_refusal refuses both switches, so such a step never reaches run_step and raises mid-forward. - ModelRunner clears _served at the top of run_model, which every forward path (RapidServeModelRunner's too) runs, not in ModelRunner.forward. - The poll-timeout check reuses one event and one pinned buffer instead of a new event every launch. - The draft temperature scale env read (ATOM_DSPARK_PROB_DRAFT_TSCALE, default 1, not a main switch) is gone: the drafts use the request's temperature. - recipes: the kernel's own copy of the routed experts (~81 GB a rank) and what it costs the KV pool (~275 -> ~199 GB a rank on MI455X TP4). Checked: ruff, black, CPU tests (aiter tests on the GPU host pending).
…s run without aiter (K9, CT11) - The MoE always reads the router's ids: step.build has no moe_ids, FPTRS no 'mroute', buffers.py no "mroute" pointer / scratch (a fixture-only buffer nothing filled). The traced IR of the 16 launches differs only in the pointer-table offsets after the removed entry. - The gfx1250 package, config, dispatch and Dims import without aiter, so the deployment / routing tests run in the non-GPU CI too (only the gfx950 mono dispatch test needs aiter). The served-step fixture patches LazyRunner as the model fixture does. Checked: the CI's non-GPU suite (Python 3.12, CPU torch, no aiter / FlyDSL): 11920 passed, 0 failed; this file 87 passed, 1 skipped.
…h's kept draft distributions When the production drafter draws the next drafts, the kernel's kept q of those requests no longer match them. The drafter wrap cleared every request's q instead, so a prefill (a new request) dropped the q of the requests still decoding and their next verification used the wrong rule. It now forgets the requests of the decode batch the kernel declined (step_override records it), and nothing on a prefill; served steps take the kernel's drafts for that batch's requests. Checked: CPU tests (88 passed, 1 skipped without aiter), ruff, black.
… in place, no second copy (K10, CT11) The step kept its own copy of every routed expert (~81 GB a rank, packed at load), which cost the KV pool ~76 GB a rank. The kernel now reads the production MoE's own tensors as aiter's gfx1250 FlyDSL fused_moe keeps them (gate / up rows interleaved, 16 x 16 tile shuffle, n32k4 scales): the layer table holds mw13 / ms13 / mw2 / ms2 instead of mplane. - kernels/moe.py: a unit's A operand and scales are addressed in that layout (lane row r' = 2 r / 2 (r - 8) + 1 for gate / up); a tile's two K steps are issued back to back (2 KB contiguous). kernels/dims.py: MOE_IPAD = 640, the padded intermediate width the strides assume (pad_align 128). - weights/planes.py: routed_experts binds the tensors, refused unless in that layout and shape; the plane packing, the MoE's load hook and the runner's plane lookup are gone (deepseek_v41/moe.py is back to main's). - recipe: no second copy. Checked: on random weights the production layout gives the kernel the same bytes as the old plane for every unit, lane and load (pad 640 and 768); the CI's non-GPU suite: 11921 passed. On MI455X TP4 (CT8): the KV pool back to ~274 GB a rank (198 GB with the copy); GSM8K first 256: 230; CT11 fixture +2.4% a layer. Whole-network A/B in progress.
…ot the padded row count run_step reads input_ids as 6 rows a request; step_override checked the possibly padded running_tokens, so a padded step could still raise mid-forward. It now serves only when input_ids.numel() == running_bs * 6 (review). A test covers padded steps.
… files outside the repository
tables.py, buffers.py and weights/{draft,layer,planes}.py named investigation
scripts (the step fixture, the plane packer) that are not in the repository;
they now say what the layout is.
…s.path changes (K11, CT11) runner.import_kernels put kernels/ first on sys.path and every kernel file inserted its own directory, then imported device / step / tail / attention / dims ... as top-level modules of the server process, where another `import step` or `import device` would get ours, or ours an earlier one. The kernels now import each other by their package path (as the gfx950 kernels do), nothing changes sys.path, and import_kernels imports <package>.kernels.<name>. Checked: traced IR of all 16 launches (widths 6..48, target / draft) identical to before; black / ruff; CI-like tests 89 passed, 1 skipped; with FlyDSL no top-level kernel module names are left in the process.
Since d9fd688 the step override decides on input_ids.numel(); the padded running_tokens ModelRunner passed was unused. StepOverrideModel.step_override is (input_ids, running_bs, batch).
24bc5e2 to
f85569f
Compare
… section - Imports at the top of every module (review rule: no function-local imports): the host code (inputs, runner, weights/planes, __init__, dispatch) and the kernels (device, moe, step, index_score; CT11). The one left is dispatch's runner import inside the LazyRunner factory: the runner needs aiter and FlyDSL, and the routing must import without them (the CPU tests run without aiter). - CLAUDE.md: the gfx1250 V4.1 single-launch decode step in the mono section. Checked: traced IR of all 16 launches identical; ruff, black; CI-like tests 89 passed / 1 skipped; serve smoke on MI455X TP4 (single launch on, c8 / c1).
|
@valar-cop can we merge this pr? |
|
@yhl-amd code side looks fine — all the earlier threads are fixed in the head (dispatch gating on but by the ci summary the e2e jobs are still running: vLLM in progress, sglang and ATOM Test queued, so hold the merge until those pass.. also the description still says head runs are TODO and lists 451, while the recipe has 453 / 455 and tpot 2.01 / 2.06 vs 2.30 / 2.15 — which pair is on this head? please align them. LGTM otherwise. |
Brings in main through 03b584b (ROCm#3109 V4.1 gfx1250 single-launch decode step, ROCm#3169, ROCm#3167, ROCm#3052, ROCm#3073, ROCm#3153, README news). One conflict, in model_runner.py: ROCm#3109's step override (`_step_override` / `_served`) and this branch's model head output (`compute_head`) sit side by side, in __init__ and as two branches of postprocess's verdict (a tuple head output; an overridden step's `verify_outputs`). The two are exclusive: an overridden step returns no head output. The hi runner test's stand-in runner gains ROCm#3109's `_served` and the stream argument of `forward_done_event.record`.
…he drafter's ids (#3177) (#3185) The gfx1250 DeepSeek-V4.1 single-launch step (#3109, step_override) runs the whole decode step, drafts included, and returns no hidden states. #3136 sliced the drafter's input ids by hidden_states.shape[0], so postprocess crashed on the first served step. ModelRunner._draft_input_ids keeps the hidden states' rows where there are any (a PCP rank's shard on a pure middle chunk) and takes the step's running_tokens otherwise, at all three call sites. #3091's served rounds return hidden states and are unaffected. Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
…frozen frame) One layer body for every target and draft layer (step.target_layer), every stage-to-stage edge an LL tile handoff (only the TP all-reduces cross devices), modelled on TileRT mega_afm: - attention: x8 made once by 32 publisher CTAs; wqkv / wqb weights issued before their waits; A -> A2 -> kv / C -> D on LL flags; the D merge stays in the 16-CTA hardware cluster (a separate A2 merge CTA set measured +2.1 us); - wo_a by (32 output rows x head) on 256 CTAs + head fold, wo_b on 160 owners x 32 cols (full K) with the TP1 slice; the FFN / attention mHC posts run in those owners; - FFN / MoE per row (40 CTAs a row): router on 128 CTAs x 3 experts (gate read once), every row's top-k, the expert union for the UG (each (expert, group) once, its rows as N), the down per (row, 128 cols) with the shared expert on wave 6 (its slice prefetched to LDS during the top-k), row partials pushed to the 4 ranks and summed by the owners; - indexer layers wait on the A2 / publisher flags (no ATTN_A handoff); the draft layers use the same body (128 experts top-3, the draft metadata, dh as the head input). Merged into the PR #3109 files (attention, device, front, moe, step); buffers: the fixture's fixed route table. Fixture (CT5 / CT17, c1 TR 6, 4 ranks): ~89.9-91.4 us a layer (was ~98.8), 0 poll timeouts, rank-0 logits and the draft outputs bit-identical to the unmerged frame. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
On gfx1250 the per-layer mono kernels are refused (gfx950 wave64 code, #2617),
so every V4.1 decode step takes the unfused path: each target layer's own
launches, then the sampler, the rejection sampler and the DSpark proposer.
This adds gfx1250's mono decode: a TP4 DSpark-5 verify step of up to 8
requests runs as two launches of one persistent FlyDSL kernel a rank, with
both TP all-reduces of every layer inside it.
its 2-level top-k selection and REINDEX, attention, routed and shared MoE),
the LM head, temperature / top-p sampling and the probabilistic block
verification.
the Markov block drafts.
The output copy waits for the target launch only; the draft launch runs while
the host schedules the next step.
Package (atom/models/deepseek_v41/mono/gfx1250/, beside the gfx950 mono):
rows, all built at startup), the per-step inputs, the kernel's pointer
tables and buffer catalog, graph capture.
index_query / index_score / topk, compressor, engram, draft_ctx; tail, head,
sample, markov). dims.py holds the model and deployment sizes, layout.py
names every hand-off counter and LDS region, index_plan.py / moe_plan.py
hold the per-step arithmetic the kernels and the CPU tests share.
Switched like the gfx950 mono decode:
with 5 speculative tokens, BF16 KV cache and FP8 index cache. Any other
deployment is refused at load (
V4.1 single-launch step off: <why>) andkeeps the production path; ATOM_MONO_ENABLE=0 turns it off.
output, logprobs, more than 8 requests, or a block table / index bound past
the kernel's.
logs the head input / logits, the index selection, the taps and the drafts
against it; ATOM_MONO_TIMELINE= logs where a served step's time
goes.
Outside the package:
atom/mono/runtime/step_override.py: a model that runs a whole decode step
itself returns its sampled ids and verdict (the rejection sampler's shapes)
and the stream they are ready on. On such a step ModelRunner computes no
logits, takes the verdict in postprocess, skips the TP id broadcast (the ids
are identical on every rank, and the broadcast would wait for the draft
launch) and records forward_done_event on that stream, so the output copy
does not wait for the draft launch either. The drafter's propose /
compute_draft_kv are wrapped by the mono package, not by ModelRunner.
process_weights_after_loading, which the loader runs before the experts'
in-place shuffle (parent modules first), and only when the step will serve
the deployment. The shared moe.py is unchanged.
uncached kind. The step's few-word peer pushes (sampling, Markov) waited
~80-100 us each to become visible in uncached memory.
Fixes found on the way (in the kernels; nothing outside this path changes):
(NaN after EOS).
its rows per cluster, and PV read past the row's key slots.
count (wrong scores past 16384 index rows).
Measured on MI455X, TP4, image gfx1250-atom-20261105-dsv41, same machine.
Native is the production path of the same tree (the recipe's probabilistic
draft + block verify).
Two runs where there are two numbers. The native long-context runs hang now
and then in warmup (NCCL broadcast timeout), so that pair has one native run;
an earlier pair measured native TPOT 2.304 / 2.154 ms.
These are from e354c8e. The later commits change the integration, the
naming and the switches but not the kernels (the traced IR of all 16 builds,
widths 6..48 target / draft, is identical); the same runs on the PR head are
TODO before review.
Docs: recipes/DeepSeek-V4.1-Flash.md gains a gfx1250 single-launch section;
docs/environment_variables.md gives the gfx1250 behaviour of ATOM_MONO_CHECK
and ATOM_MONO_TIMELINE.
Tests: tests/test_deepseek_v41_gfx1250.py, CPU only. Without aiter (as in
CI): the index and MoE plans on Python ints (selection tree coverage, task
counts, a regression for REINDEX key rows past 16384 rows, the score split)
and the layout. With aiter: deployment refusals, width routing, which steps
are served, the step override and the drafter wrap.