Skip to content

feat(dsv41 mono): serve the widths that spill by dropping K2's index stage, not the width - #3221

Open
zhiding512 wants to merge 7 commits into
mainfrom
zhiding512/dsv41-mono-intree
Open

zhiding512 wants to merge 7 commits into
mainfrom
zhiding512/dsv41-mono-intree

Conversation

@zhiding512

@zhiding512 zhiding512 commented Nov 11, 2026 •

Copy link
Copy Markdown

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 lean seam WidthBuilds already has.

Why they are declined. WidthBuilds._scratch_free refuses a width if any kernel needs a byte of scratch — a
persistent 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:

build scratch a lane vgpr
K2, fused indexer (index_bound_max > 0), MFMA FP4 index layout 8 B 256
K2, fused indexer, non-MFMA FP4 layout 0 255
K2, index_bound_max = 0 (indexer outside) 0 250
K1, every variant 0 <= 242

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 x tok_par on/off, 120
builds: the nine that need scratch are all fused (index_bound_max > 0); every unfused build compiles with
zero scratch (probe/unfused_confirm.json). That is the whole of what this change relies on.

tok_par makes no difference and K1 never spills: it is the MFMA FP4 indexer stage, one VGPR over the 255 that
fit. 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 lean callback, as kimi_k3 already does. When a width's K2 spills, that
width alone is rebuilt with index_bound_max = 0 and tok_par = False, and ATOM's own indexer runs between K1
and K2, publishing the selection into the irow / icand peer-buffer regions K2a reads. Widths that fit keep
today's fused build, which is faster. Nothing else changes: ATOM_MONO_ENABLE=0, the hi path, gfx1250, the
DSpark 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 a
silent fall-back, so an A/B cannot measure the original path by accident.

Related

PR What State
coderfeli/plugin_poc#67 The same kernels driven unfused, as an out-of-tree AITER provider merged (11d9076)
ROCm/aiter#7809 The two routed decode-layer contracts the provider serves draft
#3215 The ATOM consumer that calls those routed ops (ATOM_DSV41_MONO_BACKEND=aiter) draft

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.

point arm result
synthetic 1K/1K, c=8, 6 reps provider vs ATOM_MONO_ENABLE=0 ITL 12.842 -> 11.933 ms (+7.1%), Welch t = 10.7
same in-tree k1k2 vs the same 12.468 ms (+2.9%) — it declines its 48-row build
AgentX tp4 c8, 3600 s, fp8 KV, FULL graphs provider vs ATOM_MONO_ENABLE=0 ITL 2.860 -> 1.968 ms (-31.2%), interactivity 314.07 -> 409.76 (+30.5%), tok/s/GPU 14,992 -> 16,367 (+9.2%), TTFT p50 0.369 -> 0.182 s

Measured 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=1 on the
mono arm so a declined width would have been a startup error rather than a silent fall-back.

arm decode step (mean ITL, ms) per-run ITL TPOT mean (ms) out tok/s accept GSM8K
ATOM_MONO_ENABLE=0 12.190 12.38 / 12.01 / 12.18 2.623 2717 4.61 0.920 +/- 0.019
this branch (default) 11.257 (+7.66%) 11.40 / 11.14 / 11.23 2.320 (+11.6%) 3243 4.70 0.945 +/- 0.016

The 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:

V4.1 mono: atom::v41_layer_post_s48_e384_k6_tl0_ix1_i41_im1_tpar1 need scratch, rebuilt lean

That is the 48-row fused, token-parallel K2 (ix1, tpar1) asking for scratch, and the lean seam rebuilding
the 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:

run strict intree ITL no-mono ITL gain
1 no 11.303 12.127 +6.79%
2 yes 11.123 (arm not run) —
3 (the table above) yes 11.257 12.190 +7.66%

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-dsv41 knobs (fp8 KV, prefix cache, level 3 FULL graphs, fixed
acceptance length 3.51), one 3600 s run an arm, both points valid with zero errors.

arm interactivity E2E-norm tok/s/GPU TTFT p50 / p90 (s) requests
ATOM_MONO_ENABLE=0 331.82 98.81 15,344 0.32 / 1.88 1389
this branch (default) 498.16 (+50.1%) 213.82 17,357 (+13.1%) 0.20 / 0.94 1465

The fallback was active here, at four widths. That run's log carries:

v41_layer_post_s30_e384_k6_tl0_ix1_i41_im1_tpar0 need scratch, rebuilt lean
v41_layer_post_s36_e384_k6_tl0_ix1_i41_im1_tpar1 need scratch, rebuilt lean
v41_layer_post_s42_e384_k6_tl0_ix1_i41_im1_tpar0 need scratch, rebuilt lean
v41_layer_post_s48_e384_k6_tl0_ix1_i41_im1_tpar1 need scratch, rebuilt lean

and the target widths it went on to build were s6/s12/s18/s24 ix1 (fused, no spill) beside s30/s36/s42/s48
rebuilt 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

  • CPU, on this branch at 47e0c7075: the four mono suites give 173 passed (19 of them new). The whole
    tests/ tree has the same failure set as base main in this GPU-less environment (199 failed / 108 errors
    on both, all pre-existing); the branch passes 24 more, which are the new tests. black clean; repo-wide
    ruff is 200 findings on base and the same 200 on this branch.
  • Compile evidence, gfx950 compile-only, 120 builds over rows 24/32/48 x bf16/fp8 KV x tok_par on/off:
    the nine builds that need scratch are all fused; every unfused build compiles with zero scratch.
  • The unfused K2's scratch layout fits the buffer sized for the fused build at MAX_ROWS — it needs about
    15% 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.
  • GPU correctness: GSM8K 3-shot 200q at c=8 — 0.945 on this branch against 0.920 with mono off, within
    about one standard error, with strict on.
  • GPU perf: the table above — 3 reps an arm in one session, strict, 48-row width confirmed served from the
    server log, and the point reproduced across three runs.
  • AgentX tp4 c8: the table above, both points valid, with the fallback confirmed active at widths 30, 36,
    42 and 48. Caveat as stated: it compares mono against no mono on this branch, not against main.

Limits

  • The fallback keeps the measured provider configuration (tok_par=False). Token-parallel rows with an
    external indexer are not covered.
  • A narrower fix would save one VGPR in the MFMA FP4 indexer stage, after which the fused build — which is
    faster — fits at every width. That is kernel work and is not attempted here.
  • Only gfx950 TP4 c=8 was measured, with a bf16 KV cache and the FP4 (opus/MFMA) index layout. TP2 and TP8
    never spill in the compile probe, so they never take this path; they were not run.
  • write_selection_rows is a torch chain, where the rest of step_meta.py is Triton precisely to avoid "~19
    elementwise 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

zhiding512 and others added 7 commits November 11, 2026 17:42
``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>
@github-actions

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 3221 --add-label <label>

@zhiding512
zhiding512 marked this pull request as ready for review November 11, 2026 19:17
Copilot AI balanced review requested due to automatic review settings November 11, 2026 19:17

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🟡 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

Comment on lines +669 to +672
query, weights = indexer.project(
self.normed, self.qr, self.qr_scale, cache, step, rope
)
indexer.score(query, weights, cache, step)
Comment on lines +291 to +294
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 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..

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,

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

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.

Comment thread atom/utils/envs.py
# 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",

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

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")

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

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(

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

narrow widths (24/32) only compile-probed, _use_rows narrowing untested on gpu. run one narrow width?

@valarLip valarLip left a comment •

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

⚠️ Action required before merge (5)

  1. Remove tests/conftest.py:21-40. At import time it rewrites aiter.dtypes.fp8 for 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.
  2. Remove ATOM_DSV41_MONO_STRICT (envs.py:633-641, builds.strict / builds.refuse, the branch in runner.prepare, and its two tests). It's a measurement-only knob with one user; WidthBuilds already logs every refused width.
  3. Guard the test module. tests/test_deepseek_v41_mono_unfused.py:17 only importorskips aiter but imports flydsl kernel modules and the triton-based step_meta. Guard it like the sibling suites (valar-cop), and move the pure builds.py / write_selection_rows tests into a CI-runnable module (Copilot).
  4. State the measurement setup (valar-cop): the ATOM_DSV41_MONO_HI setting 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.
  5. Perf follow-ups (Copilot; fix here or list them in the PR): the unfused K1 still projects iq/iw that _score discards before calling indexer.project again (runner.py:666-672); write_selection_rows allocates torch.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.

@valar-cop

Copy link
Copy Markdown

This branch has not been deployed

No deployments
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.

5 participants