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/engine_core.py b/atom/model_engine/engine_core.py
index 319dbc55aa..31668e0b08 100644
--- a/atom/model_engine/engine_core.py
+++ b/atom/model_engine/engine_core.py
@@ -19,6 +19,7 @@
from atom.model_engine.sequence import Sequence, SequenceStatus, get_exit_sequence
from atom.utils import (
envs,
+ get_hf_text_config,
init_exit_handler,
make_zmq_socket,
set_process_title,
@@ -120,6 +121,20 @@ def __init__(self, config: Config, input_address: str, output_address: str):
if not good:
self._finalizer()
+ hf_text_config = get_hf_text_config(config.hf_config)
+ arches = getattr(hf_text_config, "architectures", None) or []
+ is_kimi_kda_hybrid = any("KimiK3" in str(a) for a in arches) or (
+ getattr(hf_text_config, "model_type", None) == "kimi_linear"
+ and hasattr(hf_text_config, "full_attn_layer_ids")
+ )
+ if is_kimi_kda_hybrid and config.enable_prefix_caching:
+ logger.info(
+ "%s: Kimi-K3 KDA hybrid disables prefix caching engine-wide "
+ "because recurrent attention state is per request.",
+ self.label,
+ )
+ config.enable_prefix_caching = False
+
self.scheduler = Scheduler(config)
self.kv_transfer_enabled = bool(config.kv_transfer_config)
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..a636269779 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",
@@ -771,6 +772,15 @@ def __init__(self, rank: int, config: Config):
f"[{self.rank_name}] Model load done: {config.model} "
f"(weights loaded in {load_elapsed:.2f}s)"
)
+ if (
+ os.getenv("ATOM_SYNC_AFTER_LOAD", "0").lower() in ("1", "true", "yes")
+ and get_tp_group().world_size > 1
+ ):
+ logger.info(
+ "Waiting for all TP ranks to finish model loading before warmup"
+ )
+ get_tp_group().barrier()
+ logger.info("All TP ranks finished model loading")
# Optional debug instrumentation; no-op when env vars unset.
# See atom/utils/debug_helper/.
@@ -899,10 +909,14 @@ def is_qwen_next(self) -> bool:
"qwen3_next_mtp",
"qwen3_5_text",
"qwen3_5_moe_text",
+ "kimi_linear",
):
return True
return False
+ def is_kimi_linear(self) -> bool:
+ return getattr(self.hf_text_config, "model_type", None) == "kimi_linear"
+
def is_deepseek_v4(self) -> bool:
# NOTE: `hf_text_config.model_type` reads "deepseek_v3" for V4 because
# `_CONFIG_REGISTRY` maps deepseek_v4 → deepseek_v3 (V4 reuses V3 schema).
@@ -1340,6 +1354,16 @@ def warmup_model(self):
if pcp_size > 1:
warmup_max_tokens = max(1, warmup_max_tokens // pcp_size)
+ # TEMPORARY (benign, warmup-only): cap the dummy warmup prefill. A large
+ # all-zero warmup batch samples garbage logits over all positions and
+ # faults the sampler/softmax on gfx1250; real inference only samples the
+ # last token per seq so it is unaffected. Proper fix = skip/greedy sampling
+ # during is_dummy_run warmup, then this cap can go. Not needed for the MoE
+ # (that large-M crash is fixed in aiter grouped_moe_gfx1250).
+ _warmup_cap = int(os.environ.get("ATOM_WARMUP_MAX_TOKENS", "0") or "0")
+ if _warmup_cap > 0:
+ warmup_max_tokens = min(warmup_max_tokens, _warmup_cap)
+
num_seqs = min(warmup_max_tokens // max_model_len, self.config.max_num_seqs)
if num_seqs == 0:
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..c2c137413b 100644
--- a/atom/model_ops/attention_mha.py
+++ b/atom/model_ops/attention_mha.py
@@ -331,6 +331,33 @@ 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)
+ # 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
@@ -794,23 +821,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 +867,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..a7a6001c29 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 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]
+ 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..05a9482022 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,10 @@ 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 +1378,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 +1407,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 +2612,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 +3054,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 +3092,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 +3106,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,
diff --git a/atom/model_ops/sampler.py b/atom/model_ops/sampler.py
index a1f52d47c4..b73f2ea972 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(
diff --git a/atom/models/kimi_k3.py b/atom/models/kimi_k3.py
new file mode 100644
index 0000000000..1429bb9bfd
--- /dev/null
+++ b/atom/models/kimi_k3.py
@@ -0,0 +1,1324 @@
+# 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
+
+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 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 = "",
+ ):
+ super().__init__()
+ self.config = config
+ self.prefix = prefix
+ 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
+ )
+
+ def 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 _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,
+ ):
+ 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",
+ )
+ 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()
+
+ 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,
+ ),
+ 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..b535cf8c60
--- /dev/null
+++ b/atom/models/kimi_k3_fused.py
@@ -0,0 +1,837 @@
+"""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, 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
+
+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/envs.py b/atom/utils/envs.py
index a19b1a9b41..a693fcd2a8 100644
--- a/atom/utils/envs.py
+++ b/atom/utils/envs.py
@@ -157,6 +157,27 @@
# 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"
+ ),
# --- 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.
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`).