fix(kimi-k3): lazy-import the Triton MLA builder; fix(atomesh-ci): raise the Mesh build open-file limit by junyyang-amd · Pull Request #2664 · ROCm/ATOM

junyyang-amd · 2026-10-21T09:43:08Z

Summary

Two fixes that unblock the MI350X Kimi-K3 ATOMesh cells.

Kimi-K3: import the Triton MLA builder only when `ATOM_USE_TRITON_MLA` is set

aiter #6633 removed aiter.ops.triton.attention.mla_decode. kimi_mla_gdn_attn imported triton_mla at module load to define KimiTritonMLAGDNMetadataBuilder, and triton_mla imports that module, so every Kimi-K3 server fails at startup with the new nightly image (rocm/atom-dev:nightly_202610201527-rocm10.0), even though Kimi only uses that builder when ATOM_USE_TRITON_MLA=1:

File "/workspace/ATOM/atom/model_ops/attentions/kimi_mla_gdn_attn.py", line 30, in <module>
  from .triton_mla import TritonMLAMetadataBuilder
File "/workspace/ATOM/atom/model_ops/attentions/triton_mla.py", line 7, in <module>
  from aiter.ops.triton.attention.mla_decode import csr_to_dense_block_table
ModuleNotFoundError: No module named 'aiter.ops.triton.attention.mla_decode'

All 20 Kimi-K3 cells of MI350X ATOMesh nightly #4768 failed this way, e.g. c=1.

This moves KimiTritonMLAGDNMetadataBuilder to kimi_triton_mla_gdn_attn.py and imports it from KimiMLAGDNBackend.get_builder_cls only when ATOM_USE_TRITON_MLA is set, the same way selector.py only resolves triton_mla for DeepSeek and GLM when the flag is on.

Flag off (default, and what CI runs): get_builder_cls returns the same KimiAiterMLAGDNMetadataBuilder; triton_mla is no longer imported.
Flag on: same class and MRO as before. It still fails against current aiter, now when the builder is selected instead of at module import, as DeepSeek and GLM do with the flag on. Restoring the Triton MLA path (csr_to_dense_block_table and the non-shuffled decode_attention_fwd branch in attention_mla.py) is left for a follow-up.

ATOMesh CI: raise the open-file limit for the Mesh build container

With the Kimi fix, both Kimi-K3 cells in run 39580104520 got past the import and decode came up, but the router never started. On rank 0, both Mesh build attempts failed with:

could not execute process `/mesh-cache/rustup-home/toolchains/1.94.0-x86_64-unknown-linux-gnu/bin/rustc --crate-name cfg_if ...` (never executed)
Caused by:
  Too many open files (os error 24)

#2625 moved the Mesh Cargo cache from /tmp to <log root>/.mesh-cache/<host>-<uid>, so every host now starts cold and Cargo compiles all 534 crates with one rustc per CPU (192 on these nodes). The build container from setup_mesh.sh had no --ulimit, so it ran with Docker's default soft nofile of 1024. The service containers in pd_slurm_job.sh already use nofile=65536:524288. This gives the build container the same limit.

Test plan

Kimi-K3 import, in rocm/atom-dev:nightly_202610201527-rocm10.0 on an MI350X node, with this branch mounted at /workspace/ATOM:

aiter.ops.triton.attention.mla_decode is absent in the image.
Flag unset: import atom.model_ops.attentions.kimi_mla_gdn_attn succeeds, get_builder_cls() returns KimiAiterMLAGDNMetadataBuilder, and atom.model_ops.attentions.triton_mla is not in sys.modules.
ATOM_USE_TRITON_MLA=1: get_builder_cls() raises the same ModuleNotFoundError (expected, see above).
tests/test_kda_layout_id.py, test_kda_checkpoint_slot_copy.py, test_h2d_attention_publication.py, test_page_unit_geometry.py, test_prepare_mtp_decode_contract.py: pass.
tests/test_gdn_state_relocation.py: 7 failures with KVTransferTensors.__init__() got an unexpected keyword argument 'block_regions'. Pre-existing and unrelated: feat(offload): support native DSv4 checkpoints with LMCache MP #2250 replaced the block_regions field with pages, and the test re-landed in feat(kimi-k3): support P/D disaggregation (re-land of #2154) #2332 still passes block_regions=. The file importorskips both aiter (so CPU CI skips it) and kimi_mla_gdn_attn (so it was also skipped in this image until this fix).

Mesh build, on an MI350X node with the same image:

Docker default inside the build container: ulimit -Sn is 1024 with nproc 192. With --ulimit nofile=65536:524288 it is 65536.
setup_mesh.sh with an empty cache directory: cold build succeeds on the first attempt in 5m52s, with no os error 24 and an atomesh binary produced.
tests/test_ci_latency_reports.py -k mesh: 5 passed, including a new assertion on the nofile flag.
Kimi-K3 ATOMesh benchmark cells on MI350X: run 39717855656 passed kimi-k3-mxfp4-1p1d-tp8-dspark7-agentic-1m-c1 and kimi-k3-mxfp4-1p1d-tp8-dcp8-agentic-lmcache-1m-c48. Both Mesh builds started from an empty cache and succeeded on the first attempt (3m05s and 8m05s).

…_MLA is set kimi_mla_gdn_attn imported triton_mla at module load to define KimiTritonMLAGDNMetadataBuilder, so Kimi-K3 pulled in aiter's Triton MLA decode even on the default aiter MLA path. aiter #6633 removed aiter.ops.triton.attention.mla_decode, and every Kimi-K3 server now fails at startup with ModuleNotFoundError. Move KimiTritonMLAGDNMetadataBuilder to kimi_triton_mla_gdn_attn and import it from get_builder_cls only when ATOM_USE_TRITON_MLA is set, as the selector already does for DeepSeek and GLM. The default path returns the same KimiAiterMLAGDNMetadataBuilder as before. Co-authored-by: Cursor <cursoragent@cursor.com>

github-actions · 2026-10-21T09:43:44Z

🏷️ 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 2664 --add-label <label>

Since #2625 moved the Mesh Cargo cache to <log root>/.mesh-cache, every host starts with an empty cache and Cargo compiles all crates at once, one rustc per CPU (192 on MI350 nodes). The build container ran with Docker's default soft limit of 1024 open files, so both build attempts failed with "Too many open files (os error 24)" and the router never started (run 39580104520, jobs 14061 and 14073). Give it the same nofile limit as the service containers. Co-authored-by: Cursor <cursoragent@cursor.com>

valar-cop

not ready yet.. un-skipping a test that then fails is not ok.

atom/model_ops/attentions/kimi_mla_gdn_attn.py:30 — this un-skips test_gdn_state_relocation and it fails 7 tests, fix the test here or keep it skipped with a reason.
atom/model_ops/attentions/kimi_triton_mla_gdn_attn.py:12 — root cause is still the module-level import in triton_mla.py, lazy import there fixes kimi/deepseek/glm at once.
.github/scripts/atomesh/setup_mesh.sh:39 — soft 65536 vs hard 524288, why not equal? and is the hard limit ok on every node?
atom/model_ops/attentions/kimi_mla_gdn_attn.py:40 — raw ModuleNotFoundError gives no hint, raise a clear message or drop the flag.
atom/model_ops/attentions/kimi_triton_mla_gdn_attn.py:15 — class moved, did you grep other importers?

any ci run showing the test green? no reproducer means nothing to me.

— valar-cop 👮⭐

Jasen2201 · 2026-10-30T07:18:32Z

@junyyang-amd thanks for tracking this down. Where the two fixes stand on current main:

Kimi-K3 import crash: no longer reproduces on main. fix(mla): migrate Triton decode to maintained aiter API #2669 (65a199e, merged 10-21) moved triton_mla.py to the maintained aiter API, so kimi_mla_gdn_attn → triton_mla no longer imports aiter.ops.triton.attention.mla_decode. Checked on rocm/atom-dev:nightly_202610291449 (ATOM 70dee2a): that module is absent from the image's aiter, and Kimi-K3 (TP8, DSpark) loads and serves normally. Atomesh Benchmark Kimi-K3 runs from 10-29/10-30 pass as well (e.g. 41004161113, 40992652198). The lazy import is still a nice tidy-up, but nothing depends on it any more.
Mesh build nofile: still missing on main. .github/scripts/atomesh/setup_mesh.sh has no --ulimit nofile, while the service containers in pd_slurm_job.sh set nofile=65536:524288. It only bites when a host's .mesh-cache is cold, so it isn't blocking anything today, but it's worth keeping.

Suggestion: the Kimi change is also what un-skips test_gdn_state_relocation (the 7 block_regions → pages failures from the review above). Could you drop it and keep only the setup_mesh.sh ulimit change here, or move that change into a small new PR and close this one? Either should be quick to merge.

Jasen2201 · 2026-10-30T07:18:48Z

@junyyang-amd follow-up: since main already covers the Kimi-K3 import fix via #2669, let's close this PR. Please move the setup_mesh.sh --ulimit nofile change into a small separate PR. Thanks!

valar-cop

not ready yet.. kimi change is redundant now, and it un-skips a failing test.

any ci run showing test_gdn_state_relocation green? and nofile hard limit verified on all node types?

