[ATOM][Kimi-K3] ATOM support Kimi-K3 by zejunchen-zejun · Pull Request #1712 · ROCm/ATOM

zejunchen-zejun · 2026-07-28T01:35:37Z

No description provided.

Port all Kimi-K3 code (incl. recipes) developed on the carlushuang/ATOM-K3 `xiaobing/k3` branch onto ROCm/ATOM. Provenance: - source branch: carlushuang/ATOM-K3 xiaobing/k3 @ cbca2f0 - fork base (shared commit): 8d87bd8 "[feat][PCP] Make PCP compatible with FP8 KVCache for DeepSeek V4 (#1652)" - range migrated: 8d87bd8..cbca2f0 (135 commits, squashed into this port) - 29 files, +5305/-120 Excluded per request: my_script/ (14 files, launch/debug scripts). Content is byte-for-byte identical to xiaobing/k3 outside my_script/. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

github-actions · 2026-07-28T01:36:07Z

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

Drop debug and unused experimental paths flagged in review as reviewer overhead. All removed flags were introduced by k3 (absent at the atom fork point), so removal keeps the default (flag-off) behavior: - ATOM_FWD_DUMP_STATS_ONLY: revert dump.py to base (stats-only dump path) - ATOM_K3_SHARED_EXPERT_OVERLAP: remove dual-stream shared-expert overlap (alt_stream plumbing, maybe_dual_stream_forward dispatch, dual_stream_moe_forward, module_dispatch_ops side-effect import); it currently hangs and is unused - AITER_USE_TORCH_TOPK: remove torch topK fallbacks in topK.py - ATOM_USE_TORCH_SAMPLER: remove torch sampler fallback in sampler.py - ATOM_USE_TORCH_CACHE: remove _torch_reshape_and_cache fallback in attention_mha.py Kept ATOM_DUAL_STREAM_MOE_TOKEN_THRESHOLD (base, shared with DeepSeek) and ATOM_DEBUG_TOPK (not in scope). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Review follow-ups on the Kimi-K3 port: - Revert all tests/ changes to the fork base (8d87bd8) and drop the new tests/entrypoints/test_serving_completion.py; no test cases ship with this PR. - moe.py: remove the `fp32_allreduce_enabled` branch in both FusedMoE all-reduce paths. The whole mechanism (_FP32_ALLREDUCE + setter + getter) was deleted from linear.py upstream, so the import raised ImportError for any FusedMoE with reduce_results=True and tp/ep > 1. Restoring the unconditional all_reduce matches linear.py's current RowParallel behaviour; re-adding the getter would only reinstate a branch nothing can enable. - Drop comment references to files excluded from this PR (code_k3/shuffle_kernel_fix.patch, my_script/optest_fused.py); the technical findings they described are kept inline. - moe.py `_load_w13`: fold the no-op shard_start/loaded_shard rename back into the original narrow() call. model_runner.py: drop a stray blank line. gfx1250 workarounds (ATOM_K3_MOE_CHUNK, ATOM_WARMUP_MAX_TOKENS, ATOM_KDA_FORCE_RECURRENT, FLASH-layout selection) are kept as-is. black still reports serving_completion.py and gdn_attn.py; deferred to a follow-up at the author's request. ruff clean; 663 tests pass. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

wuhuikx · 2026-07-28T02:54:08Z

            "config": config,
            "input_address": input_address,
            "output_address": output_address,
+            "env_snapshot": dict(os.environ),


do we need this change?

wuhuikx · 2026-07-28T02:54:28Z

+            f"(weights loaded in {load_elapsed:.2f}s)"
+        )
+        if (
+            os.getenv("ATOM_SYNC_AFTER_LOAD", "0").lower() in ("1", "true", "yes")


Please remove this

wuhuikx · 2026-07-28T02:54:47Z

+        # (that large-M crash is fixed in aiter grouped_moe_gfx1250).
+        _warmup_cap = int(os.environ.get("ATOM_WARMUP_MAX_TOKENS", "0") or "0")
+        if _warmup_cap > 0:
+            warmup_max_tokens = min(warmup_max_tokens, _warmup_cap)


do we need this?

wuhuikx · 2026-07-28T02:55:32Z

+                    time.perf_counter() - shard_start,
+                    os.path.basename(st_file),
+                )
+    pbar.close()


Please check the above code.

Both process levels (server -> EngineCore -> ModelRunner workers) are spawned, and spawn children inherit the parent's environ through exec -- including vars set at runtime via os.environ, since CPython writes those through to the C environ. Verified on this stack (Python 3.12, two-level spawn): shell-exported, top-parent-set and EngineCore-set vars all reach the workers without any snapshot. So `os.environ.update(env_snapshot)` was a no-op: it can only add keys to an environment that already matches. It was also mildly hazardous -- the worker snapshot is taken before the spawn loop, so a future per-rank env var set inside the loop would be silently overwritten with the stale value instead of inherited correctly. async_proc.py and engine_core_mgr.py carried no other changes on this branch and are restored to the fork base; engine_core.py keeps its Kimi-K3 prefix-caching hunk. The `import os` each file gained for the snapshot goes with it. ruff clean; 663 tests pass. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

valarLip · 2026-07-28T07:53:25Z

+            # still in flight. Do not try to backfill deferred placeholders for
+            # that seq: aborted partial-prefill requests may not have any
+            # output_tokens placeholder to overwrite.
+            if seq.status == SequenceStatus.ABORTED:


ABORTED should already supported? @yhl-amd any comments?

valarLip · 2026-07-28T07:53:56Z

 from dataclasses import dataclass
 from typing import Optional, Tuple

+KIMI_THINK_END = "<|close|>think<|sep|>"


no model specific stuff should be here

valarLip · 2026-07-28T07:56:31Z

+from transformers import PretrainedConfig

 logger = logging.getLogger("atom")
+_MOE_LOAD_COPY_LOG_SECONDS = float(os.getenv("ATOM_MOE_LOAD_COPY_LOG_SECONDS", "0"))


valarLip · 2026-07-28T07:56:53Z

+            src = src.contiguous()
+        start = time.perf_counter() if _MOE_LOAD_COPY_LOG_SECONDS > 0 else None
+
+        def log_slow_copy() -> None:


valarLip · 2026-07-28T07:57:15Z

        Returns:
            Sampled token IDs (num_tokens,)
        """
+        if envs.ATOM_DEBUG_TOPK > 0:


valarLip · 2026-07-28T07:58:17Z

most of kernels are duplicated with existing one..

valarLip · 2026-07-28T07:59:57Z

@@ -634,9 +646,15 @@ def _get_hf_token() -> str | None:
        ):
            text_config_dict["quantization_config"] = config_dict["quantization_config"]
        text_model_type = text_config_dict.get("model_type", "deepseek_v3")


text_model_type = text_config_dict.get("model_type", "deepseek_v3") ->text_model_type = text_config_dict.get("model_type", "text_config")

zejunchen-zejun · 2026-07-28T14:03:48Z

use PR #1718

wuhuikx requested a review from XiaobingSuper July 28, 2026 01:57

zejunchen-zejun and others added 2 commits July 28, 2026 10:19

wuhuikx reviewed Jul 28, 2026

View reviewed changes

valarLip reviewed Jul 28, 2026

View reviewed changes

zejunchen-zejun closed this Jul 28, 2026

Uh oh!

Conversation

zejunchen-zejun commented Jul 28, 2026

Uh oh!

github-actions Bot commented Jul 28, 2026

🏷️ CI Guide

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zejunchen-zejun commented Jul 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants