Repository navigation
feat(dsv41 mono): serve the widths that spill by dropping K2's index stage, not the width - #3221
zhiding512 wants to merge 7 commits into
Conversation
``RowsBuild`` moves to its own module and carries the K2 build key its launchers came from. The forward's token-parallel row split and row gather, the debug wait report's region names and check mode's MoE regions now read ``tok_par`` and the K2 scratch layout off that key instead of deriving them again from the step's row count. No behaviour change: ``_post_key(s)`` is exactly what each of those call sites recomputed, and ``_tok_par(s)`` is what it was built from. That stops being true in the next commit, where a width can be rebuilt with token parallel off while its row count still says on. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
At TP 4 a 32- or 48-row K2 built with the fused indexer on the MFMA FP4 index layout comes out at 256 VGPRs and needs 8 B of scratch a lane; a persistent kernel whose waves need scratch cannot be dispatched whole, so ``WidthBuilds`` refuses the width and its steps take the original model's path. With DSpark's 48-row verify that is mono off at TP 4 c=8. A compile probe over rows 24/32/48, both KV dtypes and both token-par settings says K1 never spills, only K2 does, and only with the index stage in it: ``index_bound_max`` 0 is scratch-free (236-250 VGPRs) at every width measured. So give the width a leaner form through ``WidthBuilds``' existing ``lean`` seam, which Kimi-K3 already uses: when a K2 (or, in check mode, a K2a) launcher spills, rebuild the whole width with ``index_bound_max`` 0 -- the shape the DSpark draft already builds -- and token parallel off, and run ATOM's own indexer between K1 and K2 for that width, publishing its selection as pool rows into ``irow`` (and a producer's blocks into ``icand``) exactly where the fused K2 would have written them itself. Widths that do not spill keep today's fused build, bit for bit. ``write_selection_rows`` is ``index_score._selected_rows``' arithmetic on the host: a column's pool row through its request's PAGE table, valid entries packed first because the attention takes their count off the bound (``index_plan.selected``) rather than off the array. This is the configuration already measured end to end driving these same kernels from the out-of-tree AITER provider (+6.1/+7.1% ITL at TP4 c=8 1K/1K, AgentX tp4 c8 ITL -31.2%); nothing here is a new kernel. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
A width mono will not serve is a silent fallback: the step runs the original model's path and nothing downstream says so, which is exactly what an A/B of mono against the original path then measures. Off by default, so production behaviour is unchanged; on, ``prepare`` raises with the width, after the warning that names why it was refused. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…tion ``unfused_keys`` turns a width's (K1 keys, K2 key) into the unfused pair, so the rule that K1 loses token parallel with K2 lives beside the rule that K2 loses its index stage -- and a test can read both off real build keys instead of off the runner, which needs a GPU to import. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…switch - ``test_deepseek_v41_mono_unfused``: the lean choice over a build record (a K2 or K2a spill rebuilds; a K1 or K2b spill, an already-unfused or indexerless build and a build with no key do not; nothing spilled is not a rebuild), the unfused keys (K2's index stage out, token parallel off on K2 and on every K1, nothing else moved), the selection mapping against a hand-worked case and against ``index_score._selected_rows``' arithmetic element by element, and ``ATOM_DSV41_MONO_STRICT``. - ``test_deepseek_v41_mono_contract``: the unfused K1 and K2 of the 48-row TP4 width hold their hand-off contracts, traced. - ``test_deepseek_v41_mono_build_key``: the unfused K2 is its own binary, so the lean rebuild cannot be served the fused one it replaces. - ``test_mono_runtime``: a width that does not spill is never asked for a lean build, and the "nothing leaner" test says whose answer that is -- a K2 spill is no longer universally unrebuildable. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
``lean`` returns a build of the whole width, not the spilled launchers with leaner ones dropped in, and what counts as leaner is each runner's own answer -- Kimi-K3 restages a K1's GEMV, V4.1's decode drops K2's index stage. The old wording described only the first of those. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Standalone and droppable: it changes no production code and nothing in this branch depends on it where a GPU exists. AITER picks ``dtypes.fp8`` off the detected arch and falls back to ``torch.uint8`` with no device to ask; ``torch.finfo`` of that raises at import of ``model_ops.blockscale_kernels.quantization``, so on a machine that has AITER but no GPU every module reaching it fails to COLLECT -- the V4.1 mono kernel suites among them, which then read as "no tests" rather than as an environment that cannot answer. CLAUDE.md says the suite needs no GPU; this is part of what that costs. Guarded on the fallback value itself, so a machine whose AITER did read an arch (gfx950 OCP E4M3, gfx942 FNUZ) is untouched, and CI, which has no AITER at all, ``importorskip``s those modules before reaching them. Note the trade on a GPU-less box: suites that used to ERROR at collection now import and then FAIL for want of a device. Neither passes; the second at least names what it wanted. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
There was a problem hiding this comment.
🟡 Changes recommended
The fallback introduces redundant projections and critical-path tensor operations, while dependency gating prevents the build-decision and strict-mode tests from running in non-GPU CI.
3 open issues
What changed in this PR
This PR keeps wider DeepSeek-V4.1 decode steps on ATOM’s in-tree monokernel path by rebuilding spilling widths without K2’s fused indexer, rather than falling back to the original model path.
Changes:
- Adds per-width unfused rebuilds and runs ATOM’s indexer between K1 and K2 while preserving fused builds that fit.
- Adds optional strict-mode errors for refused builds and uses the actual build configuration in runtime and diagnostic paths.
- Adds build-policy, selection-mapping, cache-key, and contract tests, plus CPU-only AITER dtype setup.
| File | Description |
|---|---|
tests/test_mono_runtime.py |
Tests that scratch-free widths retain their original build. |
tests/test_deepseek_v41_mono_unfused.py |
Tests unfused keys, rebuild decisions, selection mapping, and strict mode. |
tests/test_deepseek_v41_mono_contract.py |
Checks unfused K1/K2 hand-off contracts. |
tests/test_deepseek_v41_mono_build_key.py |
Checks that unfused K2 gets a distinct cached binary. |
tests/conftest.py |
Adjusts AITER’s fallback FP8 dtype for CPU-only collection. |
atom/utils/envs.py |
Registers ATOM_DSV41_MONO_STRICT. |
atom/mono/runtime/widths.py |
Clarifies the existing lean-rebuild contract in documentation. |
atom/models/deepseek_v41/mono/step_meta.py |
Adds selection compaction and pool-row translation. |
atom/models/deepseek_v41/mono/runner.py |
Integrates lean rebuilds, external scoring, and selection publication. |
atom/models/deepseek_v41/mono/check.py |
Uses the actual K2 build key for scratch-layout checks. |
atom/models/deepseek_v41/mono/builds.py |
Defines build records, unfused variants, rebuild policy, and strict refusal. |
Review effort: Balanced
| query, weights = indexer.project( | ||
| self.normed, self.qr, self.qr_scale, cache, step, rope | ||
| ) | ||
| indexer.score(query, weights, cache, step) |
| dest = torch.where(valid, valid.to(torch.int64).cumsum(1) - 1, k) | ||
| packed = torch.zeros(s, k + 1, dtype=torch.int64, device=sel.device) | ||
| packed.scatter_(1, dest, torch.where(valid, rows, 0)) | ||
| out.copy_(packed[:, :k]) |
| import pytest | ||
| import torch | ||
|
|
||
| pytest.importorskip("aiter") |
valar-cop
left a comment
There was a problem hiding this comment.
some changes needed..
strict-mode rerun of the tp4 c=8 point, and an in-tree number where hi refuses (tp4 fp8 kv)?
ci Run non-GPU unit tests failed — not sure yet (low): The job's log file (.valar/ci/job_130749758471.log) has 0 lines, so neither a PR cause nor an infra cause can be confirmed; the same job passes on main and in 5 of 6 other recent runs.
next: Re-run the job or attach its log; the new mono tests only importorskip in this job, so a PR-side cause is unlikely but unproven. https://github.com/ROCm/ATOM/actions/runs/43085809325/job/130749758471
| self._row_kernels, | ||
| group, | ||
| "V4.1 mono", | ||
| lean=self._lean_width, |
There was a problem hiding this comment.
this point is served by hi by default, so say the hi setting in the measurement? and add a number where hi refuses, e.g. TP4 fp8 kv.
| # deployment that loses one width still serves the rest -- but an A/B that | ||
| # means to measure mono measures the fallback instead and cannot tell from | ||
| # its numbers, so turn it on there. | ||
| "ATOM_DSV41_MONO_STRICT": lambda: os.getenv("ATOM_DSV41_MONO_STRICT", "0") == "1", |
There was a problem hiding this comment.
missing a row in docs/environment_variables.md for ATOM_DSV41_MONO_STRICT? all the other ATOM_DSV41_* have one.
| import pytest | ||
| import torch | ||
|
|
||
| pytest.importorskip("aiter") |
There was a problem hiding this comment.
only importorskip aiter, but this imports flydsl and triton too.. guard like the sibling tests.
| return | ||
| published.add(owner) | ||
| geometry = cache.geometry | ||
| write_selection_rows( |
There was a problem hiding this comment.
narrow widths (24/32) only compile-probed, _use_rows narrowing untested on gpu. run one narrow width?
There was a problem hiding this comment.
⚠️ Action required before merge (5)
- Remove
tests/conftest.py:21-40. At import time it rewritesaiter.dtypes.fp8for the whole suite, only to serve a GPU-less dev box with aiter installed. It's unrelated to this change; drop it or send it separately. - Remove
ATOM_DSV41_MONO_STRICT(envs.py:633-641,builds.strict/builds.refuse, the branch inrunner.prepare, and its two tests). It's a measurement-only knob with one user;WidthBuildsalready logs every refused width. - Guard the test module.
tests/test_deepseek_v41_mono_unfused.py:17only importorskips aiter but imports flydsl kernel modules and the triton-basedstep_meta. Guard it like the sibling suites (valar-cop), and move the purebuilds.py/write_selection_rowstests into a CI-runnable module (Copilot). - State the measurement setup (valar-cop): the
ATOM_DSV41_MONO_HIsetting used, plus one in-tree number at a point hi refuses (TP4 fp8 KV). The AgentX table compares mono vs no-mono on this branch, not this branch vs main. - Perf follow-ups (Copilot; fix here or list them in the PR): the unfused K1 still projects
iq/iwthat_scorediscards before callingindexer.projectagain (runner.py:666-672);write_selection_rowsallocatestorch.zeros(s, k+1)per selection owner on the K1->K2 critical path (step_meta.py:292).
Reviewed at 29ef52f. The width-fallback design itself is sound.
|
ci triage for 29ef52f:
|

Summary
At TP4 the V4.1 decode monokernel declines its widest step widths, so the steps that matter most at
concurrency run ATOM's original decode path instead. This teaches the in-tree runner to rebuild those widths in
a form that fits, through the
leanseamWidthBuildsalready has.Why they are declined.
WidthBuilds._scratch_freerefuses a width if any kernel needs a byte of scratch — apersistent kernel needs every CTA resident, and one whose waves need scratch can be dispatched only in part.
Compiling every in-tree build for gfx950 compile-only shows exactly one offender at TP4:
index_bound_max > 0), MFMA FP4 index layoutindex_bound_max = 0(indexer outside)The fused K2a (check mode's split half) sits at exactly 256 VGPRs from 24 rows up and spills at 48, so the
build is one register short across the whole upper width range rather than at one width. Unfused it drops to
205-249. Evidence:
probe/k1k2_scratch.json,probe/k2_sens.json,probe/k2a_scratch.json.Re-run against this branch's base (
47e0c7075) over rows 24/32/48 x bf16/fp8 KV xtok_paron/off, 120builds: the nine that need scratch are all fused (
index_bound_max > 0); every unfused build compiles withzero scratch (
probe/unfused_confirm.json). That is the whole of what this change relies on.tok_parmakes no difference and K1 never spills: it is the MFMA FP4 indexer stage, one VGPR over the 255 thatfit. The probe sampled 24/32/48 and found 32 and 48 spilling, and 24 too with a bf16 KV cache on current main
(#3128 widened the indexer's block-maxima grid). The real set is wider than that sample: in the AgentX run
below, which exercises many request counts, four widths spilled and were rebuilt — 30, 36, 42 and 48 (30
and 42 are not multiples of the TP size, so they are not even token-parallel). With DSpark-5, concurrency 8 is
a 48-row step, so today mono is off exactly there, and off across that whole upper band whenever the traffic
reaches it.
The change. dsv41 now passes a
leancallback, askimi_k3already does. When a width's K2 spills, thatwidth alone is rebuilt with
index_bound_max = 0andtok_par = False, and ATOM's own indexer runs between K1and K2, publishing the selection into the
irow/icandpeer-buffer regions K2a reads. Widths that fit keeptoday's fused build, which is faster. Nothing else changes:
ATOM_MONO_ENABLE=0, the hi path, gfx1250, theDSpark draft runner and every other model are untouched.
ATOM_DSV41_MONO_STRICT(default off) turns a width still refused after the rebuild into an error instead of asilent fall-back, so an A/B cannot measure the original path by accident.
Related
11d9076)ATOM_DSV41_MONO_BACKEND=aiter)This PR supersedes the plugin route in #3215 for shipping the win: it needs no AITER plugin mechanism, no
routed ops and no out-of-tree provider. #3215 stays open as the reference implementation and as the AITER
contract's consumer; it is not closed by this PR.
The provider measurements below are what motivated this change. They were taken with the plugin form, whose
kernels are byte-identical to these; only the driving differs.
Measured as a provider (the plugin form, for motivation)
DeepSeek-V4.1-Flash, gfx950 (MI355X), TP4, DSpark-5.
ATOM_MONO_ENABLE=0k1k2vs the sameATOM_MONO_ENABLE=0Measured on this branch: TP4, c=8 — the in-tree fallback beats no monokernel
DeepSeek-V4.1-Flash, gfx950 (MI355X), TP4, DSpark-5, ISL/OSL 1024/1024, concurrency 8, 64 prompts after 16
warm-ups, 3 reps an arm, both arms in one session on the same four GPUs,
ATOM_DSV41_MONO_STRICT=1on themono arm so a declined width would have been a startup error rather than a silent fall-back.
ATOM_MONO_ENABLE=0The six runs' ranges are disjoint (branch 11.14-11.40, no-mono 12.01-12.38). GSM8K is within about one standard
error, so accuracy is a pass rather than a difference. ITL rather than TPOT is the honest metric with DSpark,
because TPOT is per output token and absorbs the acceptance length.
The fallback is what did it, and the mono arm really was on mono — from that run's own server log:
That is the 48-row fused, token-parallel K2 (
ix1,tpar1) asking for scratch, and the lean seam rebuildingthe width. The run logged no "declined the 48-row build" line, and the widths actually built were
s48 ix0(the unfused rebuild) alongside
s6/s12/s24 ix1— the lower widths kept the fused build, exactly as intended.Note 24 rows stays fused here because the runner builds it token-parallel, which the compile probe also shows
not spilling; only 48 needed the rebuild, and 48 is the c=8 bucket.
The point was measured three times over the evening, twice with strict on, and the branch lands in the same
place each time:
The gain lands where the provider's did (+7.1% ITL on the same point), which is the expected result: same
kernels, same build choice, driven in-tree instead of through a plugin.
AgentX tp4 c8 (InferenceX agentic replay, 3600 s a point)
Same box and branch, the
atom-agentx-dsv41knobs (fp8 KV, prefix cache, level 3 FULL graphs, fixedacceptance length 3.51), one 3600 s run an arm, both points
validwith zero errors.ATOM_MONO_ENABLE=0The fallback was active here, at four widths. That run's log carries:
and the target widths it went on to build were
s6/s12/s18/s24 ix1(fused, no spill) besides30/s36/s42/s48rebuilt unfused. On main those four would have declined and their steps would have run the original path.
Read the number carefully even so. Both arms are this branch, so what is measured is the monokernel
against no monokernel. It is not this change against main: main would still have served every width below 30.
A main-with-mono arm would isolate the change exactly, and I did not run one. The claim this PR rests on is
the 1K/1K table above, where the arm ran strict and the 48-row rebuild is in the log.
Test plan
47e0c7075: the four mono suites give 173 passed (19 of them new). The wholetests/tree has the same failure set as base main in this GPU-less environment (199 failed / 108 errorson both, all pre-existing); the branch passes 24 more, which are the new tests.
blackclean; repo-wideruffis 200 findings on base and the same 200 on this branch.tok_paron/off:the nine builds that need scratch are all fused; every unfused build compiles with zero scratch.
MAX_ROWS— it needs about15% of it at every width (48 rows: 15,297,280 B against 103,377,664 B), so the existing allocation covers
the fallback with room to spare.
about one standard error, with strict on.
server log, and the point reproduced across three runs.
42 and 48. Caveat as stated: it compares mono against no mono on this branch, not against main.
Limits
tok_par=False). Token-parallel rows with anexternal indexer are not covered.
faster — fits at every width. That is kernel work and is not attempted here.
opus/MFMA) index layout. TP2 and TP8never spill in the compile probe, so they never take this path; they were not run.
write_selection_rowsis a torch chain, where the rest ofstep_meta.pyis Triton precisely to avoid "~19elementwise kernels a step". It runs once per selection owner per step on the unfused path, so it is a
handful of small launches on the critical path. The configuration measured as a provider carried the same
cost and still won, but writing it as a Triton kernel beside its neighbours is the obvious follow-up. It is
capture-safe as it stands: no host sync, static shapes.
🤖 Generated with Claude Code