valar-cop · 2026-10-30T07:27:48Z

    @staticmethod
    def get_builder_cls() -> type["_KimiMLAGDNCommon"]:
        if envs.ATOM_USE_TRITON_MLA:
+            from .kimi_triton_mla_gdn_attn import KimiTritonMLAGDNMetadataBuilder


this un-skips the test and it fails 7 cases, ci will be red.. fix block_regions -> pages or keep it skipped?

valar-cop · 2026-10-30T07:27:48Z

+
+"""Kimi's Triton MLA builder, kept apart from `kimi_mla_gdn_attn`.
+
+`triton_mla` imports aiter's Triton MLA decode, which the default aiter MLA path


#2669 already fixed the root cause on main, so this is a no-op refactor.. drop it?

valar-cop · 2026-10-30T07:27:48Z

+  # A cold build runs one rustc per CPU (192 on MI350 nodes), which exhausts
+  # Docker's default soft limit of 1024 open files.
  docker run --rm --user "$(id -u):$(id -g)" \
+    --ulimit nofile=65536:524288 \


hard limit 524288 only tested on MI350X? other node types may have lower host nofile and docker run will fail..

valar-cop · 2026-10-30T07:27:48Z

+  # A cold build runs one rustc per CPU (192 on MI350 nodes), which exhausts
+  # Docker's default soft limit of 1024 open files.
  docker run --rm --user "$(id -u):$(id -g)" \
+    --ulimit nofile=65536:524288 \


two unrelated changes in one pr.. please split.

valar-cop · 2026-10-30T07:27:48Z

+from .triton_mla import TritonMLAMetadataBuilder
+
+
+class KimiTritonMLAGDNMetadataBuilder(_KimiMLAGDNCommon, TritonMLAMetadataBuilder):


any other importer of KimiTritonMLAGDNMetadataBuilder? grep it or re-export for compat.

junyyang-amd requested a review from Jasen2201 October 21, 2026 09:57

junyyang-amd changed the title ~~fix(kimi-k3): import the Triton MLA builder only when ATOM_USE_TRITON_MLA is set~~ fix(kimi-k3): lazy-import the Triton MLA builder; fix(atomesh-ci): raise the Mesh build open-file limit Oct 22, 2026

junyyang-amd requested a review from gbyu-amd October 30, 2026 06:45

valar-cop reviewed Oct 30, 2026

View reviewed changes

junyyang-amd closed this Oct 30, 2026

Navigation Menu

Repository navigation

fix(kimi-k3): lazy-import the Triton MLA builder; fix(atomesh-ci): raise the Mesh build open-file limit - #2664

fix(kimi-k3): lazy-import the Triton MLA builder; fix(atomesh-ci): raise the Mesh build open-file limit#2664
junyyang-amd wants to merge 2 commits into
mainfrom
junyyang/kimi-lazy-triton-mla

junyyang-amd commented Oct 21, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Oct 21, 2026

Uh oh!

valar-cop left a comment

Uh oh!

Jasen2201 commented Oct 30, 2026

Uh oh!

Jasen2201 commented Oct 30, 2026

Uh oh!

valar-cop left a comment

Uh oh!

valar-cop Oct 30, 2026

Uh oh!

valar-cop Oct 30, 2026

Uh oh!

valar-cop Oct 30, 2026

Uh oh!

valar-cop Oct 30, 2026

Uh oh!

valar-cop Oct 30, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants


		"""Kimi's Triton MLA builder, kept apart from `kimi_mla_gdn_attn`.

		`triton_mla` imports aiter's Triton MLA decode, which the default aiter MLA path

		from .triton_mla import TritonMLAMetadataBuilder


		class KimiTritonMLAGDNMetadataBuilder(_KimiMLAGDNCommon, TritonMLAMetadataBuilder):

Conversation

junyyang-amd commented Oct 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Kimi-K3: import the Triton MLA builder only when ATOM_USE_TRITON_MLA is set

ATOMesh CI: raise the open-file limit for the Mesh build container

Test plan

Uh oh!

github-actions Bot commented Oct 21, 2026

🏷️ CI Guide

Uh oh!

valar-cop left a comment

Choose a reason for hiding this comment

Uh oh!

Jasen2201 commented Oct 30, 2026

Uh oh!

Jasen2201 commented Oct 30, 2026

Uh oh!

valar-cop left a comment

Choose a reason for hiding this comment

Uh oh!

valar-cop Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

valar-cop Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

valar-cop Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

valar-cop Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

valar-cop Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

junyyang-amd commented Oct 21, 2026 •

edited

Loading

Kimi-K3: import the Triton MLA builder only when `ATOM_USE_TRITON_MLA` is set