Skip to content

[DSV4.1][gfx1250] Single-launch decode step (FlyDSL) - #3109

Merged
yifehuan merged 32 commits into
mainfrom
yifehuan/gfx1250-single-refactor
Nov 10, 2026
Merged

yifehuan merged 32 commits into
mainfrom
yifehuan/gfx1250-single-refactor

Conversation

@yifehuan

@yifehuan yifehuan commented Nov 9, 2026 •

Copy link
Copy Markdown
Contributor

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.

  • Target launch: the 40 target layers (compressor, Engram, the indexer with
    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.
  • Draft launch: the draft context KV, the 3 draft stages, the draft head and
    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):

  • config.py: the deployment it is written for, and why another is refused.
  • dispatch.py: which steps are served (SingleStepModel) and the drafter wrap.
  • runner.py / inputs.py / tables.py / buffers.py: per-width builds (6..48
    rows, all built at startup), the per-step inputs, the kernel's pointer
    tables and buffer catalog, graph capture.
  • weights/: the kernel's weight layouts, including the routed experts' planes.
  • kernels/: the FlyDSL stages (step.py is the launch; attention, front, moe,
    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.
  • debug.py: ATOM_MONO_CHECK and ATOM_MONO_TIMELINE (below).

Switched like the gfx950 mono decode:

  • On with ATOM_MONO_ENABLE (default 1) on a gfx1250 device, for TP4, DSpark
    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>) and
    keeps the production path; ATOM_MONO_ENABLE=0 turns it off.
  • A step it does not serve takes the production path: prefill, structured
    output, logprobs, more than 8 requests, or a block table / index bound past
    the kernel's.
  • ATOM_MONO_CHECK (eager) runs the launch beside the production forward and
    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:

  • ModelRunner (+40 / -3) knows only a generic base,
    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.
  • The V4.1 MoE packs its routed experts into the kernel's planes in its own
    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.
  • atom/mono/runtime: fine-grained (coherent, cached) peer memory next to the
    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):

  • MoE route-weight renormalization floors the sum at 1e-20, as AITER does
    (NaN after EOS).
  • Draft attention computed no last row for a width that is not a multiple of
    its rows per cluster, and PV read past the row's key slots.
  • REINDEX key rows clamp to the token's bound, not to the compacted column
    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).

mono native
GSM8K, T=1, 512 questions 451 451
GSM8K draft acceptance 0.367 / 0.360 0.361 / 0.372
decode step, 8 requests, short prompts 7.05 ms 8.11 ms
decode step, 1 request 5.68 ms 6.76 ms
TPOT, 8 requests, ~26.6K-token Agentic first turns 1.974 / 2.001 ms 2.179 ms
throughput, same 2155 / 2304 tok/s 2105 tok/s
acceptance, same 0.507 / 0.519 0.519

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.

@github-actions

github-actions Bot commented Nov 9, 2026

Copy link
Copy Markdown
Contributor

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

  • ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

  • ✅ Run after the PR is approved and Pre Checkin passes
  • ✅ Run immediately when an approval review is submitted
  • ✅ Can be requested before approval with labels
Label Tests
ci:full Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom Run native ATOM model accuracy tests
ci:vllm Run ATOM vLLM OOT model accuracy tests
ci:sglang Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 3109 --add-label <label>

@valar-cop valar-cop left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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] Unnecessary dict() call (rewrite as a literal)
⚠️ [ruff] Unused noqa directive (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

Comment thread atom/model_engine/model_runner.py Resolved
Comment thread atom/utils/envs.py Resolved

@valar-cop valar-cop left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

yifehuan pushed a commit that referenced this pull request Nov 9, 2026
…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>
yifehuan added a commit that referenced this pull request Nov 9, 2026
…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).
@yifehuan
yifehuan force-pushed the yifehuan/gfx1250-single-refactor branch from 2613756 to 930889a Compare November 9, 2026 08:17
@yifehuan

yifehuan commented Nov 9, 2026

Copy link
Copy Markdown
Contributor Author

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.

@valar-cop valar-cop left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

some changes needed..

ci green.

…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).
@yifehuan
yifehuan force-pushed the yifehuan/gfx1250-single-refactor branch from 24bc5e2 to f85569f Compare November 9, 2026 10:02
… 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).
@yhl-amd

yhl-amd commented Nov 10, 2026

Copy link
Copy Markdown
Collaborator

@valar-cop can we merge this pr?

@valar-cop

Copy link
Copy Markdown

@yhl-amd code side looks fine — all the earlier threads are fixed in the head (dispatch gating on input_ids.numel(), _served per run_model, no second expert copy).

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.

@yifehuan
yifehuan merged commit 03b584b into main Nov 10, 2026
54 of 63 checks passed
@yifehuan
yifehuan deleted the yifehuan/gfx1250-single-refactor branch November 10, 2026 15:29
rdoshi-amd added a commit to rdoshi-amd/ATOM that referenced this pull request Nov 10, 2026
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`.
valarLip pushed a commit that referenced this pull request Nov 11, 2026
…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>
yifehuan added a commit that referenced this pull request Nov 12, 2026
…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants