Repository navigation
feat(mono): shared DSA indexer mono launch; Kimi-K3 mono to 64 rows, plain decode and ReplaySSM - #2871
Conversation
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
valar-cop
left a comment
valar-cop
left a comment
There was a problem hiding this comment.
some changes needed..
ci Run non-GPU unit tests failed — not sure yet (low): The job's log provided in the context is 0 bytes, so neither the failing test nor the failing step can be read. The same job succeeds on the 8 most recent main runs while one other recent run also failed, so a PR-caused failure and pre-existing flakiness are both possible; nothing can be attributed.
.valar/ci/job_124584532524.log is 0 bytes
same job on main: 8x success; on other recent runs: {'success': 5, 'cancelled': 1, 'failure': 1}
next: Re-run the job or fetch its log and identify the failing test; the PR description already reports 42 CPUs failures on clean main, so check whether this run hit one of those before treating it as caused by this PR. https://github.com/ROCm/ATOM/actions/runs/41217341026/job/124584532524
| """A loader thread's own OpenMP team: one thread (``omp_set_num_threads`` | ||
| sets the calling thread's count only; torch's OpenMP runtime is loaded, | ||
| so the process resolves it).""" | ||
| ctypes.CDLL(None).omp_set_num_threads(1) |
There was a problem hiding this comment.
if omp_set_num_threads is missing the worker dies and future.result() hangs forever, guard it?
| def weight_loader(self, param, loaded_weight, loaded_shard_id=None): | ||
| if param is not self.weight or loaded_shard_id not in (0, 1, 2, 3): | ||
| raise ValueError("Unexpected fused indexer projection checkpoint slice") | ||
| if loaded_weight.dtype != torch.bfloat16 or param.dtype != torch.bfloat16: |
There was a problem hiding this comment.
fp8 or fp32 indexer.wk now fails to load here, fall back to the unfused path instead?
| if param is not self.weight or loaded_shard_id not in (0, 1, 2, 3): | ||
| raise ValueError("Unexpected fused indexer projection checkpoint slice") | ||
| if loaded_weight.dtype != torch.bfloat16 or param.dtype != torch.bfloat16: | ||
| raise ValueError("Fused indexer projection requires original BF16 weights") |
There was a problem hiding this comment.
quant config is the only signal here? if cfg says No but tensor is fp8 we raise mid-load.. can we decide before allocating the module?
| """A loader thread's own OpenMP team: one thread (``omp_set_num_threads`` | ||
| sets the calling thread's count only; torch's OpenMP runtime is loaded, | ||
| so the process resolves it).""" | ||
| ctypes.CDLL(None).omp_set_num_threads(1) |
There was a problem hiding this comment.
no omp symbol -> AttributeError and the whole load fails over an optional opt.. catch it and no-op?
8441e34 to
33162a4
Compare
valar-cop
left a comment
valar-cop
left a comment
There was a problem hiding this comment.
a few small things..
kimi-k3 gsm8k or lm_eval for the changed decode and ReplaySSM paths?
ci ATOM Test failed — not sure yet (low): The only failed job in the summary; no job log or job URL was provided under .valar/ci (only summary.md), so neither the failing test nor the failing step can be read and the failure cannot be attributed to this PR or the CI.
- ATOM Test: failure https://github.com/ROCm/ATOM/actions/runs/41316310797
.valar/ci/ contains no log for this run
next: Fetch the job log (or re-run) and identify the failing test; compare it against the description's list of pre-existing CPU failures on clean main before treating it as PR-caused. https://github.com/ROCm/ATOM/actions/runs/41316310797
| """A loader thread's own OpenMP team: one thread (``omp_set_num_threads`` | ||
| sets the calling thread's count only; torch's OpenMP runtime is loaded, | ||
| so the process resolves it).""" | ||
| ctypes.CDLL(None).omp_set_num_threads(1) |
There was a problem hiding this comment.
still open from before: if omp_set_num_threads is missing the initializer raises and every weight copy fails. wrap in try/except and no-op?
| return f"wq_b {shape.qrows} rows: the q GEMV is one 16-row task a CTA" | ||
| if shape.block_rows % TILE: | ||
| return f"{shape.block_rows}-row blocks: a score wave reads 16 keys of a block" | ||
| if shape.topk % 4: |
There was a problem hiding this comment.
this check also refuses DCP/TP-split where _copy_winners is never reached. move it into unfit() gated on key.replicated?
…plain decode and ReplaySSM Shared DSA indexer (atom/mono/dsa_indexer/): one FlyDSL launch per DSA layer at decode for wq_b, the index-cache writer (NoPE or the fused RoPE writer), FP8 scoring and the exact top-K selection to the MLA's KV slots, with a distributed radix selection over the TP ranks under DCP or without it (replicated cache). Models are routed through one opaque op by a wrapper, refused before routing when the kernel cannot serve them. GLM-5.2 is registered; deepseek_v2 folds the indexer's wk / weights_proj into fused_qkv_a_proj when the attention is not quantized. Kimi-K3 mono: K1 / K2 up to 64 rows (16-row token groups), plain decode and uniform verify, with or without ReplaySSM; the bf16 GEMV moves to atom/mono/device/gemv.py. Shared runtime: uncached step scratch and peer buffers, lean rebuild of a spilled width, shared device helpers. Also: get_dcp_group reuses the TP coordinator when the groups match, one-thread OpenMP teams for loader threads, quant_spec.effective_layer_quant, a dense DSpark draft owns KV rows but no index-key slices.
33162a4 to
69a741b
Compare
Conflict in atom/models/kimi_k3/mono/plan.py (#2871 moved K3 mono to plain decode and ReplaySSM): take main's ReplaySSM support and keep this branch's state-fork refusal, now ahead of the spec/plain split so a forked plain decode is refused too -- this branch can put a fork on a decode, which #2871's state_slots docstring assumed never happens. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Brings #2365's review fixes since a128906 (6808659, acdf74d, f34cb8a, a0b8ef0) and its merge of main 99ebf63 (#2871 among it). f34cb8a's `_defer_prefill_hashes` takes over both no-output branches of the PP head loop, the ack branch included. No textual conflicts. f34cb8a's test bounds the head's in-flight window through `pp_size`, which this branch replaced with `pp_max_inflight` (ATOM_PP_MAX_INFLIGHT); the test now sets both. Co-authored-by: Cursor <cursoragent@cursor.com>
Follows #2709 onto #2365's head a0b8ef0, which merged main 99ebf63. No conflicts: none of the files this branch changes moved. The mono log lines the C1 cell checks are still printed after #2871 (`route.py`'s "%s on" / "%s off" / "declined the %d-row build"). Co-authored-by: Cursor <cursoragent@cursor.com>
ATOM 974414f16 carries ROCm/ATOM#2871 (Kimi-K3 mono up to 64-row steps), so concurrency 4 now runs mono too: max-num-seqs 8 (8 x 8 rows at DSpark k=7), fp32 SSM state and the C1 mono exclude list. C1 keeps the PrefillDelayer. Recipe frontend pinned to the tag's digest. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
… to replace main's review-pr skill (#2871) and this branch descend from the same July ancestor and were edited in parallel, so SKILL.md conflicted. Nothing of #2871 is discarded or left orphaned here: - `references/maintainer-lens.md` is that SKILL.md, verbatim, moved to a reference path when the bot pipeline took SKILL.md. Its section 3 is the maintainer's lens -- the same ten themes this branch's N1-N6 rules were distilled from, which makes it the upstream of the bot's judgement, not a competing design. A human reviewing a PR by hand reads that file. - `references/rule-catalog.md` and `scripts/affected_tests.py` merged cleanly and are now cited by SKILL.md and rules.md rather than sitting unreferenced. rules.md records the descent and both deliberate divergences, so neither reads as drift: 22 rules here that the catalog does not carry (the H/I/T families, the N family from the lens, and eight extensions), and P1-P5, which were removed here on evidence -- no ATOM citation for any of the five, a targeted search that found no instance, and the only ids claiming "always check" while being UNREACHABLE_BY_DESIGN. Re-importing them from the catalog would undo a decision that had a search behind it, so rules.md says to re-add one only with an ATOM citation. scripts/affected_tests.py is wired into validation triage, which is a real capability this branch lacked. Target selection considered only tests the PR touches, so a PR changing runtime code without shipping a test got "ships no test target" and nothing was validated. affected_tests computes the tests that import a changed atom/ module, with hub modules left unexpanded. Checked against ATOM#2893, the PR that produced that dead end: it changes atom/model_ops/moe.py alone and 15 existing tests reach it. Fifteen still routes to "name the target explicitly" rather than guessing, but "fifteen candidates, pick one" is actionable where "no target" was not. The import is guarded: a helper that cannot run must not take triage down with it. selftest 98 green / 0 red, 0 drift, 36 pytest cases, ruff and black clean. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Summary
Shared DSA indexer mono launch (
atom/mono/dsa_indexer/)One FlyDSL launch per DSA layer at decode replaces the indexer chain:
wq_b,the index-cache writer (NoPE or the fused RoPE writer), scoring against the FP8
index cache, and the exact top-K selection to the MLA's KV slots.
Indexeruses it; the layer'sIndexerShapekeys and names the build,
plan.unsupportedlists what it is not written for.keys), or without DCP (
IxBuild.replicated: every rank's cache whole, keyssplit by 128-key runs, every rank writes the whole selection at the MLA's
sparse_kv_indptr).batch_id_per_q_tokenon a verify; padded tokens write nothing.any context), row groups of up to 8 rows per key load, score tasks sized to one
CTA round.
dsa_indexer.dispatchwraps the model and routes itsindexers through one opaque op, so a compiled layer stack routes too. A
deployment the kernel cannot serve is refused before any layer is routed.
GLM-5.2 is registered for it (
mono_decode_installers).deepseek_v2folds theindexer's
wk/weights_projintofused_qkv_a_projwhen the attention isnot quantized (one GEMM), and the GLM-5.2 mono runner reads them from there.
Kimi-K3 mono decode
well as a uniform speculative verify, with or without ReplaySSM.
atom/mono/device/gemv.py, shared by K1 and K2 andalso used by the shared indexer.
Shared mono runtime
(
runtime/uncached.py).WidthBuildsrebuilds a width that spills with a lean variant.select.f32_key,ops.hw_log2/ld_f32and others are shared; V4 / V4.1 /GLM-5.2 / M3 call sites follow.
Other
get_dcp_groupreturns the TP coordinator when the DCP group is the TP groupin the same order (one communicator over the same GPUs).
x 8 ranks had become ~17k OS threads).
quant_spec.effective_layer_quant(moved from the Kimi model).Measurements
GLM-5.2-MXFP4, DCP4/TP4, MTP3, 16K ISL / c8, rank-0 trace:
GLM-5.2-MXFP4, TP4 without DCP: GSM8K 0.9310, step time even with production
(the cross-rank selection costs what the scorer saves there).
Harness (
dsa_ix_bench.py): every case bit-equal to the original ops (cachebytes, FP8 q, folded weights, scores, selection) at 128K and 1M, both modes,
with padded verify tokens.
Testing
pytest tests/ --ignore=tests/plugin --ignore=tests/entrypoints/test_openai_server.py(CPU): 14261 passed, 42 failed. Every failure also fails on a clean
mainworktree in this environment (aiter JIT
mega_mhcbuild, torch.compile,atomesh_spur_dispatch,minimax_m31_shared_fusion,sparse_attn_prefill_composition,deepseek_v41_mono_build_key,multiprocess_metrics) or is timing-flaky under a loaded run(
test_upload_numpy, passes alone).test_dsa_indexer_{plan,hook,dispatch}.py,test_kimi_k3_mono_peer_regions.py,test_dcp_device_group.py,test_loader_threads.py.