Repository navigation
[DeepSeek-V4.1] High-interactivity monokernel decode for DeepSeek-V4.1-Flash on MI355X (TP4, DSpark) - #3091
Conversation
…ey serve A monokernel decode path for DeepSeek-V4.1-Flash on MI355X (gfx950) at TP4 with DSpark. Each decoder layer runs as one persistent FlyDSL launch, and consecutive layers are chained: a layer pushes its MoE partials and the next one sums them. The verify head (with sampling), the DSpark draft's context write and its block draft also run on the path. The kernels use dual-copy mailboxes per I/O die, uncached scratch and symmetric peer memory on atom.mono's HIP wrapper. They run on ATOM's own storage, in place: - the KV pool: the compressed KV, the FP4 index keys and the compressor state are read and written where ATOM keeps them, through the step's block table. Every store is masked by its row's liveness, so a padded request or a band's dead slot writes nothing. Only the window rings are imported and exported each step. - the weights: packed from ATOM's loaded tensors. The routed experts, their E8M0 scales (aiter's shuffle_scale order) and the plain tensors are read in place. - the token embedding: each id's row is read from the rank that holds its shard (exported once at load, checked against ATOM's lookup). - RoPE: ATOM's bf16 cos / sin caches. atom/models/deepseek_v41/mono/hi holds the hi runtime: - the V4.1 recipe and schedules (frozen MI355X / TP4 tables, the Linear prefetch fits included); - the gfx950 blocks and device code; - the planner, emitter and runtime. It builds at load, compile-only (no device needed to build), from FlyDSL's disk cache. Like the other vendored FlyDSL kernel sources, the package is exempt from ruff / black (pyproject.toml). atom/mono/runtime/srcloc.py captures FlyDSL's source locations at a fraction of FlyDSL's own cost, around a build's trace only. The hook (install_mono_decode) is spread over these modules: - hi_runner: the layers and a step: the request words, the ring bridge (hi_bridge), the pointers bound to ATOM's storage. - hi_draft: the DSpark context write and the block draft. - hi_weights: the packs from ATOM's tensors. - hi_compose: the programs. - round/: the decode-round hooks: the lean host prep, the device-side round advance with its fused prologue, the verdict / Engram commit / DSpark anchors inside the target graph, and the remote embedding reads. A decode step the hook does not serve runs on ATOM's K1 / K2, built on the first such step. ATOM_DSV41_MONO_HI: - "auto" (default): on for exactly the supported deployment, which is - MI355X with every CU, TP4, no DP attention; - bf16 KV with the FP4 index cache, ATOM's 128-token V4.1 KV page; - DSpark with 5 drafts; - ATOM's Engram runtime on UVA host shards with overlap staging; - the packed H2D publications, no KV transfer, FlyDSL. Any other deployment keeps ATOM's mono decode and logs one info line. - "1": forces it. - "0": turns it off. The programs' context bucket is 1 Mi tokens; a longer --max-model-len keeps ATOM's mono decode. The decision reads only the settled configuration and device properties; nothing of the hi path is imported otherwise. The builds use FlyDSL's disk cache, which also holds the Linear fits derived where the frozen table does not apply. hi_contract checks the programs' hand-offs and memory model without a GPU. Engine hooks: opt-in extension points a model can serve parts of a step from. A model that offers none runs exactly as before. - ModelRunner: a model-side verify head. A speculative verify's ids come from the model's own head instead of logits (greedy, sampled rows it draws itself, forced acceptance), its verdict in one launch, with no TP broadcast of the ids. Steps that read logits keep the logits path. - ModelRunner: a model's decode round (decode_round). It is called - in prepare_input_ids; - at the deferred ids fill and at the staging mark; - after the target's launch and after the rejection sampler; - once the H2D publications are set up. - DeepseekV41MetadataBuilder.decode_round: a lean step's layout (no compress plans, no step planners), and committed before the speculative commit. - PagedAttentionCache.prepare_state: takes a device-round step without its probe. - DSparkProposer: skips its context-KV write and its anchor staging when the model did them (anchors_staged). - PackedCopy: a host-arena read and a deferred scatter. - EngramStaging.lookup_served: a step whose model looks the rows up itself only advances the cursor. - V4.1 RoPE: rope_frequencies. Tests are CPU unless noted: - the gate on ATOM's own Config of the served command line, and every refusal; - the decode-round hooks in forward order; - the weights packed from ATOM's tensors against the checkpoint's packs; - the programs' contracts (compile-only); - the RoPE caches; - bands, routing, scratch, and the package's self-containment; - the cache bridge and round kernels (test_mono_hi_bridge: Triton interpreter cases, plus GPU cases).
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
|
|
||
| return logits, hidden_states | ||
|
|
||
| # ---- a model-side verify head (DeepSeek-V4.1's hi hook, `ATOM_DSV41_MONO_HI`) |
There was a problem hiding this comment.
this file is for common path!! move the v4.1 eligibility policy next to its impl, keep one generic seam here.
There was a problem hiding this comment.
Done. model_runner.py now has one model-agnostic seam: the runner looks up an optional verify_head object on the model (next to the existing decode_round) and asks it for a step's ids (None → the logits path). The V4.1 policy (which verify steps the head serves: greedy, top-p sampling, forced acceptance; logprobs, structured output, a required channel, relaxed acceptance and PCP decline) moved next to its implementation in atom/models/deepseek_v41/mono/round/verify_head.py. The runner's diff vs main went from +251 to +154 lines, with no V4.1-specific code left in it.
|
|
||
| import pytest | ||
|
|
||
| pytest.importorskip("triton") |
There was a problem hiding this comment.
importorskip at module scope skips the pure-python plan_widths/bands and HiLayers._key logic in ci. split them out and test those?
There was a problem hiding this comment.
Done. The pure-Python logic (band plan, widths, program keys, the verify-head policy, lean-host predicates, RoPE tables, weight aliasing) now imports without Triton, AITER or FlyDSL (the layer plan moved to hi_plan.py; Triton parts load lazily), and skips sit only on tests that need those. With all three absent, 78 of the hi tests run (was 47); the rest need real Triton kernels, FlyDSL or a GPU.
| nothing (a few µs of host). -> False when a graph capture is recording (nothing staged). | ||
| """ | ||
| # the step's rows' (T, top_p); past them greedy | ||
| rows = np.zeros((self.par_rows, 2), dtype=np.float32) |
There was a problem hiding this comment.
why np.zeros every sampled step? keep one persistent rows buffer and fill in place.
There was a problem hiding this comment.
Done. set_head_sampling fills persistent buffers created once at install and stages a copy only when the values change; nothing is allocated per step. Covered by a CPU test.
|
|
||
| The layer recipe and schedules (``v41``), the gfx950 blocks and device code (``blocks.gfx950``, | ||
| ``device.gfx950``), the planner (``plan``), the emitter and compiler driver (``emit``) and the runtime (``runtime``: | ||
| kernels, peer memory, dual-copy and uncached scratch). Imports within the package are relative; it imports nothing |
There was a problem hiding this comment.
docstring says imports nothing from atom, but runtime/peer.py and device/gfx950/ops.py do. fix the doc?
There was a problem hiding this comment.
Fixed: it now lists the atom.mono modules the package imports (device.ops / device.mx, runtime.hip_ipc, runtime.srcloc); a test holds the package to that list.
|
|
||
|
|
||
| @dataclass(frozen=True) | ||
| class Search: |
There was a problem hiding this comment.
Search and SEARCH are never read, and the docstrings cite modules not in this repo. remove?
There was a problem hiding this comment.
Removed, together with the other tuner remnants and modules that V4.1 never reaches (conv.py, stamps.py, lds_dma.py, plan/items.py, unused stages/modes), and the docstrings now cite only code in this repo. The generated kernels are unchanged (LLVM IR identical for all served programs).
| from flydsl.expr import meta | ||
|
|
||
| # sha256 (first 16 hex) of the FlyDSL functions this module reproduces (``_fingerprint``) | ||
| FINGERPRINT = ( |
There was a problem hiding this comment.
fingerprint miss silently no-ops and doubles build time. can this go upstream? at least log one line on miss.
There was a problem hiding this comment.
It now logs one warning per process on a miss (what's skipped and that builds take ~2x as long). Agreed it belongs upstream; we'll propose the equivalent change to FlyDSL (per-code-object position caching in meta.capture_user_location plus a scoped location depth), after which this module reduces to a context manager.
- ModelRunner keeps one generic, model-agnostic seam: a model's `verify_head` (before_forward / ids / captured / verdict / ids_identical_on_tp_ranks) beside its `decode_round`. Which verify steps the V4.1 hi head serves (greedy, sampled top-p, forced acceptance; refused with logprobs, structured output, a required channel, relaxed acceptance or context parallelism) moves next to its implementation, `round/verify_head.py` (`HiVerifyHead`). Behaviour is unchanged. - The pure-Python logic imports without Triton, AITER or FlyDSL: the layer plan is `hi_plan.py` (out of `hi_bridge`), `hi_compose` no longer imports the bridge, the round package imports its Triton parts lazily, and `rope_args` lives in `hi_weights`. The tests import these directly and skip only where a test needs the dependency; the band plan, program keys, verify-head policy, lean host predicates and RoPE tables now run in a CPU job without them. - `set_head_sampling` fills persistent host buffers in place (each request's temperature / top-p on its rows) instead of allocating arrays every sampled step. - The hi package docstring names the atom.mono modules it imports. - Tuner remnants (`Search`, the `SEARCH` dicts), the DFlash2 conv stage (`blocks/gfx950/conv.py`), unused modules (`stamps`, `lds_dma`, `plan/items`) and unreachable stages, modes and helpers are removed, as is the bridge's compiled-out MXFP8 ring copy (`RING8`). Comments no longer refer to other models, measurements or files outside this repository. The served programs' LLVM IR is unchanged (110 of 110). - `srcloc` logs one warning a process when FlyDSL's location capture is not the one it reproduces: the fast capture is skipped and builds take about twice as long; the programs are the same.
|
Thanks for the review. All six comments are addressed in f13e0e9 ("refactor(deepseek-v4.1): address review of the hi monokernels"); replies inline. Validated on MI355X (TP4) after the change: hi path on all ranks, KV 145.8 GB, c1 smoke 0 errors, sampled (T 0.7 / top-p 0.95) requests at c1/c4 correct, kernel step 3.16 ms (unchanged within noise). |
valarLip
left a comment
valarLip
left a comment
There was a problem hiding this comment.
Reviewed at f13e0e9. This review covers only the PR's changes to existing files (12 files, +446 / −48), i.e. how the hi path plugs into ATOM's engine; the new mono/hi* code is out of scope here. I read the code and probed the routing predicates; no GPU. rotary.py's rope_frequencies extraction, the env var and test_packed_h2d.py are fine.
The common problem is one pattern repeated across five shared modules. Each grows a class attribute that defaults to None and is set from outside by the hi path, plus one undeclared field read via getattr. All of them encode the same fact: "this step's round is served by hi". Because they are separate predicates, they already disagree, and two of the bugs below come from that disagreement. Each item ends with the change I'm asking for.
1. One declared per-step ownership instead of five injected hooks. Today the fact is spread over:
DeepseekV41MetadataBuilder.decode_round = None(attentions/deepseek_v41/backend.py:256);getattr(step, "device_round", False)on aBatchStepthat doesn't declare it (attentions/deepseek_v41/cache.py:494);EngramStaging.lookup_served = None/_served(engram/device/staging.py:185-195);DSparkProposer.anchors_staged = None(spec_decode/dspark_proposer.py:50) andgetattr(type(self.model), "context_kv_served", None)(:642), sniffing the draft model's class for a hi-only method;PackedCopy.defer/host_arena(utils/packed_h2d.py:117-125, item 4).
They have already drifted:
- Engram:
staging.lookup_servedchecks onlyserves_step(width, step), whileHiDecodeModel.forwardalso requiresstep_supported(noinputs_embeds, noimage_mask, noubatch_slices). On a served-width decode carryinginputs_embeds/image_mask,EngramStaging.startskips the snapshot and side-stream gather while the forward falls back to ATOM's path.EngramStagedRowsthen reads the previous step'shost.buffers[layer]. - DSpark:
ctx_served(hi_draft.py:246-252) keys on(width, max_q_len)and never checks that the step is a decode. A 6-token prefill (a short prompt, a chunk tail, or a prefix hit leaving 6 new tokens) has the key(6, 6), the same as a served bs=1 verify.write_context_kv(dspark_proposer.py:642) then returns early, CTX didn't run either, and the draft reads the slot's previous occupant's KV. Acceptance drops silently, and prefix caching is on in the PR's own command.
→ Change: declare one field, e.g. BatchStep.served_by_round: bool (or on the forward context), set once by the routing decision HiDecodeModel.forward uses. Have the metadata builder, the cache, Engram staging and the DSpark proposer read that field. Delete the four injected class attributes and the getattr(step, "device_round") / getattr(type(self.model), ...) sniffing.
2. The runner's model hooks should be a declared interface (model_engine/model_runner.py:3426-3470). verify_head (5 members) and decode_round (7 members) are duck-typed protocols, documented only in a comment and called from 10 points of the generic step loop. Specifically:
_model_hookcaches lookups viaself.__dict__.setdefault("_model_hooks", {}), so:1015has topopit by hand after mono install._decode_round/_verify_headre-checkgetattr(self, "model", None)on every call, andpostprocessreadsgetattr(self, "verify_token_ids", None).run_model→_run_modelsplits the method only soafter_forwardruns afterwards, and the two names don't tell them apart.fill_decode_idsreturns a bool meaning "I filled them".
→ Change:
- Define the two protocols in core (
typing.Protocolor small base classes), and resolve them once into plain attributes when the model is installed (Nonefor every other model); drop_model_hooksand the pop. - Initialise
verify_token_idsin__init__. - Name the split for what it does (e.g. keep
run_modeland calldecode_round.after_forwardat its single call site). - Make
fill_decode_idsan explicit capability rather than a bool return.
3. Two runner-side behaviours to justify or remove (model_runner.py:561-565, :3685-3700):
num_rejectedon a device-round step.spec_status()returns(None, None), soprev_rejected_numis None and the remap is skipped.num_rejectedthen takes the scheduler'sbatch.num_rejected, one verdict older, not 0, which contradictsdevice_round's stated invariant that the host builds the step "as if no draft had been rejected". The device advance fixesstep.positionsonly, so spans,cached_until/ floors,token_maskand the Engram host rows disagree with the device positions. I found no hi reader of those fields yet, so this is at least a broken invariant.- Skipping the TP id broadcast. It's skipped on
ids_identical_on_tp_ranks, i.e. on the model's own declaration.
→ Change: on a device-round step, have the runner pass num_rejected = 0 explicitly (or assert nothing reads the host-side fields). Keep the broadcast unless a debug check verifies equality across ranks, since drafter inputs differing across TP ranks would be very hard to diagnose.
4. PackedCopy gains hi-only state and a second runtime handle (utils/packed_h2d.py):
defer/host_arenaare set from outside:round/__init__.py:104-105andround/prologue.py:276reach into the runner's privategroup._backend.- The
host_arena is Nonesub-branch of the defer path is unreachable on the hi path. If it ever ran, a refusedtakewould issue the same H2D twice. is_current_stream_capturing()is evaluated twice.read_host_arenaswitches every publication of the runner, prefill included, to a zero-copy PCIe read, with no measurement._device_addresscallsctypes.CDLL("libamdhip64.so")on every call.engram/host_memory.py:100-110(_hip_runtime) explains why that's wrong: the unversioned dev link may be missing on a wheel install, or may be a second HIP runtime, different from torch's, with its own state. A device pointer obtained there need not be valid for torch's runtime._Addressis defined three times (here,round/prologue.py:154,hi_runner.py:1064).
→ Change:
- Make deferral and zero-copy an explicit API on
h2d_owner/PublicationGroupthat a model opts into, with a size threshold and a measured A/B for prefill-sized publications. - Get the HIP handle from
engram.host_memory._hip_runtime()(or move it to a shared util), computed once. - Keep one
_Address. - Delete the dead branch.
5. The DeepSeek-V4.1 metadata builder's lean path (attentions/deepseek_v41/backend.py:725-835, :1345). A lean step sets plans = {} because None "would build private ones", so {} vs None now carry two meanings for begin_step. commit_speculative_state returns early when the hook says the forward already committed.
→ Change: with item 1's field, branch on step.served_by_round. Pass an explicit "no plans" value (or skip begin_step's plan handling), not an empty dict that differs from None.
6. auto should fall back, not fail (models/deepseek_v41/mono/dispatch.py:158-175). The local import of hi_dispatch is fine (the hi package isn't imported when off). But install_hi checks many deployment assumptions with ValueError, and nothing catches them:
- the expert layout (
hi_weights._expert_layout,:154,:167); - head-vocab divisibility,
eps,sliding_window,MAX_TAPS, EngramNGRAM.
For example, gfx950 + TP4 + ATOM_USE_TRITON_MOE=1 (a backend recipes/DeepSeek-V4.md documents) passes hi_gate, then fails in _expert_layout with "not in the gfx950 FlyDSL layout", and the server doesn't start. ATOM's own mono uses MonoUnsupported here and falls back.
→ Change: raise MonoUnsupported for deployment mismatches in install_hi and catch it in install_mono_decode: under auto, log one line and keep ATOM's mono; under 1, re-raise. Better still, move those checks into hi_gate so decide answers them before anything is built.
7. No lint exemption for new code (pyproject.toml:69-87). atom/models/deepseek_v41/mono/hi/**/*.py gets ruff ALL and black force-exclude, "vendored the same way". But this is new code written for ATOM (it imports atom.mono.* and has its own tests), and its same-lineage sibling qwen38/mono/kernels has no exemption. About 23k lines would leave CI's black . && ruff check ..
→ Change: remove both entries and run black / ruff on the package.
8. Docs. The env doc entry is good. CLAUDE.md's architecture index and docs/architecture_guide.md don't mention the new runner hook protocols or the hi path (Docs-match-code).
→ Change: add them, with item 2's protocol names, once those settle.
… hooks
- One declared per-step ownership: AttentionMetaData.served_by_round
(RoundService: NONE / FORWARD / LEAN / DEVICE), decided once by the
V4.1 metadata builder through the model's decode round
(DecodeRound.service: a decode at a served width and bucket, no image
rows). The builder, its cache, Engram staging
(EngramStep.served_by_round), the DSpark proposer and the hi forward
read it; the injected DeepseekV41MetadataBuilder.decode_round,
EngramStaging.lookup_served, DSparkProposer.anchors_staged and the
context_kv_served / step.device_round sniffing are gone. This fixes two
disagreements: Engram staging skipping its lookup on a served-width
step whose forward took ATOM's path (image rows, inputs_embeds), and the
DSpark context write skipped on a 6-token prefill (the draft read the
slot's previous occupant). A routed step whose forward the hi layers
cannot run now raises.
- Declared runner hooks in atom/model_engine/model_hooks.py (VerifyHead,
DecodeRound), resolved once at model install into
ModelRunner.verify_head / decode_round; verify_token_ids is initialised
in __init__; run_model is unsplit, the decode round's after_forward
sits right after it in forward; the decode-ids fill is a capability
(decode_ids_filler) instead of a bool return.
- A device-round step's host status is zeros (built as if no draft had
been rejected, as its positions are); the ids' TP broadcast is kept on
head-served steps.
- Zero-copy and deferral of packed publications are an explicit opt-in
per publication (h2d.HostArenaReader, PublicationGroup.read_host_arena):
the hi path opts in only its decode publications, decode steps only,
up to 1 MiB. One DeviceAddress / device_address helper on the HIP
runtime torch loaded (engram.host_memory), resolved once; the dead
device-arena deferral branch is gone; capture state is read once.
- A lean step skips its compress plans and step planners on
served_by_round in begin_step (no {} vs None).
- "auto" falls back instead of failing: hi_gate refuses Triton MoE,
two-batch overlap and a geometry other than V4.1-Flash's (vocabulary,
norm eps, window, draft taps and stages, Engram n-grams) before
anything is built, and a mismatch only the loaded model shows raises
MonoUnsupported, which install_mono_decode turns into ATOM's mono
decode under "auto" (re-raised under "1").
- No lint exemption for the hi package: black and ruff run on it as on
the rest of the tree.
- CLAUDE.md and docs/architecture_guide.md describe the model hooks and
the hi path.
The lines every TP rank would log alike (the gate's and the install fallback's "V4.1 mono hi off", and the install's own info lines) go through hi_gate.log_once: rank 0 only. Per-rank lines (the build summary with its rank, the Engram host shard rows and checks, a build that does not fit on one rank) stay per rank.
- A head-served verify takes each request's next id from the TP-broadcast output ids, as the logits path does: ModelRunner.postprocess gathers them at the head verdict's accepted counts, so the next forward's ids (prev_token_ids) and the drafter's input are rank 0's on every rank. VerifyHead.verdict returns (output ids, verdict); the head's own next ids are gone. The in-graph post-verify no longer stages the DSpark anchors from this rank's picks: the proposer stages them from the broadcast ids on every step (DecodeRound.anchors_staged and its skip in DSparkProposer._stage_anchors are gone). The post-verify keeps the verdict and the Engram cursor commit; committed() now takes the step's in-graph results. - Zero-copy decode publications only on steps the hi path serves: the token inputs when the step's requests fit a served width (device_round.may_serve, before the routing decision), the V4.1 step metadata on the routing decision itself (HiRound.service). Decode steps at widths the hi path does not serve copy as before. - One HIP loader: atom/utils/hip_runtime.py (the runtime torch loaded, moved from engram.host_memory) is imported by packed_h2d at module level and is hip_ipc.HipRuntime's default library; the hi peer groups and uncached scratch use hip_ipc's module runtime instead of a second HipRuntime. - The PackedCopy opt-in tests load packed_h2d under the repo's stubs and run without Triton. - hi_gate.log_once logs each line once a process, on rank 0. The checkpoint-shaped refusals in hi_weights (FP8 dense dtype, padding scale code 255, RoPE cache layout) raise MonoUnsupported, so "auto" falls back. DecodeRound.service is black-formatted. - Tests: the runner test checks the next ids are the broadcast ones; the lean attributes test is behavioural; the old-name scan skips a worktree's .git file; the GPU head-verdict check hands ATOM's rejection sampler logits (it no longer accepts None) and the post-verify check tests the commit's skip, not a removed installer.
|
Thanks for the careful review. All eight items are addressed in 7bf3f5b, bc360a2 and 0eb7bfe: 1. One per-step ownership. The 6-token case was a real hole in the predicate. On the previous head we couldn't trigger it through the server: 2. Declared hooks. 3. Runner behaviours. A device-round step now reports 4. PackedCopy. Zero-copy is an explicit opt-in: a model passes an 5. Lean path. 6. 7. Lint exemption. Both 8. Docs. CLAUDE.md's architecture index and Validated on MI355X (TP4) at 027489698:
|
valarLip
left a comment
valarLip
left a comment
There was a problem hiding this comment.
Re-read the changes to existing files at 0eb7bfe (same scope as last time), by reading the code and running the affected unit tests on the PR tree; no GPU. Thanks for addressing all eight items, and for the measurements. served_by_round, the declared hooks, the explicit zero-copy opt-in and the single HIP loader read right. Four things are left; 1–3 I'd like fixed before approving. Each ends with the change I'm asking for.
1. A CI test file now fails (atom/model_engine/model_runner.py:545, :730). tokenIDProcessor reads self.runner.decode_round unconditionally. tests/test_spec_deferred_status_queue.py builds its runner as a SimpleNamespace without that attribute, so on the PR tree it gives 5 failed / 28 passed, all AttributeError: 'types.SimpleNamespace' object has no attribute 'decode_round'. The same file passes 33/33 at the merge base, and it is in the CI unit-test set.
→ Change: give the test's runner double decode_round = None (or have tokenIDProcessor take the hook at construction). Then run the CI unit command over the PR once more to catch any other doubles.
2. The MonoUnsupported fallback must be agreed across ranks (models/deepseek_v41/mono/dispatch.py:177-182). This completes my last item 6: I said "raise MonoUnsupported and fall back" but didn't say the decision has to be collective, as ATOM's own mono does with consensus.bind_agreed. Some refusals are rank-local data checks. For example, hi_weights.py:216 refuses when this rank's expert-scale shard holds a padding code 255. If one rank raises and falls back while the others continue, the others block in the next collective of install_hi (embed_full.install's all_gather_object, the peer groups), and loading hangs. When ranks do fall back together, whatever _install_hi already built stays allocated: packed weights, peer groups, the uncached arenas, and the embedding moved into uncached memory by embed_full.
→ Change:
- Run the hi install's refusals through
consensus.bind_agreed: each rank checks, all agree, then all build or all fall back. - Do every rank-local check (weights, embedding) before the first collective and before anything is allocated, so a refusal leaves nothing behind.
- Have
embed_full's own check raiseMonoUnsupported, notRuntimeError.
3. Device-round steps skip the stale-slot check (model_ops/attentions/deepseek_v41/cache.py:482-494). prepare_device_round_state does no _judge_previous_probe, no fresh / resume slot reset and no stale-state probe. Its safety rests on the routing predicate alone. In steady decode on the served deployment every step is DEVICE, so the hard check that used to raise on a recycled slot never runs there. If the predicate ever admits a recycled slot, the step silently reads the previous tenant's state instead of failing.
→ Change: in prepare_device_round_state, assert the step has no fresh slots and no resumes, i.e. the same requests as the previous step. That keeps the hard failure without the probe's cost.
4. Docs and small leftovers:
CLAUDE.md:109-111still saysruntime/peer_memory.pyholds "BOTH peer-buffer implementations" and that every model's step scratch goes throughruntime/uncached.py. The hi path now has its own peer group and arenas.- Docstrings cite things that don't exist:
send_to_cpu_async_draft(round/post_verify.py:27),tests/gpu/test_topk.py(hi/device/gfx950/topk.py:52),tests/cpu/test_source_digest.py(hi/emit/kernel.py:96),tests/gpu/test_indexer_v41.py(hi/blocks/gfx950/attention/dsa_indexer.py:61). packed_h2d.device_addressduplicatesqwen38/mono/ple.py:device_address, and the two differ when HIP is absent (one raises, one returns the host address).
→ Change:
- Update the CLAUDE.md paragraph.
- Point those docstrings at what exists, or drop the references.
- Have
ple.pyuse the shareddevice_address.
- The hi install's refusals are agreed over the TP group: every check that reads what a rank loaded (hi_runner.check: the geometry and plan, each target layer's and draft stage's FP8 linears, the routed experts' layout and padding codes, the RoPE caches, the verify head, the draft's stages and norms, the token embedding, the Engram host shards and their device read) runs first, before any collective or allocation, and its verdict goes through consensus.tp_uniform_local_validation, so every rank builds or every rank falls back with nothing built, the refusal naming the rank. The checks moved out of the build (HiLayers, _head_tables, bind_draft, use_host_tables, expert_packs -> hi_weights.check_block / check_experts / check_linears). embed_full's checks raise MonoUnsupported; its install runs first in the build, and a read that differs from ATOM's lookup refuses on every rank with the shard moved back into ordinary memory. - prepare_device_round_state raises unless the step continues the previous prepared step: the same requests in the same order, none fresh or resuming (a recycled slot fails instead of reading its last tenant's state); it judges the previous step's deferred probe too. - tokenIDProcessor holds the decode round (set at install) instead of reading the runner's, so runner doubles need no hook attributes. - packed_h2d.device_address serves qwen38's PLE table as well (identity_without_hip keeps its host address where there is no HIP runtime); CLAUDE.md names the hi path's own peer groups and arenas; docstrings point at what exists or drop the reference. - Tests: one rank refusing (its own check, another's, any exception) refuses on every rank before anything is built, and through install_mono_decode keeps ATOM's mono path; embed_full.check; check_block refuses each FP8 linear and the padding code; the device-round check in steady state and on recycled, resumed, missing or reordered requests; device_address without a HIP runtime. The EngramStaging source pin of the vLLM copy test moves to round 2's start / join (their served-step skip is not ported: the plugin's runner has no decode round). Test doubles the CI run with AITER found: the V4.1 metadata test's stand-in metadata carries served_by_round, the output-transfer test's runner the runner's hook defaults; the hi routing test stubs ATOM's mono runners only where they are not imported yet (an imported mono.dispatch holds their classes).
|
Thanks for the re-review. All four items are addressed in 8657c15: 1. CI test failure. 2. Agreed fallback. Verified on MI355X with rank 2 forced to refuse: all four ranks fell back to ATOM's mono path without hanging, and 3. Device-round slot check. 4. Docs and leftovers. The CLAUDE.md paragraph now covers the hi path's own peer groups and uncached arenas. The Validated on MI355X (TP4):
|
Brings in main through bf65756 (ROCm#3136, attention metadata staged, published and finished once per micro-batch). Conflicts resolved on main's structure: - model_runner.py: the verify head's ids follow main's running-token hidden states (the LM head runs the running tokens, postprocess cuts). - deepseek_v41/backend.py: the routing decision (served_by_round) is taken before the step is staged into main's step_metadata group; lean steps stage no compress plans; the decode round sees the step after its publication and micro-batches. Dummy steps stay unserved. - deepseek_v41/cache.py: main's begin_step signature; a lean step (the builder's lean=) keeps no plans or planners. - utils/h2d.py: HostArenaReader and read_host_arena beside main's after_publish and count-driven publish. - tests/test_packed_h2d.py, tests/attentions/deepseek_v41/test_metadata.py: main's helpers (_publish, use_packed, stage_query_layout); the host-arena tests on them. Adapted to main: - The hi prologue reads the step_metadata publication (was v41_metadata) and the step's running_tokens (was width). - The hi gate no longer reads ATOM_H2D_BACKEND (removed by ROCm#3136: every eligible group packs) and refuses an FP4 index layout the hi programs do not address. - The hi index sources read and write ATOM's FP4 index plane in either gfx950 layout: preshuffle (unchanged programs) or mfma, the OPUS scorer's 64-row groups that ROCm#2809 makes the default; the layout comes from Config.index_cache_layout and is checked against the cache's.
|
Merged main (bf65756) to resolve the conflicts with #3136, #2809 and #3098. Main's structure is kept, and our changes
Revalidated on MI355X (TP4, nightly 202611091503), with both index layouts:
The CI unit command shows 0 tests failing only on this branch versus main, both CPU-only and with aiter. |
valarLip
left a comment
valarLip
left a comment
There was a problem hiding this comment.
Re-read at d618872. Thanks for the last round: all four items are addressed, and the forced-refusal run on rank 2 is exactly the evidence I wanted. The merge with #3136 / #2809 also reads right.
Having now traced everything the hi path takes over in a served round, I'd like the engine side to be narrower than the protocols I asked for in the first round. Today ModelRunner carries two model-specific protocols (VerifyHead: 4 methods, DecodeRound: 7 runner-side methods) called from 10 points, all implemented only by the hi path. What the runner actually needs reduces to two capabilities that aren't specific to hi, and that ATOM's own path could use too. This is a further narrowing after seeing the full scope, not a reversal of the last round's fixes. Each item ends with the change I'm asking for.
1. Fold the verify head into the in-graph head output (logits_in_graph) (model_engine/model_runner.py:4178, :4339-4346, :3219-3227, :3256-3264, :3428). The runner already records the head inside a decode graph: with logits_in_graph, the recorded full_forward returns self.model.compute_logits(...) and req_graph.replay hands those logits back. The verify head is the same idea one step further: the graph's tail produces the rows' token ids (and the verdict) instead of logits. Making that one concept would:
- let ATOM's own path put its greedy argmax / rejection verdict in the graph as well, so a TP=1 or non-hi decode also returns ids straight from the replay, not only the hi path;
- give hi's head a single model method (e.g. a
compute_head(hidden_states)that returns logits or(ids, verdict), recorded the same waycompute_logitsis today); - remove
VerifyHead(before_forward/ids/captured/verdict),_verify_head_ids, theverify_token_idsattribute (three initialisations: class,__init__,run_model), the duplicated "head orcompute_logits" blocks in the eager and replay branches, andnote_verdict.postprocessthen branches once, on whether the head output carries ids.
The sampling rows the head needs (temperature / top-p per row) can be published by the builder with the step's metadata rather than handed over through before_forward. captured isn't needed once the head is part of the recorded forward.
→ Change: generalise logits_in_graph into an in-graph head output (logits, or ids + verdict), have hi supply its head through that model method, and delete the VerifyHead protocol and its runner call sites.
2. Make "the step doesn't need the host-side status" a builder capability (model_engine/model_runner.py:535-553, :3689, :4763). The device round needs one thing from the runner: don't block in prepare_input_ids on recv_spec_status before building the step; take that wait after the target launch. That's useful for any spec-decode builder that can advance positions from the device-resident verdict, not only hi.
→ Change: replace before_input_ids / spec_status / after_forward with one builder-declared capability (e.g. the builder says it consumes the device verdict for this step). The runner then defers its status wait generically. Keep the decision of which steps qualify (device_round_step / first_step) inside V4.1's builder, where service / laid_out / committed already live.
3. Move the prologue's hooks out of the runner, or drop them (model_engine/model_runner.py:714-715, :1676-1679, :1743-1744). decode_ids_filler, host_arena_reader and before_staging_mark exist so the prologue can fuse the step's four pre-target launches and read the pinned arenas zero-copy.
→ Change: if the fusion is worth keeping, express it in the H2D layer (PublicationGroup / h2d) as an opt-in any publisher can use. Otherwise remove these three hooks and let the runner launch as it does on main.
4. Ablation for each piece. To decide what stays, please post the decode step time (and c1 X) with each one disabled in turn on the current head:
- the in-graph head (ids instead of logits);
- the device round;
- the prologue fusion and zero-copy publications;
- the in-place embedding read.
Pieces that measure within noise should go rather than move.
With 1–3, the runner keeps two generic switches (an in-graph head output, a deferrable status wait) and no model-specific protocol. model_hooks.py can then shrink to what V4.1's builder needs, or disappear.
Review round 5: the runner keeps an in-graph head output and a deferrable status wait, and no model-specific protocol. - In-graph head output (generalises logits_in_graph): a model's compute_head(hidden_states) gives the forward's logits, or a verify's output ids and verdict. ModelRunner._head_output takes it after an eager forward and records it at the tail of a decode graph's capture, at any TP. _step_head keeps the ids and verdict where they stand for the step (_head_ids_serve), else computes the logits; postprocess branches once on whether the head output carries ids. The ids are still broadcast over TP and the next ids, drafter input and DSpark anchors gathered from the broadcast ones. The hi path supplies its verify head and in-graph post-verify through HiDecodeModel.compute_head; each served step's (T, top_p) rows are published by V4.1's builder with its routing decision. Removed: VerifyHead, _verify_head_ids, verify_token_ids, captured, note_verdict, round/verify_head.py and round/head_verdict.py (a step the in-graph verdict does not stand for reads the logits). - Deferrable status wait: AttentionMetadataBuilder.consumes_device_verdict. On such a step the runner builds from zero counts and takes the spec status right after the target's launch. V4.1's builder answers from the device round's decision (DecodeRound.consumes_device_verdict); the verdict reaches it through commit_speculative_state's anchors. Removed: DecodeRound.before_input_ids / spec_status / after_forward. The device-round slot check stays. - The prologue fusion and the zero-copy publications are removed (round/prologue.py, decode_ids_filler, host_arena_reader, before_staging_mark, h2d.HostArenaReader, read_host_arena, packed_h2d.DeviceAddress / device_address). h2d.py, packed_h2d.py and hip_ipc.py are back to main's, and atom/utils/hip_runtime.py is gone. device_address now sits in engram.host_memory beside its HIP runtime; the hi path's Engram host shards and Qwen3.8's PLE table use it. - The in-place embedding read is removed (round/embed_full.py): the hi begin-step and block draft expand ATOM's own embedding output. - model_hooks.py keeps DecodeRound only, what V4.1's builder asks (DeepseekV41MetadataBuilder.use_decode_round).
Brings in main through cd5fa97 (ROCm#3083 per-row sampling noise opt-in, ROCm#3105 delayed P/D KV completion, ROCm#2668 LMCache LRU refresh, ROCm#2709 Kimi-K3 PP state checkpoints). No conflicts. The sampler opt-in (--sampler-independent-noise) changes nothing the hi verify head serves: its sampled rows draw per-row Philox noise either way.
|
Thanks. We ran the ablation first (item 4), and it decided items 1-3. 4. Ablation (d618872, MI355X TP4, c1, forced acceptance 3.51, rank-0 profile over ~680 steps; each piece
1. In-graph head output. 2. Deferrable status wait. 3. Prologue hooks. Deleted, together with the zero-copy publications:
Validated at 18e1805 (MI355X TP4, nightly 202611091503):
82789bb then merges main up to b5d1b96 (no conflicts; CPU CI and the affected Engram / PLE GPU tests re-run clean). |
valarLip
left a comment
valarLip
left a comment
There was a problem hiding this comment.
Approving at 82789bb. Thanks for running the ablation first and letting it decide items 1–3. The runner no longer carries a model-specific protocol, the deleted pieces were the ones that measured within noise, and the validation (fallback with rank 2 forced to refuse, GSM8K, CI vs main) covers what I'd want. Three small leftovers in the existing files, all from that narrowing, for a follow-up or before merge as you prefer:
1. A subclass the postprocess signature change missed (atom/rollout/model_runner_ext.py:121). postprocess now takes logits: torch.Tensor | tuple[...], but the rollout override still does logits.shape[-1] > self._true_vocab_size. A rollout runner on a model with compute_head (today only V4.1 hi) would hit AttributeError on the tuple when _true_vocab_size > 0, or skip the vocab-tail mask when it's 0.
→ Change: mask only when the head output is logits, and sweep the other postprocess overrides.
2. Leftovers in the runner (atom/model_engine/model_runner.py:335-339, :790, :2913, :3347, :3371):
ModelRunner.spec_status_deferredandtokenIDProcessor.spec_status_owedencode one fact.take_deferred_spec_status()already checksspec_status_owed, so callers can call it unconditionally and the runner flag can go._head_ids_serve(~35 lines of eligibility: logprobs, structured output, channels, relaxed top-N, PCP, temperature, top-k, noise) sits in the generic runner and must stay in step with the model-side rule inHiRound._publish_samplingby convention. Let the head decide once and return logits when the step doesn't qualify, so the runner keeps only the "does the output carry ids" branch.logits(and_head_output(..., logits: bool)) now often holds ids and a verdict.
→ Change:
- Drop
spec_status_deferred. - Move the eligibility into the head.
- Rename
logitsto e.g.head_output.
3. DecodeRound / RoundService still live in generic modules (atom/model_engine/model_hooks.py, atom/utils/forward_context.py:50, :727). Their only user is V4.1's builder. AttentionMetaData.served_by_round is the four-level enum, while EngramStep.served_by_round (engram/device/staging.py:52) is a bool copy with the same name.
→ Change:
- Move
DecodeRoundandRoundServiceinto V4.1's package. - Keep a single bool on the generic metadata, the shape I suggested in the first round.
- Give the Engram field a name that says it's the collapsed bool.
valar-cop
left a comment
valar-cop
left a comment
There was a problem hiding this comment.
LGTM
ci ATOM Test failed — not sure yet (low): The job failed at the PR head but no log is present under .valar/ci/, so the failure cannot be attributed to the PR or to CI.
ATOM Test: failure https://github.com/ROCm/ATOM/actions/runs/42814039587
next: Fetch the job log (or re-run) and compare the failing test against the merge base b5d1b96 to decide pr_bug vs main_broken/flaky. https://github.com/ROCm/ATOM/actions/runs/42814039587
ci Pre Checkin failed — not sure yet (low): Pre Checkin is red at the head with no log archived in .valar/ci/; a lint/format failure cannot be distinguished from an infrastructure failure.
Pre Checkin: failure https://github.com/ROCm/ATOM/actions/runs/42777329390
next: Fetch the Pre Checkin log; if it is a formatting/lint diff, run black/ruff locally (CI enforces them) rather than waiting on a re-run. https://github.com/ROCm/ATOM/actions/runs/42814039587
CI's ruff job installs the latest ruff (0.17.4), which flags 188 findings on the lines this PR adds. The safe fixes: dict() calls as literals (C408), sorted imports (I001), unused noqa directives (RUF100), functools.cache (UP033), a dict comprehension (C402), re.MULTILINE (FURB167), no bare return None (RET501 / PLR1711). By hand: unused unpacked names underscore-prefixed (RUF059); nested ifs merged where both are host-side, kept with a noqa where a traced kernel or Triton nests a compile-time and a runtime if (SIM102); the nested with merged (SIM117); fmt: skip moved to statement level (RUF028); check=False made explicit (PLW1510); a broad except kept with a noqa where any error is a build refusal, the rank's install refusal or an aiter-free image (BLE001); the ceil-div without double negation (B002); a bound loop variable (B023); class constants as a tuple / ClassVar (RUF012); the exec's noqa on its line (S102); the FlyDSL launch default stream noqa'd (B008); a parenthesized implicit concatenation (ISC004). No behaviour change.
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`.
ROCm#3091's compute_draft_kv reads attn_metadata.served_by_round first; the stub forward context now carries RoundService.NONE, as main's own tests do. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
…p test (#3182) #2582 made DSparkProposer._stage_anchors take the block pass as its first argument, and #3091, merged five minutes later, added a test that calls it without one. main's non-GPU unit tests fail on the three anchor cases since 5b4fbfd. Hand the test's fake block in, as both production callers do. Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Resolve dspark_proposer.compute_draft_kv against #3091's served_by_round early return, and fix three tests this branch broke: - Drafter._aux_blocks becomes the module-level aux_blocks(drafter), so SimpleNamespace stubs with only _aux_buffers/_relayed_aux keep working (test_pp_aux_relay). - compute_draft_kv keys the non-PCP path on aux_for() returning None, as main does; only the PCP split path reads _captures_aux (test_deepseek_v41_mono_hi_served_step). - test_sp_moe_dispatch patches tbo_on_comm instead of the removed moe.tbo_active. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…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>
… branch's interfaces The stub drafter carries `aux_extract` (read by bind_aux_taps), and the stub draft model `kv_pool_spec` and a `write_context_kv` taking the slots, with the target's `slot_mapping` on the forward context, as compute_draft_kv reads them. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Summary
This PR adds a high-interactivity monokernel decode path for DeepSeek-V4.1-Flash on MI355X at TP4 with DSpark
speculative decoding. On a decode step of 1-8 requests, every decoder layer runs as one persistent FlyDSL kernel, the
verify head's ids and verdict come out of the decode graph, and the next step is laid out from the device-resident
verdict: 51 target + 8 draft kernels per step instead of 129 + 41.
cache with the FP4 index cache, DSpark with 5 draft tokens, DeepSeek-V4.1-Flash geometry, no KV transfer, context
up to 1Mi tokens. Any other deployment keeps ATOM's existing V4.1 mono decode,
unchanged, with one info log line saying why (
V4.1 mono hi off: <reason>; ATOM's mono path).ATOM_DSV41_MONO_HI=0turns it off;ATOM_DSV41_MONO_HI=1forces it and raises if the deployment is not served.atom/models/deepseek_v41/mono/hi/(the same shape as the Qwen3.8 mono kernels)and share the device primitives in
atom/mono/. The engine integration is two generic runner switches, an in-graph headoutput (a model's
compute_head, generalisinglogits_in_graphto return logits or a verify's ids and verdict)and a builder capability that defers the spec-status wait (
consumes_device_verdict), plus one per-step routingdecision (
AttentionMetaData.served_by_round) inside the V4.1 builder. With the path off nothing changes.in ATOM's paged pool (every store masked by its row's liveness, so padded requests and dead slots write nothing),
the routed experts, their E8M0 scales and the RoPE caches are read from ATOM's tensors, and the token embedding is
ATOM's own. The kernels compile at server start, as ATOM's V4.1 mono kernels do, and reuse FlyDSL's disk cache on later starts.
Kernel speedup
MI355X, TP4 on GPUs 0-3 of an 8-GPU node, ATOM nightly 202611091503, c1, DSpark 5 draft tokens, forced acceptance 3.51,
both columns from this branch (ATOM main's path =
ATOM_DSV41_MONO_HI=0). GPU timeper decode step from ATOM's own torch profiler (
decode[...]range = target verify forward,propose_dspark[...]range= draft step), rank 0, ~680 steps, median (P90):
On the same servers, a 300 s AgentX c1 window gives X 561.7 vs 755.7 (1.35x) and median ITL 1.470 vs 1.090 ms.
End to end: AgentX, c1
InferenceX's AgentX scenario through AIPerf (c1, 1,200 s window after a 600 s warm pass, forced acceptance 3.51;
X = 1000 / P90 inter-token latency), measured on the first revision of this PR (nightly 202611060646):
Accuracy: GSM8K
InferenceX's GSM8K eval (1,319 questions, 5-shot, chat template, greedy), c4 server, real (not forced) acceptance:
Paired over questions (McNemar, exact): main vs first revision 4 vs 6 right-only, p = 0.754; current head vs the
previous head 5 vs 8, p = 0.58 (no significant difference).
Memory and startup
--gpu-memory-utilization 0.9: 148.4 GB available on main, 145.8 GB with this path on. Thepath reads and writes ATOM's own KV pool, FP4 index planes and compressor state in place, and reads the routed
experts, their scales and the RoPE caches from ATOM's tensors; what it adds is its packed dense
weights, verify head, peer buffers and scratch.
How to reproduce
Image: the ATOM nightly built from this branch. Checkpoint:
deepseek-ai/DeepSeek-V4.1-Flash-1017(
1bae0e1088993b9b82de08f59de5be5a34e42170). No extra environment is needed: the path turns on by itself in thisdeployment. For the baseline, set
ATOM_DSV41_MONO_HI=0.Server (the InferenceX
dsv41flash/atom/mi355x-fp4-mtp/agentic.yamlarguments at TP4):On start, each rank logs
V4.1 mono hi: 40 layers on the hi path (TP4 rank r) ....Kernel timing: add
--spec-decode-acceptance-length 3.51 --torch-profiler-dir <dir>to the server. Then:POST /start_profile.ignore_eos).POST /stop_profile.In rank 0's trace, per step, take the span of the GPU kernels launched inside each
decode[bs=1 ...]andpropose_dspark[...]range (matched by correlation id; a graph launch carries its kernels). Report the median andP90 over the steps.
GSM8K: run the server without
--spec-decode-acceptance-length, then InferenceX's GSM8K task through lm-eval:Tests
CPU tests (need FlyDSL 0.4.2; the aiter-dependent ones skip without aiter):
python -m pytest tests/test_deepseek_v41_mono_hi_*.py tests/models/deepseek_v41/test_mono_hi_bridge.pytest_deepseek_v41_mono_hi_gate.py: the on/off decision on the realConfigbuilt from the served command line (on),and every refused deployment as a one-field change of it (off, or an error with
ATOM_DSV41_MONO_HI=1).test_deepseek_v41_mono_hi_contract.py: the served programs' contract and memory-model checks (CPU compile-only).test_deepseek_v41_mono_hi_weights.py: the weights packed from ATOM's loaded tensors are byte-equal to a checkpoint pack.test_deepseek_v41_mono_hi_runner.py,_host.py,_routing.py,_served_step.py: the in-graph head output, thedeferred status wait, the per-step routing decision (a short prefill or a step with image rows is never routed), that the next ids are the TP-broadcast
ones, and that both switches are inert when the path is off.
test_mono_hi_bridge.py: the cache bridge kernels against a reference.test_deepseek_v41_mono_hi_package.py,_rope.py,_bands.py,_band_plan.py,_aux_tap_key.py,_experts_alias.py,_scratch_uncached.py: package imports, RoPE tables, batch bands, aliasing and scratch.On GPU (MI355X, TP4), validated in this deployment:
ATOM_DSV41_MONO_HI=0and an unsupported deployment (--kv_cache_dtype fp8) take ATOM's mono path.