[ATOM][Kimi-K3] ATOM support Kimi-K3

atom/config.py

@@ -266,6 +266,17 @@ class QuantizationConfig:
     - ``quant_type``, ``quant_dtype``, ``is_dynamic`` convenience properties
     """
 
+    _ONLINE_QUANT_SOURCE_METHODS = frozenset(
+        {
+            "",
+            "fp8",
+            "mxfp4",
+            "mxfp8",
+            "quark",
+            "compressed-tensors",
+        }
+    )
+
     def __init__(
         self,
         config: PretrainedConfig = None,
@@ -306,14 +317,15 @@ def __init__(
         self.online_global_spec: LayerQuantConfig = LayerQuantConfig()
         self.online_layer_pattern_specs: list[tuple[str, LayerQuantConfig]] = []
         self.online_exclude_layers: list[str] = []
-        if online_quant_config and self.quant_method in [
-            "",
-            "fp8",
-            "mxfp4",
-            "mxfp8",
-            "quark",
-        ]:
+        online_source_supported = self.quant_method in self._ONLINE_QUANT_SOURCE_METHODS
+        if online_quant_config and online_source_supported:
             self.online_quant = True
+            if self.quant_method == "compressed-tensors":
+                logger.warning(
+                    "Online quantization for compressed-tensors checkpoints "
+                    "relies on the caller's exclude_layer configuration. Ensure "
+                    "unsupported or already-quantized layers are excluded."
+                )
             online_parser = get_quant_parser("online_quant")
             online_parsed_quant_config = online_parser.parse(online_quant_config)
             self.online_global_spec = online_parsed_quant_config.global_spec
@@ -584,6 +596,7 @@ def _remap_layer_name(name: str) -> list[str]:
 
 _MULTIMODAL_MODEL_TYPES: dict[str, str] = {
     # Maps multimodal model_type -> key in config_dict for the text sub-config
+    "kimi_k3": "text_config",
     "kimi_k25": "text_config",
     "qwen3_5": "text_config",
     "qwen3_5_moe": "text_config",
@@ -633,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")
-        mapped_type = _CONFIG_REGISTRY.get(text_model_type, text_model_type)
-        config_class = AutoConfig.for_model(mapped_type)
-        hf_config = config_class.from_dict(text_config_dict)
+        if text_model_type == "kimi_linear":
+            # Transformers does not ship KimiLinearConfig yet in this image.
+            # Keep the remote-code fields as plain PretrainedConfig attrs; the
+            # ATOM model normalizes the aliases it needs at construction time.
+            hf_config = PretrainedConfig.from_dict(text_config_dict)
+        else:
+            mapped_type = _CONFIG_REGISTRY.get(text_model_type, text_model_type)
+            config_class = AutoConfig.for_model(mapped_type)
+            hf_config = config_class.from_dict(text_config_dict)
         # Override architectures so that ATOM selects the correct model class
         # which can handle the multimodal weight prefix during loading.
         original_arch = config_dict.get("architectures", [])

atom/entrypoints/openai/api_server.py

@@ -200,19 +200,33 @@ def _validate_context_length(
     )
 
 
-def _get_engine_max_model_len() -> Optional[int]:
+def _get_engine_config():
     config = getattr(engine, "config", None)
     if config is None:
         config = getattr(getattr(engine, "io_processor", None), "config", None)
-    return getattr(config, "max_model_len", None)
+    return config
 
 
 def _validate_sequence_context_length(seq) -> None:
+    config = _get_engine_config()
     _validate_context_length(
         seq.num_prompt_tokens,
         seq.max_tokens,
-        _get_engine_max_model_len(),
+        getattr(config, "max_model_len", None),
     )
+    max_num_batched_tokens = getattr(config, "max_num_batched_tokens", None)
+    enable_chunked_prefill = bool(getattr(config, "enable_chunked_prefill", False))
+    if (
+        not enable_chunked_prefill
+        and max_num_batched_tokens is not None
+        and seq.num_prompt_tokens > max_num_batched_tokens
+    ):
+        raise ValueError(
+            f"Prompt contains {seq.num_prompt_tokens} input tokens, which exceeds "
+            f"max_num_batched_tokens={max_num_batched_tokens} while chunked prefill "
+            f"is disabled. Increase --max-num-batched-tokens, enable chunked "
+            f"prefill, or shorten the prompt."
+        )
 
 
 def _has_multimodal_content(messages: List[Any]) -> bool:
@@ -1320,7 +1334,9 @@ async def completions(request: CompletionRequest, raw_request: Request):
             )
             if not outputs:
                 raise RuntimeError("No output generated")
-            resp = build_completion_response_multi(request_id, model_name, outputs)
+            resp = build_completion_response_multi(
+                request_id, model_name, outputs, request.stop
+            )
         else:
             final_output = await _run_nonstream_with_disconnect(
                 generate_async(
@@ -1337,7 +1353,9 @@ async def completions(request: CompletionRequest, raw_request: Request):
             if final_output is None:
                 raise RuntimeError("No output generated")
 
-            resp = build_completion_response(request_id, model_name, final_output)
+            resp = build_completion_response(
+                request_id, model_name, final_output, request.stop
+            )
         _log_request_event("response", request_id, resp.model_dump())
         return resp
 

atom/entrypoints/openai/reasoning.py

@@ -12,6 +12,22 @@
 from dataclasses import dataclass
 from typing import Optional, Tuple
 
+KIMI_THINK_END = "<|close|>think<|sep|>"
+KIMI_RESPONSE_START = "<|open|>response<|sep|>"
+KIMI_RESPONSE_END = "<|close|>response<|sep|>"
+KIMI_MESSAGE_END = "<|close|>message<|sep|>"
+KIMI_END_OF_MSG = "<|end_of_msg|>"
+
+
+def _strip_kimi_response_markers(text: str) -> str:
+    if text.startswith(KIMI_RESPONSE_START):
+        text = text[len(KIMI_RESPONSE_START) :]
+
+    for marker in (KIMI_RESPONSE_END, KIMI_MESSAGE_END, KIMI_END_OF_MSG):
+        if marker in text:
+            text = text.partition(marker)[0]
+    return text.strip()
+
 
 def separate_reasoning(text: str) -> Tuple[Optional[str], str]:
     """Separate reasoning content from the final answer.
