Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
junyyang-amd wants to merge 2 commits into from
Closed

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8a332db
fix(kimi-k3): import the Triton MLA builder only when ATOM_USE_TRITON…
junyyang-amd Oct 21, 2026
bda6b02
fix(atomesh-ci): raise the open-file limit for the Mesh build container
junyyang-amd Oct 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions .github/scripts/atomesh/setup_mesh.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,10 @@ if [[ ! -x "${artifact_dir}/atomesh" || ! -f "${artifact_dir}/mesh-build.json" ]
source_dirty=true
fi
echo "[setup] Building Mesh before launching services" >&2
# 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 \

@valar-cop valar-cop Oct 30, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

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

@valar-cop valar-cop Oct 30, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

two unrelated changes in one pr.. please split.

--env-file "${env_file}" \
-e ATOMESH_MESH_SOURCE_COMMIT="${source_commit}" \
-e ATOMESH_MESH_SOURCE_DIRTY="${source_dirty}" \
Expand Down
7 changes: 2 additions & 5 deletions atom/model_ops/attentions/kimi_mla_gdn_attn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,6 @@
from .gdn_attn import LINEAR_STATE_ROWS, GDNStateMixin
from .pool_layout.page_unit_geometry import PageUnitGeometryMixin
from .pool_layout.sub_pool_spec import SubPoolSpec, page_pool, state_pool
from .triton_mla import TritonMLAMetadataBuilder


class KimiMLAGDNBackend(AttentionBackend):
Expand All @@ -38,6 +37,8 @@ def get_name() -> str:
@staticmethod
def get_builder_cls() -> type["_KimiMLAGDNCommon"]:
if envs.ATOM_USE_TRITON_MLA:
from .kimi_triton_mla_gdn_attn import KimiTritonMLAGDNMetadataBuilder

@valar-cop valar-cop Oct 30, 2026

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


return KimiTritonMLAGDNMetadataBuilder
return KimiAiterMLAGDNMetadataBuilder

Expand Down Expand Up @@ -540,7 +541,3 @@ def build_for_cudagraph_capture(self, bs: int):
class KimiAiterMLAGDNMetadataBuilder(_KimiMLAGDNCommon, AiterMLAMetadataBuilder):
def _supports_dcp_index_staging(self) -> bool:
return False


class KimiTritonMLAGDNMetadataBuilder(_KimiMLAGDNCommon, TritonMLAMetadataBuilder):
pass
16 changes: 16 additions & 0 deletions atom/model_ops/attentions/kimi_triton_mla_gdn_attn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
# SPDX-License-Identifier: MIT
# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.

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

@valar-cop valar-cop Oct 30, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

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

does not need, so only `KimiMLAGDNBackend.get_builder_cls` imports this module,
and only when `ATOM_USE_TRITON_MLA` is set.
"""

from .kimi_mla_gdn_attn import _KimiMLAGDNCommon
from .triton_mla import TritonMLAMetadataBuilder


class KimiTritonMLAGDNMetadataBuilder(_KimiMLAGDNCommon, TritonMLAMetadataBuilder):

@valar-cop valar-cop Oct 30, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

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

pass
1 change: 1 addition & 0 deletions tests/test_ci_latency_reports.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -613,6 +613,7 @@ def test_mesh_setup_builds_once_and_passes_a_reusable_binary_path(tmp_path):
recorded = [json.loads(line) for line in calls.read_text().splitlines()]
assert len(recorded) == 1
assert "--user" in recorded[0] and "--env-file" in recorded[0]
assert "nofile=65536:524288" in recorded[0]
assert any(arg.startswith("ATOMESH_MESH_SOURCE_COMMIT=") for arg in recorded[0])
assert f"{tmp_path / 'cache'}:/mesh-cache" in recorded[0]
for overrides, expected in (
Expand Down
Loading