From c84f297565926f1eed1f5f94288a61ae9d68ab50 Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Tue, 28 Jul 2026 09:32:55 +0800 Subject: [PATCH 1/4] Migrate Kimi-K3 support from carlushuang/ATOM-K3 xiaobing/k3 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 --- atom/config.py | 39 +- atom/entrypoints/openai/api_server.py | 28 +- atom/entrypoints/openai/reasoning.py | 33 + atom/entrypoints/openai/serving_completion.py | 22 +- atom/examples/simple_inference.py | 4 +- atom/model_engine/async_proc.py | 7 + atom/model_engine/engine_core.py | 26 +- atom/model_engine/engine_core_mgr.py | 2 + atom/model_engine/llm_engine.py | 1 + atom/model_engine/model_runner.py | 25 + atom/model_engine/scheduler.py | 9 + atom/model_loader/loader.py | 220 ++- atom/model_ops/attention_mha.py | 149 +- atom/model_ops/attentions/aiter_attention.py | 35 +- atom/model_ops/attentions/gdn_attn.py | 55 +- .../model_ops/fla_ops/fused_sigmoid_gating.py | 34 +- atom/model_ops/fused_moe_triton.py | 35 +- atom/model_ops/moe.py | 85 +- atom/model_ops/sampler.py | 21 +- atom/model_ops/topK.py | 31 + atom/models/kimi_k3.py | 1420 +++++++++++++++++ atom/models/kimi_k3_fused.py | 836 ++++++++++ atom/utils/debug_helper/dump.py | 25 + atom/utils/envs.py | 39 + recipes/Kimi-K3.md | 78 + tests/entrypoints/test_api_server_helpers.py | 38 + tests/entrypoints/test_reasoning.py | 17 + tests/entrypoints/test_serving_completion.py | 51 + tests/test_scheduler.py | 24 + 29 files changed, 3269 insertions(+), 120 deletions(-) create mode 100644 atom/models/kimi_k3.py create mode 100644 atom/models/kimi_k3_fused.py create mode 100644 recipes/Kimi-K3.md create mode 100644 tests/entrypoints/test_serving_completion.py diff --git a/atom/config.py b/atom/config.py index 2c35848a5e..6a37c82821 100644 --- a/atom/config.py +++ b/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", []) diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index 41d0b5e81b..97b6a3c55b 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/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 diff --git a/atom/entrypoints/openai/reasoning.py b/atom/entrypoints/openai/reasoning.py index 6fb8a8e004..53aa47c0b7 100644 --- a/atom/entrypoints/openai/reasoning.py +++ b/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: ... match = re.match(r"(.*?)\s*(.*)", text, flags=re.DOTALL) if match: diff --git a/atom/entrypoints/openai/serving_completion.py b/atom/entrypoints/openai/serving_completion.py index a240c0b7ce..50f3a8aa18 100644 --- a/atom/entrypoints/openai/serving_completion.py +++ b/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) diff --git a/atom/examples/simple_inference.py b/atom/examples/simple_inference.py index 73d51c2e69..040a9093b9 100644 --- a/atom/examples/simple_inference.py +++ b/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 diff --git a/atom/model_engine/async_proc.py b/atom/model_engine/async_proc.py index 7169108637..ba63c0973f 100644 --- a/atom/model_engine/async_proc.py +++ b/atom/model_engine/async_proc.py @@ -15,6 +15,7 @@ import logging import multiprocessing +import os import pickle import queue import threading @@ -69,9 +70,13 @@ def __init__( rank: int, kv_output_addr: str | None = None, all_ranks_barrier=None, + env_snapshot: dict[str, str] | None = None, *args, **kwargs, ): + if env_snapshot: + os.environ.update(env_snapshot) + # Bind this worker's lifetime to its parent EngineCore: if the parent # exits for any reason, have the kernel reap this process immediately # instead of leaving it orphaned. A ModelRunner worker holds a large GPU @@ -275,6 +280,7 @@ def __init__(self, finalizer, proc_num: int, runner: str, *args): atexit.register(self._cleanup_shared_memory) self.all_ranks_barrier = ctx.Barrier(proc_num) init_exit_handler(self) + env_snapshot = dict(os.environ) # KV output aggregation infrastructure self.kv_output_aggregator: KVOutputAggregator | None = None @@ -300,6 +306,7 @@ def __init__(self, finalizer, proc_num: int, runner: str, *args): i, self.kv_output_addrs[i], self.all_ranks_barrier, + env_snapshot, *args, ), ) diff --git a/atom/model_engine/engine_core.py b/atom/model_engine/engine_core.py index 319dbc55aa..6680d234dd 100644 --- a/atom/model_engine/engine_core.py +++ b/atom/model_engine/engine_core.py @@ -2,6 +2,7 @@ # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. import logging +import os import pickle import queue import threading @@ -19,6 +20,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 +122,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) @@ -172,7 +188,15 @@ def _send_engine_dead(self): self.output_thread.join(timeout=0.5) @staticmethod - def run_engine(config: Config, input_address: str, output_address: str): + def run_engine( + config: Config, + input_address: str, + output_address: str, + env_snapshot: dict[str, str] | None = None, + ): + if env_snapshot: + os.environ.update(env_snapshot) + # Bind this EngineCore's lifetime to its parent (the server / # CoreManager): if the parent exits, have the kernel reap this process — # and, transitively, the ModelRunner workers it spawns — instead of diff --git a/atom/model_engine/engine_core_mgr.py b/atom/model_engine/engine_core_mgr.py index c399645b00..6121870437 100644 --- a/atom/model_engine/engine_core_mgr.py +++ b/atom/model_engine/engine_core_mgr.py @@ -4,6 +4,7 @@ import asyncio import logging import multiprocessing +import os import pickle import queue import weakref @@ -791,6 +792,7 @@ def launch_engine_core(config: Config, dp_rank: int = 0): "config": config, "input_address": input_address, "output_address": output_address, + "env_snapshot": dict(os.environ), }, ) diff --git a/atom/model_engine/llm_engine.py b/atom/model_engine/llm_engine.py index c1263b9859..b4eb3c48e7 100644 --- a/atom/model_engine/llm_engine.py +++ b/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", } ) diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index 4f8b989cf5..a243dc233b 100644 --- a/atom/model_engine/model_runner.py +++ b/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", @@ -679,6 +680,7 @@ def __init__(self, rank: int, config: Config): self.use_mla = self.is_deepseek_mla() self.use_gdn = self.is_qwen_next() self.use_v4 = self.is_deepseek_v4() + rope_parameters = getattr(self.hf_text_config, "rope_parameters", None) or {} self.use_mrope = "mrope_section" in rope_parameters self.is_deepseek_v32 = ( @@ -771,6 +773,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 +910,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 +1355,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: diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index 978777a1fc..0c18f3cfd2 100644 --- a/atom/model_engine/scheduler.py +++ b/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: diff --git a/atom/model_loader/loader.py b/atom/model_loader/loader.py index 81222eab05..bf0b31e8f5 100644 --- a/atom/model_loader/loader.py +++ b/atom/model_loader/loader.py @@ -6,10 +6,11 @@ import os import logging import re +import sys import threading import time from glob import glob -from typing import Generator, Tuple +from typing import Callable, Generator, Tuple from collections.abc import Iterable, Mapping from dataclasses import dataclass, field from typing import Any, Optional @@ -149,9 +150,42 @@ def default_weight_loader(param: nn.Parameter, loaded_weight: torch.Tensor): def safetensors_weights_iterator( model_name_or_path: str, disable_mmap: bool = False, -) -> Generator[Tuple[str, torch.Tensor], None, None]: +) -> Generator[Tuple[str, Callable[[], torch.Tensor]], None, None]: """Iterate over the weights in the model safetensor files.""" logger.info(f"disable_mmap: {disable_mmap}") + use_fastsafetensors = envs.ATOM_USE_FASTSAFETENSORS and not disable_mmap + fastsafe_open = None + fastsafetensors_device = "cpu" + if use_fastsafetensors: + try: + from fastsafetensors import fastsafe_open as _fastsafe_open + + fastsafe_open = _fastsafe_open + requested_device = envs.ATOM_FASTSAFETENSORS_DEVICE.lower() + if requested_device == "auto": + requested_device = ( + "cuda" if not envs.ATOM_FASTSAFETENSORS_NOGDS else "cpu" + ) + if requested_device == "cuda": + if not torch.cuda.is_available(): + raise RuntimeError( + "ATOM_FASTSAFETENSORS_DEVICE=cuda requested but CUDA/HIP " + "is not available" + ) + fastsafetensors_device = f"cuda:{torch.cuda.current_device()}" + elif requested_device != "cpu": + fastsafetensors_device = requested_device + logger.info( + "Using fastsafetensors for safetensors reads " + f"(nogds={envs.ATOM_FASTSAFETENSORS_NOGDS}, " + f"device={fastsafetensors_device})" + ) + except ImportError: + use_fastsafetensors = False + logger.warning( + "ATOM_USE_FASTSAFETENSORS=1 but fastsafetensors is not installed; " + "falling back to safetensors.safe_open" + ) path = ( model_name_or_path if os.path.isdir(model_name_or_path) @@ -165,37 +199,121 @@ def safetensors_weights_iterator( enable_tqdm = ( not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 ) + tp_group = None + tp_world_size = 1 + tp_rank = 0 + fastsafe_dist_load = False + if use_fastsafetensors and envs.ATOM_FASTSAFETENSORS_DIST_LOAD: + try: + tp_group = get_tp_group() + tp_world_size = tp_group.world_size + tp_rank = tp_group.rank_in_group + fastsafe_dist_load = tp_world_size > 1 + except Exception: + logger.exception( + "ATOM_FASTSAFETENSORS_DIST_LOAD=1 but TP group is unavailable; " + "falling back to per-rank fastsafetensors loading" + ) + if fastsafe_dist_load and tp_rank == 0: + logger.info( + "Using distributed fastsafetensors shard loading across %d TP ranks", + tp_world_size, + ) - iters = tqdm( - hf_weights_files, + pbar = tqdm( + total=len(hf_weights_files), desc=f"Loading safetensors shards[{model_name_or_path}]", - disable=not enable_tqdm, + disable=not enable_tqdm or not sys.stderr.isatty(), ) - for st_file in iters: - # Advise kernel for sequential read-ahead (mmap optimization) - if not disable_mmap and hasattr(os, "posix_fadvise"): + + def lazy_tensor_getter(reader, name: str) -> Callable[[], torch.Tensor]: + tensor: torch.Tensor | None = None + + def get_tensor() -> torch.Tensor: + nonlocal tensor + if tensor is None: + tensor = reader.get_tensor(name) + return tensor + + return get_tensor + + batch_size = tp_world_size if fastsafe_dist_load else 1 + for batch_start in range(0, len(hf_weights_files), batch_size): + st_files = hf_weights_files[batch_start : batch_start + batch_size] + if use_fastsafetensors and fastsafe_open is not None: try: - fd = os.open(st_file, os.O_RDONLY) - file_size = os.fstat(fd).st_size - os.posix_fadvise( - fd, - 0, - file_size, - os.POSIX_FADV_SEQUENTIAL | os.POSIX_FADV_WILLNEED, + fastsafe_filenames = st_files + fastsafe_pg = None + if fastsafe_dist_load: + fastsafe_filenames = { + rank: ([st_files[rank]] if rank < len(st_files) else []) + for rank in range(tp_world_size) + } + fastsafe_pg = tp_group.device_group + shard_start = time.perf_counter() + with fastsafe_open( + filenames=fastsafe_filenames, + framework="pt", + pg=fastsafe_pg, + device=fastsafetensors_device, + nogds=envs.ATOM_FASTSAFETENSORS_NOGDS, + debug_log=envs.ATOM_FASTSAFETENSORS_DEBUG, + ) as f: + for name in f.keys(): + yield name, lazy_tensor_getter(f, name) + pbar.update(len(st_files)) + if enable_tqdm: + logger.info( + "Finished safetensors shards %d-%d/%d in %.2fs: %s", + batch_start + 1, + batch_start + len(st_files), + len(hf_weights_files), + time.perf_counter() - shard_start, + ", ".join(os.path.basename(st_file) for st_file in st_files), + ) + continue + except Exception: + logger.exception( + "fastsafetensors failed for %s; falling back to safe_open", + ", ".join(st_files), ) - os.close(fd) - except OSError: - pass - - if disable_mmap: - with open(st_file, "rb") as f: - result = safetensors.torch.load(f.read()) - for name, param in result.items(): - yield name, param - else: - with safetensors.safe_open(st_file, framework="pt", device="cpu") as f: - for name in f.keys(): - yield name, f.get_tensor(name) + + for st_file in st_files: + shard_start = time.perf_counter() + # Advise kernel for sequential read-ahead (mmap optimization) + if not disable_mmap and hasattr(os, "posix_fadvise"): + try: + fd = os.open(st_file, os.O_RDONLY) + file_size = os.fstat(fd).st_size + os.posix_fadvise( + fd, + 0, + file_size, + os.POSIX_FADV_SEQUENTIAL | os.POSIX_FADV_WILLNEED, + ) + os.close(fd) + except OSError: + pass + + if disable_mmap: + with open(st_file, "rb") as f: + result = safetensors.torch.load(f.read()) + for name, param in result.items(): + yield name, lambda param=param: param + else: + with safetensors.safe_open(st_file, framework="pt", device="cpu") as f: + for name in f.keys(): + yield name, lazy_tensor_getter(f, name) + pbar.update(1) + if enable_tqdm: + logger.info( + "Finished safetensors shard %d/%d in %.2fs: %s", + batch_start + 1, + len(hf_weights_files), + time.perf_counter() - shard_start, + os.path.basename(st_file), + ) + pbar.close() # when plugin mode, model loader method is bind to model implementation @@ -600,6 +718,12 @@ def _stage_task(param, full_param_name, shard_id, global_expert_id, loaded_weigh del staging_map[pid] num_threads = envs.ATOM_LOADER_NUM_THREADS + if envs.ATOM_USE_FASTSAFETENSORS and num_threads > 1: + logger.warning( + "Disabling parallel weight loading because fastsafetensors tensors " + "must be consumed before their file context closes." + ) + num_threads = 1 if num_threads > 1: executor = concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) else: @@ -614,17 +738,21 @@ def _submit(fn, *args): try: disable_mmap = envs.ATOM_DISABLE_MMAP - for name, weight_tensor in safetensors_weights_iterator( - model_name_or_path, disable_mmap=disable_mmap - ): + if load_dummy: + logger.info("load_dummy=True: skipping checkpoint weight iteration") + loaded_weights_record.update(prefix + name for name in params_dict) + weight_iter = () + else: + weight_iter = safetensors_weights_iterator( + model_name_or_path, disable_mmap=disable_mmap + ) + for name, get_weight_tensor in weight_iter: _orig_ckpt_name = name # preserve for ckpt-side coverage report if weights_mapper is not None: mapped_name = weights_mapper._map_name(name) if mapped_name is None: continue name = mapped_name - if load_dummy: - continue # Draft models may remap ckpt-side `mtp.*` entries into params # whose names do not themselves contain `mtp` (e.g. Qwen3.5 MTP # rewrites `mtp.*` -> `model.*`). Gate only on `spec_decode`, @@ -713,7 +841,12 @@ def _submit(fn, *args): ) continue weight_loader = getattr(param, "weight_loader") - _submit(weight_loader, param, weight_tensor, shard_idx) + _submit( + weight_loader, + param, + get_weight_tensor(), + shard_idx, + ) loaded_weights_record.add(prefix + param_name) else: # Checkpoint has separate weights, load into fused param @@ -727,7 +860,12 @@ def _submit(fn, *args): dropped_ckpt_keys.append((_orig_ckpt_name, param_name)) break weight_loader = getattr(param, "weight_loader") - _submit(weight_loader, param, weight_tensor, shard_id) + _submit( + weight_loader, + param, + get_weight_tensor(), + shard_id, + ) loaded_weights_record.add(prefix + param_name) break else: @@ -763,7 +901,7 @@ def _submit(fn, *args): name, # Original checkpoint name name_mapped, # Mapped parameter name params_dict, - weight_tensor, + get_weight_tensor(), shard_id, num_experts, ) @@ -802,7 +940,7 @@ def _submit(fn, *args): name, shard_id, expert_id, - weight_tensor, + get_weight_tensor(), ) loaded_weights_record.add(prefix + name) matched = True @@ -811,7 +949,7 @@ def _submit(fn, *args): _submit( weight_loader, param, - weight_tensor, + get_weight_tensor(), name, shard_id, expert_id, @@ -835,7 +973,7 @@ def _submit(fn, *args): _submit( weight_loader, param, - weight_tensor, + get_weight_tensor(), "", # use merged moe loader "", expert_id, @@ -850,7 +988,7 @@ def _submit(fn, *args): weight_loader = getattr( param, "weight_loader", default_weight_loader ) - _submit(weight_loader, param, weight_tensor) + _submit(weight_loader, param, get_weight_tensor()) loaded_weights_record.add(prefix + name) else: # Model doesn't have expert mapping, use generic loading @@ -862,7 +1000,7 @@ def _submit(fn, *args): weight_loader = getattr( param, "weight_loader", default_weight_loader ) - _submit(weight_loader, param, weight_tensor) + _submit(weight_loader, param, get_weight_tensor()) loaded_weights_record.add(prefix + name) if executor is not None: diff --git a/atom/model_ops/attention_mha.py b/atom/model_ops/attention_mha.py index 80f349cba3..96a4a240ec 100644 --- a/atom/model_ops/attention_mha.py +++ b/atom/model_ops/attention_mha.py @@ -26,6 +26,54 @@ ) +def _torch_reshape_and_cache( + k: torch.Tensor, + v: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + slot_mapping: torch.Tensor, +) -> None: + slots = slot_mapping.flatten().to(torch.long) + # The valid-slot filter needs a host sync (valid.all()/boolean indexing) that + # is illegal during CUDA-graph capture. Decode batches (the only ones captured) + # have no padded (-1) slots, so skip the data-dependent filter while capturing. + if not torch.cuda.is_current_stream_capturing(): + valid = slots >= 0 + if not valid.all(): + slots = slots[valid] + k = k[valid] + v = v[valid] + if slots.numel() == 0: + return + + if k_cache.dim() == 5: + x = k_cache.shape[-1] + block_size = k_cache.shape[-2] + block_idx = slots // block_size + block_offset = slots % block_size + k_cache[block_idx, :, :, block_offset, :] = k.reshape( + k.shape[0], k.shape[1], k.shape[2] // x, x + ) + if v_cache.dim() == 5: + v_x = v_cache.shape[-1] + v_cache[ + block_idx, + :, + block_offset // v_x, + :, + block_offset % v_x, + ] = v + else: + v_cache[block_idx, :, :, block_offset] = v + return + + block_size = k_cache.shape[1] + block_idx = slots // block_size + block_offset = slots % block_size + k_cache[block_idx, block_offset] = k + v_cache[block_idx, block_offset] = v + + @cache def use_pa_decode_bf16_asm() -> bool: return ( @@ -331,6 +379,38 @@ def rope_cache(self, q, k, v, qkv, position, fwd_ctx: ForwardContext): output_zeros=False, ) self._cache_format = "NHD" + elif self.use_flash_layout: + # 4D FLASH layout [num_blocks, block_size, num_kv_heads, head_dim]. + # Used by Kimi-K3 full-attn (MLA-as-MHA, head_dim=192) so decode + # reads via unified_attention(shuffled_kv_cache=False), which reads + # the non-power-of-2 head_dim correctly (SHUFFLE read is broken at + # 192). Write the new (decode) token(s) with reshape_and_cache_flash. + if self.rotary_emb is not None: + assert position is not None + q, k = self.rotary_emb(position, q, k) + if self.q_norm is not None: + q = self.q_norm(q) + if self.k_norm is not None: + k = self.k_norm(k) + if envs.ATOM_USE_TORCH_CACHE: + _torch_reshape_and_cache( + k, v, k_cache, v_cache, attn_metadata.slot_mapping + ) + else: + # reshape_and_cache_flash dereferences k_scale/v_scale even for the + # "auto" bf16 path; pass the pre-allocated on-device 1.0 scale. + flash_scale = self._pa_decode_bf16_asm_scale + aiter.reshape_and_cache_flash( + k, + v, + k_cache, + v_cache, + attn_metadata.slot_mapping, + ("auto" if self.kv_cache_dtype == "bf16" else self.kv_cache_dtype), + flash_scale, + flash_scale, + ) + self._cache_format = "NHD" else: # for asm paged attention asm_layout = True @@ -355,17 +435,22 @@ def rope_cache(self, q, k, v, qkv, position, fwd_ctx: ForwardContext): asm_layout=asm_layout, ) else: - aiter.reshape_and_cache( - k, - v, - k_cache, - v_cache, - attn_metadata.slot_mapping, - kv_cache_dtype="auto", - k_scale=None, - v_scale=None, - asm_layout=asm_layout, - ) + if envs.ATOM_USE_TORCH_CACHE: + _torch_reshape_and_cache( + k, v, k_cache, v_cache, attn_metadata.slot_mapping + ) + else: + aiter.reshape_and_cache( + k, + v, + k_cache, + v_cache, + attn_metadata.slot_mapping, + kv_cache_dtype="auto", + k_scale=None, + v_scale=None, + asm_layout=asm_layout, + ) self._cache_format = "SHUFFLE" if asm_layout else "NHD" # NOTE: on a prefix-cache hit the cached+new KV is gathered into a dense @@ -794,23 +879,31 @@ def prefill_attention_triton( (self.sliding_window - 1, 0) if self.sliding_window > 0 else (-1, -1) ) - # `block_tables` is always populated by TritonMHAMetadataBuilder. - # For pure prefill (no cached tokens) it is, by default, the fake table - # built in prepare_prefill that maps seq i to token indices - # [cu_seqlens_k[i], ..., cu_seqlens_k[i+1]-1], paired with raw K/V - # treated as kv_cache with block_size=1. - # - # Under ATOM_USE_UNIFIED_ATTN, prepare_prefill instead uploads the real - # per-seq block_table and reads from KV cache, the new tokens - # already written into the paged flash-layout cache during rope_cache - # are read straight from `k_cache`/`v_cache`, identical to the - # prefix-cache-hit path. + # Prefix-cache hits use real paged block tables. Pure prefill may not + # have paged KV metadata; in that case, feed raw K/V as a block_size=1 + # flash-layout cache and synthesize a table that maps each sequence's + # logical positions to dense token offsets. + block_table = attn_metadata.block_tables + if attn_metadata.block_tables is None and not attn_metadata.has_cached: + offsets = torch.arange( + attn_metadata.max_seqlen_k, dtype=torch.int32, device=q.device + ) + seq_starts = attn_metadata.cu_seqlens_k[:-1].to(torch.int32) + block_table = seq_starts[:, None] + offsets[None, :] + if envs.ATOM_USE_UNIFIED_ATTN or attn_metadata.has_cached: - k_for_attn = k_cache - v_for_attn = v_cache - # Reads the paged KV cache, which is 5D SHUFFLE unless the (default) - # 4D flash layout is in use. - shuffled_kv_cache = not self.use_flash_layout + if attn_metadata.block_tables is None and not attn_metadata.has_cached: + # k: [total_tokens, num_kv_heads, head_size] + # -> [total_tokens, 1, num_kv_heads, head_size] + k_for_attn = k.unsqueeze(1) + v_for_attn = v.unsqueeze(1) + shuffled_kv_cache = False + else: + k_for_attn = k_cache + v_for_attn = v_cache + # Reads the paged KV cache, which is 5D SHUFFLE unless the + # (default) 4D flash layout is in use. + shuffled_kv_cache = not self.use_flash_layout else: # k: [total_tokens, num_kv_heads, head_size] # -> [total_tokens, 1, num_kv_heads, head_size] @@ -832,7 +925,7 @@ def prefill_attention_triton( causal=True, alibi_slopes=None, window_size=sliding_window, - block_table=attn_metadata.block_tables, + block_table=block_table, softcap=0, q_descale=None, k_descale=self.kv_scale, diff --git a/atom/model_ops/attentions/aiter_attention.py b/atom/model_ops/attentions/aiter_attention.py index bd22713269..89d02b18a2 100644 --- a/atom/model_ops/attentions/aiter_attention.py +++ b/atom/model_ops/attentions/aiter_attention.py @@ -617,7 +617,13 @@ def build_kv_cache_tensor(self, layer_id: int, module): # attn_idx: hybrid models (Qwen3-Next) skip linear-attention layers # in the kv_cache slot ordering; non-hybrid models use layer_id 1:1. - if runner.is_qwen_next(): + if getattr(runner, "is_kimi_linear", lambda: False)(): + mtp_start = runner.mtp_start_layer_idx + if layer_id < mtp_start: + attn_idx = runner.full_attention_layers.index(layer_id) + else: + attn_idx = runner.num_full_attn + (layer_id - mtp_start) + elif runner.is_qwen_next(): mtp_start = runner.mtp_start_layer_idx if layer_id < mtp_start: attn_idx = layer_id // runner.full_attention_interval @@ -668,6 +674,33 @@ def build_kv_cache_tensor(self, layer_id: int, module): runner._kv_layer_cache_store.append( (k_cache, v_cache, module.k_scale, module.v_scale) ) + elif getattr(runner, "is_kimi_linear", lambda: False)(): + # Kimi-K3 full-attn (MLA-as-MHA) layers run at head_dim=192 + # (non-power-of-2) and are bound in the 4D FLASH layout + # [num_blocks, block_size, num_kv_heads, head_dim]: + # unified_attention(shuffled_kv_cache=False) reads it correctly at + # 192, writes go through reshape_and_cache_flash, and the backing + # storage runner.kv_cache[k/v, attn_idx] is already exactly this 4D + # shape, so no reshape is needed. Validated: gsm8k-1319 graph = + # 0.9431 flex / 0.9424 strict (> baseline 0.9378). + # + # The 5D SHUFFLE decode path is deliberately NOT used at 192: + # (1) the stock aiter SHUFFLE read mis-indexes the padded head_dim + # at 192 (~100% error); + # (2) a candidate kernel fix for (1) (kept in + # code_k3/shuffle_kernel_fix.patch) is bit-accurate in op-tests + # and eager server decode, but it still (a) regresses under + # CUDA-graph capture (full-1319 graph: SHUFFLE 0.9204 vs FLASH + # 0.9431) and (b) destabilises the SHARED unified_attention + # FLASH path in-server (non-deterministic garbage) even though + # the edit is SHUFFLED_KV_CACHE-guarded and op-test-clean. + # So SHUFFLE decode at 192 stays a follow-up; keep FLASH. + k_cache = runner.kv_cache[0, attn_idx] + v_cache = runner.kv_cache[1, attn_idx] + module.impl.use_flash_layout = True + if config.kv_cache_dtype == "fp8": + module.k_scale = runner.kv_scale[0, attn_idx] + module.v_scale = runner.kv_scale[1, attn_idx] else: x = 16 // runner.kv_cache.element_size() k_cache = runner.kv_cache[0, attn_idx].view( diff --git a/atom/model_ops/attentions/gdn_attn.py b/atom/model_ops/attentions/gdn_attn.py index 4e2cb9602f..d0cc923534 100644 --- a/atom/model_ops/attentions/gdn_attn.py +++ b/atom/model_ops/attentions/gdn_attn.py @@ -84,13 +84,42 @@ def __init__( # a hidden ordering dependency on _compute_block_bytes being # called first. hf = model_runner.config.hf_config - model_runner.full_attention_interval = hf.full_attention_interval - model_runner.num_full_attn = ( - hf.num_hidden_layers // model_runner.full_attention_interval - ) - model_runner.num_gdn_attn_state = ( - hf.num_hidden_layers - model_runner.num_full_attn - ) + if getattr(hf, "model_type", None) == "kimi_linear": + lin = getattr(hf, "linear_attn_config", {}) or {} + model_runner.full_attention_layers = [ + int(i) - 1 for i in lin.get("full_attn_layers", []) + ] + model_runner.kda_attention_layers = [ + int(i) - 1 for i in lin.get("kda_layers", []) + ] + model_runner.num_full_attn = len(model_runner.full_attention_layers) + model_runner.num_gdn_attn_state = len(model_runner.kda_attention_layers) + hf.linear_num_key_heads = getattr( + hf, "linear_num_key_heads", lin.get("num_heads", hf.num_attention_heads) + ) + hf.linear_num_value_heads = getattr( + hf, "linear_num_value_heads", lin.get("num_heads", hf.num_attention_heads) + ) + hf.linear_key_head_dim = getattr( + hf, "linear_key_head_dim", lin.get("head_dim", hf.qk_nope_head_dim) + ) + hf.linear_value_head_dim = getattr( + hf, "linear_value_head_dim", lin.get("head_dim", hf.v_head_dim) + ) + hf.linear_conv_kernel_dim = getattr( + hf, + "linear_conv_kernel_dim", + lin.get("short_conv_kernel_size", 4), + ) + hf.head_dim = hf.qk_nope_head_dim + hf.qk_rope_head_dim + else: + model_runner.full_attention_interval = hf.full_attention_interval + model_runner.num_full_attn = ( + hf.num_hidden_layers // model_runner.full_attention_interval + ) + model_runner.num_gdn_attn_state = ( + hf.num_hidden_layers - model_runner.num_full_attn + ) self.num_spec = 0 if hasattr(model_runner, "drafter"): @@ -186,6 +215,11 @@ def _state_shape( return conv_state_shape, temporal_state_shape def _state_dtypes(self) -> tuple[torch.dtype, torch.dtype]: + if getattr(self.model_runner.config.hf_config, "model_type", None) == "kimi_linear": + return ( + self.model_runner.config.torch_dtype, + torch.float32, + ) return ( self.model_runner.config.torch_dtype, self.model_runner.config.torch_dtype, @@ -316,8 +350,11 @@ def build_kv_cache_tensor(self, layer_id: int, module): from atom.config import KVCacheTensor runner = self.model_runner - interval = runner.full_attention_interval - gdn_idx = (layer_id // interval) * (interval - 1) + (layer_id % interval) + if getattr(runner, "is_kimi_linear", lambda: False)(): + gdn_idx = runner.kda_attention_layers.index(layer_id) + else: + interval = runner.full_attention_interval + gdn_idx = (layer_id // interval) * (interval - 1) + (layer_id % interval) return KVCacheTensor( layer_num=layer_id, k_cache=runner.mamba_k_cache[gdn_idx], diff --git a/atom/model_ops/fla_ops/fused_sigmoid_gating.py b/atom/model_ops/fla_ops/fused_sigmoid_gating.py index 236504ddc1..c143fc7703 100644 --- a/atom/model_ops/fla_ops/fused_sigmoid_gating.py +++ b/atom/model_ops/fla_ops/fused_sigmoid_gating.py @@ -19,6 +19,7 @@ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, "IS_CONTINUOUS_BATCHING": lambda args: args["ssm_state_indices"] is not None, "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None, + "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None, } ) @triton.jit(do_not_specialize=["N", "T"]) @@ -29,6 +30,7 @@ def fused_sigmoid_gating_delta_rule_update_kernel( dt_bias, beta, threshold, + lower_bound, q, k, v, @@ -61,6 +63,7 @@ def fused_sigmoid_gating_delta_rule_update_kernel( IS_CONTINUOUS_BATCHING: tl.constexpr, IS_SPEC_DECODING: tl.constexpr, IS_KDA: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, ): i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_n, i_hv = i_nh // HV, i_nh % HV @@ -130,10 +133,19 @@ def fused_sigmoid_gating_delta_rule_update_kernel( # If the model is loaded in fp16, without the .float() here, A might be -inf x = tl.load(p_a).to(tl.float32) + tl.load(p_dt_bias).to(tl.float32) - softplus_x = tl.where( - beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x - ) - b_g = -tl.exp(tl.load(p_A_log).to(tl.float32)) * softplus_x + b_A = tl.load(p_A_log).to(tl.float32) + if USE_LOWER_BOUND: + # Kimi-KDA lower-bounded sigmoid gate (matches fla + # fused_recurrent_kda's USE_LOWER_BOUND branch): the per-channel + # decay is `lower_bound * sigmoid(exp(A_log) * (a + dt_bias))`, + # keeping the log-decay bounded in `(lower_bound, 0)` rather than + # the unbounded `-exp(A) * softplus(.)` used by GDN. + b_g = lower_bound * tl.sigmoid(tl.exp(b_A) * x) + else: + softplus_x = tl.where( + beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x + ) + b_g = -tl.exp(b_A) * softplus_x # compute beta_output = sigmoid(b) b_beta = tl.sigmoid(b_b.to(tl.float32)) @@ -200,11 +212,20 @@ def fused_sigmoid_gating_delta_rule_update( num_accepted_tokens: torch.Tensor | None = None, use_qk_l2norm_in_kernel: bool = False, is_kda: bool = False, + lower_bound: float | None = None, ): """ Fused triton implementation of sigmoid gating delta rule update. This function uses a single fused kernel that combines both sigmoid gating computation and the recurrent delta rule update for better performance. + + Gate variants: + * lower_bound is None (default, GDN / Qwen3-Next): the log-decay is the + unbounded `-exp(A_log) * softplus(a + dt_bias)`. + * lower_bound set (Kimi-KDA): the log-decay is the lower-bounded + `lower_bound * sigmoid(exp(A_log) * (a + dt_bias))`, matching fla's + `fused_recurrent_kda` USE_LOWER_BOUND branch. Pair with is_kda=True so + `a`/`dt_bias` are read as per-K-channel vectors (KDA's diagonal decay). """ B, T, H, K, V = *k.shape, v.shape[-1] HV = v.shape[2] @@ -253,6 +274,7 @@ def fused_sigmoid_gating_delta_rule_update( dt_bias=dt_bias, beta=beta, threshold=threshold, + lower_bound=lower_bound, q=q.contiguous(), k=k.contiguous(), v=v.contiguous(), @@ -272,7 +294,9 @@ def fused_sigmoid_gating_delta_rule_update( V=V, BK=BK, BV=BV, - stride_a_token=a.stride(-2), + # Per-token stride of `a`: GDN gate is [B, T, HV] (T stride = stride(-2)); + # KDA gate is per-K-channel [B, T, HV, K], so the token stride is stride(-3). + stride_a_token=a.stride(-3) if is_kda else a.stride(-2), stride_b_token=b.stride(-2), stride_init_state_token=stride_init_state_token, stride_final_state_token=stride_final_state_token, diff --git a/atom/model_ops/fused_moe_triton.py b/atom/model_ops/fused_moe_triton.py index ac5d5b0f7c..1a1f30d075 100644 --- a/atom/model_ops/fused_moe_triton.py +++ b/atom/model_ops/fused_moe_triton.py @@ -129,6 +129,8 @@ def triton_kernel_moe_forward( global_num_experts: int = -1, expert_map: torch.Tensor | None = None, act_quant: MoEActivationQuant = MoEActivationQuant.BF16, + situ_beta: float | None = None, + situ_linear_beta: float | None = None, ) -> torch.Tensor: routing_data, gather_idx, scatter_idx = routing( gating_output, topk, sm_first=not renormalize @@ -159,6 +161,8 @@ def triton_kernel_moe_forward( global_num_experts=global_num_experts, expert_map=expert_map, act_quant=act_quant, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, ) @@ -188,6 +192,8 @@ def triton_kernel_fused_experts( expert_map: torch.Tensor | None = None, intermediate_cache: torch.Tensor | None = None, act_quant: MoEActivationQuant = MoEActivationQuant.BF16, + situ_beta: float | None = None, + situ_linear_beta: float | None = None, ) -> torch.Tensor: # type check, uint8 means mxfp4 assert hidden_states.dtype == torch.bfloat16 @@ -338,13 +344,28 @@ def triton_kernel_fused_experts( raw_2d = raw_intermediate.view(M * topk, N) intermediate_cache = intermediate_cache.view(M * topk, half_N) - fused_clamp_act_mul( - raw_2d, - out=intermediate_cache, - swiglu_limit=swiglu_limit, - activation="silu", - dtype_quant=None, - ) + if activation == ActivationType.Situv2: + # Kimi situ: beta*tanh(gate/beta)*sigmoid(gate) * linear_beta*tanh(up/linear_beta). + # fused_clamp_act_mul only implements silu/gelu (act(gate)*up), which + # is NOT situ, so compute it in fp32 here. situ ~= silu for small + # activations but diverges as tanh saturates, so the silu shortcut + # corrupts large-magnitude expert outputs. + gate = raw_2d[:, :half_N].float() + up = raw_2d[:, half_N:].float() + beta = float(situ_beta) if situ_beta is not None else 1.0 + situ_gate = beta * torch.tanh(gate / beta) * torch.sigmoid(gate) + if situ_linear_beta is not None: + lbeta = float(situ_linear_beta) + up = lbeta * torch.tanh(up / lbeta) + intermediate_cache.copy_((situ_gate * up).to(intermediate_cache.dtype)) + else: + fused_clamp_act_mul( + raw_2d, + out=intermediate_cache, + swiglu_limit=swiglu_limit, + activation="silu", + dtype_quant=None, + ) if act_quant == MoEActivationQuant.FP4: intermediate_fp4, intermediate_mx_scale = mxfp4_quant(intermediate_cache) diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index 4a27882689..e082cbc56d 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -2,6 +2,8 @@ # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. import logging +import os +import time from abc import abstractmethod from dataclasses import dataclass from enum import Enum @@ -72,6 +74,7 @@ from transformers import PretrainedConfig logger = logging.getLogger("atom") +_MOE_LOAD_COPY_LOG_SECONDS = float(os.getenv("ATOM_MOE_LOAD_COPY_LOG_SECONDS", "0")) class MoEActivationQuant(Enum): @@ -1215,6 +1218,13 @@ def apply( ) -> torch.Tensor: if self.use_triton_decode and not get_forward_context().context.is_prefill: # Triton decode is GGUU-only; GUGU uses the FlyDSL path. + # NOTE: the a8w4 GGUU decode kernel hardcodes SiLU activation and does + # not implement Kimi situ — guard so it fails loudly instead of + # silently degrading accuracy if enabled for a situ model. + assert activation != ActivationType.Situv2, ( + "ATOM_USE_TRITON_MOE_DECODE (a8w4 GGUU) does not support situ " + "activation; disable it for Kimi-K3." + ) from aiter.ops.triton.moe.moe_routing.routing import routing from atom.model_ops.fused_moe_triton import ( triton_kernel_fused_experts_a8w4_silu_gguu, @@ -1328,6 +1338,8 @@ def apply( global_num_experts=n_expts_tot, expert_map=expert_map, act_quant=self.act_quant, + situ_beta=getattr(layer, "activation_situ_beta", None), + situ_linear_beta=getattr(layer, "activation_situ_linear_beta", None), ) # Always-on shared expert(s) via a standalone dense GEMM, @@ -1364,6 +1376,8 @@ def apply( apply_router_weight_on_input=apply_router_weight_on_input, global_num_experts=global_num_experts, act_quant=self.act_quant, + situ_beta=getattr(layer, "activation_situ_beta", None), + situ_linear_beta=getattr(layer, "activation_situ_linear_beta", None), ) topk_weights, topk_ids = self.select_experts_with_record( @@ -1391,6 +1405,11 @@ def apply( ), "swiglu_limit": getattr(layer, "swiglu_limit", 0.0), } + if activation == ActivationType.Situv2: + moe_extra_args["beta"] = getattr(layer, "activation_situ_beta", None) + moe_extra_args["linear_beta"] = getattr( + layer, "activation_situ_linear_beta", None + ) if self.fused_experts is None: return fused_moe( x, @@ -2591,6 +2610,14 @@ def __init__( self.scoring_func = scoring_func self.e_score_correction_bias = e_score_correction_bias self.activation = activation + if config is not None: + self.activation_situ_beta = getattr(config, "activation_situ_beta", None) + self.activation_situ_linear_beta = getattr( + config, "activation_situ_linear_beta", None + ) + else: + self.activation_situ_beta = None + self.activation_situ_linear_beta = None self.use_chunked = get_dp_group().world_size > 1 @@ -3025,8 +3052,31 @@ def _load_per_channel_weight_scale( @staticmethod def _copy_quant_storage(dst: torch.Tensor, src: torch.Tensor) -> None: """Copy quantized tensors without numeric conversion of byte formats.""" + if src.device.type == "cpu" and not src.is_contiguous(): + src = src.contiguous() + start = time.perf_counter() if _MOE_LOAD_COPY_LOG_SECONDS > 0 else None + + def log_slow_copy() -> None: + if start is None: + return + if dst.device.type == "cuda": + torch.cuda.synchronize(dst.device) + elapsed = time.perf_counter() - start + if elapsed >= _MOE_LOAD_COPY_LOG_SECONDS: + logger.info( + "Slow MoE quant copy: %.3fs dst_shape=%s dst_dtype=%s " + "src_shape=%s src_dtype=%s src_contiguous=%s", + elapsed, + tuple(dst.shape), + dst.dtype, + tuple(src.shape), + src.dtype, + src.is_contiguous(), + ) + if dst.dtype == dtypes.fp4x2: dst.view(torch.uint8).copy_(src.view(torch.uint8)) + log_slow_copy() return fp8_storage_dtypes = ( torch.float8_e4m3fn, @@ -3040,11 +3090,13 @@ def _copy_quant_storage(dst: torch.Tensor, src: torch.Tensor) -> None: # numeric conversion when destination storage uses a different FP8 # variant; later scale fixups expect the original bytes. dst.view(torch.uint8).copy_(src.view(torch.uint8)) + log_slow_copy() return if dst.dtype == dtypes.fp8_e8m0 and src.dtype == torch.uint8: # e8m0 microscale tensors are byte-encoded; copy_ would convert the # uint8 values numerically instead of preserving the scale bits. dst.view(torch.uint8).copy_(src) + log_slow_copy() return if dst.dtype == torch.uint8 and src.dtype in ( torch.float8_e8m0fnu, @@ -3052,6 +3104,7 @@ def _copy_quant_storage(dst: torch.Tensor, src: torch.Tensor) -> None: ): src = src.view(torch.uint8) dst.copy_(src) + log_slow_copy() def _load_w13( self, @@ -3125,9 +3178,9 @@ def _load_w2( # Narrow parameter and load. shard_size = expert_data.shape[shard_dim] load_shard_size = loaded_weight.shape[shard_dim] // self.tp_size - loaded_weight = loaded_weight.narrow( - shard_dim, load_shard_size * tp_rank, load_shard_size - ) + shard_start = load_shard_size * tp_rank + loaded_shard = loaded_weight.narrow(shard_dim, shard_start, load_shard_size) + loaded_weight = loaded_shard if load_shard_size != shard_size: expert_data = expert_data.narrow(shard_dim, 0, load_shard_size) # w2, down_proj: Load into only logical weight of w2. @@ -3734,9 +3787,16 @@ def forward_impl_graph( if self.reduce_results and (self.tp_size > 1 or self.ep_size > 1): # Default set to False. (May have to add shared expert outputs.) - final_hidden_states = get_tp_group().all_reduce( - final_hidden_states, ca_fp8_quant=False - ) + from atom.model_ops.linear import fp32_allreduce_enabled + + if fp32_allreduce_enabled(): + hs32 = final_hidden_states.float() + torch.distributed.all_reduce(hs32, group=get_tp_group().device_group) + final_hidden_states = hs32.to(final_hidden_states.dtype) + else: + final_hidden_states = get_tp_group().all_reduce( + final_hidden_states, ca_fp8_quant=False + ) return final_hidden_states @@ -3787,9 +3847,16 @@ def forward_impl(self, hidden_states: torch.Tensor, router_logits: torch.Tensor) if self.reduce_results and (self.tp_size > 1 or self.ep_size > 1): # Default set to False. (May have to add shared expert outputs.) - final_hidden_states = get_tp_group().all_reduce( - final_hidden_states, ca_fp8_quant=False - ) + from atom.model_ops.linear import fp32_allreduce_enabled + + if fp32_allreduce_enabled(): + hs32 = final_hidden_states.float() + torch.distributed.all_reduce(hs32, group=get_tp_group().device_group) + final_hidden_states = hs32.to(final_hidden_states.dtype) + else: + final_hidden_states = get_tp_group().all_reduce( + final_hidden_states, ca_fp8_quant=False + ) return final_hidden_states diff --git a/atom/model_ops/sampler.py b/atom/model_ops/sampler.py index a1f52d47c4..6d64a94670 100644 --- a/atom/model_ops/sampler.py +++ b/atom/model_ops/sampler.py @@ -7,6 +7,7 @@ from aiter import mixed_sample_outer_exponential from aiter.ops.triton.softmax import softmax from aiter.ops.triton.topk import topk +from atom.utils import envs from torch import nn # Try to import aiter top-k/top-p sampling ops @@ -73,6 +74,11 @@ def forward( Returns: Sampled token IDs (num_tokens,) """ + if envs.ATOM_DEBUG_TOPK > 0: + from atom.utils.debug_helper import maybe_log_topk + + maybe_log_topk(logits, prefix="sampler ") + # No Top-K Top-P parameters, perform temperature-based sampling if not self._needs_filtering(top_ks, top_ps): return self._temperature_sample( @@ -126,6 +132,19 @@ def _temperature_sample( exponential = get_per_token_exponential(vocab_size, logits.device).expand( num_tokens, vocab_size ) + if envs.ATOM_USE_TORCH_SAMPLER: + greedy_mask = temperatures <= self.eps + sampled = torch.empty(num_tokens, dtype=torch.int, device=logits.device) + if greedy_mask.any(): + sampled[greedy_mask] = logits[greedy_mask].argmax(dim=-1).to(torch.int) + if (~greedy_mask).any(): + rows = ~greedy_mask + scaled_logits = logits[rows] / temperatures[rows].unsqueeze(-1) + probs = torch.softmax(scaled_logits, dim=-1, dtype=torch.float32) + sampled[rows] = torch.multinomial(probs, num_samples=1).squeeze(-1).to( + torch.int + ) + return sampled mixed_sample_outer_exponential( sampled_tokens, logits, exponential, temperatures, eps=self.eps ) @@ -163,7 +182,7 @@ def _topk_topp_sample( has_topk = top_ks is not None has_topp = top_ps is not None - if AITER_TOPK_TOPP_AVAILABLE: + if AITER_TOPK_TOPP_AVAILABLE and not envs.ATOM_USE_TORCH_SAMPLER: return self._aiter_sample( probs, top_ks, top_ps, has_topk, has_topp, temperatures ) diff --git a/atom/model_ops/topK.py b/atom/model_ops/topK.py index 2c1a3599e5..b66caf826d 100644 --- a/atom/model_ops/topK.py +++ b/atom/model_ops/topK.py @@ -8,6 +8,7 @@ from aiter.jit.utils.torch_guard import torch_compile_guard from atom.config import get_current_atom_config from atom.model_ops.utils import _has_module +from atom.utils import envs from atom.utils.custom_register import direct_register_custom_op @@ -213,6 +214,7 @@ def rocm_aiter_biased_grouped_topk_impl( ) -> tuple[torch.Tensor, torch.Tensor]: from aiter import biased_grouped_topk + from aiter.ops.topk import biased_grouped_topk_torch token = gating_output.shape[0] device = gating_output.device @@ -237,6 +239,20 @@ def rocm_aiter_biased_grouped_topk_impl( else: topk_ids = torch.empty((token, topk), dtype=torch.int32, device=device) topk_weights = torch.empty((token, topk), dtype=torch.float32, device=device) + if envs.AITER_USE_TORCH_TOPK: + torch_weights, torch_ids = biased_grouped_topk_torch( + gating_output, + correction_bias, + topk, + need_renorm, + num_expert_group, + topk_group, + ) + topk_weights.copy_(torch_weights * routed_scaling_factor) + topk_ids.copy_(torch_ids) + if num_fused_shared_experts > 0: + return total_topk_weights, total_topk_ids + return topk_weights, topk_ids biased_grouped_topk( gating_output, correction_bias, @@ -314,6 +330,7 @@ def rocm_aiter_grouped_topk_impl( ) -> tuple[torch.Tensor, torch.Tensor]: from aiter import grouped_topk + from aiter.ops.topk import grouped_topk_torch token = gating_output.shape[0] device = gating_output.device @@ -338,6 +355,20 @@ def rocm_aiter_grouped_topk_impl( else: topk_ids = torch.empty((token, topk), dtype=torch.int32, device=device) topk_weights = torch.empty((token, topk), dtype=torch.float32, device=device) + if envs.AITER_USE_TORCH_TOPK: + torch_weights, torch_ids = grouped_topk_torch( + gating_output, + topk, + need_renorm, + num_expert_group, + topk_group, + scoring_func, + ) + topk_weights.copy_(torch_weights * routed_scaling_factor) + topk_ids.copy_(torch_ids) + if num_fused_shared_experts > 0: + return total_topk_weights, total_topk_ids + return topk_weights, topk_ids grouped_topk( gating_output, topk_weights, diff --git a/atom/models/kimi_k3.py b/atom/models/kimi_k3.py new file mode 100644 index 0000000000..43c12f5eda --- /dev/null +++ b/atom/models/kimi_k3.py @@ -0,0 +1,1420 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +"""Inference-only Kimi-K3 text model. + +The checkpoint is multimodal, but ATOM serves the text path here. The language +weights live under ``language_model.*`` in the checkpoint, so this module keeps +the same object hierarchy and skips the vision tower/projector tensors. +""" + +import os +from typing import Optional, Union + +import torch +from aiter import ActivationType, QuantType +from aiter.dist.communication_op import tensor_model_parallel_all_reduce +from aiter.dist.parallel_state import ( + get_pp_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from einops import rearrange +from torch import nn + +from atom.config import Config, QuantizationConfig, get_current_atom_config +from atom.model_ops.base_attention import Attention + +# Side-effect import: registers `torch.ops.aiter.maybe_dual_stream_forward`, +# the Dynamo-opaque dispatcher used to overlap shared experts with routed +# experts on a separate CUDA stream (shared with DeepSeek V2/V4). +from atom.model_ops import module_dispatch_ops as _module_dispatch_ops # noqa: F401 +from atom.model_ops.embed_head import ParallelLMHead, VocabParallelEmbedding +from atom.model_ops.layernorm import RMSNorm +from atom.model_ops.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from atom.model_ops.mamba_ops.causal_conv1d import ( + causal_conv1d_fn, + causal_conv1d_update, +) +from atom.model_ops.fla_ops.fused_sigmoid_gating import ( + fused_sigmoid_gating_delta_rule_update, +) +from atom.model_ops.moe import FusedMoE +from atom.model_ops.utils import atom_parameter +from atom.models.utils import ( + IntermediateTensors, + PPMissingLayer, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_prefix, +) +from atom.utils import envs, mark_spliting_op +from atom.utils.forward_context import get_forward_context +from atom.utils.decorators import support_torch_compile + + +def _text_config(config): + return getattr(config, "text_config", config) + + +def _normalize_kimi_config(config) -> None: + """Fill the aliases expected by shared ATOM MoE/GDN infrastructure.""" + + config.n_routed_experts = getattr(config, "n_routed_experts", config.num_experts) + config.num_experts_per_tok = getattr( + config, "num_experts_per_tok", config.num_experts_per_token + ) + config.n_shared_experts = getattr( + config, "n_shared_experts", getattr(config, "num_shared_experts", 0) + ) + config.norm_topk_prob = getattr( + config, "norm_topk_prob", getattr(config, "moe_renormalize", True) + ) + config.scoring_func = getattr( + config, "scoring_func", getattr(config, "moe_router_activation_func", "sigmoid") + ) + config.n_group = getattr(config, "n_group", getattr(config, "num_expert_group", 1)) + + lin = getattr(config, "linear_attn_config", {}) or {} + config.linear_num_key_heads = getattr( + config, "linear_num_key_heads", lin.get("num_heads", config.num_attention_heads) + ) + config.linear_num_value_heads = getattr( + config, + "linear_num_value_heads", + lin.get("num_heads", config.num_attention_heads), + ) + config.linear_key_head_dim = getattr( + config, "linear_key_head_dim", lin.get("head_dim", config.qk_nope_head_dim) + ) + config.linear_value_head_dim = getattr( + config, "linear_value_head_dim", lin.get("head_dim", config.v_head_dim) + ) + config.linear_conv_kernel_dim = getattr( + config, "linear_conv_kernel_dim", lin.get("short_conv_kernel_size", 4) + ) + config.kimi_full_attn_layers = [int(i) - 1 for i in lin.get("full_attn_layers", [])] + config.kimi_kda_layers = [int(i) - 1 for i in lin.get("kda_layers", [])] + config.num_gdn_attn_state = len(config.kimi_kda_layers) + config.num_full_attn = len(config.kimi_full_attn_layers) + + # Kimi full-attention layers run MLA math but are stored in the standard + # paged-MHA cache by padding V to q_head_dim. + config.head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim + if getattr(config, "rope_parameters", None) is None: + config.rope_parameters = { + "rope_theta": getattr(config, "rope_theta", 10000.0), + "rope_type": "default", + } + + +def _kda_packed_modules_mapping( + kda_layer_indices: list[int], +) -> dict[str, tuple[str, int]]: + mapping = { + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + } + projection_names = ("q_proj", "k_proj", "v_proj", "g_proj") + for layer_idx in kda_layer_indices: + prefix = f".layers.{layer_idx}.self_attn." + for shard_id, projection_name in enumerate(projection_names): + mapping[f"{prefix}{projection_name}"] = (f"{prefix}in_proj", shard_id) + return mapping + + +def _extract_layer_idx(prefix: str) -> int: + for part in reversed(prefix.split(".")): + if part.isdigit(): + return int(part) + return 0 + + +class SituAndMul(nn.Module): + def __init__(self, beta: float = 1.0, linear_beta: float | None = None): + super().__init__() + self.beta = beta + self.linear_beta = linear_beta + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if os.environ.get("ATOM_K3_FUSED", "1") == "1": + from atom.models.kimi_k3_fused import situ_and_mul + + return situ_and_mul(x, self.beta, self.linear_beta) + gate, up = x.chunk(2, dim=-1) + gate_f = gate.float() + up_f = up.float() + out = self.beta * torch.tanh(gate_f / self.beta) * torch.sigmoid(gate_f) + if self.linear_beta is not None: + up_f = self.linear_beta * torch.tanh(up_f / self.linear_beta) + return (out * up_f).to(x.dtype) + + +class KimiRMSNormGated(nn.Module): + def __init__(self, hidden_size: int, eps: float): + super().__init__() + self.weight = atom_parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + if os.environ.get("ATOM_K3_FUSED", "1") == "1": + from atom.models.kimi_k3_fused import rmsnorm_gated + + return rmsnorm_gated(x, self.weight, gate, self.variance_epsilon) + dtype = x.dtype + x_f = x.float() + var = x_f.pow(2).mean(dim=-1, keepdim=True) + x = x_f * torch.rsqrt(var + self.variance_epsilon) + return (x.to(dtype) * self.weight.to(dtype)) * torch.sigmoid(gate) + + +def _sharded_vector_loader(tp_rank: int, tp_size: int): + def loader(param: nn.Parameter, loaded_weight: torch.Tensor): + shard = loaded_weight.narrow(0, tp_rank * param.numel(), param.numel()) + param.data.copy_(shard.to(param.dtype).view_as(param)) + + return loader + + +class KimiMLP(nn.Module): + def __init__( + self, + config, + hidden_size: int | None = None, + intermediate_size: int | None = None, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + prefix: str = "", + ): + super().__init__() + hidden_size = hidden_size or config.hidden_size + intermediate_size = intermediate_size or config.intermediate_size + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size, intermediate_size], + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=f"{prefix}.down_proj", + ) + if config.hidden_act != "situ": + raise ValueError(f"Unsupported Kimi-K3 activation: {config.hidden_act}") + self.act_fn = SituAndMul( + beta=getattr(config, "activation_situ_beta", None) or 1.0, + linear_beta=getattr(config, "activation_situ_linear_beta", None), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(self.act_fn(self.gate_up_proj(x))) + + +class KimiSparseMoeBlock(nn.Module): + def __init__( + self, + config, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + alt_stream: Optional["torch.cuda.Stream"] = None, + ): + super().__init__() + self.config = config + self.prefix = prefix + self.alt_stream = alt_stream + self.hidden_dim = config.hidden_size + self.num_experts = config.num_experts + self.top_k = config.num_experts_per_token + self.tp_size = get_tensor_model_parallel_world_size() + self.use_latent_moe = ( + getattr(config, "routed_expert_hidden_size", None) is not None + ) + self.moe_hidden_size = ( + config.routed_expert_hidden_size + if self.use_latent_moe + else config.hidden_size + ) + + self.gate = ReplicatedLinear( + config.hidden_size, + config.num_experts, + bias=False, + quant_config=None, + prefix=f"{prefix}.gate", + ) + self.gate.e_score_correction_bias = atom_parameter( + torch.empty(config.num_experts, dtype=torch.bfloat16) + ) + self.experts = FusedMoE( + num_experts=config.num_experts, + top_k=config.num_experts_per_token, + hidden_size=self.moe_hidden_size, + intermediate_size=config.moe_intermediate_size, + reduce_results=False, + renormalize=config.moe_renormalize, + quant_config=quant_config, + use_grouped_topk=getattr(config, "use_grouped_topk", True), + num_expert_group=getattr(config, "num_expert_group", 1), + topk_group=getattr(config, "topk_group", 1), + scoring_func=config.moe_router_activation_func, + e_score_correction_bias=self.gate.e_score_correction_bias, + activation=ActivationType.Situv2, + config=config, + prefix=f"{prefix}.experts", + # inter=3072/TP8=384 is a 128-multiple; pad to 128 (not the 256 + # default) to avoid padding the MXFP4 MoE intermediate up to 512. + pad_align=128, + ) + if getattr(config, "num_shared_experts", 0): + self.shared_experts = KimiMLP( + config, + intermediate_size=config.moe_intermediate_size + * config.num_shared_experts, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.shared_experts", + ) + else: + self.shared_experts = None + + if self.use_latent_moe: + + def _routed_source_quant_dtype(layer_prefix: str) -> torch.dtype | None: + if quant_config is None: + return None + layer_quant_config = quant_config.get_layer_quant_config(layer_prefix) + if ( + layer_quant_config.quant_type == QuantType.per_1x32 + and layer_quant_config.quant_dtype + == getattr(torch, "float4_e2m1fn_x2", None) + ): + return torch.bfloat16 + return None + + down_proj_prefix = f"{prefix}.routed_expert_down_proj" + up_proj_prefix = f"{prefix}.routed_expert_up_proj" + self.routed_expert_down_proj = ReplicatedLinear( + config.hidden_size, + self.moe_hidden_size, + bias=False, + quant_config=quant_config, + source_quant_dtype=_routed_source_quant_dtype(down_proj_prefix), + prefix=down_proj_prefix, + ) + self.routed_expert_up_proj = ReplicatedLinear( + self.moe_hidden_size, + config.hidden_size, + bias=False, + quant_config=quant_config, + source_quant_dtype=_routed_source_quant_dtype(up_proj_prefix), + prefix=up_proj_prefix, + ) + self.routed_expert_norm = ( + RMSNorm(self.moe_hidden_size, eps=config.rms_norm_eps) + if getattr(config, "latent_moe_use_norm", False) + else None + ) + + # Shared-expert overlap (dual-stream): run the shared-expert MLP on + # `alt_stream` concurrently with the routed experts. Only meaningful + # when there ARE shared experts and the model handed us an alt_stream + # (created iff ATOM_K3_SHARED_EXPERT_OVERLAP is set). Per-call + # decode/prefill gating happens in the maybe_dual_stream_forward + # dispatcher (ATOM_DUAL_STREAM_MOE_TOKEN_THRESHOLD). + self._use_dual_stream = ( + self.shared_experts is not None and self.alt_stream is not None + ) + if self._use_dual_stream: + # Register so the maybe_dual_stream_forward custom op can look us + # up by layer_name (== self.prefix) without putting sub-modules / + # per-fwd state in the op signature. + get_current_atom_config().compilation_config.static_forward_context[ + prefix + ] = self + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self._use_dual_stream: + # Custom op = Dynamo barrier, so the alt_stream context inside + # dual_stream_moe_forward stays opaque to torch.compile. The + # dispatcher picks dual vs single per call (num_tokens vs + # ATOM_DUAL_STREAM_MOE_TOKEN_THRESHOLD; forces single under + # piecewise cudagraph / TBO). + return torch.ops.aiter.maybe_dual_stream_forward(hidden_states, self.prefix) + return self.single_stream_moe_forward(hidden_states) + + def single_stream_moe_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # gfx1250 work-around (REQUIRED for correctness, not just stability): the + # grouped MoE + MXFP4 latent projections are only numerically correct at + # small M on gfx1250 — at large prefill M the contiguous-M path OOB-faults + # AND the non-contiguous path returns coherent-but-wrong values (gsm8k + # 0.0). The MoE block is per-token independent, so split a large prefill + # into <=chunk sub-batches (numerically identical) so every kernel sees a + # small, correct M. Keeps whole-prompt prefill (correct MLA). Gated by + # ATOM_K3_MOE_CHUNK; remove once the gfx1250 MoE kernel is fixed at large M. + chunk = int(os.environ.get("ATOM_K3_MOE_CHUNK", "0") or "0") + n = hidden_states.shape[0] + if chunk > 0 and n > chunk: + outs = [ + self._forward_impl(hidden_states[i : i + chunk]) + for i in range(0, n, chunk) + ] + return torch.cat(outs, dim=0) + return self._forward_impl(hidden_states) + + def dual_stream_moe_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Overlap the shared-expert MLP with the routed experts. + + The shared branch (``shared_experts``, a plain MLP producing a TP-partial + output — no collective inside) runs on ``alt_stream`` while the routed + experts (+ latent all-reduce / norm / up_proj) run on the main stream. + All collectives stay on the main stream and keep their original order, so + this is numerically identical to ``_forward_impl`` — only the shared + GEMMs move off the critical path. Used for small (decode) batches only, + so the single-stream chunking work-around is unnecessary here. + """ + identity = hidden_states + current_stream = get_forward_context().main_stream + + # Fork shared-expert compute onto alt_stream (waits for the main stream + # to have produced `hidden_states` first). + self.alt_stream.wait_stream(current_stream) + with torch.cuda.stream(self.alt_stream): + shared_partial = self.shared_experts(identity) + + # Routed experts on the main stream. + router_logits = self.gate(hidden_states) + routed_input = ( + self.routed_expert_down_proj(hidden_states) + if self.use_latent_moe + else hidden_states + ) + routed_output = self.experts(routed_input, router_logits) + + if self.use_latent_moe: + if self.tp_size > 1: + routed_output = tensor_model_parallel_all_reduce(routed_output) + if self.routed_expert_norm is not None: + routed_output = self.routed_expert_norm(routed_output) + routed_output = self.routed_expert_up_proj(routed_output) + # Join alt_stream, then all-reduce the shared partial on the main + # stream (ordered after the routed all-reduce above). + current_stream.wait_stream(self.alt_stream) + if self.tp_size > 1: + shared_partial = tensor_model_parallel_all_reduce(shared_partial) + return routed_output + shared_partial + + # Non-latent: routed + shared are both TP-partial; a single deferred + # all-reduce over their sum is correct. + current_stream.wait_stream(self.alt_stream) + routed_output = routed_output + shared_partial + if self.tp_size > 1: + routed_output = tensor_model_parallel_all_reduce(routed_output) + return routed_output + + def _forward_impl(self, hidden_states: torch.Tensor) -> torch.Tensor: + identity = hidden_states + router_logits = self.gate(hidden_states) + routed_input = ( + self.routed_expert_down_proj(hidden_states) + if self.use_latent_moe + else hidden_states + ) + routed_output = self.experts(routed_input, router_logits) + if self.use_latent_moe: + # self.experts runs with reduce_results=False, so routed_output is a + # TP-partial sum over the sharded expert intermediate. routed_expert_norm + # is a (nonlinear) RMSNorm, so it must operate on the FULL sum: + # sum_r norm(partial_r) != norm(sum_r partial_r). All-reduce here first; + # routed_expert_norm/up_proj are replicated, so the result stays full. + if self.tp_size > 1: + routed_output = tensor_model_parallel_all_reduce(routed_output) + if self.routed_expert_norm is not None: + routed_output = self.routed_expert_norm(routed_output) + routed_output = self.routed_expert_up_proj(routed_output) + if self.shared_experts is not None: + # Shared branch is TP-partial (down_proj is row-parallel); reduce + # it separately and add to the already-full routed output. + shared_output = self.shared_experts(identity) + if self.tp_size > 1: + shared_output = tensor_model_parallel_all_reduce(shared_output) + routed_output = routed_output + shared_output + return routed_output + # Non-latent path: routed experts and shared experts are both TP-partial + # and everything after them is linear, so a single deferred all-reduce + # over their sum is correct. + if self.shared_experts is not None: + routed_output = routed_output + self.shared_experts(identity) + if self.tp_size > 1: + routed_output = tensor_model_parallel_all_reduce(routed_output) + return routed_output + + +class KimiFullAttention(nn.Module): + def __init__( + self, + atom_config: Config, + quant_config: QuantizationConfig | None, + prefix: str = "", + ): + super().__init__() + config = _text_config(atom_config.hf_config) + self.config = config + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.qk_nope_head_dim = config.qk_nope_head_dim + self.qk_rope_head_dim = config.qk_rope_head_dim + self.q_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + self.v_head_dim = config.v_head_dim + self.q_lora_rank = config.q_lora_rank + self.kv_lora_rank = config.kv_lora_rank + self.scaling = self.q_head_dim**-0.5 + self.tp_size = get_tensor_model_parallel_world_size() + self.num_local_heads = self.num_heads // self.tp_size + self.local_q_size = self.num_local_heads * self.q_head_dim + self.local_v_size = self.num_local_heads * self.v_head_dim + + self.q_a_proj = ReplicatedLinear( + self.hidden_size, + self.q_lora_rank, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_a_proj", + ) + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=1e-6) + self.q_b_proj = ColumnParallelLinear( + self.q_lora_rank, + self.num_heads * self.q_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_b_proj", + ) + self.kv_a_proj_with_mqa = ReplicatedLinear( + self.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_a_proj_with_mqa", + ) + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=1e-6) + self.kv_b_proj = ColumnParallelLinear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_b_proj", + ) + self.g_proj = ColumnParallelLinear( + self.hidden_size, + self.num_heads * self.v_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.g_proj", + ) + self.o_proj = RowParallelLinear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + self.layer_num = _extract_layer_idx(prefix) + self.attn = Attention( + self.num_local_heads, + self.q_head_dim, + self.scaling, + num_kv_heads=self.num_local_heads, + kv_cache_dtype=atom_config.kv_cache_dtype, + quant_config=quant_config, + use_mla=False, + layer_num=self.layer_num, + config=atom_config, + prefix=prefix, + ) + + def forward( + self, positions: torch.Tensor, hidden_states: torch.Tensor + ) -> torch.Tensor: + # q already has the [nope | rope] head layout from q_b_proj; RoPE is + # applied inside self.attn, so there is no split/re-cat here (that would + # be an identity round-trip that just copies q). + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + q = q.view(-1, self.num_local_heads, self.q_head_dim) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + k_latent, k_rope = torch.split( + compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + kv = self.kv_b_proj(self.kv_a_layernorm(k_latent)) + kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim) + + # Fused assembly (one kernel instead of split + cat + pad): builds + # k = [k_nope (from kv) | k_rope (MQA-shared)] and v_padded = [v | 0]. + # Kimi MLA is stored in the standard paged-MHA cache, so V is padded to + # the query head dim to share a cache entry with K (sliced back after). + from atom.models.kimi_k3_fused import fuse_mla_kv + + k, v_padded = fuse_mla_kv( + kv, + k_rope, + self.qk_nope_head_dim, + self.v_head_dim, + self.qk_rope_head_dim, + ) + + # MLA-as-MHA at head_dim=192. self.attn both writes the paged KV cache + # (4D FLASH layout, reshape_and_cache_flash) and runs aiter unified_attention + # for prefill (varlen) AND decode: the standard Triton kernel handles the + # non-power-of-two head dim via HEAD_SIZE_PADDED, so prefill needs no torch + # SDPA fallback. This routing also lets chunked prefill read the paged cache. + attn_out = self.attn( + q.reshape(-1, self.local_q_size), + k.reshape(-1, self.local_q_size), + v_padded.reshape(-1, self.local_q_size), + ) + attn_out = attn_out.view(-1, self.num_local_heads, self.q_head_dim)[ + :, :, : self.v_head_dim + ].reshape(-1, self.local_v_size) + attn_out = attn_out * torch.sigmoid(self.g_proj(hidden_states)) + return self.o_proj(attn_out) + + +def _kda_attention_with_output_fake( + hidden_states: torch.Tensor, layer_name: str +) -> torch.Tensor: + return torch.empty_like(hidden_states) + + +@mark_spliting_op( + is_custom=True, + gen_fake=_kda_attention_with_output_fake, + mutates_args=[], +) +def kda_attention_with_output( + hidden_states: torch.Tensor, layer_name: str +) -> torch.Tensor: + """Opaque splitting-op boundary for the KDA mixer. + + The KDA recurrence reads the forward context, calls fla causal-conv/kda + kernels and mutates the per-request conv/ssm cache in place. torch.compile + (level 3) mis-compiles that stateful path into garbage if it is allowed to + trace through it, so the whole mixer is wrapped in a custom op — inductor + treats it as opaque and the piecewise backend splits the graph here, + exactly as the GDN path does via aiter.linear_attention_with_output_base. + """ + self = get_current_atom_config().compilation_config.static_forward_context[ + layer_name + ] + return self._forward_impl(hidden_states) + + +class KimiKDAAttention(nn.Module): + @property + def mamba_type(self) -> str: + return "kimi_kda" + + def __init__( + self, + atom_config: Config, + quant_config: QuantizationConfig | None, + prefix: str = "", + ): + super().__init__() + config = _text_config(atom_config.hf_config) + self.config = config + self.hidden_size = config.hidden_size + self.num_heads = config.linear_num_key_heads + self.head_dim = config.linear_key_head_dim + self.tp_size = get_tensor_model_parallel_world_size() + self.tp_rank = get_tensor_model_parallel_rank() + self.num_local_heads = self.num_heads // self.tp_size + self.proj_size = self.num_heads * self.head_dim + self.local_proj_size = self.num_local_heads * self.head_dim + self.conv_kernel_size = config.linear_conv_kernel_dim + self.prefix = prefix + self.layer_num = _extract_layer_idx(prefix) + self.activation = "silu" + self.base_linear_attention = True + + # Register under a stable name so the kda_attention_with_output custom op + # can recover this module from the forward context. The op is the + # graph-split boundary that keeps torch.compile from tracing (and + # mis-compiling) the stateful KDA recurrence. + self.layer_name = prefix + compilation_config = atom_config.compilation_config + if self.layer_name in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer: {self.layer_name}") + compilation_config.static_forward_context[self.layer_name] = self + + # The top-level model maps four separate checkpoint projections + # directly into this fused [q | k | v | g] parameter. Mapping keys + # include the KDA layer index so KimiFullAttention.g_proj is untouched. + self.in_proj = MergedColumnParallelLinear( + self.hidden_size, + [ + self.proj_size, + self.proj_size, + self.proj_size, + self.proj_size, + ], + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.in_proj", + ) + # Keep beta separate so the fused in-proj output width remains the + # tile-aligned 4 * local_proj_size. Beta is widened to fp32 in _run_kda. + self.b_proj = ColumnParallelLinear( + self.hidden_size, + self.num_heads, + bias=False, + quant_config=None, + prefix=f"{prefix}.b_proj", + ) + + self.q_conv1d = ColumnParallelLinear( + self.conv_kernel_size, + self.proj_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_conv1d", + ) + self.k_conv1d = ColumnParallelLinear( + self.conv_kernel_size, + self.proj_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.k_conv1d", + ) + self.v_conv1d = ColumnParallelLinear( + self.conv_kernel_size, + self.proj_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.v_conv1d", + ) + for conv in (self.q_conv1d, self.k_conv1d, self.v_conv1d): + conv.weight.data = conv.weight.data.unsqueeze(1) + + self.A_log = atom_parameter(torch.empty(self.num_local_heads)) + self.dt_bias = atom_parameter(torch.empty(self.local_proj_size)) + # Lower bound of the KDA forget gate (Kimi uses -5.0). Consumed by both + # the fla prefill path (_run_kda) and the fused decode kernel. + self._kda_gate_lower_bound = ( + getattr(config, "linear_attn_config", {}) or {} + ).get("gate_lower_bound", None) + loader = _sharded_vector_loader(self.tp_rank, self.tp_size) + self.A_log.weight_loader = loader + self.dt_bias.weight_loader = loader + + self.f_a_proj = ReplicatedLinear( + self.hidden_size, + self.head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.f_a_proj", + ) + self.f_b_proj = ColumnParallelLinear( + self.head_dim, + self.proj_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.f_b_proj", + ) + self.o_norm = KimiRMSNormGated(self.head_dim, eps=config.rms_norm_eps) + self.o_proj = RowParallelLinear( + self.proj_size, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + def process_weights_after_loading(self) -> None: + """Fuse all hidden-input projections into the single in-proj (one GEMM). + + Upstream already loads q/k/v/g into a single ``in_proj`` + (``MergedColumnParallelLinear``) via ``packed_modules_mapping``. Here we + extend that fused weight in place with the two remaining projections that + also consume ``hidden_states`` -- ``b_proj`` (beta) and ``f_a_proj`` -- + so ``forward`` runs one ``self.in_proj(...)`` producing + ``[q | k | v | g | b | f_a]`` instead of three separate launches. The + tail storage is then released (the modules stay as empty shells; their + bf16 post-load hooks are no-ops and never re-run). Runs once; idempotent. + + f_b_proj is NOT fused: it consumes ``f_a_proj``'s output, not + ``hidden_states``, so it is a data-dependent second GEMM and cannot ride + the same launch. + + Fused output width is ``4*local_proj + num_local_heads + head_dim``. The + two small tails (``b`` = num_local_heads, ``f_a`` = head_dim) make N a + non-multiple of the GEMM tile, so the fused shape may fall back to an + untuned tgemm config until one is tuned for it; the saved launches + dominate on the launch-bound decode path. + + Assumes bf16 (unquantized) attention weights, which the Kimi-K3 + checkpoint guarantees (``re:.*self_attn.*`` is in the quant ignore + list). A quantized-attention checkpoint would need per-shard scale + handling and is rejected loudly rather than silently mis-fused. + """ + if getattr(self, "_in_proj_fused", False): + return + # Order defines the forward-time slice boundaries below; keep in sync. + # in_proj already holds the fused [q | k | v | g] (4 * local_proj_size); + # b_proj and f_a_proj are appended as the two tails. + tails = (self.b_proj, self.f_a_proj) + assert all(m.quant_type == QuantType.No for m in (self.in_proj, *tails)), ( + "KDA in-proj fusion assumes unquantized (bf16) attention weights; " + "this checkpoint quantizes self_attn projections." + ) + fused = torch.cat( + [self.in_proj.weight.data, *[m.weight.data for m in tails]], dim=0 + ).contiguous() + # Grow in_proj's weight in place so the existing module (and its + # unquantized tgemm.mm path) produces the wide fused output directly. + self.in_proj.weight = nn.Parameter(fused, requires_grad=False) + # Release the tail weight storage. The modules stay as empty shells; + # their bf16 post-load hooks are no-ops and are never re-run. + for m in tails: + m.weight.data = m.weight.data.new_empty(0) + + # Pre-concatenate the static q/k/v causal-conv weights once here instead + # of rebuilding the [3*local_proj_size, K] tensor on every forward. + self.conv_weight = torch.cat( + [ + self.q_conv1d.weight.view(self.local_proj_size, self.conv_kernel_size), + self.k_conv1d.weight.view(self.local_proj_size, self.conv_kernel_size), + self.v_conv1d.weight.view(self.local_proj_size, self.conv_kernel_size), + ], + dim=0, + ).contiguous() + self._in_proj_fused = True + + def _conv_weights(self) -> torch.Tensor: + return self.conv_weight + + def _run_kda( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, + output_final_state: bool, + recurrent: bool, + ): + from fla.ops.kda import chunk_kda, fused_recurrent_kda + + kwargs = dict( + q=q, + k=k, + v=v, + g=g, + # Keep beta in fp32: fla computes b = sigmoid(beta) in-kernel with + # use_beta_sigmoid_in_kernel, and triton's sigmoid follows the input + # dtype -- a bf16 beta yields a bf16 write strength, which erodes the + # delta-rule state update across the 71 KDA layers (measured gsm8k + # regression). b_proj stays bf16; only this reduction is widened. + beta=beta.float(), + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=initial_state, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + safe_gate=self._kda_gate_lower_bound is not None, + lower_bound=self._kda_gate_lower_bound, + transpose_state_layout=True, + cu_seqlens=cu_seqlens, + ) + if recurrent: + kwargs.pop("safe_gate", None) + return fused_recurrent_kda(**kwargs) + return chunk_kda(**kwargs) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # Route through the opaque custom op so torch.compile splits the graph + # here instead of tracing the stateful recurrence in _forward_impl. + return torch.ops.aiter.kda_attention_with_output(hidden_states, self.layer_name) + + def _forward_impl(self, hidden_states: torch.Tensor) -> torch.Tensor: + fwd_ctx = get_forward_context() + gdn_metadata = getattr(fwd_ctx.attn_metadata, "gdn_metadata", None) + if gdn_metadata is None: + return hidden_states.new_zeros(hidden_states.shape) + + cache = fwd_ctx.kv_cache_data[f"layer_{self.layer_num}"] + conv_state = cache.k_cache + ssm_state = cache.v_cache + if conv_state.size(1) != self.local_proj_size * 3: + conv_state = conv_state.transpose(-1, -2) + + num_actual_tokens = gdn_metadata.num_actual_tokens + hidden_states = hidden_states[:num_actual_tokens] + # Single fused in-proj GEMM producing [q | k | v | g]; slice out each + # part. `out_gate` is the KDA output gate consumed at o_norm below + # (computed here so it rides the same GEMM instead of a separate one + # after the recurrence). in_proj's weight was grown in + # process_weights_after_loading to the fused [q | k | v | g | b | f_a], + # so this single unquantized call emits all six; f_b_proj stays a + # separate GEMM because it consumes f_a's output, not hidden_states. + lp = self.local_proj_size + nlh = self.num_local_heads + hd = self.head_dim + fused_in = self.in_proj(hidden_states) + # No .contiguous() needed: mixed_qkv is a column slice (feature stride 1, + # row stride N_fused). Both causal-conv consumers read the token stride + # from the tensor itself — causal_conv1d_fn uses x.stride(1) after + # transpose (channel-last: stride(0)==1), and causal_conv1d_update only + # requires x.stride(1)==1 (feature-contiguous, which the slice preserves). + mixed_qkv = fused_in[..., : 3 * lp] + out_gate = fused_in[..., 3 * lp : 4 * lp] + # beta is widened to fp32 inside _run_kda (see the note there): the KDA + # delta-rule write strength must stay fp32 for accuracy. + beta = fused_in[..., 4 * lp : 4 * lp + nlh].unsqueeze(0) + # f_a feeds a second GEMM (f_b_proj); make it contiguous so tgemm sees a + # unit row stride rather than the fused output's N_fused stride. + f_a = fused_in[..., 4 * lp + nlh : 4 * lp + nlh + hd].contiguous() + gate = self.f_b_proj(f_a) + gate = rearrange(gate, "t (h d) -> 1 t h d", d=self.head_dim) + out = hidden_states.new_empty( + (num_actual_tokens, self.num_local_heads, self.head_dim) + ) + + conv_weights = self._conv_weights() + state_indices = gdn_metadata.non_spec_state_indices_tensor + query_start_loc = gdn_metadata.non_spec_query_start_loc + + if gdn_metadata.num_prefills > 0: + q, k, v = causal_conv1d_fn( + mixed_qkv.transpose(0, 1), + conv_weights, + None, + activation=self.activation, + conv_states=conv_state, + has_initial_state=gdn_metadata.has_initial_state, + cache_indices=state_indices, + query_start_loc=query_start_loc, + k_dim_size=self.local_proj_size, + v_dim_size=self.local_proj_size, + metadata=gdn_metadata, + ) + q = rearrange(q, "t (h d) -> 1 t h d", d=self.head_dim) + k = rearrange(k, "t (h d) -> 1 t h d", d=self.head_dim) + v = rearrange(v, "t (h d) -> 1 t h d", d=self.head_dim) + # Fused masked gather: ssm_state[state_indices] with fresh + # sequences (~has_initial_state) written as zeros in one pass, + # replacing the gather + separate zero-write. + from atom.models.kimi_k3_fused import gather_kda_initial_state + + initial = gather_kda_initial_state( + ssm_state, state_indices, gdn_metadata.has_initial_state + ) + # gfx1250 workaround: chunk_kda NaNs on short prompts (seq < chunk + # size) and its `transpose_state_layout` output can mismatch the + # decode-time fused_recurrent_kda reader, producing NaN on the first + # decode step. Forcing the recurrent path for prefill keeps the KDA + # state layout consistent across prefill/decode. Env-gated so the + # fast chunk path stays default on archs where it is correct. + _kda_force_recurrent = os.getenv("ATOM_KDA_FORCE_RECURRENT", "0") == "1" + kda_out, last_state = self._run_kda( + q, + k, + v, + gate, + beta, + initial, + query_start_loc, + True, + recurrent=_kda_force_recurrent, + ) + # last_state already has ssm_state's dtype (fla preserves the + # initial_state dtype; the gathered initial is allocated as such), + # so no .to() cast is needed. + ssm_state[state_indices] = last_state + out.copy_(kda_out.squeeze(0)) + elif gdn_metadata.num_decodes > 0: + # Slice the per-token cache-slot indices once (used for both the + # conv update and the fused recurrence below). + decode_state_indices = state_indices[:num_actual_tokens] + q, k, v = causal_conv1d_update( + mixed_qkv, + conv_state, + conv_weights, + self.local_proj_size, + self.local_proj_size, + None, + self.activation, + conv_state_indices=decode_state_indices, + validate_data=True, + ) + q = rearrange(q, "t (h d) -> 1 t h d", d=self.head_dim) + k = rearrange(k, "t (h d) -> 1 t h d", d=self.head_dim) + v = rearrange(v, "t (h d) -> 1 t h d", d=self.head_dim) + # Fused KDA decode: the kernel gathers the initial state from + # ssm_state[decode_state_indices], writes the final state back to + # the same slots inplace (inplace_final_state), and writes the + # recurrence output straight into `out`. This folds the manual + # gather / scatter-back / out.copy_ that the fla path required into + # one kernel. is_kda + lower_bound select the per-K-channel, + # lower-bounded sigmoid gate that Kimi-KDA uses (beta stays raw + # logits; the kernel applies sigmoid in fp32 internally). + fused_sigmoid_gating_delta_rule_update( + A_log=self.A_log, + a=gate, + b=beta, + dt_bias=self.dt_bias, + q=q, + k=k, + v=v, + o=out, + initial_state=ssm_state, + inplace_final_state=True, + cu_seqlens=query_start_loc[: gdn_metadata.num_decodes + 1], + ssm_state_indices=decode_state_indices, + use_qk_l2norm_in_kernel=True, + is_kda=True, + lower_bound=self._kda_gate_lower_bound, + ) + else: + out.zero_() + + out = self.o_norm(out, rearrange(out_gate, "t (h d) -> t h d", d=self.head_dim)) + return self.o_proj(rearrange(out, "t h d -> t (h d)")) + + +class KimiDecoderLayer(nn.Module): + def __init__( + self, + atom_config: Config, + prefix: str, + layer_num: int = 0, + alt_stream: Optional["torch.cuda.Stream"] = None, + ): + super().__init__() + config = _text_config(atom_config.hf_config) + quant_config = atom_config.quant_config + self.config = config + self.layer_idx = layer_num + self.hidden_size = config.hidden_size + if layer_num in config.kimi_kda_layers: + self.self_attn = KimiKDAAttention( + atom_config, quant_config, prefix=f"{prefix}.self_attn" + ) + self.is_linear_attn = True + else: + self.self_attn = KimiFullAttention( + atom_config, quant_config, prefix=f"{prefix}.self_attn" + ) + self.is_linear_attn = False + + if ( + config.num_experts is not None + and layer_num >= config.first_k_dense_replace + and layer_num % getattr(config, "moe_layer_freq", 1) == 0 + ): + self.block_sparse_moe = KimiSparseMoeBlock( + config, + quant_config=quant_config, + prefix=f"{prefix}.block_sparse_moe", + alt_stream=alt_stream, + ) + else: + self.mlp = KimiMLP( + config, quant_config=quant_config, prefix=f"{prefix}.mlp" + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + self.use_attn_residuals = ( + getattr(config, "attn_res_block_size", None) is not None + ) + if self.use_attn_residuals: + self.attn_res_block_size = config.attn_res_block_size + self.self_attention_res_norm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.mlp_res_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.self_attention_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.self_attention_res_proj", + ) + self.mlp_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.mlp_res_proj", + ) + + def _ffn(self, hidden_states: torch.Tensor) -> torch.Tensor: + if hasattr(self, "block_sparse_moe"): + return self.block_sparse_moe(hidden_states) + return self.mlp(hidden_states) + + def process_weights_after_loading(self) -> None: + # Fold each attn-residual (norm.weight * proj.weight) into a single + # static score vector consumed by apply_attn_res (see + # _attn_res_score_weight). Both operands are load-time constants. + if not self.use_attn_residuals: + return + for proj, norm in ( + (self.self_attention_res_proj, self.self_attention_res_norm), + (self.mlp_res_proj, self.mlp_res_norm), + ): + proj.score_weight = _attn_res_score_weight(proj, norm) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + block_residual: torch.Tensor | None = None, + pending_add: torch.Tensor | None = None, + ): + if not self.use_attn_residuals: + if pending_add is not None: + hidden_states = hidden_states + pending_add + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + if self.is_linear_attn: + hidden_states = self.self_attn(hidden_states) + else: + hidden_states = self.self_attn(positions, hidden_states) + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self._ffn(hidden_states) + return residual + hidden_states, None, block_residual + + prefix_sum = hidden_states + if block_residual is not None and block_residual.shape[1] > 0: + hidden_states, prefix_sum = _apply_attn_res( + prefix_sum, + block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + add_hidden=pending_add, + ) + elif pending_add is not None: + prefix_sum = prefix_sum + pending_add + hidden_states = prefix_sum + if self.layer_idx % self.attn_res_block_size == 0: + assert block_residual is not None + block_residual = torch.cat([block_residual, prefix_sum.unsqueeze(1)], dim=1) + prefix_sum = None + + hidden_states = self.input_layernorm(hidden_states) + if self.is_linear_attn: + hidden_states = self.self_attn(hidden_states) + else: + hidden_states = self.self_attn(positions, hidden_states) + + if prefix_sum is None: + prefix_sum = hidden_states + hidden_states, prefix_sum = _apply_attn_res( + prefix_sum, block_residual, self.mlp_res_proj, self.mlp_res_norm + ) + else: + # Fold prefix_sum = prefix_sum + hidden_states into the fused kernel. + hidden_states, prefix_sum = _apply_attn_res( + prefix_sum, + block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + add_hidden=hidden_states, + ) + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self._ffn(hidden_states) + if prefix_sum is None: + return hidden_states, None, block_residual + return prefix_sum, hidden_states, block_residual + + +def _attn_res_score_weight(proj: ReplicatedLinear, norm: RMSNorm) -> torch.Tensor: + """Fold the static rmsnorm gain and projection into one [H] score vector. + + Both operands are load-time constants, so this is precomputed once in + ``process_weights_after_loading`` and cached on ``proj.score_weight``; the + apply_attn_res kernel then reads a single vector per H-chunk instead of + reloading norm.weight and proj.weight and multiplying them every forward. + """ + return (norm.weight.float() * proj.weight.squeeze(0).float()).contiguous() + + +def _apply_attn_res( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj: ReplicatedLinear, + norm: RMSNorm, + add_hidden: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + # Returns (mixed_output, prefix_out). When add_hidden is given the fused path + # folds ``prefix_sum = prefix_sum + add_hidden`` into the kernel and returns + # the summed prefix; otherwise prefix_out is prefix_sum unchanged. + eps = getattr(norm, "variance_epsilon", getattr(norm, "eps", 1e-6)) + score_weight = getattr(proj, "score_weight", None) + if score_weight is None: + score_weight = _attn_res_score_weight(proj, norm) + from atom.models.kimi_k3_fused import apply_attn_res + + return apply_attn_res(prefix_sum, block_residual, score_weight, eps, add_hidden) + + +@support_torch_compile +class KimiLinearModel(nn.Module): + def __init__(self, atom_config: Config, prefix: str = ""): + super().__init__() + config = _text_config(atom_config.hf_config) + _normalize_kimi_config(config) + self.config = config + self.vocab_size = config.vocab_size + + if get_pp_group().is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, config.hidden_size + ) + else: + self.embed_tokens = PPMissingLayer() + + # Single alt_stream shared across layers for shared-expert overlap. + # Created only when the feature is enabled (default off); when None, + # every MoE block falls back to single-stream. Attention/MoE never + # contend since each block runs its MoE after its attention. + self.alt_stream: Optional[torch.cuda.Stream] = ( + torch.cuda.Stream() + if envs.ATOM_K3_SHARED_EXPERT_OVERLAP and torch.cuda.is_available() + else None + ) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix, layer_num=None: KimiDecoderLayer( + atom_config, + prefix=prefix, + layer_num=layer_num or 0, + alt_stream=self.alt_stream, + ), + prefix=f"{prefix}.layers", + layer_num_offset=0, + ) + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if getattr(config, "attn_res_block_size", None) is not None: + self.output_attn_res_norm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.output_attn_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.output_attn_res_proj", + ) + else: + self.norm = PPMissingLayer() + + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "block_residual"], config.hidden_size + ) + + def process_weights_after_loading(self) -> None: + # Fold the final output attn-residual (norm.weight * proj.weight) into a + # single static score vector for apply_attn_res. Present only on the last + # PP rank when attn residuals are enabled. + if hasattr(self, "output_attn_res_proj"): + self.output_attn_res_proj.score_weight = _attn_res_score_weight( + self.output_attn_res_proj, self.output_attn_res_norm + ) + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: Optional[IntermediateTensors] = None, + inputs_embeds: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, IntermediateTensors]: + if get_pp_group().is_first_rank: + hidden_states = ( + inputs_embeds + if inputs_embeds is not None + else self.embed_tokens(input_ids) + ) + block_residual = ( + hidden_states.new_zeros( + hidden_states.shape[0], 0, hidden_states.shape[1] + ) + if getattr(self.config, "attn_res_block_size", None) is not None + else None + ) + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + block_residual = intermediate_tensors["block_residual"] + + pending_add = None + for layer in self.layers[self.start_layer : self.end_layer]: + hidden_states, pending_add, block_residual = layer( + positions, + hidden_states, + block_residual, + pending_add=pending_add, + ) + + if not get_pp_group().is_last_rank: + if pending_add is not None: + hidden_states = hidden_states + pending_add + return IntermediateTensors( + {"hidden_states": hidden_states, "block_residual": block_residual} + ) + if getattr(self.config, "attn_res_block_size", None) is not None: + hidden_states, _ = _apply_attn_res( + hidden_states, + block_residual, + self.output_attn_res_proj, + self.output_attn_res_norm, + add_hidden=pending_add, + ) + elif pending_add is not None: + hidden_states = hidden_states + pending_add + return self.norm(hidden_states) + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.config.num_experts + (self.config.num_shared_experts or 0), + ) + + +class KimiLinearForCausalLM(nn.Module): + packed_modules_mapping = _kda_packed_modules_mapping([]) + weights_mapping = { + "weight_packed": "weight", + } + + def __init__(self, atom_config: Config, prefix: str = ""): + super().__init__() + config = _text_config(atom_config.hf_config) + _normalize_kimi_config(config) + self.config = config + self.quant_config = atom_config.quant_config + self.packed_modules_mapping = _kda_packed_modules_mapping( + config.kimi_kda_layers + ) + self.model = KimiLinearModel(atom_config, prefix=maybe_prefix(prefix, "model")) + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.get_input_embeddings(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: Optional[IntermediateTensors] = None, + inputs_embeds: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, IntermediateTensors]: + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds) + + def compute_logits(self, hidden_states: torch.Tensor) -> Optional[torch.Tensor]: + return self.lm_head(hidden_states) + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return self.model.get_expert_mapping() + + +class KimiK3ForCausalLM(nn.Module): + skip_weight_prefixes = ["vision_tower.", "mm_projector."] + quant_exclude_name_mapping = { + "language_model.model.": "language_model.model.", + "language_model.lm_head": "language_model.lm_head", + } + packed_modules_mapping = KimiLinearForCausalLM.packed_modules_mapping + weights_mapping = KimiLinearForCausalLM.weights_mapping + + def __init__(self, atom_config: Config, prefix: str = ""): + super().__init__() + root_config = atom_config.hf_config + if ( + hasattr(root_config, "text_config") + and root_config.text_config is not root_config + ): + _normalize_kimi_config(root_config.text_config) + if ( + getattr(root_config, "quantization_config", None) is None + and getattr(root_config.text_config, "quantization_config", None) + is not None + ): + atom_config.quant_config = QuantizationConfig( + root_config.text_config, + atom_config.online_quant_config, + ) + else: + _normalize_kimi_config(root_config) + self.config = _text_config(root_config) + self.quant_config = atom_config.quant_config + self.packed_modules_mapping = _kda_packed_modules_mapping( + self.config.kimi_kda_layers + ) + self.language_model = KimiLinearForCausalLM( + atom_config=atom_config, + prefix=maybe_prefix(prefix, "language_model"), + ) + self.make_empty_intermediate_tensors = ( + self.language_model.make_empty_intermediate_tensors + ) + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.language_model.get_input_embeddings(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: Optional[IntermediateTensors] = None, + inputs_embeds: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, IntermediateTensors]: + return self.language_model( + input_ids, positions, intermediate_tensors, inputs_embeds + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> Optional[torch.Tensor]: + return self.language_model.compute_logits(hidden_states) + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + # The loader matches expert entries as substrings of full checkpoint + # names, so keep these generic enough to match each layer's + # `block_sparse_moe.experts.{id}.w*.weight` entries. + return self.language_model.get_expert_mapping() diff --git a/atom/models/kimi_k3_fused.py b/atom/models/kimi_k3_fused.py new file mode 100644 index 0000000000..15935bdb2c --- /dev/null +++ b/atom/models/kimi_k3_fused.py @@ -0,0 +1,836 @@ +"""Fused Triton kernels for Kimi-K3 small ops (gfx1250 perf). + +Drop-in replacements for the elementwise/reduction glue that otherwise runs as +several separate torch ops (each a kernel launch + full HBM round-trip): + +- ``situ_and_mul`` : SiTU(gate) * linear(up) (dense/shared-expert MLP act) +- ``rmsnorm_gated`` : rmsnorm(x) * weight * sigmoid(gate) (KDA o_norm) +- ``apply_attn_res`` : block-residual soft-attention mix (attn_res boundaries) + +The activation and gated-rmsnorm kernels remain gated behind ``ATOM_K3_FUSED`` +in kimi_k3.py. ``apply_attn_res`` is the production path; correctness is checked +against a local torch reference in my_script/optest_fused.py. +""" + +from __future__ import annotations + +import torch + +try: + import triton + import triton.language as tl + + _HAS_TRITON = True +except Exception: # pragma: no cover + _HAS_TRITON = False + + +if _HAS_TRITON: + + @triton.jit + def _situ_and_mul_kernel( + x_ptr, + y_ptr, + M, + D, + stride_xm, + stride_ym, + beta, + inv_beta, + linear_beta, + inv_linear_beta, + HAS_LINEAR: tl.constexpr, + BLOCK: tl.constexpr, + ): + row = tl.program_id(0) + col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + mask = col < D + g = tl.load(x_ptr + row * stride_xm + col, mask=mask, other=0.0).to(tl.float32) + u = tl.load(x_ptr + row * stride_xm + D + col, mask=mask, other=0.0).to( + tl.float32 + ) + # SiTUv2 gate: beta * tanh(gate/beta) * sigmoid(gate); tanh via sigmoid + # identity (tanh(z) = 2*sigmoid(2z) - 1) for portability across triton. + out = beta * (2.0 * tl.sigmoid(2.0 * g * inv_beta) - 1.0) * tl.sigmoid(g) + if HAS_LINEAR: + u = linear_beta * (2.0 * tl.sigmoid(2.0 * u * inv_linear_beta) - 1.0) + y = out * u + tl.store(y_ptr + row * stride_ym + col, y.to(y_ptr.dtype.element_ty), mask=mask) + + @triton.jit + def _rmsnorm_gated_kernel( + x_ptr, + w_ptr, + g_ptr, + y_ptr, + H, + eps, + stride_xm, + stride_ym, + stride_g_outer, + stride_g_head, + HEADS: tl.constexpr, + BLOCK: tl.constexpr, + ): + row = tl.program_id(0) + cols = tl.arange(0, BLOCK) + mask = cols < H + # x / y are row-contiguous [M, H]; gate may be strided. Its logical row + # `row` decomposes into (outer, head) so the token-boundary jump + # (stride_g_outer) and per-head step (stride_g_head) are read directly, + # avoiding a contiguous copy of the strided gate slice. + g_off = (row // HEADS) * stride_g_outer + (row % HEADS) * stride_g_head + cols + x = tl.load(x_ptr + row * stride_xm + cols, mask=mask, other=0.0).to(tl.float32) + var = tl.sum(x * x, axis=0) / H + rstd = 1.0 / tl.sqrt(var + eps) + w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32) + gate = tl.load(g_ptr + g_off, mask=mask, other=0.0).to(tl.float32) + y = (x * rstd * w) * tl.sigmoid(gate) + tl.store( + y_ptr + row * stride_ym + cols, y.to(y_ptr.dtype.element_ty), mask=mask + ) + + +if _HAS_TRITON: + + @triton.jit + def _attn_res_fused_kernel( + br_ptr, + ps_ptr, + sw_ptr, + y_ptr, + hs_ptr, + pref_ptr, + B, + Bp, + H, + eps, + stride_br_t, + stride_br_b, + stride_ps_t, + stride_yt, + stride_hs_t, + stride_pref_t, + BP: tl.constexpr, # Bp padded to a power of 2 (vectorized candidate axis) + BLOCK_H: tl.constexpr, + NS: tl.constexpr, # num_stages for the H-loop software pipeline + DO_ADD: tl.constexpr, # fold prefix += hidden_states on-load + WRITE_PREF: tl.constexpr, # write the (summed) prefix back to pref_ptr + ): + # One program per row t: rmsnorm each of the Bp = B+1 candidates, score = + # , softmax over Bp, then weighted sum -> y[t]. + # Candidates 0..B-1 are block_residual rows; candidate B is prefix_sum. + # Read both source tensors directly (no torch.cat materialization); the + # Bp axis is vectorized, so scores/probs stay in registers and softmax + + # weighted-sum never touch HBM. + # + # DO_ADD folds the caller's ``prefix_sum = prefix_sum + hidden_states`` + # elementwise add into the last-candidate on-load (saving a separate + # kernel launch + HBM round-trip); WRITE_PREF then stores that summed + # prefix once (first pass) so downstream layers reuse it. + # + # Two HBM passes over H: the softmax + weighted-sum combine is over the + # small Bp axis, so probs isn't known until the whole H-reduction is done. + # The second pass re-reads br/ps -- but for one row that footprint is only + # ~Bp*H*2B (L2-resident), so the reload is served from cache, not HBM. + # (Holding the [BP, H] tile in registers to avoid the reload was measured + # slower: it blows the VGPR file and collapses occupancy.) num_stages + # pipelines each pass so the next chunk's load overlaps the current + # reduce -- the only win at small T where occupancy alone can't hide it. + t = tl.program_id(0) + b_idx = tl.arange(0, BP) + b_mask = b_idx < Bp + is_last = b_idx == B # prefix_sum candidate + br_base = t * stride_br_t + b_idx * stride_br_b # [BP] + ps_base = t * stride_ps_t + + acc_sq = tl.zeros((BP,), dtype=tl.float32) + acc_dot = tl.zeros((BP,), dtype=tl.float32) + for h0 in tl.range(0, H, BLOCK_H, num_stages=NS): + cols = h0 + tl.arange(0, BLOCK_H) + h_mask = cols < H + br = tl.load( + br_ptr + br_base[:, None] + cols[None, :], + mask=(b_idx < B)[:, None] & h_mask[None, :], + other=0.0, + ).to(tl.float32) + ps = tl.load(ps_ptr + ps_base + cols, mask=h_mask, other=0.0).to( + tl.float32 + ) # [BLOCK_H] + if DO_ADD: + ps += tl.load( + hs_ptr + t * stride_hs_t + cols, mask=h_mask, other=0.0 + ).to(tl.float32) + if WRITE_PREF: + tl.store( + pref_ptr + t * stride_pref_t + cols, + ps.to(pref_ptr.dtype.element_ty), + mask=h_mask, + ) + v = tl.where( + is_last[:, None], ps[None, :], br + ) # [BP, BLOCK_H], ps broadcast in-reg + # score_weight = norm_weight * proj_weight, precomputed at load time + sw = tl.load(sw_ptr + cols, mask=h_mask, other=0.0).to(tl.float32) + acc_sq += tl.sum(v * v, axis=1) # [BP] + acc_dot += tl.sum(v * sw[None, :], axis=1) # [BP] + + rstd = 1.0 / tl.sqrt(acc_sq / H + eps) + scores = tl.where(b_mask, rstd * acc_dot, float("-inf")) + scores = scores - tl.max(scores, axis=0) + probs = tl.exp(scores) + probs = probs / tl.sum(probs, axis=0) # [BP], softmax over Bp + + for h0 in tl.range(0, H, BLOCK_H, num_stages=NS): + cols = h0 + tl.arange(0, BLOCK_H) + h_mask = cols < H + br = tl.load( + br_ptr + br_base[:, None] + cols[None, :], + mask=(b_idx < B)[:, None] & h_mask[None, :], + other=0.0, + ).to(tl.float32) + ps = tl.load(ps_ptr + ps_base + cols, mask=h_mask, other=0.0).to(tl.float32) + if DO_ADD: + ps += tl.load( + hs_ptr + t * stride_hs_t + cols, mask=h_mask, other=0.0 + ).to(tl.float32) + v = tl.where(is_last[:, None], ps[None, :], br) + out = tl.sum(probs[:, None] * v, axis=0) # [BLOCK_H] + tl.store( + y_ptr + t * stride_yt + cols, + out.to(y_ptr.dtype.element_ty), + mask=h_mask, + ) + + @triton.jit + def _attn_res_reduce_kernel( + br_ptr, + ps_ptr, + sw_ptr, + psq_ptr, + pdot_ptr, + hs_ptr, + pref_ptr, + B, + H, + S, + stride_br_t, + stride_br_b, + stride_ps_t, + stride_o_t, + stride_o_s, + stride_hs_t, + stride_pref_t, + BP: tl.constexpr, + BLOCK_H: tl.constexpr, + NS: tl.constexpr, + DO_ADD: tl.constexpr, + WRITE_PREF: tl.constexpr, + ): + # Split-H stage 1, grid=(T, S): S workgroups cooperate on one row t, each + # owning a block-cyclic subset of the H-chunks. They emit PARTIAL sums + # psq/pdot[t, s] -- no softmax here, the reduction axis (H) is orthogonal + # to the softmax axis (Bp), so we stop strictly before the softmax fence. + # This multiplies the grid by S to fill the GPU at small T (where the + # grid=(T,) kernel launches too few workgroups to reach full occupancy). + # + # DO_ADD folds prefix += hidden_states on-load; each (t, s) owns a disjoint + # block-cyclic slice of H, so WRITE_PREF here stores that slice of the + # summed prefix with no overlap -- together the S programs cover all of H. + t = tl.program_id(0) + s = tl.program_id(1) + b_idx = tl.arange(0, BP) + is_last = b_idx == B + br_base = t * stride_br_t + b_idx * stride_br_b + ps_base = t * stride_ps_t + acc_sq = tl.zeros((BP,), dtype=tl.float32) + acc_dot = tl.zeros((BP,), dtype=tl.float32) + for h0 in tl.range(s * BLOCK_H, H, S * BLOCK_H, num_stages=NS): + cols = h0 + tl.arange(0, BLOCK_H) + h_mask = cols < H + br = tl.load( + br_ptr + br_base[:, None] + cols[None, :], + mask=(b_idx < B)[:, None] & h_mask[None, :], + other=0.0, + ).to(tl.float32) + ps = tl.load(ps_ptr + ps_base + cols, mask=h_mask, other=0.0).to(tl.float32) + if DO_ADD: + ps += tl.load( + hs_ptr + t * stride_hs_t + cols, mask=h_mask, other=0.0 + ).to(tl.float32) + if WRITE_PREF: + tl.store( + pref_ptr + t * stride_pref_t + cols, + ps.to(pref_ptr.dtype.element_ty), + mask=h_mask, + ) + v = tl.where(is_last[:, None], ps[None, :], br) + sw = tl.load(sw_ptr + cols, mask=h_mask, other=0.0).to(tl.float32) + acc_sq += tl.sum(v * v, axis=1) + acc_dot += tl.sum(v * sw[None, :], axis=1) + o = t * stride_o_t + s * stride_o_s + b_idx + tl.store(psq_ptr + o, acc_sq) + tl.store(pdot_ptr + o, acc_dot) + + @triton.jit + def _attn_res_combine_kernel( + br_ptr, + ps_ptr, + psq_ptr, + pdot_ptr, + y_ptr, + hs_ptr, + B, + Bp, + H, + S, + eps, + stride_br_t, + stride_br_b, + stride_ps_t, + stride_i_t, + stride_i_s, + stride_yt, + stride_hs_t, + BP: tl.constexpr, + BLOCK_H: tl.constexpr, + NS: tl.constexpr, + DO_ADD: tl.constexpr, + ): + # Split-H stage 2, grid=(T,): sum the S partials back to the full + # H-reduction (associative, so exact), then the identical softmax + + # weighted-sum tail as the single-kernel path. + t = tl.program_id(0) + b_idx = tl.arange(0, BP) + b_mask = b_idx < Bp + is_last = b_idx == B + acc_sq = tl.zeros((BP,), dtype=tl.float32) + acc_dot = tl.zeros((BP,), dtype=tl.float32) + for s in range(S): + o = t * stride_i_t + s * stride_i_s + b_idx + acc_sq += tl.load(psq_ptr + o) + acc_dot += tl.load(pdot_ptr + o) + rstd = 1.0 / tl.sqrt(acc_sq / H + eps) + scores = tl.where(b_mask, rstd * acc_dot, float("-inf")) + scores = scores - tl.max(scores, axis=0) + probs = tl.exp(scores) + probs = probs / tl.sum(probs, axis=0) + br_base = t * stride_br_t + b_idx * stride_br_b + ps_base = t * stride_ps_t + for h0 in tl.range(0, H, BLOCK_H, num_stages=NS): + cols = h0 + tl.arange(0, BLOCK_H) + h_mask = cols < H + br = tl.load( + br_ptr + br_base[:, None] + cols[None, :], + mask=(b_idx < B)[:, None] & h_mask[None, :], + other=0.0, + ).to(tl.float32) + ps = tl.load(ps_ptr + ps_base + cols, mask=h_mask, other=0.0).to(tl.float32) + if DO_ADD: + ps += tl.load( + hs_ptr + t * stride_hs_t + cols, mask=h_mask, other=0.0 + ).to(tl.float32) + v = tl.where(is_last[:, None], ps[None, :], br) + out = tl.sum(probs[:, None] * v, axis=0) + tl.store( + y_ptr + t * stride_yt + cols, + out.to(y_ptr.dtype.element_ty), + mask=h_mask, + ) + + +# Per-token-count tuning for apply_attn_res at H=7168 (gfx1250). Each bucket maps +# an upper-bound token count to (split, S, num_stages, num_warps): +# split=True -> two-kernel split-H (S workgroups per row); wins at small T by +# filling the GPU that grid=(T,) leaves idle (~1.3-1.55x). +# split=False -> single grid=(T,) two-pass kernel with num_stages pipelining; +# wins once T alone saturates the machine (the split-H partial +# round-trip then becomes pure overhead). +# Dispatch rounds T UP to the smallest bucket >= T (ceil-to-bucket), matching how +# CUDAGraph captures a handful of fixed batch sizes; T above the largest bucket +# falls through to the catch-all two-pass path. +_ATTN_RES_CONFIGS = ( + # (max_tokens, split, S, num_stages, num_warps) + (8, True, 7, 1, 2), + (16, True, 7, 1, 4), + (32, True, 7, 1, 4), + (64, True, 7, 1, 4), + (128, True, 6, 1, 4), + (256, False, 1, 2, 4), +) +_ATTN_RES_CATCHALL = (False, 1, 2, 4) # T > largest bucket +_ATTN_RES_BLOCK_H = 1024 + + +def _pick_attn_res_config(tokens: int): + for max_tokens, split, s, ns, nw in _ATTN_RES_CONFIGS: + if tokens <= max_tokens: + return split, s, ns, nw + return _ATTN_RES_CATCHALL + + +def _apply_attn_res_impl( + prefix_sum: torch.Tensor, # [T, H] + block_residual: torch.Tensor, # [T, B, H] + score_weight: torch.Tensor, # [H] (norm_weight * proj_weight, precomputed) + eps: float, + add_hidden: torch.Tensor | None = None, # [T, H], folded: prefix += add_hidden +) -> tuple[torch.Tensor, torch.Tensor]: + """Block-residual soft-attention mix: rmsnorm each of the B+1 candidates, + score = , softmax over B+1, weighted sum. + + ``score_weight`` folds ``norm_weight * proj_weight`` once at load time (both + are static loaded weights), so the kernel loads a single vector instead of + two and drops the per-forward multiply. + + When ``add_hidden`` is given, the caller's ``prefix_sum = prefix_sum + + hidden_states`` elementwise add is folded into the kernel (added on-load to + the prefix candidate, one HBM read + one launch saved) and the summed prefix + is written back. Returns ``(y, prefix_out)`` where ``prefix_out`` is the + summed prefix (or the unchanged ``prefix_sum`` when ``add_hidden`` is None), + so downstream layers reuse it. + + Dispatches by token count (see ``_ATTN_RES_CONFIGS``): split-H at small T to + fill the GPU, the single-pass pipelined kernel once T saturates it.""" + T, B, H = block_residual.shape + Bp = B + 1 + do_add = add_hidden is not None + br = block_residual.contiguous() + ps = prefix_sum.contiguous() + sw = score_weight.contiguous() + y = torch.empty((T, H), device=block_residual.device, dtype=prefix_sum.dtype) + # hs/pref pointers are always passed (triton needs a tensor); when not adding + # they alias ps and are never dereferenced (DO_ADD / WRITE_PREF are False). + if do_add: + hs = add_hidden.contiguous() + pref = torch.empty((T, H), device=block_residual.device, dtype=prefix_sum.dtype) + else: + hs = ps + pref = ps + BP = triton.next_power_of_2(Bp) + BLOCK_H = _ATTN_RES_BLOCK_H + nchunk = triton.cdiv(H, BLOCK_H) + + split, s, ns, nw = _pick_attn_res_config(T) + # S can't exceed the chunk count (a workgroup with no chunk to own is wasted). + s = min(s, nchunk) + if split and s > 1: + psq = torch.empty((T, s, BP), device=br.device, dtype=torch.float32) + pdot = torch.empty((T, s, BP), device=br.device, dtype=torch.float32) + _attn_res_reduce_kernel[(T, s)]( + br, + ps, + sw, + psq, + pdot, + hs, + pref, + B, + H, + s, + br.stride(0), + br.stride(1), + ps.stride(0), + psq.stride(0), + psq.stride(1), + hs.stride(0), + pref.stride(0), + BP=BP, + BLOCK_H=BLOCK_H, + NS=ns, + num_warps=nw, + DO_ADD=do_add, + WRITE_PREF=do_add, + ) + _attn_res_combine_kernel[(T,)]( + br, + ps, + psq, + pdot, + y, + hs, + B, + Bp, + H, + s, + eps, + br.stride(0), + br.stride(1), + ps.stride(0), + psq.stride(0), + psq.stride(1), + y.stride(0), + hs.stride(0), + BP=BP, + BLOCK_H=BLOCK_H, + NS=ns, + num_warps=nw, + DO_ADD=do_add, + ) + return y, (pref if do_add else prefix_sum) + + _attn_res_fused_kernel[(T,)]( + br, + ps, + sw, + y, + hs, + pref, + B, + Bp, + H, + float(eps), + br.stride(0), + br.stride(1), + ps.stride(0), + y.stride(0), + hs.stride(0), + pref.stride(0), + BP=BP, + BLOCK_H=BLOCK_H, + NS=ns, + num_warps=nw, + DO_ADD=do_add, + WRITE_PREF=do_add, + ) + return y, (pref if do_add else prefix_sum) + + +@torch.library.custom_op("atom::apply_attn_res", mutates_args=(), device_types="cuda") +def _apply_attn_res_op( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + score_weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + mixed_output, _ = _apply_attn_res_impl( + prefix_sum, block_residual, score_weight, eps + ) + return mixed_output + + +@_apply_attn_res_op.register_fake +def _apply_attn_res_op_fake( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + score_weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + return torch.empty_like(prefix_sum) + + +@torch.library.custom_op( + "atom::apply_attn_res_add", mutates_args=(), device_types="cuda" +) +def _apply_attn_res_add_op( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + score_weight: torch.Tensor, + eps: float, + add_hidden: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + return _apply_attn_res_impl( + prefix_sum, block_residual, score_weight, eps, add_hidden + ) + + +@_apply_attn_res_add_op.register_fake +def _apply_attn_res_add_op_fake( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + score_weight: torch.Tensor, + eps: float, + add_hidden: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + return torch.empty_like(prefix_sum), torch.empty_like(prefix_sum) + + +def apply_attn_res( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + score_weight: torch.Tensor, + eps: float, + add_hidden: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Dispatch an opaque custom op whose CUDA implementation selects by concrete T.""" + if add_hidden is None: + return ( + _apply_attn_res_op(prefix_sum, block_residual, score_weight, eps), + prefix_sum, + ) + return _apply_attn_res_add_op( + prefix_sum, block_residual, score_weight, eps, add_hidden + ) + + +def situ_and_mul( + x: torch.Tensor, beta: float, linear_beta: float | None +) -> torch.Tensor: + """SiTUv2 gated activation over the last dim (x[..., :D] gate, x[..., D:] up).""" + *lead, two_d = x.shape + assert two_d % 2 == 0 + d = two_d // 2 + x2 = x.reshape(-1, two_d) + m = x2.shape[0] + y = torch.empty((m, d), dtype=x.dtype, device=x.device) + if not _HAS_TRITON or m == 0: + return _situ_and_mul_torch(x, beta, linear_beta) + BLOCK = 1024 + grid = (m, triton.cdiv(d, BLOCK)) + has_linear = linear_beta is not None + _situ_and_mul_kernel[grid]( + x2, + y, + m, + d, + x2.stride(0), + y.stride(0), + float(beta), + 1.0 / float(beta), + float(linear_beta) if has_linear else 0.0, + (1.0 / float(linear_beta)) if has_linear else 0.0, + HAS_LINEAR=has_linear, + BLOCK=BLOCK, + ) + return y.reshape(*lead, d) + + +def rmsnorm_gated( + x: torch.Tensor, weight: torch.Tensor, gate: torch.Tensor, eps: float +) -> torch.Tensor: + """rmsnorm(x) over last dim * weight * sigmoid(gate). + + ``gate`` may be strided (e.g. a column slice of a fused GEMM output): the + kernel reads it via (outer, head) strides so no contiguous copy is needed. + ``x`` is normed row-wise and is made contiguous (cheap; the caller's ``out`` + already is). Supports a 2D ``[M, H]`` or 3D ``[outer, heads, H]`` gate. + """ + h = x.shape[-1] + x2 = x.reshape(-1, h) + m = x2.shape[0] + if not _HAS_TRITON or m == 0 or h > 8192: + return _rmsnorm_gated_torch(x, weight, gate, eps) + if gate.ndim == 3: + heads = gate.shape[1] + stride_g_outer, stride_g_head = gate.stride(0), gate.stride(1) + else: + # 2D: one logical head per row; the head term drops out (row % 1 == 0). + heads = 1 + stride_g_outer, stride_g_head = gate.stride(0), 0 + x2 = x2.contiguous() + y = torch.empty_like(x2) + BLOCK = triton.next_power_of_2(h) + _rmsnorm_gated_kernel[(m,)]( + x2, + weight, + gate, + y, + h, + float(eps), + x2.stride(0), + y.stride(0), + stride_g_outer, + stride_g_head, + HEADS=heads, + BLOCK=BLOCK, + ) + return y.reshape_as(x) + + +# --------------------------------------------------------------------------- # +# torch references (also the fallback when triton is unavailable) +# --------------------------------------------------------------------------- # +def _situ_and_mul_torch( + x: torch.Tensor, beta: float, linear_beta: float | None +) -> torch.Tensor: + gate, up = x.chunk(2, dim=-1) + gate_f = gate.float() + up_f = up.float() + out = beta * torch.tanh(gate_f / beta) * torch.sigmoid(gate_f) + if linear_beta is not None: + up_f = linear_beta * torch.tanh(up_f / linear_beta) + return (out * up_f).to(x.dtype) + + +def _rmsnorm_gated_torch( + x: torch.Tensor, weight: torch.Tensor, gate: torch.Tensor, eps: float +) -> torch.Tensor: + dtype = x.dtype + x_f = x.float() + var = x_f.pow(2).mean(dim=-1, keepdim=True) + xn = x_f * torch.rsqrt(var + eps) + return (xn.to(dtype) * weight.to(dtype)) * torch.sigmoid(gate) + + +# --------------------------------------------------------------------------- # +# KDA initial-state gather (masked) # +# --------------------------------------------------------------------------- # +if _HAS_TRITON: + + @triton.jit + def _gather_kda_state_kernel( + src_ptr, # ssm_state viewed as [num_lines, S] + idx_ptr, # state_indices [num_seqs] + mask_ptr, # has_initial_state [num_seqs] int8 (unused when HAS_MASK=False) + dst_ptr, # initial viewed as [num_seqs, S] + S, + stride_src, + stride_dst, + HAS_MASK: tl.constexpr, + BLOCK: tl.constexpr, + ): + seq = tl.program_id(0) + offs = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + col_mask = offs < S + if HAS_MASK: + keep = tl.load(mask_ptr + seq) != 0 + load_mask = col_mask & keep + else: + load_mask = col_mask + # int64 offsets: line * S can exceed int32 for large state caches. + line = tl.load(idx_ptr + seq).to(tl.int64) + vals = tl.load(src_ptr + line * stride_src + offs, mask=load_mask, other=0.0) + tl.store(dst_ptr + seq.to(tl.int64) * stride_dst + offs, vals, mask=col_mask) + + +def gather_kda_initial_state( + ssm_state: torch.Tensor, + state_indices: torch.Tensor, + has_initial_state: torch.Tensor | None = None, +) -> torch.Tensor: + """Gather ``ssm_state[state_indices]`` into a packed ``[num_seqs, ...]`` + initial-state tensor, zeroing sequences whose ``has_initial_state`` is + False -- in a single kernel. + + Fuses the ``ssm_state[idx].contiguous()`` gather with the + ``initial[~has_initial_state] = 0`` masking so fresh sequences are written + as zeros in the same pass instead of a gather followed by a separate + zero-write pass. Falls back to the torch path when triton is unavailable or + there are no sequences. + """ + num_seqs = int(state_indices.shape[0]) + if not _HAS_TRITON or num_seqs == 0: + initial = ssm_state[state_indices].contiguous() + if has_initial_state is not None: + initial[~has_initial_state] = 0 + return initial + tail = ssm_state.shape[1:] + src = ssm_state.reshape(ssm_state.shape[0], -1) + S = src.shape[1] + initial = torch.empty( + (num_seqs, *tail), dtype=ssm_state.dtype, device=ssm_state.device + ) + dst = initial.reshape(num_seqs, -1) + has_mask = has_initial_state is not None + mask_arg = has_initial_state.to(torch.int8) if has_mask else src + BLOCK = 1024 + grid = (num_seqs, triton.cdiv(S, BLOCK)) + _gather_kda_state_kernel[grid]( + src, + state_indices, + mask_arg, + dst, + S, + src.stride(0), + dst.stride(0), + HAS_MASK=has_mask, + BLOCK=BLOCK, + ) + return initial + + +# --------------------------------------------------------------------------- # +# MLA k/v assembly (fused cat + pad) # +# --------------------------------------------------------------------------- # +if _HAS_TRITON: + + @triton.jit + def _fuse_mla_kv_kernel( + kv_ptr, # [T, NH, QK_NOPE + V_HEAD] contiguous ([k_nope | v] per head) + krope_ptr, # [T, QK_ROPE] (MQA-shared rope, same for every head) + k_ptr, # [T, NH, QHD] out ([k_nope | k_rope]) + vpad_ptr, # [T, NH, QHD] out ([v | 0]) + stride_kv_t, + stride_kv_h, + stride_kr_t, + stride_out_t, + stride_out_h, + NH, + QK_NOPE: tl.constexpr, + V_HEAD: tl.constexpr, + QHD: tl.constexpr, + BLOCK: tl.constexpr, + ): + pid = tl.program_id(0) # one program per (t, h) + t = pid // NH + h = pid % NH + d = tl.arange(0, BLOCK) + dmask = d < QHD + is_nope = d < QK_NOPE + + kv_base = kv_ptr + t * stride_kv_t + h * stride_kv_h + out_base = t * stride_out_t + h * stride_out_h + + # k = [k_nope (kv[:QK_NOPE]) | k_rope (shared)] + k_nope = tl.load(kv_base + d, mask=dmask & is_nope, other=0.0) + krope = tl.load( + krope_ptr + t * stride_kr_t + (d - QK_NOPE), + mask=dmask & (~is_nope), + other=0.0, + ) + tl.store(k_ptr + out_base + d, tl.where(is_nope, k_nope, krope), mask=dmask) + + # v_padded = [v (kv[QK_NOPE:]) | 0] + is_v = d < V_HEAD + v_val = tl.load(kv_base + QK_NOPE + d, mask=dmask & is_v, other=0.0) + tl.store(vpad_ptr + out_base + d, v_val, mask=dmask) + + +def fuse_mla_kv( + kv: torch.Tensor, + k_rope: torch.Tensor, + qk_nope_head_dim: int, + v_head_dim: int, + qk_rope_head_dim: int, +): + """Assemble the MLA-as-MHA k and padded-v in one kernel. + + ``kv`` is ``[T, NH, qk_nope_head_dim + v_head_dim]`` (kv_b output, laid out + as ``[k_nope | v]`` per head); ``k_rope`` is ``[T, qk_rope_head_dim]`` shared + across heads. Returns:: + + k = [T, NH, q_head_dim] == cat(k_nope, k_rope) + v_padded = [T, NH, q_head_dim] == pad(v, 0..q_head_dim) + + Fuses the ``cat`` + ``F.pad`` (which also reads kv twice via the split) into + a single pass. Torch fallback when triton is unavailable / T == 0. + """ + T, nh, _ = kv.shape + qhd = qk_nope_head_dim + qk_rope_head_dim + if not _HAS_TRITON or T == 0: + k_nope, v = torch.split(kv, [qk_nope_head_dim, v_head_dim], dim=-1) + k_rope_e = k_rope.unsqueeze(1).expand(-1, nh, -1) + k = torch.cat((k_nope, k_rope_e), dim=-1) + v_padded = torch.nn.functional.pad(v, (0, qhd - v_head_dim)) + return k, v_padded + kv = kv.contiguous() + k = torch.empty((T, nh, qhd), dtype=kv.dtype, device=kv.device) + v_padded = torch.empty((T, nh, qhd), dtype=kv.dtype, device=kv.device) + BLOCK = triton.next_power_of_2(qhd) + _fuse_mla_kv_kernel[(T * nh,)]( + kv, + k_rope, + k, + v_padded, + kv.stride(0), + kv.stride(1), + k_rope.stride(0), + k.stride(0), + k.stride(1), + nh, + QK_NOPE=qk_nope_head_dim, + V_HEAD=v_head_dim, + QHD=qhd, + BLOCK=BLOCK, + ) + return k, v_padded diff --git a/atom/utils/debug_helper/dump.py b/atom/utils/debug_helper/dump.py index 1f18714514..4975947b8e 100644 --- a/atom/utils/debug_helper/dump.py +++ b/atom/utils/debug_helper/dump.py @@ -116,6 +116,31 @@ def _hook(_mod, _args, output): t = output[0] if isinstance(output, tuple) else output if not isinstance(t, torch.Tensor): return + if envs.ATOM_FWD_DUMP_STATS_ONLY: + t_float = t.detach().float() + finite = torch.isfinite(t_float) + finite_vals = t_float[finite] + if finite_vals.numel() > 0: + stats = { + "mean": float(finite_vals.mean().item()), + "std": float(finite_vals.std(unbiased=False).item()), + "min": float(finite_vals.min().item()), + "max": float(finite_vals.max().item()), + } + else: + stats = {"mean": None, "std": None, "min": None, "max": None} + torch.save( + { + "shape": tuple(t.shape), + "dtype": str(t.dtype), + "numel": int(t.numel()), + "finite_ratio": float(finite.float().mean().item()), + "zero_ratio": float((t_float == 0).float().mean().item()), + **stats, + }, + fname, + ) + return torch.save( {"hidden": t.detach().cpu(), "shape": tuple(t.shape)}, fname, diff --git a/atom/utils/envs.py b/atom/utils/envs.py index a19b1a9b41..c19d20647a 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -157,6 +157,34 @@ # with that many threads; set to 1 to fall back to the original sequential # per-expert path. "ATOM_LOADER_NUM_THREADS": lambda: int(os.getenv("ATOM_LOADER_NUM_THREADS", "16")), + # Use the optional fastsafetensors package for safetensors file reads. + "ATOM_USE_FASTSAFETENSORS": lambda: ( + os.getenv("ATOM_USE_FASTSAFETENSORS", "0") == "1" + ), + # Split safetensors shard ownership across TP ranks when using + # fastsafetensors. This reduces disk/GDS read amplification while keeping + # the existing ATOM weight-loader semantics. + "ATOM_FASTSAFETENSORS_DIST_LOAD": lambda: ( + os.getenv("ATOM_FASTSAFETENSORS_DIST_LOAD", "0") == "1" + ), + "ATOM_FASTSAFETENSORS_NOGDS": lambda: ( + os.getenv("ATOM_FASTSAFETENSORS_NOGDS", "1") == "1" + ), + # fastsafetensors target device. Keep "cpu" as the conservative default; + # set to "cuda" to read tensors directly onto the current rank's GPU. + "ATOM_FASTSAFETENSORS_DEVICE": lambda: os.getenv( + "ATOM_FASTSAFETENSORS_DEVICE", "cpu" + ), + "ATOM_FASTSAFETENSORS_DEBUG": lambda: ( + os.getenv("ATOM_FASTSAFETENSORS_DEBUG", "0") == "1" + ), + # Force MoE grouped topK routing through the torch reference path. Useful + # when aiter.jit.module_moe_asm cannot compile for the current ROCm/CK arch. + "AITER_USE_TORCH_TOPK": lambda: os.getenv("AITER_USE_TORCH_TOPK", "0") == "1", + # Force sampling through PyTorch instead of aiter.jit.module_sample. + "ATOM_USE_TORCH_SAMPLER": lambda: os.getenv("ATOM_USE_TORCH_SAMPLER", "0") == "1", + # Force bf16 KV cache insert through PyTorch instead of aiter.jit.module_cache. + "ATOM_USE_TORCH_CACHE": lambda: os.getenv("ATOM_USE_TORCH_CACHE", "0") == "1", # --- Attention Backend --- # Use unified_attention (flash-style) for MHA paged/prefill attention instead # of pa_decode_gluon. Set to 1 to enable the unified_attention path. @@ -181,6 +209,14 @@ "ATOM_DUAL_STREAM_MOE_TOKEN_THRESHOLD": lambda: int( os.getenv("ATOM_DUAL_STREAM_MOE_TOKEN_THRESHOLD", "1024") ), + # Kimi-K3: overlap the shared-expert MLP with the routed experts on a + # separate CUDA stream (shared compute runs on alt_stream while routed + # experts + latent post-proc run on the main stream). Default OFF; set to + # "1" to enable. Per-call decode/prefill gating still uses + # ATOM_DUAL_STREAM_MOE_TOKEN_THRESHOLD. + "ATOM_K3_SHARED_EXPERT_OVERLAP": lambda: ( + os.getenv("ATOM_K3_SHARED_EXPERT_OVERLAP", "0") == "1" + ), # Gate/Up interleave mode for MoE weight preshuffle and kernel gate_mode. # "0" (default) = SEPARATED layout; "1" = INTERLEAVE layout. "ATOM_MOE_GU_ITLV": lambda: os.getenv("ATOM_MOE_GU_ITLV", "0") == "1", @@ -234,6 +270,9 @@ "ATOM_FWD_DUMP_LAYER_ATTR", "layer_id" ), "ATOM_FWD_DUMP_ONE_SHOT": lambda: os.getenv("ATOM_FWD_DUMP_ONE_SHOT", "1") == "1", + "ATOM_FWD_DUMP_STATS_ONLY": lambda: ( + os.getenv("ATOM_FWD_DUMP_STATS_ONLY", "0") == "1" + ), # Per-rank weight dump + sys.exit(0) — for byte-equal weight comparison. "ATOM_WEIGHT_DUMP_DIR": lambda: os.getenv("ATOM_WEIGHT_DUMP_DIR", ""), "ATOM_WEIGHT_DUMP_LAYERS": lambda: os.getenv("ATOM_WEIGHT_DUMP_LAYERS", "0"), diff --git a/recipes/Kimi-K3.md b/recipes/Kimi-K3.md new file mode 100644 index 0000000000..03f2330bb0 --- /dev/null +++ b/recipes/Kimi-K3.md @@ -0,0 +1,78 @@ +# Kimi-K3 Usage Guide (gfx950 / MI355) + +Kimi-K3 is a **KimiLinear hybrid-attention MoE** model (`KimiLinearForCausalLM`). Each decoder layer is either a **KDA linear-attention** layer or an **MLA full-attention** layer, on top of a large MXFP4 latent MoE. ATOM serves the **text-only** backbone. + +This guide targets **AMD MI355 (gfx950) only**, `-tp 8`. + +| Variant | Quantization | Description | +|---------|-------------|-------------| +| **MXFP4** | MXFP4 (w4a4, e8m0 scales, group_size=32) | Routed MoE expert weights in microscale FP4. On gfx950 the SiTU experts run the FlyDSL **native SiTUv2** grouped-MoE path. Attention / shared experts / dense MLP stay BF16. | + +**Validated (full 1319, GSM8K 5-shot, base completions, tp8):** + +- **flexible-extract 0.9621 / strict-match 0.9621** — above Kimi-K2-Thinking (0.9363). +- Decode ~20 tok/s. + +--- + +## Launching server + +### MXFP4 on 8×MI355 GPUs (TP8) + +```bash +#!/bin/bash +# ---- load (load-only) ---- +export ATOM_LOADER_USE_THREADPOOL=1 +export ATOM_LOADER_THREADPOOL_WORKERS=16 +export ATOM_SYNC_AFTER_LOAD=1 # one-off TP barrier after load +export ATOM_DIST_TIMEOUT_SECONDS=3600 + +# ---- gfx950 MoE: FlyDSL native SiTUv2 ---- +export ATOM_USE_TRITON_GEMM=1 +export AITER_USE_GROUPED_GEMM=0 +export ATOM_USE_TRITON_MOE=0 # use FlyDSL native SiTUv2 (in-kernel), not the torch-fp32 triton path +export AITER_FLYDSL_FORCE=1 +export AITER_FORCE_GFX1250=0 + +# ---- attention: aiter Triton unified_attention (head_dim=192) ---- +export ATOM_USE_UNIFIED_ATTN=1 +export ATOM_FORCE_ATTN_TRITON=1 + +python -m atom.entrypoints.openai_server \ + --model Kimi-K3 \ + --kv_cache_dtype fp8 -tp 8 \ + --trust-remote-code \ + --max-model-len 16384 \ + --max-num-seqs 64 \ + --max-num-batched-tokens 10240 \ + --gpu-memory-utilization 0.93 \ + --block-size 128 \ + --no-enable_prefix_caching +``` + +MLA layers use the flash KV-cache layout to support `head_dim = 192` (dodges the aiter SHUFFLE-read bug at non-power-of-two head dim), and MLA prefill runs torch SDPA, so prefix caching and chunked prefill stay off — a cached or chunked prefix would be missed by the in-batch SDPA prefill. `-tp 8` is required (tp4 OOMs: MoE weights ~175 GB/GPU). `gpu-memory-utilization 0.93` (not 0.90) so the CUDA-graph pool fits alongside the KDA per-request state cache; at 0.90 startup fails with "Per-request cache tensor exceeds available KV budget". + +--- + +## Accuracy test + +Start the server as above, then run the full 1319-question GSM8K eval: + +```bash +lm_eval \ + --model local-completions \ + --model_args "model=Kimi-K3,base_url=http://localhost:8000/v1/completions,num_concurrent=64,max_retries=3,tokenized_requests=False,trust_remote_code=True" \ + --tasks gsm8k \ + --num_fewshot 5 +``` + +Validated GSM8K result (gfx950 tp8): + +```text +|Tasks|Version| Filter |n-shot| Metric | |Value | |Stderr| +|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:| +|gsm8k| 3|flexible-extract| 5|exact_match|↑ |0.9621|± |0.0067| +| | |strict-match | 5|exact_match|↑ |0.9621|± |0.0067| +``` + +Run on a **clean GPU set**. A competing job stealing VRAM mid-run corrupts the score (a one-off 0.53 was traced to exactly that). Verify no other container holds VRAM and the server stays up (grep the eval output for `ServerDisconnected`). diff --git a/tests/entrypoints/test_api_server_helpers.py b/tests/entrypoints/test_api_server_helpers.py index 6d3ed48a6a..5c62917241 100644 --- a/tests/entrypoints/test_api_server_helpers.py +++ b/tests/entrypoints/test_api_server_helpers.py @@ -15,6 +15,7 @@ import sys import types +from types import SimpleNamespace import pytest @@ -205,3 +206,40 @@ def test_missing_max_model_len_skips_validation(self): max_tokens=8, max_model_len=None, ) + + def test_prompt_over_batch_tokens_without_chunked_prefill_is_rejected( + self, monkeypatch + ): + monkeypatch.setattr( + api_server, + "engine", + SimpleNamespace( + config=SimpleNamespace( + max_model_len=8192, + max_num_batched_tokens=128, + enable_chunked_prefill=False, + ) + ), + ) + seq = SimpleNamespace(num_prompt_tokens=685, max_tokens=256) + + with pytest.raises(ValueError, match="max_num_batched_tokens=128"): + api_server._validate_sequence_context_length(seq) + + def test_prompt_over_batch_tokens_with_chunked_prefill_is_allowed( + self, monkeypatch + ): + monkeypatch.setattr( + api_server, + "engine", + SimpleNamespace( + config=SimpleNamespace( + max_model_len=8192, + max_num_batched_tokens=128, + enable_chunked_prefill=True, + ) + ), + ) + seq = SimpleNamespace(num_prompt_tokens=685, max_tokens=256) + + api_server._validate_sequence_context_length(seq) diff --git a/tests/entrypoints/test_reasoning.py b/tests/entrypoints/test_reasoning.py index 0961fe488e..c7e5bb2a86 100644 --- a/tests/entrypoints/test_reasoning.py +++ b/tests/entrypoints/test_reasoning.py @@ -96,6 +96,23 @@ def test_no_think_start_tag_empty_content(self): assert reasoning == "Reasoning only" assert content == "" + def test_kimi_token_reasoning_and_response(self): + text = ( + "The user asks for the capital." + "<|close|>think<|sep|><|open|>response<|sep|>" + "Paris" + "<|close|>response<|sep|><|close|>message<|sep|><|end_of_msg|>" + ) + reasoning, content = separate_reasoning(text) + assert reasoning == "The user asks for the capital." + assert content == "Paris" + + def test_kimi_response_without_reasoning(self): + text = "<|open|>response<|sep|>42<|close|>response<|sep|>" + reasoning, content = separate_reasoning(text) + assert reasoning is None + assert content == "42" + # ============================================================================ # ReasoningFilter (Streaming) Tests diff --git a/tests/entrypoints/test_serving_completion.py b/tests/entrypoints/test_serving_completion.py new file mode 100644 index 0000000000..f489aa1334 --- /dev/null +++ b/tests/entrypoints/test_serving_completion.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: MIT + +from atom.entrypoints.openai.serving_completion import ( + _strip_stop_strings, + build_completion_response, + build_completion_response_multi, +) + + +def _final_output(text: str) -> dict: + return { + "text": text, + "finish_reason": "stop", + "num_tokens_input": 3, + "num_tokens_output": 5, + "ttft": 0.0, + "tpot": 0.0, + "latency": 0.0, + } + + +def test_strip_stop_strings_removes_earliest_match(): + text = "answer #### 18\nQuestion: next\n" + + assert _strip_stop_strings(text, ["", "Question:"]) == "answer #### 18\n" + + +def test_build_completion_response_strips_stop_string(): + resp = build_completion_response( + "cmpl-test", + "test-model", + _final_output("18\nQuestion: leaked prompt"), + ["Question:"], + ) + + assert resp.choices[0]["text"] == "18\n" + + +def test_build_completion_response_multi_strips_stop_string_per_choice(): + resp = build_completion_response_multi( + "cmpl-test", + "test-model", + [ + _final_output("first\nQuestion: leaked"), + _final_output("second leaked"), + ], + ["Question:", ""], + ) + + assert resp.choices[0]["text"] == "first\n" + assert resp.choices[1]["text"] == "second" diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index cbdf1d54d0..2360a687ed 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -752,6 +752,30 @@ def test_finished_removed_from_running(self, scheduler, seq_factory): scheduler.postprocess(list(scheduler.running), self._output(seq.id, [2])) assert scheduler.get_request_counts() == (0, 0) + def test_aborted_deferred_output_finishes_without_placeholder( + self, scheduler, seq_factory + ): + seq = self._prefill(scheduler, seq_factory([1, 2, 3, 4])) + seq.status = SequenceStatus.ABORTED + assert seq.output_tokens == [] + + finished = scheduler.postprocess( + list(scheduler.running), + ScheduledBatchOutput( + req_ids=[seq.id], + token_ids=[(10,)], + num_rejected=np.array([0], dtype=np.int32), + num_bonus=np.array([0], dtype=np.int32), + draft_token_ids=None, + is_deferred_out=True, + ), + ) + + assert finished == [seq] + assert seq.status == SequenceStatus.FINISHED + assert seq.leave_reason == "aborted" + assert scheduler.get_request_counts() == (0, 0) + # ── get_next_batch_info ──────────────────────────────────────────────────── From 9841bc5d80cfa5f0f1cd91b3a562e44d21b5cd65 Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Tue, 28 Jul 2026 10:19:28 +0800 Subject: [PATCH 2/4] kimi-k3: remove debug/experimental code from PR 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 --- atom/model_ops/attention_mha.py | 106 ++++++++------------------------ atom/model_ops/sampler.py | 15 +---- atom/model_ops/topK.py | 31 ---------- atom/models/kimi_k3.py | 98 +---------------------------- atom/utils/debug_helper/dump.py | 25 -------- atom/utils/envs.py | 18 ------ 6 files changed, 26 insertions(+), 267 deletions(-) diff --git a/atom/model_ops/attention_mha.py b/atom/model_ops/attention_mha.py index 96a4a240ec..c2c137413b 100644 --- a/atom/model_ops/attention_mha.py +++ b/atom/model_ops/attention_mha.py @@ -26,54 +26,6 @@ ) -def _torch_reshape_and_cache( - k: torch.Tensor, - v: torch.Tensor, - k_cache: torch.Tensor, - v_cache: torch.Tensor, - slot_mapping: torch.Tensor, -) -> None: - slots = slot_mapping.flatten().to(torch.long) - # The valid-slot filter needs a host sync (valid.all()/boolean indexing) that - # is illegal during CUDA-graph capture. Decode batches (the only ones captured) - # have no padded (-1) slots, so skip the data-dependent filter while capturing. - if not torch.cuda.is_current_stream_capturing(): - valid = slots >= 0 - if not valid.all(): - slots = slots[valid] - k = k[valid] - v = v[valid] - if slots.numel() == 0: - return - - if k_cache.dim() == 5: - x = k_cache.shape[-1] - block_size = k_cache.shape[-2] - block_idx = slots // block_size - block_offset = slots % block_size - k_cache[block_idx, :, :, block_offset, :] = k.reshape( - k.shape[0], k.shape[1], k.shape[2] // x, x - ) - if v_cache.dim() == 5: - v_x = v_cache.shape[-1] - v_cache[ - block_idx, - :, - block_offset // v_x, - :, - block_offset % v_x, - ] = v - else: - v_cache[block_idx, :, :, block_offset] = v - return - - block_size = k_cache.shape[1] - block_idx = slots // block_size - block_offset = slots % block_size - k_cache[block_idx, block_offset] = k - v_cache[block_idx, block_offset] = v - - @cache def use_pa_decode_bf16_asm() -> bool: return ( @@ -392,24 +344,19 @@ def rope_cache(self, q, k, v, qkv, position, fwd_ctx: ForwardContext): q = self.q_norm(q) if self.k_norm is not None: k = self.k_norm(k) - if envs.ATOM_USE_TORCH_CACHE: - _torch_reshape_and_cache( - k, v, k_cache, v_cache, attn_metadata.slot_mapping - ) - else: - # reshape_and_cache_flash dereferences k_scale/v_scale even for the - # "auto" bf16 path; pass the pre-allocated on-device 1.0 scale. - flash_scale = self._pa_decode_bf16_asm_scale - aiter.reshape_and_cache_flash( - k, - v, - k_cache, - v_cache, - attn_metadata.slot_mapping, - ("auto" if self.kv_cache_dtype == "bf16" else self.kv_cache_dtype), - flash_scale, - flash_scale, - ) + # reshape_and_cache_flash dereferences k_scale/v_scale even for the + # "auto" bf16 path; pass the pre-allocated on-device 1.0 scale. + flash_scale = self._pa_decode_bf16_asm_scale + aiter.reshape_and_cache_flash( + k, + v, + k_cache, + v_cache, + attn_metadata.slot_mapping, + ("auto" if self.kv_cache_dtype == "bf16" else self.kv_cache_dtype), + flash_scale, + flash_scale, + ) self._cache_format = "NHD" else: # for asm paged attention @@ -435,22 +382,17 @@ def rope_cache(self, q, k, v, qkv, position, fwd_ctx: ForwardContext): asm_layout=asm_layout, ) else: - if envs.ATOM_USE_TORCH_CACHE: - _torch_reshape_and_cache( - k, v, k_cache, v_cache, attn_metadata.slot_mapping - ) - else: - aiter.reshape_and_cache( - k, - v, - k_cache, - v_cache, - attn_metadata.slot_mapping, - kv_cache_dtype="auto", - k_scale=None, - v_scale=None, - asm_layout=asm_layout, - ) + aiter.reshape_and_cache( + k, + v, + k_cache, + v_cache, + attn_metadata.slot_mapping, + kv_cache_dtype="auto", + k_scale=None, + v_scale=None, + asm_layout=asm_layout, + ) self._cache_format = "SHUFFLE" if asm_layout else "NHD" # NOTE: on a prefix-cache hit the cached+new KV is gathered into a dense diff --git a/atom/model_ops/sampler.py b/atom/model_ops/sampler.py index 6d64a94670..b73f2ea972 100644 --- a/atom/model_ops/sampler.py +++ b/atom/model_ops/sampler.py @@ -132,19 +132,6 @@ def _temperature_sample( exponential = get_per_token_exponential(vocab_size, logits.device).expand( num_tokens, vocab_size ) - if envs.ATOM_USE_TORCH_SAMPLER: - greedy_mask = temperatures <= self.eps - sampled = torch.empty(num_tokens, dtype=torch.int, device=logits.device) - if greedy_mask.any(): - sampled[greedy_mask] = logits[greedy_mask].argmax(dim=-1).to(torch.int) - if (~greedy_mask).any(): - rows = ~greedy_mask - scaled_logits = logits[rows] / temperatures[rows].unsqueeze(-1) - probs = torch.softmax(scaled_logits, dim=-1, dtype=torch.float32) - sampled[rows] = torch.multinomial(probs, num_samples=1).squeeze(-1).to( - torch.int - ) - return sampled mixed_sample_outer_exponential( sampled_tokens, logits, exponential, temperatures, eps=self.eps ) @@ -182,7 +169,7 @@ def _topk_topp_sample( has_topk = top_ks is not None has_topp = top_ps is not None - if AITER_TOPK_TOPP_AVAILABLE and not envs.ATOM_USE_TORCH_SAMPLER: + if AITER_TOPK_TOPP_AVAILABLE: return self._aiter_sample( probs, top_ks, top_ps, has_topk, has_topp, temperatures ) diff --git a/atom/model_ops/topK.py b/atom/model_ops/topK.py index b66caf826d..2c1a3599e5 100644 --- a/atom/model_ops/topK.py +++ b/atom/model_ops/topK.py @@ -8,7 +8,6 @@ from aiter.jit.utils.torch_guard import torch_compile_guard from atom.config import get_current_atom_config from atom.model_ops.utils import _has_module -from atom.utils import envs from atom.utils.custom_register import direct_register_custom_op @@ -214,7 +213,6 @@ def rocm_aiter_biased_grouped_topk_impl( ) -> tuple[torch.Tensor, torch.Tensor]: from aiter import biased_grouped_topk - from aiter.ops.topk import biased_grouped_topk_torch token = gating_output.shape[0] device = gating_output.device @@ -239,20 +237,6 @@ def rocm_aiter_biased_grouped_topk_impl( else: topk_ids = torch.empty((token, topk), dtype=torch.int32, device=device) topk_weights = torch.empty((token, topk), dtype=torch.float32, device=device) - if envs.AITER_USE_TORCH_TOPK: - torch_weights, torch_ids = biased_grouped_topk_torch( - gating_output, - correction_bias, - topk, - need_renorm, - num_expert_group, - topk_group, - ) - topk_weights.copy_(torch_weights * routed_scaling_factor) - topk_ids.copy_(torch_ids) - if num_fused_shared_experts > 0: - return total_topk_weights, total_topk_ids - return topk_weights, topk_ids biased_grouped_topk( gating_output, correction_bias, @@ -330,7 +314,6 @@ def rocm_aiter_grouped_topk_impl( ) -> tuple[torch.Tensor, torch.Tensor]: from aiter import grouped_topk - from aiter.ops.topk import grouped_topk_torch token = gating_output.shape[0] device = gating_output.device @@ -355,20 +338,6 @@ def rocm_aiter_grouped_topk_impl( else: topk_ids = torch.empty((token, topk), dtype=torch.int32, device=device) topk_weights = torch.empty((token, topk), dtype=torch.float32, device=device) - if envs.AITER_USE_TORCH_TOPK: - torch_weights, torch_ids = grouped_topk_torch( - gating_output, - topk, - need_renorm, - num_expert_group, - topk_group, - scoring_func, - ) - topk_weights.copy_(torch_weights * routed_scaling_factor) - topk_ids.copy_(torch_ids) - if num_fused_shared_experts > 0: - return total_topk_weights, total_topk_ids - return topk_weights, topk_ids grouped_topk( gating_output, topk_weights, diff --git a/atom/models/kimi_k3.py b/atom/models/kimi_k3.py index 43c12f5eda..1429bb9bfd 100644 --- a/atom/models/kimi_k3.py +++ b/atom/models/kimi_k3.py @@ -25,10 +25,6 @@ from atom.config import Config, QuantizationConfig, get_current_atom_config from atom.model_ops.base_attention import Attention -# Side-effect import: registers `torch.ops.aiter.maybe_dual_stream_forward`, -# the Dynamo-opaque dispatcher used to overlap shared experts with routed -# experts on a separate CUDA stream (shared with DeepSeek V2/V4). -from atom.model_ops import module_dispatch_ops as _module_dispatch_ops # noqa: F401 from atom.model_ops.embed_head import ParallelLMHead, VocabParallelEmbedding from atom.model_ops.layernorm import RMSNorm from atom.model_ops.linear import ( @@ -53,7 +49,7 @@ make_layers, maybe_prefix, ) -from atom.utils import envs, mark_spliting_op +from atom.utils import mark_spliting_op from atom.utils.forward_context import get_forward_context from atom.utils.decorators import support_torch_compile @@ -226,12 +222,10 @@ def __init__( config, quant_config: QuantizationConfig | None = None, prefix: str = "", - alt_stream: Optional["torch.cuda.Stream"] = None, ): super().__init__() self.config = config self.prefix = prefix - self.alt_stream = alt_stream self.hidden_dim = config.hidden_size self.num_experts = config.num_experts self.top_k = config.num_experts_per_token @@ -325,34 +319,7 @@ def _routed_source_quant_dtype(layer_prefix: str) -> torch.dtype | None: else None ) - # Shared-expert overlap (dual-stream): run the shared-expert MLP on - # `alt_stream` concurrently with the routed experts. Only meaningful - # when there ARE shared experts and the model handed us an alt_stream - # (created iff ATOM_K3_SHARED_EXPERT_OVERLAP is set). Per-call - # decode/prefill gating happens in the maybe_dual_stream_forward - # dispatcher (ATOM_DUAL_STREAM_MOE_TOKEN_THRESHOLD). - self._use_dual_stream = ( - self.shared_experts is not None and self.alt_stream is not None - ) - if self._use_dual_stream: - # Register so the maybe_dual_stream_forward custom op can look us - # up by layer_name (== self.prefix) without putting sub-modules / - # per-fwd state in the op signature. - get_current_atom_config().compilation_config.static_forward_context[ - prefix - ] = self - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if self._use_dual_stream: - # Custom op = Dynamo barrier, so the alt_stream context inside - # dual_stream_moe_forward stays opaque to torch.compile. The - # dispatcher picks dual vs single per call (num_tokens vs - # ATOM_DUAL_STREAM_MOE_TOKEN_THRESHOLD; forces single under - # piecewise cudagraph / TBO). - return torch.ops.aiter.maybe_dual_stream_forward(hidden_states, self.prefix) - return self.single_stream_moe_forward(hidden_states) - - def single_stream_moe_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # gfx1250 work-around (REQUIRED for correctness, not just stability): the # grouped MoE + MXFP4 latent projections are only numerically correct at # small M on gfx1250 — at large prefill M the contiguous-M path OOB-faults @@ -371,56 +338,6 @@ def single_stream_moe_forward(self, hidden_states: torch.Tensor) -> torch.Tensor return torch.cat(outs, dim=0) return self._forward_impl(hidden_states) - def dual_stream_moe_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - """Overlap the shared-expert MLP with the routed experts. - - The shared branch (``shared_experts``, a plain MLP producing a TP-partial - output — no collective inside) runs on ``alt_stream`` while the routed - experts (+ latent all-reduce / norm / up_proj) run on the main stream. - All collectives stay on the main stream and keep their original order, so - this is numerically identical to ``_forward_impl`` — only the shared - GEMMs move off the critical path. Used for small (decode) batches only, - so the single-stream chunking work-around is unnecessary here. - """ - identity = hidden_states - current_stream = get_forward_context().main_stream - - # Fork shared-expert compute onto alt_stream (waits for the main stream - # to have produced `hidden_states` first). - self.alt_stream.wait_stream(current_stream) - with torch.cuda.stream(self.alt_stream): - shared_partial = self.shared_experts(identity) - - # Routed experts on the main stream. - router_logits = self.gate(hidden_states) - routed_input = ( - self.routed_expert_down_proj(hidden_states) - if self.use_latent_moe - else hidden_states - ) - routed_output = self.experts(routed_input, router_logits) - - if self.use_latent_moe: - if self.tp_size > 1: - routed_output = tensor_model_parallel_all_reduce(routed_output) - if self.routed_expert_norm is not None: - routed_output = self.routed_expert_norm(routed_output) - routed_output = self.routed_expert_up_proj(routed_output) - # Join alt_stream, then all-reduce the shared partial on the main - # stream (ordered after the routed all-reduce above). - current_stream.wait_stream(self.alt_stream) - if self.tp_size > 1: - shared_partial = tensor_model_parallel_all_reduce(shared_partial) - return routed_output + shared_partial - - # Non-latent: routed + shared are both TP-partial; a single deferred - # all-reduce over their sum is correct. - current_stream.wait_stream(self.alt_stream) - routed_output = routed_output + shared_partial - if self.tp_size > 1: - routed_output = tensor_model_parallel_all_reduce(routed_output) - return routed_output - def _forward_impl(self, hidden_states: torch.Tensor) -> torch.Tensor: identity = hidden_states router_logits = self.gate(hidden_states) @@ -1001,7 +918,6 @@ def __init__( atom_config: Config, prefix: str, layer_num: int = 0, - alt_stream: Optional["torch.cuda.Stream"] = None, ): super().__init__() config = _text_config(atom_config.hf_config) @@ -1029,7 +945,6 @@ def __init__( config, quant_config=quant_config, prefix=f"{prefix}.block_sparse_moe", - alt_stream=alt_stream, ) else: self.mlp = KimiMLP( @@ -1193,23 +1108,12 @@ def __init__(self, atom_config: Config, prefix: str = ""): else: self.embed_tokens = PPMissingLayer() - # Single alt_stream shared across layers for shared-expert overlap. - # Created only when the feature is enabled (default off); when None, - # every MoE block falls back to single-stream. Attention/MoE never - # contend since each block runs its MoE after its attention. - self.alt_stream: Optional[torch.cuda.Stream] = ( - torch.cuda.Stream() - if envs.ATOM_K3_SHARED_EXPERT_OVERLAP and torch.cuda.is_available() - else None - ) - self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, lambda prefix, layer_num=None: KimiDecoderLayer( atom_config, prefix=prefix, layer_num=layer_num or 0, - alt_stream=self.alt_stream, ), prefix=f"{prefix}.layers", layer_num_offset=0, diff --git a/atom/utils/debug_helper/dump.py b/atom/utils/debug_helper/dump.py index 4975947b8e..1f18714514 100644 --- a/atom/utils/debug_helper/dump.py +++ b/atom/utils/debug_helper/dump.py @@ -116,31 +116,6 @@ def _hook(_mod, _args, output): t = output[0] if isinstance(output, tuple) else output if not isinstance(t, torch.Tensor): return - if envs.ATOM_FWD_DUMP_STATS_ONLY: - t_float = t.detach().float() - finite = torch.isfinite(t_float) - finite_vals = t_float[finite] - if finite_vals.numel() > 0: - stats = { - "mean": float(finite_vals.mean().item()), - "std": float(finite_vals.std(unbiased=False).item()), - "min": float(finite_vals.min().item()), - "max": float(finite_vals.max().item()), - } - else: - stats = {"mean": None, "std": None, "min": None, "max": None} - torch.save( - { - "shape": tuple(t.shape), - "dtype": str(t.dtype), - "numel": int(t.numel()), - "finite_ratio": float(finite.float().mean().item()), - "zero_ratio": float((t_float == 0).float().mean().item()), - **stats, - }, - fname, - ) - return torch.save( {"hidden": t.detach().cpu(), "shape": tuple(t.shape)}, fname, diff --git a/atom/utils/envs.py b/atom/utils/envs.py index c19d20647a..a693fcd2a8 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -178,13 +178,6 @@ "ATOM_FASTSAFETENSORS_DEBUG": lambda: ( os.getenv("ATOM_FASTSAFETENSORS_DEBUG", "0") == "1" ), - # Force MoE grouped topK routing through the torch reference path. Useful - # when aiter.jit.module_moe_asm cannot compile for the current ROCm/CK arch. - "AITER_USE_TORCH_TOPK": lambda: os.getenv("AITER_USE_TORCH_TOPK", "0") == "1", - # Force sampling through PyTorch instead of aiter.jit.module_sample. - "ATOM_USE_TORCH_SAMPLER": lambda: os.getenv("ATOM_USE_TORCH_SAMPLER", "0") == "1", - # Force bf16 KV cache insert through PyTorch instead of aiter.jit.module_cache. - "ATOM_USE_TORCH_CACHE": lambda: os.getenv("ATOM_USE_TORCH_CACHE", "0") == "1", # --- Attention Backend --- # Use unified_attention (flash-style) for MHA paged/prefill attention instead # of pa_decode_gluon. Set to 1 to enable the unified_attention path. @@ -209,14 +202,6 @@ "ATOM_DUAL_STREAM_MOE_TOKEN_THRESHOLD": lambda: int( os.getenv("ATOM_DUAL_STREAM_MOE_TOKEN_THRESHOLD", "1024") ), - # Kimi-K3: overlap the shared-expert MLP with the routed experts on a - # separate CUDA stream (shared compute runs on alt_stream while routed - # experts + latent post-proc run on the main stream). Default OFF; set to - # "1" to enable. Per-call decode/prefill gating still uses - # ATOM_DUAL_STREAM_MOE_TOKEN_THRESHOLD. - "ATOM_K3_SHARED_EXPERT_OVERLAP": lambda: ( - os.getenv("ATOM_K3_SHARED_EXPERT_OVERLAP", "0") == "1" - ), # Gate/Up interleave mode for MoE weight preshuffle and kernel gate_mode. # "0" (default) = SEPARATED layout; "1" = INTERLEAVE layout. "ATOM_MOE_GU_ITLV": lambda: os.getenv("ATOM_MOE_GU_ITLV", "0") == "1", @@ -270,9 +255,6 @@ "ATOM_FWD_DUMP_LAYER_ATTR", "layer_id" ), "ATOM_FWD_DUMP_ONE_SHOT": lambda: os.getenv("ATOM_FWD_DUMP_ONE_SHOT", "1") == "1", - "ATOM_FWD_DUMP_STATS_ONLY": lambda: ( - os.getenv("ATOM_FWD_DUMP_STATS_ONLY", "0") == "1" - ), # Per-rank weight dump + sys.exit(0) — for byte-equal weight comparison. "ATOM_WEIGHT_DUMP_DIR": lambda: os.getenv("ATOM_WEIGHT_DUMP_DIR", ""), "ATOM_WEIGHT_DUMP_LAYERS": lambda: os.getenv("ATOM_WEIGHT_DUMP_LAYERS", "0"), From 91f71fc03ceea7effc7203a0395547a25f4f03f2 Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Tue, 28 Jul 2026 10:53:12 +0800 Subject: [PATCH 3/4] kimi-k3: drop tests, fix dead moe import, clean stale refs 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) --- atom/model_engine/model_runner.py | 1 - atom/model_ops/attentions/aiter_attention.py | 14 +++--- atom/model_ops/moe.py | 36 +++++--------- atom/models/kimi_k3_fused.py | 5 +- tests/entrypoints/test_api_server_helpers.py | 38 --------------- tests/entrypoints/test_reasoning.py | 17 ------- tests/entrypoints/test_serving_completion.py | 51 -------------------- tests/test_scheduler.py | 24 --------- 8 files changed, 22 insertions(+), 164 deletions(-) delete mode 100644 tests/entrypoints/test_serving_completion.py diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index a243dc233b..a636269779 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -680,7 +680,6 @@ def __init__(self, rank: int, config: Config): self.use_mla = self.is_deepseek_mla() self.use_gdn = self.is_qwen_next() self.use_v4 = self.is_deepseek_v4() - rope_parameters = getattr(self.hf_text_config, "rope_parameters", None) or {} self.use_mrope = "mrope_section" in rope_parameters self.is_deepseek_v32 = ( diff --git a/atom/model_ops/attentions/aiter_attention.py b/atom/model_ops/attentions/aiter_attention.py index 89d02b18a2..a7a6001c29 100644 --- a/atom/model_ops/attentions/aiter_attention.py +++ b/atom/model_ops/attentions/aiter_attention.py @@ -687,13 +687,13 @@ def build_kv_cache_tensor(self, layer_id: int, module): # The 5D SHUFFLE decode path is deliberately NOT used at 192: # (1) the stock aiter SHUFFLE read mis-indexes the padded head_dim # at 192 (~100% error); - # (2) a candidate kernel fix for (1) (kept in - # code_k3/shuffle_kernel_fix.patch) is bit-accurate in op-tests - # and eager server decode, but it still (a) regresses under - # CUDA-graph capture (full-1319 graph: SHUFFLE 0.9204 vs FLASH - # 0.9431) and (b) destabilises the SHARED unified_attention - # FLASH path in-server (non-deterministic garbage) even though - # the edit is SHUFFLED_KV_CACHE-guarded and op-test-clean. + # (2) a candidate aiter-side kernel fix for (1) is bit-accurate in + # op-tests and eager server decode, but it still (a) regresses + # under CUDA-graph capture (full-1319 graph: SHUFFLE 0.9204 vs + # FLASH 0.9431) and (b) destabilises the SHARED + # unified_attention FLASH path in-server (non-deterministic + # garbage) even though the edit is SHUFFLED_KV_CACHE-guarded + # and op-test-clean. # So SHUFFLE decode at 192 stays a follow-up; keep FLASH. k_cache = runner.kv_cache[0, attn_idx] v_cache = runner.kv_cache[1, attn_idx] diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index e082cbc56d..05a9482022 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -1339,7 +1339,9 @@ def apply( expert_map=expert_map, act_quant=self.act_quant, situ_beta=getattr(layer, "activation_situ_beta", None), - situ_linear_beta=getattr(layer, "activation_situ_linear_beta", None), + situ_linear_beta=getattr( + layer, "activation_situ_linear_beta", None + ), ) # Always-on shared expert(s) via a standalone dense GEMM, @@ -3178,9 +3180,9 @@ def _load_w2( # Narrow parameter and load. shard_size = expert_data.shape[shard_dim] load_shard_size = loaded_weight.shape[shard_dim] // self.tp_size - shard_start = load_shard_size * tp_rank - loaded_shard = loaded_weight.narrow(shard_dim, shard_start, load_shard_size) - loaded_weight = loaded_shard + loaded_weight = loaded_weight.narrow( + shard_dim, load_shard_size * tp_rank, load_shard_size + ) if load_shard_size != shard_size: expert_data = expert_data.narrow(shard_dim, 0, load_shard_size) # w2, down_proj: Load into only logical weight of w2. @@ -3787,16 +3789,9 @@ def forward_impl_graph( if self.reduce_results and (self.tp_size > 1 or self.ep_size > 1): # Default set to False. (May have to add shared expert outputs.) - from atom.model_ops.linear import fp32_allreduce_enabled - - if fp32_allreduce_enabled(): - hs32 = final_hidden_states.float() - torch.distributed.all_reduce(hs32, group=get_tp_group().device_group) - final_hidden_states = hs32.to(final_hidden_states.dtype) - else: - final_hidden_states = get_tp_group().all_reduce( - final_hidden_states, ca_fp8_quant=False - ) + final_hidden_states = get_tp_group().all_reduce( + final_hidden_states, ca_fp8_quant=False + ) return final_hidden_states @@ -3847,16 +3842,9 @@ def forward_impl(self, hidden_states: torch.Tensor, router_logits: torch.Tensor) if self.reduce_results and (self.tp_size > 1 or self.ep_size > 1): # Default set to False. (May have to add shared expert outputs.) - from atom.model_ops.linear import fp32_allreduce_enabled - - if fp32_allreduce_enabled(): - hs32 = final_hidden_states.float() - torch.distributed.all_reduce(hs32, group=get_tp_group().device_group) - final_hidden_states = hs32.to(final_hidden_states.dtype) - else: - final_hidden_states = get_tp_group().all_reduce( - final_hidden_states, ca_fp8_quant=False - ) + final_hidden_states = get_tp_group().all_reduce( + final_hidden_states, ca_fp8_quant=False + ) return final_hidden_states diff --git a/atom/models/kimi_k3_fused.py b/atom/models/kimi_k3_fused.py index 15935bdb2c..b535cf8c60 100644 --- a/atom/models/kimi_k3_fused.py +++ b/atom/models/kimi_k3_fused.py @@ -8,8 +8,9 @@ - ``apply_attn_res`` : block-residual soft-attention mix (attn_res boundaries) The activation and gated-rmsnorm kernels remain gated behind ``ATOM_K3_FUSED`` -in kimi_k3.py. ``apply_attn_res`` is the production path; correctness is checked -against a local torch reference in my_script/optest_fused.py. +in kimi_k3.py, with the torch reference implementations at the bottom of this +module doubling as the no-triton fallback. ``apply_attn_res`` is the production +path and is not env-gated. """ from __future__ import annotations diff --git a/tests/entrypoints/test_api_server_helpers.py b/tests/entrypoints/test_api_server_helpers.py index 5c62917241..6d3ed48a6a 100644 --- a/tests/entrypoints/test_api_server_helpers.py +++ b/tests/entrypoints/test_api_server_helpers.py @@ -15,7 +15,6 @@ import sys import types -from types import SimpleNamespace import pytest @@ -206,40 +205,3 @@ def test_missing_max_model_len_skips_validation(self): max_tokens=8, max_model_len=None, ) - - def test_prompt_over_batch_tokens_without_chunked_prefill_is_rejected( - self, monkeypatch - ): - monkeypatch.setattr( - api_server, - "engine", - SimpleNamespace( - config=SimpleNamespace( - max_model_len=8192, - max_num_batched_tokens=128, - enable_chunked_prefill=False, - ) - ), - ) - seq = SimpleNamespace(num_prompt_tokens=685, max_tokens=256) - - with pytest.raises(ValueError, match="max_num_batched_tokens=128"): - api_server._validate_sequence_context_length(seq) - - def test_prompt_over_batch_tokens_with_chunked_prefill_is_allowed( - self, monkeypatch - ): - monkeypatch.setattr( - api_server, - "engine", - SimpleNamespace( - config=SimpleNamespace( - max_model_len=8192, - max_num_batched_tokens=128, - enable_chunked_prefill=True, - ) - ), - ) - seq = SimpleNamespace(num_prompt_tokens=685, max_tokens=256) - - api_server._validate_sequence_context_length(seq) diff --git a/tests/entrypoints/test_reasoning.py b/tests/entrypoints/test_reasoning.py index c7e5bb2a86..0961fe488e 100644 --- a/tests/entrypoints/test_reasoning.py +++ b/tests/entrypoints/test_reasoning.py @@ -96,23 +96,6 @@ def test_no_think_start_tag_empty_content(self): assert reasoning == "Reasoning only" assert content == "" - def test_kimi_token_reasoning_and_response(self): - text = ( - "The user asks for the capital." - "<|close|>think<|sep|><|open|>response<|sep|>" - "Paris" - "<|close|>response<|sep|><|close|>message<|sep|><|end_of_msg|>" - ) - reasoning, content = separate_reasoning(text) - assert reasoning == "The user asks for the capital." - assert content == "Paris" - - def test_kimi_response_without_reasoning(self): - text = "<|open|>response<|sep|>42<|close|>response<|sep|>" - reasoning, content = separate_reasoning(text) - assert reasoning is None - assert content == "42" - # ============================================================================ # ReasoningFilter (Streaming) Tests diff --git a/tests/entrypoints/test_serving_completion.py b/tests/entrypoints/test_serving_completion.py deleted file mode 100644 index f489aa1334..0000000000 --- a/tests/entrypoints/test_serving_completion.py +++ /dev/null @@ -1,51 +0,0 @@ -# SPDX-License-Identifier: MIT - -from atom.entrypoints.openai.serving_completion import ( - _strip_stop_strings, - build_completion_response, - build_completion_response_multi, -) - - -def _final_output(text: str) -> dict: - return { - "text": text, - "finish_reason": "stop", - "num_tokens_input": 3, - "num_tokens_output": 5, - "ttft": 0.0, - "tpot": 0.0, - "latency": 0.0, - } - - -def test_strip_stop_strings_removes_earliest_match(): - text = "answer #### 18\nQuestion: next\n" - - assert _strip_stop_strings(text, ["", "Question:"]) == "answer #### 18\n" - - -def test_build_completion_response_strips_stop_string(): - resp = build_completion_response( - "cmpl-test", - "test-model", - _final_output("18\nQuestion: leaked prompt"), - ["Question:"], - ) - - assert resp.choices[0]["text"] == "18\n" - - -def test_build_completion_response_multi_strips_stop_string_per_choice(): - resp = build_completion_response_multi( - "cmpl-test", - "test-model", - [ - _final_output("first\nQuestion: leaked"), - _final_output("second leaked"), - ], - ["Question:", ""], - ) - - assert resp.choices[0]["text"] == "first\n" - assert resp.choices[1]["text"] == "second" diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 2360a687ed..cbdf1d54d0 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -752,30 +752,6 @@ def test_finished_removed_from_running(self, scheduler, seq_factory): scheduler.postprocess(list(scheduler.running), self._output(seq.id, [2])) assert scheduler.get_request_counts() == (0, 0) - def test_aborted_deferred_output_finishes_without_placeholder( - self, scheduler, seq_factory - ): - seq = self._prefill(scheduler, seq_factory([1, 2, 3, 4])) - seq.status = SequenceStatus.ABORTED - assert seq.output_tokens == [] - - finished = scheduler.postprocess( - list(scheduler.running), - ScheduledBatchOutput( - req_ids=[seq.id], - token_ids=[(10,)], - num_rejected=np.array([0], dtype=np.int32), - num_bonus=np.array([0], dtype=np.int32), - draft_token_ids=None, - is_deferred_out=True, - ), - ) - - assert finished == [seq] - assert seq.status == SequenceStatus.FINISHED - assert seq.leave_reason == "aborted" - assert scheduler.get_request_counts() == (0, 0) - # ── get_next_batch_info ──────────────────────────────────────────────────── From fdcafd041e8196701f65df7a418ac3ce97f47513 Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Tue, 28 Jul 2026 11:03:59 +0800 Subject: [PATCH 4/4] kimi-k3: drop redundant env_snapshot plumbing 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) --- atom/model_engine/async_proc.py | 7 ------- atom/model_engine/engine_core.py | 11 +---------- atom/model_engine/engine_core_mgr.py | 2 -- 3 files changed, 1 insertion(+), 19 deletions(-) diff --git a/atom/model_engine/async_proc.py b/atom/model_engine/async_proc.py index ba63c0973f..7169108637 100644 --- a/atom/model_engine/async_proc.py +++ b/atom/model_engine/async_proc.py @@ -15,7 +15,6 @@ import logging import multiprocessing -import os import pickle import queue import threading @@ -70,13 +69,9 @@ def __init__( rank: int, kv_output_addr: str | None = None, all_ranks_barrier=None, - env_snapshot: dict[str, str] | None = None, *args, **kwargs, ): - if env_snapshot: - os.environ.update(env_snapshot) - # Bind this worker's lifetime to its parent EngineCore: if the parent # exits for any reason, have the kernel reap this process immediately # instead of leaving it orphaned. A ModelRunner worker holds a large GPU @@ -280,7 +275,6 @@ def __init__(self, finalizer, proc_num: int, runner: str, *args): atexit.register(self._cleanup_shared_memory) self.all_ranks_barrier = ctx.Barrier(proc_num) init_exit_handler(self) - env_snapshot = dict(os.environ) # KV output aggregation infrastructure self.kv_output_aggregator: KVOutputAggregator | None = None @@ -306,7 +300,6 @@ def __init__(self, finalizer, proc_num: int, runner: str, *args): i, self.kv_output_addrs[i], self.all_ranks_barrier, - env_snapshot, *args, ), ) diff --git a/atom/model_engine/engine_core.py b/atom/model_engine/engine_core.py index 6680d234dd..31668e0b08 100644 --- a/atom/model_engine/engine_core.py +++ b/atom/model_engine/engine_core.py @@ -2,7 +2,6 @@ # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. import logging -import os import pickle import queue import threading @@ -188,15 +187,7 @@ def _send_engine_dead(self): self.output_thread.join(timeout=0.5) @staticmethod - def run_engine( - config: Config, - input_address: str, - output_address: str, - env_snapshot: dict[str, str] | None = None, - ): - if env_snapshot: - os.environ.update(env_snapshot) - + def run_engine(config: Config, input_address: str, output_address: str): # Bind this EngineCore's lifetime to its parent (the server / # CoreManager): if the parent exits, have the kernel reap this process — # and, transitively, the ModelRunner workers it spawns — instead of diff --git a/atom/model_engine/engine_core_mgr.py b/atom/model_engine/engine_core_mgr.py index 6121870437..c399645b00 100644 --- a/atom/model_engine/engine_core_mgr.py +++ b/atom/model_engine/engine_core_mgr.py @@ -4,7 +4,6 @@ import asyncio import logging import multiprocessing -import os import pickle import queue import weakref @@ -792,7 +791,6 @@ def launch_engine_core(config: Config, dp_rank: int = 0): "config": config, "input_address": input_address, "output_address": output_address, - "env_snapshot": dict(os.environ), }, )