@@ -23,6 +39,23 @@ def separate_reasoning(text: str) -> Tuple[Optional[str], str]:
         Tuple of (reasoning_content, content). reasoning_content is None if
         no thinking block was found.
     """
+    kimi_response_delim = KIMI_THINK_END + KIMI_RESPONSE_START
+    if kimi_response_delim in text:
+        reasoning, _, content = text.partition(kimi_response_delim)
+        reasoning = reasoning.strip()
+        content = _strip_kimi_response_markers(content)
+        return (reasoning if reasoning else None, content)
+
+    if KIMI_RESPONSE_START in text:
+        _, _, content = text.partition(KIMI_RESPONSE_START)
+        return (None, _strip_kimi_response_markers(content))
+
+    if KIMI_THINK_END in text:
+        reasoning, _, content = text.partition(KIMI_THINK_END)
+        reasoning = reasoning.strip()
+        content = _strip_kimi_response_markers(content)
+        return (reasoning if reasoning else None, content)
+
     # Check for closed thinking block: <think>...</think>
     match = re.match(r"<think>(.*?)</think>\s*(.*)", text, flags=re.DOTALL)
     if match:

atom/entrypoints/openai/serving_completion.py

@@ -18,6 +18,22 @@
 logger = logging.getLogger("atom")
 
 
+def _strip_stop_strings(
+    text: str, stop_strings: Optional[List[str]]
+) -> str:
+    """OpenAI responses must not include the matched stop string."""
+    if not stop_strings:
+        return text
+
+    stop_pos = min(
+        (idx for stop in stop_strings if stop and (idx := text.find(stop)) != -1),
+        default=-1,
+    )
+    if stop_pos == -1:
+        return text
+    return text[:stop_pos]
+
+
 def create_completion_chunk(
     request_id: str,
     model: str,
@@ -126,6 +142,7 @@ def build_completion_response(
     request_id: str,
     model: str,
     final_output: Dict[str, Any],
+    stop_strings: Optional[List[str]] = None,
 ) -> CompletionResponse:
     """Build a non-streaming text completion response (single choice)."""
     response = CompletionResponse(
@@ -135,7 +152,7 @@ def build_completion_response(
         choices=[
             {
                 "index": 0,
-                "text": final_output["text"],
+                "text": _strip_stop_strings(final_output["text"], stop_strings),
                 "finish_reason": final_output["finish_reason"],
             }
         ],
@@ -162,13 +179,14 @@ def build_completion_response_multi(
     request_id: str,
     model: str,
     final_outputs: List[Dict[str, Any]],
+    stop_strings: Optional[List[str]] = None,
 ) -> CompletionResponse:
     """Build a non-streaming response with one choice per fan-out sibling."""
     assert final_outputs, "build_completion_response_multi requires at least one output"
     choices = [
         {
             "index": i,
-            "text": out["text"],
+            "text": _strip_stop_strings(out["text"], stop_strings),
             "finish_reason": out["finish_reason"],
         }
         for i, out in enumerate(final_outputs)

atom/examples/simple_inference.py

@@ -58,7 +58,9 @@ def main():
     engine_args = EngineArgs.from_cli_args(args)
     llm = engine_args.create_engine()
 
-    tokenizer = AutoTokenizer.from_pretrained(args.model)
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model, trust_remote_code=args.trust_remote_code
+    )
 
     sampling_params = SamplingParams(
         temperature=args.temperature, max_tokens=args.max_tokens

atom/model_engine/engine_core.py

@@ -19,6 +19,7 @@
 from atom.model_engine.sequence import Sequence, SequenceStatus, get_exit_sequence
 from atom.utils import (
     envs,
+    get_hf_text_config,
     init_exit_handler,
     make_zmq_socket,
     set_process_title,
@@ -120,6 +121,20 @@ def __init__(self, config: Config, input_address: str, output_address: str):
             if not good:
                 self._finalizer()
 
+        hf_text_config = get_hf_text_config(config.hf_config)
+        arches = getattr(hf_text_config, "architectures", None) or []
+        is_kimi_kda_hybrid = any("KimiK3" in str(a) for a in arches) or (
+            getattr(hf_text_config, "model_type", None) == "kimi_linear"
+            and hasattr(hf_text_config, "full_attn_layer_ids")
+        )
+        if is_kimi_kda_hybrid and config.enable_prefix_caching:
+            logger.info(
+                "%s: Kimi-K3 KDA hybrid disables prefix caching engine-wide "
+                "because recurrent attention state is per request.",
+                self.label,
+            )
+            config.enable_prefix_caching = False
+
         self.scheduler = Scheduler(config)
 
         self.kv_transfer_enabled = bool(config.kv_transfer_config)

atom/model_engine/llm_engine.py

@@ -377,6 +377,7 @@ def _per_req_cache_model_types() -> frozenset[str]:
                 "qwen3_next",
                 "qwen3_5_text",
                 "qwen3_5_moe_text",
+                "kimi_linear",
                 "deepseek_v4",
             }
         )

atom/model_engine/model_runner.py

@@ -87,6 +87,7 @@
     "Qwen3_5ForConditionalGeneration": "atom.models.qwen3_5.Qwen3_5MultimodalModel",
     "Qwen3_5MoeForConditionalGeneration": "atom.models.qwen3_5.Qwen3_5MoeMultimodalModel",
     "KimiK25ForConditionalGeneration": "atom.models.kimi_k25.KimiK25ForCausalLM",
+    "KimiK3ForConditionalGeneration": "atom.models.kimi_k3.KimiK3ForCausalLM",
     "MiniMaxM2ForCausalLM": "atom.models.minimax_m2.MiniMaxM2ForCausalLM",
     "MiMoV2ForCausalLM": "atom.models.mimo_v2.MiMoV2ForCausalLM",
     "MiMoV2FlashForCausalLM": "atom.models.mimo_v2.MiMoV2ForCausalLM",
@@ -771,6 +772,15 @@ def __init__(self, rank: int, config: Config):
             f"[{self.rank_name}] Model load done: {config.model} "
             f"(weights loaded in {load_elapsed:.2f}s)"
         )
+        if (
+            os.getenv("ATOM_SYNC_AFTER_LOAD", "0").lower() in ("1", "true", "yes")
+            and get_tp_group().world_size > 1
+        ):
+            logger.info(
+                "Waiting for all TP ranks to finish model loading before warmup"
+            )
+            get_tp_group().barrier()
+            logger.info("All TP ranks finished model loading")
 
         # Optional debug instrumentation; no-op when env vars unset.
         # See atom/utils/debug_helper/.
@@ -899,10 +909,14 @@ def is_qwen_next(self) -> bool:
             "qwen3_next_mtp",
             "qwen3_5_text",
             "qwen3_5_moe_text",
+            "kimi_linear",
         ):
             return True
         return False
 
+    def is_kimi_linear(self) -> bool:
+        return getattr(self.hf_text_config, "model_type", None) == "kimi_linear"
+
     def is_deepseek_v4(self) -> bool:
         # NOTE: `hf_text_config.model_type` reads "deepseek_v3" for V4 because
         # `_CONFIG_REGISTRY` maps deepseek_v4 → deepseek_v3 (V4 reuses V3 schema).
@@ -1340,6 +1354,16 @@ def warmup_model(self):
         if pcp_size > 1:
             warmup_max_tokens = max(1, warmup_max_tokens // pcp_size)
 
+        # TEMPORARY (benign, warmup-only): cap the dummy warmup prefill. A large
+        # all-zero warmup batch samples garbage logits over all positions and
+        # faults the sampler/softmax on gfx1250; real inference only samples the
+        # last token per seq so it is unaffected. Proper fix = skip/greedy sampling
+        # during is_dummy_run warmup, then this cap can go. Not needed for the MoE
+        # (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)
+
         num_seqs = min(warmup_max_tokens // max_model_len, self.config.max_num_seqs)
 
         if num_seqs == 0:

atom/model_engine/scheduler.py

@@ -1455,6 +1455,15 @@ def postprocess(
             num_placeholder += 1
 
         for seq in self.running:
+            # A disconnected client may abort a seq while its current forward is
+            # 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:
+                seq.leave_reason = "aborted"
+                seq.status = SequenceStatus.FINISHED
+                finished_seqs.append(seq)
+                continue
             # Update the running status
             idx = fwd_output.get_idx(seq.id)
             if idx is None: