From 5f624a0a4dabfa2ee5e06b3379385ab2023f03d3 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Thu, 27 Aug 2026 18:52:01 -0700 Subject: [PATCH 1/2] [TRTLLM-15159][feat] K3 checkpoint loading: generalize lazy-safetensors and decouple from HfWeightLoader Follow-up to the Kimi K3 / KimiLinear support in #17269. That PR shipped a load path validated end-to-end on a real 1.5 TB K3 checkpoint, but it reached the lazy-loading behavior through model-name conditionals in shared code and configured K3 loading through a block of environment variables. This removes both, without changing runtime behavior. 1. Generalize lazy-safetensors into a selectable LoadFormat. Opening safetensors shards and streaming only the rank-local slices, instead of materializing the full checkpoint in host RAM, becomes LoadFormat.LAZY_SAFETENSORS, reachable via load_format="lazy_safetensors" for any HF checkpoint. The sharded-vs-consolidated file selection is handled inside the format. K3 opts in by declaring it as its default in get_model_defaults; an explicit user load_format still wins. 2. Decouple the shared HfWeightLoader from K3. The Kimi/K3 model-name branch (_is_kimi_k3_checkpoint and its dispatch) is removed; the shared loader now dispatches on the resolved LoadFormat, not on any model name. No kimi/k3 string checks remain in weight_loader.py. 3. Move the K3 FP8 weight-read env vars onto the config surface. These 5 knobs decide whether a replicated K3 projection is read from an FP8 (e4m3, 128x128 block-scale) copy of its weights instead of BF16, and they are consumed on the checkpoint-loading path: load_weights keeps the FP8_PB_WO checkpoint pairs only when the read is enabled, and the post-load conversion swaps the modules. They therefore move to QuantConfig. Resolution lives in the new, CPU-unit-testable kimi_k3_knobs.py (config value wins, else the deprecated env var honored with a one-time warning, else the historical default). Deprecated env var -> config equivalent: KIMI_K3_FP8_WEIGHT_READ -> quant_config.kimi_k3_fp8_weight_read KIMI_K3_FP8_WEIGHT_READ_KDA -> quant_config.kimi_k3_fp8_weight_read_kda KIMI_K3_FP8_WEIGHT_READ_MLA -> quant_config.kimi_k3_fp8_weight_read_mla KIMI_K3_FP8_WEIGHT_READ_GATE_UP -> quant_config.kimi_k3_fp8_weight_read_gate_up KIMI_K3_KDA_GLUE_FP8 -> quant_config.kimi_k3_kda_glue_fp8 The env vars are still honored for back-compat, emitting a one-time deprecation warning; an explicit config value takes precedence when both are set. The four remaining K3 env vars (TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ, KIMI_K3_MLA_MAX_POSITIONS, KIMI_K3_FUSED_ATTN_RES, KIMI_K3_ROUTER_BF16) are left as they are. They configure the forward path -- the latent MoE GEMM selection, the identity-RoPE table size, the fused attn_res op, and the router GEMM dtype -- not loading or quantization, and none of them is required for K3 to load. Migrating them, and deciding which of them belong on the public API surface at all, is tracked separately. Test coverage: - tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py exercises the LAZY_SAFETENSORS format on a synthetic multi-shard fixture (sharded and consolidated selection, lazy slice materialization, dtype/shape/values, FP8 weights and scale tensors), no large checkpoint required. - tests/unittest/_torch/modeling/test_kimi_k3_fp8_weight_read_gates.py covers the 5 gates: config defaults, sub-gate narrowing, the is_sm_100f() arch gate, deprecated-env back-compat, config-over-env precedence, warn-once, and a zero-env/zero-config completeness check. - API-stability reference updated for the new QuantConfig fields. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../checkpoints/hf/checkpoint_loader.py | 1 + .../models/checkpoints/hf/weight_loader.py | 128 ++++++---- tensorrt_llm/_torch/models/kimi_k3_knobs.py | 176 ++++++++++++++ .../_torch/models/modeling_kimi_k3_vl.py | 17 +- .../_torch/models/modeling_kimi_linear.py | 114 +++------ .../_torch/pyexecutor/model_loader.py | 20 +- tensorrt_llm/llmapi/llm_args.py | 5 + tensorrt_llm/models/modeling_utils.py | 34 +++ .../test_kimi_k3_fp8_weight_read_gates.py | 218 ++++++++++++++++-- .../checkpoints/hf/test_weight_loader.py | 144 +++++++++--- .../references/quant_config.yaml | 15 ++ 11 files changed, 683 insertions(+), 189 deletions(-) create mode 100644 tensorrt_llm/_torch/models/kimi_k3_knobs.py diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py index c495230779cc..0e626bd2ccd3 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py @@ -91,6 +91,7 @@ def build_checkpoint_catalog(self, checkpoint_dir: str, return self.weight_loader.build_checkpoint_catalog( checkpoint_dir, use_consolidated=kwargs.get("use_consolidated", False), + load_lazily=kwargs.get("load_lazily", False), ) @property diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py index 3eeae6ceb74a..ea7b247d50ac 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. import glob -import json import multiprocessing import os import threading @@ -56,9 +55,6 @@ _RANK_STRIPED_IO_POLICY = "rank_striped_read_ahead" _SUPPORTED_IO_POLICIES = (_NATIVE_IO_POLICY, _RANK_STRIPED_IO_POLICY) _SUPPORTED_REQUESTED_IO_POLICIES = (_AUTO_IO_POLICY, ) + _SUPPORTED_IO_POLICIES -# Model families whose checkpoints are too large to materialize in host RAM; -# their models stream rank-local slices out of the lazy mmapped handles. -_LAZY_SAFETENSORS_MODEL_TYPES = ("kimi_k3", "kimi_linear") # Default to a single cached checkpoint: each entry pins a full copy of the # raw weights in CPU RAM, so callers wanting cross-model caching must opt in # via TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES. @@ -360,27 +356,6 @@ def _with_weight_cache(self, self._cache_loaded_weights(cache_key, weights) return weights - @staticmethod - def _requires_lazy_safetensors(checkpoint_dir: str) -> bool: - """Whether this checkpoint must stay mmapped instead of being read - into host RAM. - - The listed model families ship checkpoints too large to materialize - in host RAM (Kimi K3 is about 1.5 TB), and their models stream - rank-local slices (expert-parallel expert ranges) out of the lazy - handles during `load_weights`. - """ - config_path = os.path.join(checkpoint_dir, "config.json") - if not os.path.isfile(config_path): - return False - # Do not swallow read/parse failures: every rank must take the same - # branch here (the eager path enqueues collectives), so a rank-local - # transient error routing one rank differently would deadlock the job. - # Propagating fails fast on all ranks instead. - with open(config_path) as f: - model_type = json.load(f).get("model_type") - return model_type in _LAZY_SAFETENSORS_MODEL_TYPES - def _load_lazy_safetensors( self, checkpoint_dir: str, @@ -390,7 +365,7 @@ def _load_lazy_safetensors( Values are ``safetensors`` PySafeSlice objects: ``v[:]`` (or any indexing) materializes only the requested bytes from the mmapped file. This lets a model's ``load_weights`` stream a huge checkpoint - and read only its rank-local shard (e.g. Kimi K3 expert-parallel + and read only its rank-local shard (e.g. expert-parallel routed expert slices) without ever holding the full checkpoint in RAM. """ weight_files = sorted(glob.glob(f"{checkpoint_dir}/*.safetensors")) @@ -421,7 +396,7 @@ def _load_lazy_safetensors( # transient source weights when ModelLoader.load() returns. lazy_weights = _LazySafetensorsWeights(weights, handles) # A lazy slice does not carry the file it came from, and a model that - # wants to re-open shards itself (Kimi K3 streams rank-local experts + # wants to re-open shards itself (e.g. streaming rank-local experts # per shard file, precisely to avoid holding this mapping open) has no # other reliable source: transformers no longer sets # ``PretrainedConfig._name_or_path``. @@ -432,8 +407,12 @@ def load_weights(self, checkpoint_dir: str, mapping: Mapping, use_consolidated: bool = False, + load_lazily: bool = False, **kwargs) -> dict[str, Any]: - """Load synchronously without activating session-scoped read-ahead.""" + """Load synchronously without activating session-scoped read-ahead. + + See :meth:`_load_weights_native` for what ``load_lazily`` selects. + """ self._reset_checkpoint_io_status() if self._checkpoint_io_policy != _NATIVE_IO_POLICY: status = self._last_checkpoint_io_status @@ -446,8 +425,11 @@ def load_weights(self, f"load: requested={status.requested}, " f"selected={status.selected}, reason={status.fallback_reason}.") logger.info(message) - weights = self._load_weights_native(checkpoint_dir, mapping, - use_consolidated, **kwargs) + weights = self._load_weights_native(checkpoint_dir, + mapping, + use_consolidated, + load_lazily=load_lazily, + **kwargs) self._last_checkpoint_io_status.effective = _NATIVE_IO_POLICY self._log_checkpoint_io_status() return weights @@ -457,12 +439,21 @@ def open_weight_session(self, checkpoint_dir: str, mapping: Mapping, use_consolidated: bool = False, + load_lazily: bool = False, **kwargs) -> Iterator[dict[str, Any]]: - """Keep opt-in read-ahead alive through model materialization.""" + """Keep opt-in read-ahead alive through model materialization. + + See :meth:`_load_weights_native` for what ``load_lazily`` selects. A + lazy load never reads whole shards up front, so it is not eligible for + read-ahead; the flag is forwarded so every path below agrees on it. + """ self._reset_checkpoint_io_status() if self._checkpoint_io_policy == _NATIVE_IO_POLICY: - weights = self._load_weights_native(checkpoint_dir, mapping, - use_consolidated, **kwargs) + weights = self._load_weights_native(checkpoint_dir, + mapping, + use_consolidated, + load_lazily=load_lazily, + **kwargs) self._last_checkpoint_io_status.effective = _NATIVE_IO_POLICY self._log_checkpoint_io_status() yield weights @@ -492,6 +483,7 @@ def open_weight_session(self, str(coordinated_mapping_error), active_communicator=active_communicator, allow_native_prefetch=False, + load_lazily=load_lazily, **kwargs, ) yield weights @@ -506,6 +498,7 @@ def open_weight_session(self, # Ray workers intentionally disable MPI. Preserve the native # path's per-process read-ahead in that supported mode. allow_native_prefetch=mpi_disabled(), + load_lazily=load_lazily, **kwargs, ) yield weights @@ -516,6 +509,7 @@ def open_weight_session(self, mapping, use_consolidated, active_communicator, + load_lazily=load_lazily, **kwargs, ) if session is None: @@ -583,9 +577,10 @@ def _selected_safetensors_files(checkpoint_dir: str, def build_checkpoint_catalog( self, checkpoint_dir: str, - use_consolidated: bool = False) -> CheckpointCatalog | None: + use_consolidated: bool = False, + load_lazily: bool = False) -> CheckpointCatalog | None: """Inspect the selected eager SafeTensors files without reading payloads.""" - if self._requires_lazy_safetensors(checkpoint_dir): + if load_lazily: return None if (self._partial_model_loading or int(os.environ.get("TLLM_OVERRIDE_LAYER_NUM", "0")) != 0): @@ -629,6 +624,7 @@ def _fallback_to_native(self, session=None, *, allow_native_prefetch: bool | None = None, + load_lazily: bool = False, **kwargs) -> dict[str, Any]: status = self._last_checkpoint_io_status status.activated = False @@ -671,6 +667,7 @@ def load_native() -> dict[str, Any]: checkpoint_dir, mapping, use_consolidated, + load_lazily=load_lazily, _local_communicator=fallback_communicator, _allow_prefetch=allow_native_prefetch, **kwargs) @@ -722,6 +719,7 @@ def _start_rank_striped_read_ahead( mapping: Mapping, use_consolidated: bool, active_communicator, + load_lazily: bool = False, **kwargs, ) -> tuple[dict[str, Any], RankStripedReadAheadSession | None]: """Load weights via rank-striped read-ahead, or fall back to native.""" @@ -742,10 +740,13 @@ def _start_rank_striped_read_ahead( None, node_communicator, active_communicator, "rank-striped node communicator cleanup") self._last_checkpoint_io_status.selected = _NATIVE_IO_POLICY - return self._fallback_to_native(checkpoint_dir, mapping, + return self._fallback_to_native(checkpoint_dir, + mapping, use_consolidated, str(coordinated_split_error), - active_communicator, **kwargs), None + active_communicator, + load_lazily=load_lazily, + **kwargs), None weight_files = [] stats = [] @@ -753,10 +754,9 @@ def _start_rank_striped_read_ahead( eligibility_reason = None preflight_error = None try: - if self._requires_lazy_safetensors(checkpoint_dir): + if load_lazily: eligibility_reason = ( - "the checkpoint requires model-specific lazy SafeTensors loading" - ) + "the checkpoint is loaded as lazy SafeTensors slices") weight_files = self._selected_safetensors_files( checkpoint_dir, use_consolidated) stats = [(path, os.stat(path)) for path in weight_files] @@ -776,11 +776,14 @@ def _start_rank_striped_read_ahead( active_communicator, "rank-striped preflight", preflight_error) if coordinated_preflight_error is not None: self._last_checkpoint_io_status.selected = _NATIVE_IO_POLICY - return self._fallback_to_native(checkpoint_dir, mapping, + return self._fallback_to_native(checkpoint_dir, + mapping, use_consolidated, str(coordinated_preflight_error), active_communicator, - node_communicator, **kwargs), None + node_communicator, + load_lazily=load_lazily, + **kwargs), None local_rank = 0 local_size = 1 @@ -810,10 +813,13 @@ def _start_rank_striped_read_ahead( None, node_communicator, active_communicator, "rank-striped node preflight cleanup") self._last_checkpoint_io_status.selected = _NATIVE_IO_POLICY - return self._fallback_to_native(checkpoint_dir, mapping, + return self._fallback_to_native(checkpoint_dir, + mapping, use_consolidated, str(coordinated_node_error), - active_communicator, **kwargs), None + active_communicator, + load_lazily=load_lazily, + **kwargs), None file_sizes = [(path, stat.st_size) for path, stat in stats] checkpoint_bytes = sum(size for _, size in file_sizes) @@ -845,11 +851,14 @@ def _start_rank_striped_read_ahead( if fallback_reasons: rank, reason = fallback_reasons[0] self._last_checkpoint_io_status.selected = _NATIVE_IO_POLICY - return self._fallback_to_native(checkpoint_dir, mapping, + return self._fallback_to_native(checkpoint_dir, + mapping, use_consolidated, f"rank {rank}: {reason}", active_communicator, - node_communicator, **kwargs), None + node_communicator, + load_lazily=load_lazily, + **kwargs), None self._last_checkpoint_io_status.selected = _RANK_STRIPED_IO_POLICY session = None @@ -865,11 +874,14 @@ def _start_rank_striped_read_ahead( "rank-striped reader setup", setup_error) if coordinated_setup_error is not None: - return self._fallback_to_native(checkpoint_dir, mapping, + return self._fallback_to_native(checkpoint_dir, + mapping, use_consolidated, str(coordinated_setup_error), active_communicator, - node_communicator, session, + node_communicator, + session, + load_lazily=load_lazily, **kwargs), None assert session is not None @@ -912,11 +924,27 @@ def _load_weights_native(self, mapping: Mapping, use_consolidated: bool = False, *, + load_lazily: bool = False, _local_communicator=None, _allow_prefetch: bool = True, **kwargs) -> dict[str, Any]: - """Load weights with the native (no read-ahead) I/O policy.""" - if self._requires_lazy_safetensors(checkpoint_dir): + """Load weights with the native (no read-ahead) I/O policy. + + When ``load_lazily`` is set (selected via ``LoadFormat.LAZY_SAFETENSORS`` + on the ``load_format`` surface, or a model's declared default), the + safetensors shards are opened lazily so only the rank-local slices each + model reads are materialized -- never the whole checkpoint in CPU + memory. This is required for checkpoints too large to fit in host RAM. + + Otherwise the checkpoint is loaded eagerly: safetensors shards may be + prefetched in parallel to warm up the OS file cache if the CPU memory is + large enough, before their tensors are loaded via mmap. When + `_WEIGHT_CACHE_ENV` is on, eager loads can also use a CPU weight cache to + accelerate repeated loading under the same process. + + Returns a `ConsumableWeightsDict` mapping checkpoint tensor names to tensors. + """ + if load_lazily: return self._load_lazy_safetensors(checkpoint_dir, use_consolidated) weight_files = glob.glob(f"{checkpoint_dir}/*.safetensors") # Some model checkpoint directories contain not only the sharded safetensors, but one diff --git a/tensorrt_llm/_torch/models/kimi_k3_knobs.py b/tensorrt_llm/_torch/models/kimi_k3_knobs.py new file mode 100644 index 000000000000..4d3ee870b638 --- /dev/null +++ b/tensorrt_llm/_torch/models/kimi_k3_knobs.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Resolution of the Kimi K3 FP8 weight-read knobs. + +These knobs decide whether a replicated K3 projection is read from an FP8 +(e4m3, 128x128 block-scale) copy of its weights instead of BF16. They are +consumed on the checkpoint-loading path — ``load_weights`` keeps the FP8 +checkpoint pairs only when the read is enabled, and the post-load conversion +swaps the modules — so they live on +:class:`~tensorrt_llm.models.modeling_utils.QuantConfig`. + +They used to be read straight from the ``KIMI_K3_*`` environment variables +inside :mod:`modeling_kimi_linear`. Resolution precedence for every knob: + +1. an explicit config value (not ``None``) wins; +2. else the deprecated environment variable, if set, is honored (emitting a + one-time deprecation warning); otherwise +3. the historical default (some are computed at runtime from the arch / + parallelism). + +This keeps the shipped env-var behavior working for back-compat while making the +config surface authoritative. Kept intentionally lightweight (no torch model +imports) so the resolution logic is unit-testable on CPU. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Any, Callable, Optional + +from ..._utils import is_sm_100f +from ...logger import logger + +# Deprecated env var -> the config path that now owns the knob. Used only to +# make the one-time deprecation warning actionable. +_ENV_TO_CONFIG_PATH = { + "KIMI_K3_FP8_WEIGHT_READ": "quant_config.kimi_k3_fp8_weight_read", + "KIMI_K3_FP8_WEIGHT_READ_KDA": "quant_config.kimi_k3_fp8_weight_read_kda", + "KIMI_K3_FP8_WEIGHT_READ_MLA": "quant_config.kimi_k3_fp8_weight_read_mla", + "KIMI_K3_FP8_WEIGHT_READ_GATE_UP": "quant_config.kimi_k3_fp8_weight_read_gate_up", + "KIMI_K3_KDA_GLUE_FP8": "quant_config.kimi_k3_kda_glue_fp8", +} + +# FP8 weight-read knob field names on ``QuantConfig``. These are carried from the +# user-facing ``llm_args.quant_config`` onto the checkpoint-derived +# ``model_config.quant_config`` by :func:`carry_user_quant_knobs`. +KIMI_K3_QUANT_KNOB_FIELDS = ( + "kimi_k3_fp8_weight_read", + "kimi_k3_fp8_weight_read_kda", + "kimi_k3_fp8_weight_read_mla", + "kimi_k3_fp8_weight_read_gate_up", + "kimi_k3_kda_glue_fp8", +) + + +def _resolve( + config_value: Optional[Any], env_name: str, env_parser: Callable[[str], Any], default: Any +) -> Any: + """Resolve one knob: config wins; else deprecated env (warn-once); else default. + + ``config_value`` is ``None`` when the knob was not set on the config surface. + Whenever the deprecated env var is present a one-time warning is emitted, + even if a config value overrides it, so users learn to migrate. + """ + env_raw = os.environ.get(env_name) + if env_raw is not None: + logger.warning_once( + f"Environment variable '{env_name}' is deprecated and will be " + f"removed; set '{_ENV_TO_CONFIG_PATH[env_name]}' on the config " + f"surface instead (via extra_llm_api_options). It is still honored " + f"for now, but the config value takes precedence when both are set.", + key=f"kimi_k3_deprecated_env::{env_name}", + ) + if config_value is not None: + return config_value + if env_raw is not None: + return env_parser(env_raw) + return default + + +def _knob(config: Optional[Any], name: str) -> Optional[Any]: + """Read ``name`` off a (possibly ``None``) config object; ``None`` if absent.""" + if config is None: + return None + return getattr(config, name, None) + + +@dataclass(frozen=True) +class Fp8WeightReadGates: + """Resolved FP8 weight-read gates. + + ``master`` folds in the ``is_sm_100f()`` arch gate: it is ``False`` off + Blackwell regardless of the requested value. The sub-gates only ever narrow + an enabled master, so all of them are ``False`` when ``master`` is ``False``. + """ + + master: bool + kda: bool + kda_glue: bool + mla: bool + gate_up: bool + + +def resolve_fp8_weight_read_gates( + quant_config: Optional[Any], *, enable_attention_dp: bool +) -> Fp8WeightReadGates: + """Resolve the 5 FP8 weight-read gates from ``quant_config`` (+ deprecated env). + + The master switch is opt-in (FP8 weight reads are lossy relative to BF16, so + a default run keeps BF16 and matches the published accuracy numbers) and is + additionally gated on ``is_sm_100f()`` — the DeepGEMM ``fp8_swap_ab_gemm`` + kernel is Blackwell-only. The KDA / KDA-glue / MLA / gate-up sub-gates only + narrow an already-enabled master and stay default-on. + """ + master_requested = _resolve( + _knob(quant_config, "kimi_k3_fp8_weight_read"), + "KIMI_K3_FP8_WEIGHT_READ", + lambda s: s not in ("", "0"), + False, + ) + master = bool(is_sm_100f() and master_requested) + + kda = master and bool( + _resolve( + _knob(quant_config, "kimi_k3_fp8_weight_read_kda"), + "KIMI_K3_FP8_WEIGHT_READ_KDA", + lambda s: s != "0", + True, + ) + ) + kda_glue = kda and bool( + _resolve( + _knob(quant_config, "kimi_k3_kda_glue_fp8"), + "KIMI_K3_KDA_GLUE_FP8", + lambda s: s != "0", + True, + ) + ) + mla = master and bool( + _resolve( + _knob(quant_config, "kimi_k3_fp8_weight_read_mla"), + "KIMI_K3_FP8_WEIGHT_READ_MLA", + lambda s: s != "0", + True, + ) + ) + gate_up = master and bool( + _resolve( + _knob(quant_config, "kimi_k3_fp8_weight_read_gate_up"), + "KIMI_K3_FP8_WEIGHT_READ_GATE_UP", + lambda s: s != "0", + enable_attention_dp, + ) + ) + return Fp8WeightReadGates(master=master, kda=kda, kda_glue=kda_glue, mla=mla, gate_up=gate_up) + + +def carry_user_quant_knobs( + src_quant_config: Optional[Any], dst_quant_config: Optional[Any] +) -> None: + """Copy the K3 FP8 weight-read knobs from the user's ``quant_config`` + (``llm_args.quant_config``) onto the checkpoint-derived + ``model_config.quant_config``. + + ``model_config.quant_config`` is built from the checkpoint's + ``hf_quant_config.json`` and does not carry the user's ``extra_llm_api_options`` + ``quant_config`` values, so the K3 FP8-read knobs are threaded across here. + A no-op for knobs the user did not set (``None``) and for non-K3 runs. + """ + if src_quant_config is None or dst_quant_config is None: + return + for name in KIMI_K3_QUANT_KNOB_FIELDS: + value = getattr(src_quant_config, name, None) + if value is not None: + setattr(dst_quant_config, name, value) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py b/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py index d00236574b83..edad8a749df2 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py @@ -37,7 +37,7 @@ """ import copy -from typing import Optional, Tuple +from typing import TYPE_CHECKING, Dict, Optional, Tuple import torch import torch.nn as nn @@ -77,6 +77,9 @@ register_vision_encoder, ) +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + # --------------------------------------------------------------------------- # Native MoonViT3d Vision Encoder Components (K3 deltas) # --------------------------------------------------------------------------- @@ -460,6 +463,18 @@ class KimiK3ForConditionalGeneration(KimiK25ForConditionalGeneration): _VISION_MODEL_CLS = KimiK3VisionModel mamba_metadata_cls = KimiLinearForCausalLM.mamba_metadata_cls + @classmethod + def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> Dict[str, str]: + """Default this model to the lazy safetensors load format. + + The K3 checkpoint (~1.5 TB) must be streamed shard-by-shard, so K3 + declares the lazy ``LoadFormat`` as its default now that the shared + loader no longer sniffs K3 by model type. This wrapper subclasses HF + ``PreTrainedModel``, not the TRT-LLM base, so there is no inherited + ``get_model_defaults`` to extend. A user-set ``load_format`` still wins. + """ + return {"load_format": "lazy_safetensors"} + def __init__( self, model_config: ModelConfig[PretrainedConfig], diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 38dad076d57f..64988e42d50d 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -103,7 +103,6 @@ from safetensors import safe_open from torch import nn -from ..._utils import is_sm_100f from ...logger import logger from ...mapping import Mapping from ...models.modeling_utils import QuantAlgo, QuantConfig @@ -121,6 +120,7 @@ from ..moe.fused_moe import ConfigurableMoE, SiTuActivation, TRTLLMGenFusedMoE, create_moe from ..moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod from ..utils import AuxStreamType +from .kimi_k3_knobs import resolve_fp8_weight_read_gates from .modeling_speculative import SpecDecOneEngineForCausalLM from .modeling_utils import DecoderModel, register_auto_model, run_concurrently @@ -142,61 +142,14 @@ # bounds. KIMI_K3_MLA_MAX_POSITIONS overrides the size for short-context # deployments. _KIMI_K3_MLA_MAX_POSITIONS_ENV = "KIMI_K3_MLA_MAX_POSITIONS" +# The FP8 block-scale weight-read knobs (master + KDA/MLA/gate_up/KDA-glue +# sub-gates) live on ``QuantConfig`` and are resolved in ``kimi_k3_knobs`` +# (config value > deprecated env, warn-once > historical default). _KIMI_K3_MLA_DERIVED_PARAM_SUFFIXES = ( ".self_attn.mixer.k_b_proj_trans", ".self_attn.mixer.v_b_proj", ) -# Serve the replicated MoE-layer MLP projections (shared-expert gate/up/down -# and the latent up/down projection) from an FP8 copy of their weights instead -# of BF16. Under attention data-parallelism every rank re-reads these dense -# weights in full on every decode step, so decode is bound by that HBM read; -# an FP8 (e4m3) weight with 128x128 block scales roughly halves those bytes. -# The MLA projections and the routed MXFP4 experts are left untouched (the KDA -# q/k/v/g/o projections have their own switch below). The FP8 weight read is -# lossy relative to BF16, so it is opt-in: set this to "1" to trade accuracy -# for decode bandwidth. Default "0" keeps BF16, which is what the published -# accuracy numbers are measured against. -_KIMI_K3_FP8_WEIGHT_READ_ENV = "KIMI_K3_FP8_WEIGHT_READ" - -# Also read the KDA linear-attention q/k/v/g/o projections at FP8 block-scale. -# These are the largest single replicated weight read (~61 GB/rank of the -# ~109 GB BF16 read per decode step). They use the same FP8 path as the MLP -# projections above but are gated separately: the recurrent linear-attention -# core is more accuracy-sensitive than the feed-forward MLPs, so set this to -# "0" to keep the KDA projections in BF16 while still reading the MLPs at FP8. -# The master KIMI_K3_FP8_WEIGHT_READ switch and the SM100 gate still apply. -_KIMI_K3_FP8_WEIGHT_READ_KDA_ENV = "KIMI_K3_FP8_WEIGHT_READ_KDA" - -# Also read the MLA (full-attention) q_a/q_b/o and output-gate projections at -# FP8 block-scale. These are the replicated attention weights the MLP pass and -# the KDA pass above leave in BF16, and they are re-read in full by every rank -# each decode step under attention data-parallelism. Two MLA projections are -# deliberately kept in BF16: kv_a_proj_with_mqa outputs kv_lora_rank + -# qk_rope_head_dim (576, not a multiple of 128, so no exact 128x128 block -# scale), and kv_b_proj's weight is consumed directly (not through its forward) -# by the absorbed-decode _kv_b_absorb_split to build the k/v absorb matrices, -# which has no FP8 dequant path. The master KIMI_K3_FP8_WEIGHT_READ switch and -# the SM100 gate still apply. -_KIMI_K3_FP8_WEIGHT_READ_MLA_ENV = "KIMI_K3_FP8_WEIGHT_READ_MLA" - -# Expert override (prototype): set to "0" to drop the -# KimiKDALinearAttention decode -# fast path — fused qkvg and [f_a | b] projections, persistent conv staging, -# and precomputed kernel-layout constants (``forward_decode``) — when the -# KDA projections are read at FP8 block-scale. With the fast path kept (the -# default on an enabled master), decode issues the loader's fused FP8 -# ``qkvg_proj`` GEMM for q/k/v/g plus one small BF16 GEMV for [f_a | b] -# (``finalize_decode_weights_fp8``), so FP8 weight storage and the decode -# glue savings coexist. Requires the FP8 KDA read to be active; no effect -# otherwise. Default on ("0" disables). -_KIMI_K3_KDA_GLUE_FP8_ENV = "KIMI_K3_KDA_GLUE_FP8" - -# FP8 read for the fused shared-expert gate_up_proj. Default follows the -# parallel layout (on under attention DP, off under TP — see the conversion -# helper's comment); set 0/1 to force either. -_KIMI_K3_FP8_WEIGHT_READ_GATE_UP_ENV = "KIMI_K3_FP8_WEIGHT_READ_GATE_UP" - class KimiK3MoEGate(nn.Module): """Kimi K3 gate weights and routing method for ``ConfigurableMoE``.""" @@ -290,23 +243,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return self.weight * hidden_states_float.to(input_dtype) -def _resolve_fp8_weight_read_gates() -> tuple[bool, bool, bool]: - """Resolve the FP8 weight-read switches into (master, kda, kda_glue). - - The master switch is opt-in: FP8 weight reads are lossy relative to BF16, - so a default run keeps BF16 and matches the published accuracy numbers. - The KDA and KDA-glue switches only narrow an already-enabled master, so - they stay default-on and are inert while the master is off. - """ - fp8_weight_read = is_sm_100f() and os.environ.get(_KIMI_K3_FP8_WEIGHT_READ_ENV, "0") not in ( - "", - "0", - ) - kda_fp8 = fp8_weight_read and os.environ.get(_KIMI_K3_FP8_WEIGHT_READ_KDA_ENV, "1") != "0" - kda_glue_fp8 = kda_fp8 and os.environ.get(_KIMI_K3_KDA_GLUE_FP8_ENV, "1") != "0" - return fp8_weight_read, kda_fp8, kda_glue_fp8 - - def _resolve_kimi_situ_betas(cfg: Any) -> tuple[float, float]: """Return the finite SiTu betas required by the routed-expert kernels.""" config_situ_beta = getattr(cfg, "activation_situ_beta", None) @@ -2076,6 +2012,13 @@ def _setup_helix_mappings( @classmethod def get_model_defaults(cls, llm_args) -> dict: + # - load_format=lazy_safetensors: the K3 checkpoint (~1.5 TB) must be + # streamed shard-by-shard (rank-local slices only); eagerly loading + # it into host RAM OOM-kills the job. Declaring the lazy LoadFormat + # as the model default is how K3 opts into that path without any + # model-name check in the shared weight loader, and without the user + # having to set load_format. A user-set load_format still wins + # (apply_model_defaults_to_llm_args honors explicit overrides). # - enable_block_reuse defaults off: reuse is supported as an # explicit opt-in (routes to CppMambaHybridCacheManager with # per-block KDA state snapshots); the default stays on the @@ -2085,10 +2028,11 @@ def get_model_defaults(cls, llm_args) -> dict: # the fallback C++ path requires num_heads % 64 == 0, which K3's # 96 query heads violate. return { + "load_format": "lazy_safetensors", "kv_cache_config": { "enable_block_reuse": False, "tokens_per_block": 64, - } + }, } @classmethod @@ -2274,7 +2218,10 @@ def _load_trunk_params( ) # Keep each FP8_PB_WO checkpoint pair alongside the BF16 # parameter only when the later weight-read conversion consumes it. - stash_ckpt_fp8 = _resolve_fp8_weight_read_gates()[0] + stash_ckpt_fp8 = resolve_fp8_weight_read_gates( + self.model_config.quant_config, + enable_attention_dp=self.model_config.mapping.enable_attention_dp, + ).master # KDA head-shard (attention-DP off): rank r loads head rows/cols # [r*local : (r+1)*local] of every head-major KDA tensor. kda_tp_size, kda_tp_rank = 1, 0 @@ -2662,18 +2609,21 @@ def _finalize_weight_load(self, num_params: int, num_moe_layers: int) -> None: # the bf16 wrapper fast path (finalize_decode_weights) is NOT built: # both fuse the same projections and the wrapper path — checked first # at decode — would bypass the FP8 modules entirely, leaving the FP8 - # copies resident but inert. KIMI_K3_FP8_WEIGHT_READ_KDA=0 restores - # the bf16 wrapper fast path; KIMI_K3_KDA_GLUE_FP8=1 instead rebuilds - # the wrapper fast path on top of the FP8 modules after the + # copies resident but inert. quant_config.kimi_k3_fp8_weight_read_kda=0 + # restores the bf16 wrapper fast path; kimi_k3_kda_glue_fp8=1 instead + # rebuilds the wrapper fast path on top of the FP8 modules after the # conversion (finalize_decode_weights_fp8), so neither is traded away. - fp8_weight_read, kda_fp8, kda_glue_fp8 = _resolve_fp8_weight_read_gates() + gates = resolve_fp8_weight_read_gates( + self.model_config.quant_config, + enable_attention_dp=self.model_config.mapping.enable_attention_dp, + ) # Build the KDA fused projection views and decode kernel constants. # This must run after every KDA parameter is loaded and sharded. num_kda_fused = 0 for layer in self.model.layers: if getattr(layer, "is_kda", False) and _has_weights(layer): - if not kda_fp8: + if not gates.kda: layer.linear_attn.finalize_decode_weights() num_kda_fused += int( layer.linear_attn._qkvg_proj_weight is not None @@ -2696,14 +2646,10 @@ def _finalize_weight_load(self, num_params: int, num_moe_layers: int) -> None: # FP8 block-scale weight read for the replicated MoE-layer MLPs. The # DeepGEMM fp8_swap_ab_gemm kernel is Blackwell-only; keep BF16 on any # other SM or when explicitly disabled. - if fp8_weight_read: - gate_up_default = "1" if self.model_config.mapping.enable_attention_dp else "0" + if gates.master: n_fp8 = _convert_moe_mlps_to_fp8_weight_read( self.model, - include_fused_gate_up=os.environ.get( - _KIMI_K3_FP8_WEIGHT_READ_GATE_UP_ENV, gate_up_default - ) - != "0", + include_fused_gate_up=gates.gate_up, ) logger.info( f"Kimi K3: reading {n_fp8} MoE-layer MLP projections " @@ -2713,14 +2659,14 @@ def _finalize_weight_load(self, num_params: int, num_moe_layers: int) -> None: # weight read; convert them to the same FP8 block-scale read unless # kept in BF16 for accuracy (their own switch — the recurrent core # is the most precision-sensitive slice). - if os.environ.get(_KIMI_K3_FP8_WEIGHT_READ_KDA_ENV, "1") != "0": + if gates.kda: n_kda = _convert_kda_projections_to_fp8_weight_read(self.model) logger.info( f"Kimi K3: reading {n_kda} KDA q/k/v/g/o projections " f"at FP8 block-scale (q/k/v/g fused into one prefill/decode/verify GEMM " f"per layer)" ) - if kda_glue_fp8: + if gates.kda_glue: # Rebuild the fused projection path on top of the FP8 # modules. This must run after the conversion above so # fused FP8 qkvg_proj exists and only [f_a | b] is fused @@ -2739,7 +2685,7 @@ def _finalize_weight_load(self, num_params: int, num_moe_layers: int) -> None: # leave in BF16; convert them to the same FP8 block-scale read # (kv_a/kv_b stay BF16 — see the switch's comment) unless kept in # BF16 for accuracy. - if os.environ.get(_KIMI_K3_FP8_WEIGHT_READ_MLA_ENV, "1") != "0": + if gates.mla: n_mla = _convert_mla_projections_to_fp8_weight_read(self.model) logger.info( f"Kimi K3: reading {n_mla} MLA q_a/q_b/o/g projections at FP8 block-scale" diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index e1ee1d5ee08f..0dac51684c91 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -971,8 +971,9 @@ def allocate_buffer_on_cuda(t: torch.Tensor): _apply_to_buffers_only(model, allocate_buffer_on_cuda) - need_initialized_weights = load_format not in (LoadFormat.AUTO, - LoadFormat.DUMMY) + need_initialized_weights = load_format not in ( + LoadFormat.AUTO, LoadFormat.DUMMY, + LoadFormat.LAZY_SAFETENSORS) def allocate_weights_on_cuda(t: torch.Tensor): if t not in memo: @@ -1033,7 +1034,13 @@ def init_meta_tensor(t: torch.Tensor): # prior to zero-copy mapping, then refreshes derived state after # real GMS tensors are bound. gms_post_load_handled = False - if load_format == LoadFormat.AUTO: + if load_format in (LoadFormat.AUTO, LoadFormat.LAZY_SAFETENSORS): + # LAZY_SAFETENSORS shares the AUTO disk-load path; it differs + # only in that the checkpoint weight loader opens shards lazily + # (streaming rank-local slices) instead of materializing the + # whole checkpoint. The load_lazily flag carries that choice + # down to HfWeightLoader so the shared loader dispatches on the + # resolved LoadFormat, not on any model name. # Pass model= so format-specific loaders (e.g. MX) can # write weights directly into parameter buffers via P2P. # Generic loaders ignore model=; loaders that can consume a @@ -1043,6 +1050,7 @@ def init_meta_tensor(t: torch.Tensor): "model": model, # Generic loaders ignore it; MXCheckpointLoader pops it. "source_identity": self._source_identity, + "load_lazily": load_format == LoadFormat.LAZY_SAFETENSORS, } if checkpoint_loader.checkpoint_format == "MX": load_weights_kwargs["model_config"] = config @@ -1858,6 +1866,12 @@ def _load_and_validate_config( 'nvfp4_gemm_allowed_backends'] = config.nvfp4_gemm_allowed_backends config.extra_attrs[ 'kv_cache_dtype'] = self.llm_args.kv_cache_config.dtype + # Thread the Kimi K3 FP8 weight-read knobs onto the checkpoint-derived + # quant_config, which is built from hf_quant_config.json and does not + # otherwise carry the user's quant_config. A no-op for non-K3 runs. + # See _torch/models/kimi_k3_knobs.py. + from ..models.kimi_k3_knobs import carry_user_quant_knobs + carry_user_quant_knobs(self.llm_args.quant_config, config.quant_config) # Store allreduce pre-allocation config for AllReduce module access. # Use get_text_config() so VLM wrapper configs (e.g. KimiK2VLConfig, # KimiK25Config) that store the text config under .text_config are diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index c47f5e107984..04b714b8f4ab 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5278,6 +5278,11 @@ class LoadFormat(Enum): VISION_ONLY = 2 # Load weights through GPU Memory Service. GMS = 3 + # Open safetensors shards lazily and stream only the rank-local slices, + # never materializing the full checkpoint in host RAM. A model opts in by + # declaring this as its default load format (get_model_defaults), or a user + # can select it via load_format="lazy_safetensors" for any HF checkpoint. + LAZY_SAFETENSORS = 4 class ModelExpressConfig(StrictBaseModel): diff --git a/tensorrt_llm/models/modeling_utils.py b/tensorrt_llm/models/modeling_utils.py index f5a9f27364f7..34dea2d784cc 100644 --- a/tensorrt_llm/models/modeling_utils.py +++ b/tensorrt_llm/models/modeling_utils.py @@ -167,6 +167,40 @@ class QuantConfig(StrictBaseModel): "Number of Philox rounds for stochastic rounding PRNG. Higher values give better randomness." ) + # --- Kimi K3 FP8 weight-read knobs (opt-in decode-bandwidth optimization) --- + # Opt-in FP8 (e4m3, 128x128 block-scale) reads for selected replicated K3 + # projections; lossy vs BF16, so off by default. ``None`` means unset; + # resolution (is_sm_100f() gate, deprecated-env fallback, per-knob defaults) + # lives in ``_torch/models/kimi_k3_knobs.py``. Per-knob detail on each field. + kimi_k3_fp8_weight_read: Optional[bool] = Field( + default=None, + description= + "Kimi K3: master FP8 block-scale weight read for the replicated MoE-layer " + "MLP/KDA/MLA projections (opt-in, additionally gated on is_sm_100f()). " + "None keeps the historical default (off).") + kimi_k3_fp8_weight_read_kda: Optional[bool] = Field( + default=None, + description= + "Kimi K3: read the KDA q/k/v/g/o projections at FP8 block-scale; narrows " + "an enabled master. None keeps the historical default (on).") + kimi_k3_fp8_weight_read_mla: Optional[bool] = Field( + default=None, + description= + "Kimi K3: read the MLA q_a/q_b/o/gate projections at FP8 block-scale; " + "narrows an enabled master. None keeps the historical default (on).") + kimi_k3_fp8_weight_read_gate_up: Optional[bool] = Field( + default=None, + description= + "Kimi K3: read the fused shared-expert gate_up_proj at FP8 block-scale; " + "narrows an enabled master. None follows the parallel layout (on under " + "attention DP, off otherwise).") + kimi_k3_kda_glue_fp8: Optional[bool] = Field( + default=None, + description= + "Kimi K3: rebuild the fused KDA decode-glue projection on top of the FP8 " + "modules; narrows an enabled FP8 KDA read. None keeps the historical " + "default (on).") + @cached_property def quant_mode(self) -> QuantModeWrapper: quant_mode_list = [ diff --git a/tests/unittest/_torch/modeling/test_kimi_k3_fp8_weight_read_gates.py b/tests/unittest/_torch/modeling/test_kimi_k3_fp8_weight_read_gates.py index d28aeb03a0ac..f3e279abc454 100644 --- a/tests/unittest/_torch/modeling/test_kimi_k3_fp8_weight_read_gates.py +++ b/tests/unittest/_torch/modeling/test_kimi_k3_fp8_weight_read_gates.py @@ -12,71 +12,245 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Pin the Kimi K3 FP8 weight-read switch defaults. +"""Pin the Kimi K3 FP8 weight-read gate resolution and its defaults. The FP8 weight read is lossy relative to BF16, so the master switch must stay opt-in — a default run keeps BF16 and matches the published accuracy numbers. These tests exist so that default cannot drift silently: it did once, and the only symptom was accuracy measured against a configuration nobody selected. +The 5 gates now resolve from ``QuantConfig`` (an explicit config value wins), +else the deprecated ``KIMI_K3_*`` env var (still honored, warn-once), else the +historical default. Both surfaces are exercised here, along with the migration's +acceptance criteria at the config/unit level: each of the 5 former env vars has +a config equivalent, a run with neither env vars nor config values resolves a +complete and valid set of gates, and a set env var is honored exactly once with +a deprecation warning. + CPU-only: ``is_sm_100f`` is patched so the gates can be exercised without Blackwell hardware. """ import pytest -from tensorrt_llm._torch.models import modeling_kimi_linear -from tensorrt_llm._torch.models.modeling_kimi_linear import _resolve_fp8_weight_read_gates +from tensorrt_llm._torch.models import kimi_k3_knobs +from tensorrt_llm._torch.models.kimi_k3_knobs import resolve_fp8_weight_read_gates +from tensorrt_llm.models.modeling_utils import QuantConfig ENVS = ( "KIMI_K3_FP8_WEIGHT_READ", "KIMI_K3_FP8_WEIGHT_READ_KDA", + "KIMI_K3_FP8_WEIGHT_READ_MLA", + "KIMI_K3_FP8_WEIGHT_READ_GATE_UP", "KIMI_K3_KDA_GLUE_FP8", ) @pytest.fixture def sm100f(monkeypatch): - """Report Blackwell so the SM gate never masks the env behavior.""" - monkeypatch.setattr(modeling_kimi_linear, "is_sm_100f", lambda: True) + """Report Blackwell so the SM gate never masks env/config behavior; clear env.""" + monkeypatch.setattr(kimi_k3_knobs, "is_sm_100f", lambda: True) for env in ENVS: monkeypatch.delenv(env, raising=False) +def _gates(quant_config=None, *, enable_attention_dp=False): + return resolve_fp8_weight_read_gates(quant_config, enable_attention_dp=enable_attention_dp) + + +def _tuple(g): + return (g.master, g.kda, g.kda_glue, g.mla, g.gate_up) + + +# --- Defaults (config surface) --------------------------------------------- + + def test_master_switch_defaults_off(sm100f): """With nothing set, no FP8 weight read anywhere.""" - assert _resolve_fp8_weight_read_gates() == (False, False, False) + assert _tuple(_gates(QuantConfig())) == (False, False, False, False, False) -def test_master_switch_opt_in(sm100f, monkeypatch): +def test_master_switch_opt_in_via_config(sm100f): """Setting the master switch enables it and the default-on sub-gates.""" + g = _gates(QuantConfig(kimi_k3_fp8_weight_read=True), enable_attention_dp=True) + assert _tuple(g) == (True, True, True, True, True) + + +def test_gate_up_default_follows_attention_dp(sm100f): + """gate_up defaults on under attention-DP, off otherwise (master on).""" + on = _gates(QuantConfig(kimi_k3_fp8_weight_read=True), enable_attention_dp=True) + off = _gates(QuantConfig(kimi_k3_fp8_weight_read=True), enable_attention_dp=False) + assert on.gate_up is True and off.gate_up is False + + +@pytest.mark.parametrize( + "field,expected", + [ + # kda off -> kda_glue collapses too; mla/gate_up independent. + ("kimi_k3_fp8_weight_read_kda", (True, False, False, True, True)), + ("kimi_k3_kda_glue_fp8", (True, True, False, True, True)), + ("kimi_k3_fp8_weight_read_mla", (True, True, True, False, True)), + ("kimi_k3_fp8_weight_read_gate_up", (True, True, True, True, False)), + ], +) +def test_sub_gates_narrow_enabled_master_via_config(sm100f, field, expected): + """Sub-gates only ever narrow an enabled master; they never enable alone.""" + g = _gates( + QuantConfig(kimi_k3_fp8_weight_read=True, **{field: False}), enable_attention_dp=True + ) + assert _tuple(g) == expected + + +@pytest.mark.parametrize( + "field", + [ + "kimi_k3_fp8_weight_read_kda", + "kimi_k3_fp8_weight_read_mla", + "kimi_k3_fp8_weight_read_gate_up", + "kimi_k3_kda_glue_fp8", + ], +) +def test_sub_gates_inert_while_master_off_via_config(sm100f, field): + """A sub-gate set on must not turn on FP8 reads by itself.""" + g = _gates(QuantConfig(**{field: True}), enable_attention_dp=True) + assert _tuple(g) == (False, False, False, False, False) + + +def test_non_blackwell_never_reads_fp8_config(monkeypatch): + """The SM gate wins even with the master switch explicitly on (config).""" + monkeypatch.setattr(kimi_k3_knobs, "is_sm_100f", lambda: False) + for env in ENVS: + monkeypatch.delenv(env, raising=False) + g = _gates(QuantConfig(kimi_k3_fp8_weight_read=True), enable_attention_dp=True) + assert _tuple(g) == (False, False, False, False, False) + + +# --- Deprecated env var back-compat (still honored) ------------------------ + + +def test_master_switch_opt_in_via_deprecated_env(sm100f, monkeypatch): + """The deprecated master env var is still honored when config is unset.""" monkeypatch.setenv("KIMI_K3_FP8_WEIGHT_READ", "1") - assert _resolve_fp8_weight_read_gates() == (True, True, True) + g = _gates(QuantConfig(), enable_attention_dp=False) + # gate_up follows attention_dp=False -> off; the rest default on. + assert _tuple(g) == (True, True, True, True, False) @pytest.mark.parametrize( "env,expected", [ - ("KIMI_K3_FP8_WEIGHT_READ_KDA", (True, False, False)), - ("KIMI_K3_KDA_GLUE_FP8", (True, True, False)), + ("KIMI_K3_FP8_WEIGHT_READ_KDA", (True, False, False, True, True)), + ("KIMI_K3_KDA_GLUE_FP8", (True, True, False, True, True)), + ("KIMI_K3_FP8_WEIGHT_READ_MLA", (True, True, True, False, True)), + ("KIMI_K3_FP8_WEIGHT_READ_GATE_UP", (True, True, True, True, False)), ], ) -def test_sub_gates_narrow_an_enabled_master(sm100f, monkeypatch, env, expected): - """Sub-gates only ever narrow; they never enable on their own.""" +def test_sub_gates_narrow_via_deprecated_env(sm100f, monkeypatch, env, expected): monkeypatch.setenv("KIMI_K3_FP8_WEIGHT_READ", "1") monkeypatch.setenv(env, "0") - assert _resolve_fp8_weight_read_gates() == expected + assert _tuple(_gates(QuantConfig(), enable_attention_dp=True)) == expected -@pytest.mark.parametrize("env", ENVS[1:]) -def test_sub_gates_are_inert_while_master_is_off(sm100f, monkeypatch, env): - """A sub-gate set to 1 must not turn on FP8 reads by itself.""" - monkeypatch.setenv(env, "1") - assert _resolve_fp8_weight_read_gates() == (False, False, False) +def test_non_blackwell_never_reads_fp8_env(monkeypatch): + """The SM gate wins even with the master env var on.""" + monkeypatch.setattr(kimi_k3_knobs, "is_sm_100f", lambda: False) + for env in ENVS: + monkeypatch.delenv(env, raising=False) + monkeypatch.setenv("KIMI_K3_FP8_WEIGHT_READ", "1") + assert _tuple(_gates(QuantConfig(), enable_attention_dp=True)) == ( + False, + False, + False, + False, + False, + ) + + +# --- Precedence: config wins over the deprecated env var ------------------- + + +def test_config_master_off_beats_env_master_on(sm100f, monkeypatch): + """Explicit config value takes precedence over the deprecated env var.""" + monkeypatch.setenv("KIMI_K3_FP8_WEIGHT_READ", "1") + g = _gates(QuantConfig(kimi_k3_fp8_weight_read=False), enable_attention_dp=True) + assert _tuple(g) == (False, False, False, False, False) + + +def test_config_subgate_off_beats_env_subgate_on(sm100f, monkeypatch): + """A config sub-gate overrides a conflicting deprecated env sub-gate.""" + monkeypatch.setenv("KIMI_K3_FP8_WEIGHT_READ", "1") + monkeypatch.setenv("KIMI_K3_FP8_WEIGHT_READ_KDA", "1") + g = _gates( + QuantConfig(kimi_k3_fp8_weight_read=True, kimi_k3_fp8_weight_read_kda=False), + enable_attention_dp=True, + ) + assert g.master is True and g.kda is False and g.kda_glue is False + + +# --- Config surface: every former env var has a config equivalent ----------- + + +def test_fp8_knobs_exist_on_quant_config(): + """Each of the 5 former KIMI_K3_* env vars has a QuantConfig field.""" + fields = QuantConfig.model_fields + for name in kimi_k3_knobs.KIMI_K3_QUANT_KNOB_FIELDS: + assert name in fields, f"QuantConfig missing {name}" + assert fields[name].default is None, f"{name} must default to None (unset)" + assert len(kimi_k3_knobs.KIMI_K3_QUANT_KNOB_FIELDS) == len(ENVS) + + +def test_knobs_settable_from_extra_llm_api_options_dict(): + """The knobs round-trip through a plain dict, as extra_llm_api_options does.""" + qc = QuantConfig.model_validate( + {"kimi_k3_fp8_weight_read": True, "kimi_k3_fp8_weight_read_kda": False} + ) + assert qc.kimi_k3_fp8_weight_read is True + assert qc.kimi_k3_fp8_weight_read_kda is False + + +def test_zero_env_zero_config_resolves_complete_valid(sm100f): + """No env vars + a default QuantConfig -> every gate resolves to a concrete + value (the historical BF16 default), with no missing-knob error.""" + g = _gates(QuantConfig(), enable_attention_dp=False) + assert _tuple(g) == (False, False, False, False, False) + assert all(isinstance(v, bool) for v in _tuple(g)) + + +# --- Deprecation warning: emitted once, and even when config overrides ------ + + +class _RecordingLogger: + """Mirrors ``logger.warning_once`` dedup so warn-once is deterministic.""" + + def __init__(self): + self._seen = set() + self.emitted = [] + + def warning_once(self, *msg, key): + if key not in self._seen: + self._seen.add(key) + self.emitted.append((key, " ".join(str(m) for m in msg))) + + +def test_deprecated_env_warns_once_and_is_honored(sm100f, monkeypatch): + rec = _RecordingLogger() + monkeypatch.setattr(kimi_k3_knobs, "logger", rec) + monkeypatch.setenv("KIMI_K3_FP8_WEIGHT_READ", "1") + # Resolve several times; the deprecation warning must appear exactly once. + for _ in range(3): + assert _gates(QuantConfig(), enable_attention_dp=True).master is True + keys = [k for k, _ in rec.emitted] + assert keys == ["kimi_k3_deprecated_env::KIMI_K3_FP8_WEIGHT_READ"] + assert "extra_llm_api_options" in rec.emitted[0][1] -def test_non_blackwell_never_reads_fp8(monkeypatch): - """The SM gate wins even with the master switch explicitly on.""" - monkeypatch.setattr(modeling_kimi_linear, "is_sm_100f", lambda: False) +def test_warn_once_even_when_config_overrides(sm100f, monkeypatch): + """The deprecation is about the env var existing; it warns even when a + config value overrides the env value.""" + rec = _RecordingLogger() + monkeypatch.setattr(kimi_k3_knobs, "logger", rec) monkeypatch.setenv("KIMI_K3_FP8_WEIGHT_READ", "1") - assert _resolve_fp8_weight_read_gates() == (False, False, False) + assert _gates(QuantConfig(kimi_k3_fp8_weight_read=False), enable_attention_dp=True).master is ( + False + ) + assert [k for k, _ in rec.emitted] == ["kimi_k3_deprecated_env::KIMI_K3_FP8_WEIGHT_READ"] diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py b/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py index ccf1973bb7f2..7043b4006797 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py @@ -13,18 +13,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json import mmap import os import threading -from pathlib import Path from unittest import mock import pytest from tensorrt_llm._torch.models.checkpoints import HfWeightLoader from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict -from tensorrt_llm._torch.models.checkpoints.hf.weight_loader import _LAZY_SAFETENSORS_MODEL_TYPES from tensorrt_llm.mapping import Mapping pytestmark = pytest.mark.cpu_only @@ -496,51 +493,140 @@ def test_prefetch_files_emits_progress_heartbeat(tmp_path, monkeypatch): assert len(progress_logs) >= 12 -def test_kimi_k3_lazy_load_records_the_checkpoint_dir(tmp_path): +def test_lazy_load_records_the_checkpoint_dir(tmp_path): """A model that re-opens shards itself needs the directory back. - Kimi K3 streams its rank-local experts per shard file to avoid holding - the whole mapping open. A lazy slice does not carry its file, and - transformers 5.x no longer sets ``PretrainedConfig._name_or_path``, so - without this the model silently fell back to the shared mapping and the - step was OOM-killed. + Expert-parallel models stream their rank-local experts per shard file to + avoid holding the whole mapping open. A lazy slice does not carry its + file, and transformers 5.x no longer sets ``PretrainedConfig._name_or_path``, + so without this the model silently fell back to the shared mapping and the + step was OOM-killed. The lazy ``LoadFormat`` is selected here via the + ``load_lazily`` flag (what ``LoadFormat.LAZY_SAFETENSORS`` resolves to in + the model loader) -- no checkpoint model-name detection is involved. """ - import safetensors.torch import torch - (tmp_path / "config.json").write_text(json.dumps({"model_type": "kimi_k3"})) safetensors.torch.save_file( {"w": torch.zeros(2, 2)}, tmp_path / "model-00001-of-00001.safetensors" ) loader = HfWeightLoader() + # ``weights`` is bound before the try so a load_weights() failure surfaces + # as itself instead of an UnboundLocalError raised from the cleanup path. + weights = None try: - weights = loader.load_weights(str(tmp_path), Mapping()) + weights = loader.load_weights(str(tmp_path), Mapping(), load_lazily=True) assert isinstance(weights, ConsumableWeightsDict) assert weights.checkpoint_dir == str(tmp_path) finally: + if weights is not None: + weights.clear() loader.cleanup() -@pytest.mark.parametrize("model_type", _LAZY_SAFETENSORS_MODEL_TYPES) -def test_requires_lazy_safetensors_for_every_listed_model_type( - tmp_path: Path, model_type: str -) -> None: - """Each model type in the table routes to the lazy (mmapped) load path.""" - (tmp_path / "config.json").write_text(json.dumps({"model_type": model_type})) - assert HfWeightLoader._requires_lazy_safetensors(str(tmp_path)) is True +def _materialize_lazy_slice(value): + """Realize a lazy safetensors slice the way a model's load_weights does. + + Mirrors ``modeling_kimi_linear._materialize``: ``[:]`` realizes a normal + slice, but a 0-dim (scalar) entry needs ``[()]``. + """ + import torch + + if isinstance(value, torch.Tensor): + return value + get_shape = getattr(value, "get_shape", None) + if get_shape is not None and len(get_shape()) == 0: + return value[()] + return value[:] + +def _write_multi_shard_fixture(checkpoint_dir): + """Write a small multi-shard safetensors fixture (no big checkpoint). -@pytest.mark.parametrize("config", [{"model_type": "llama"}, {}]) -def test_requires_lazy_safetensors_is_false_for_other_checkpoints( - tmp_path: Path, config: dict[str, str] -) -> None: - """Any other model type, or no model type at all, takes the eager path.""" - (tmp_path / "config.json").write_text(json.dumps(config)) - assert HfWeightLoader._requires_lazy_safetensors(str(tmp_path)) is False + Two sharded files carry distinct keys; a consolidated file carries a + different key so which file set was selected is directly observable from + the loaded key set. Includes a quantized (fp8) weight and a companion + fp32 block-scale tensor so dtype and scale tensors are exercised. + Returns (sharded_keys, consolidated_keys). + """ + import safetensors.torch + import torch + + sharded_a = { + "layer.0.weight": torch.arange(16, dtype=torch.bfloat16).reshape(4, 4), + # FP8 block-scale companion (kept BF16-free on purpose: a real scale). + "layer.0.weight_scale_inv": torch.rand(2, 2, dtype=torch.float32), + } + sharded_b = { + # A genuinely quantized tensor: proves the fp8 dtype survives the + # lazy path byte-for-byte. + "layer.1.weight": torch.arange(16).reshape(4, 4).to(torch.float8_e4m3fn), + } + consolidated = { + "consolidated.only.weight": torch.ones(2, 2, dtype=torch.bfloat16), + } + safetensors.torch.save_file(sharded_a, str(checkpoint_dir / "model-00001-of-00002.safetensors")) + safetensors.torch.save_file(sharded_b, str(checkpoint_dir / "model-00002-of-00002.safetensors")) + safetensors.torch.save_file(consolidated, str(checkpoint_dir / "consolidated.safetensors")) + return (set(sharded_a) | set(sharded_b)), set(consolidated) + + +@pytest.mark.parametrize("use_consolidated", [False, True]) +def test_lazy_format_load_matches_eager_load(tmp_path, use_consolidated): + """The lazy ``LoadFormat`` equals an eager load on the same fixture. + + Covers criterion s1.g1: a general lazy format whose sharded-vs-consolidated + selection matches the eager path, whose materialized values/dtype/scale + tensors match the eager load exactly, and whose rank-local sliced reads do + not materialize the full tensor set. Runs on CPU, no 1.5 TB checkpoint. + """ + import torch -def test_requires_lazy_safetensors_is_false_without_a_config(tmp_path: Path) -> None: - """A directory without config.json cannot opt in to lazy loading.""" - assert HfWeightLoader._requires_lazy_safetensors(str(tmp_path)) is False + checkpoint_dir = tmp_path / "ckpt" + checkpoint_dir.mkdir() + sharded_keys, consolidated_keys = _write_multi_shard_fixture(checkpoint_dir) + expected_keys = consolidated_keys if use_consolidated else sharded_keys + + loader = HfWeightLoader() + + # Eager reference load. Mock only the OS-cache prefetch; the real file + # selection and the real safetensors load both run. + with mock.patch.object(loader, "prefetch_files"): + eager = loader.load_weights( + str(checkpoint_dir), Mapping(), use_consolidated=use_consolidated + ) + eager_values = {key: eager[key] for key in eager.keys()} + + # Lazy-format load of the same fixture. + lazy = loader.load_weights( + str(checkpoint_dir), Mapping(), use_consolidated=use_consolidated, load_lazily=True + ) + try: + # (1) use_consolidated selection matches the eager path: same key set. + assert set(eager.keys()) == expected_keys + assert set(lazy.keys()) == expected_keys + + # (2) Lazy values are slice handles, not materialized tensors: merely + # opening the checkpoint never realizes the full tensor set. + for key in lazy.keys(): + assert not isinstance(lazy[key], torch.Tensor), key + + # (3) A rank-local sliced read materializes ONLY the requested bytes. + row_key = "consolidated.only.weight" if use_consolidated else "layer.0.weight" + partial = lazy[row_key][0:1] + assert isinstance(partial, torch.Tensor) + assert partial.shape[0] == 1 # only the requested row, not all rows + + # (4) Materialized lazy values equal the eager values in dtype, shape, + # and bytes -- scale and fp8 tensors included. + for key in expected_keys: + lazy_val = _materialize_lazy_slice(lazy[key]) + eager_val = eager_values[key] + assert lazy_val.dtype == eager_val.dtype, key + assert lazy_val.shape == eager_val.shape, key + assert torch.equal(lazy_val.to(torch.float32), eager_val.to(torch.float32)), key + finally: + lazy.clear() + loader.cleanup() diff --git a/tests/unittest/api_stability/references/quant_config.yaml b/tests/unittest/api_stability/references/quant_config.yaml index 9aa0f6b5366f..36c1bf150734 100644 --- a/tests/unittest/api_stability/references/quant_config.yaml +++ b/tests/unittest/api_stability/references/quant_config.yaml @@ -13,6 +13,21 @@ methods: has_zero_point: annotation: bool default: false + kimi_k3_fp8_weight_read: + annotation: Optional[bool] + default: null + kimi_k3_fp8_weight_read_gate_up: + annotation: Optional[bool] + default: null + kimi_k3_fp8_weight_read_kda: + annotation: Optional[bool] + default: null + kimi_k3_fp8_weight_read_mla: + annotation: Optional[bool] + default: null + kimi_k3_kda_glue_fp8: + annotation: Optional[bool] + default: null kv_cache_quant_algo: annotation: Optional[tensorrt_llm.quantization.mode.QuantAlgo] default: null From 2a7210cfb2c5ec8044bf07bcb0ebdebe8f380fc9 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Mon, 14 Sep 2026 07:24:13 -0700 Subject: [PATCH 2/2] [TRTLLM-15159][fix] Register lazy_safetensors in telemetry and the standard-HF-loader gate Two places still enumerate the LoadFormat members by hand, so the new LAZY_SAFETENSORS member was invisible to one and rejected by the other. - TorchLlmArgs.load_format is annotated Union[str, LoadFormat]. The bare str arm means telemetry capture cannot derive the domain from the enum and requires an explicit allowlist, so adding the member alone left every lazy_safetensors run reporting no load_format at all: _sanitize_allowlist finds no matching candidate, drops the value and flags it unsafe_excluded. Add the member to the allowlist and refresh the two artifacts derived from it, the golden manifest row and the rendered telemetry reference table. - ModelLoader.uses_standard_hf_loader tested == LoadFormat.AUTO, so a model that has _requires_standard_hf_loading modules (today only Qwen4ExpPinnedHostEmbedding) and loads lazily would have been rejected with the "requires the standard Hugging Face AUTO loader" error. LAZY_SAFETENSORS is the AUTO disk path with lazily opened shards and goes through the same HF mapper, which is the property the gate is guarding. No model combines the two today, so this is latent. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- docs/source/developer-guide/telemetry.md | 2 +- tensorrt_llm/_torch/pyexecutor/model_loader.py | 12 +++++++----- tensorrt_llm/llmapi/llm_args.py | 2 +- tensorrt_llm/usage/llm_args_golden_manifest.json | 3 ++- 4 files changed, 11 insertions(+), 8 deletions(-) diff --git a/docs/source/developer-guide/telemetry.md b/docs/source/developer-guide/telemetry.md index 9a67bfa2c5dd..ad822e040273 100644 --- a/docs/source/developer-guide/telemetry.md +++ b/docs/source/developer-guide/telemetry.md @@ -153,7 +153,7 @@ unset or when the safety sanitizer rejects the runtime value. | `kv_connector_config.connector` | `Optional[str]` | `categorical` | allowlist | `lmcache`, `lmcache-mp`, `kvbm` | | `layer_wise_benchmarks_config.calibration_layer_indices` | `Optional[List[int]]` | `value` | | | | `layer_wise_benchmarks_config.calibration_mode` | `Literal['NONE', 'MARK', 'COLLECT']` | `categorical` | | `NONE`, `MARK`, `COLLECT` | -| `load_format` | `Union[str, tensorrt_llm.llmapi.llm_args.LoadFormat]` | `categorical` | allowlist | `auto`, `dummy`, `vision_only`, `gms` | +| `load_format` | `Union[str, tensorrt_llm.llmapi.llm_args.LoadFormat]` | `categorical` | allowlist | `auto`, `dummy`, `vision_only`, `gms`, `lazy_safetensors` | | `lora_config.lora_ckpt_source` | `Literal['hf', 'nemo']` | `categorical` | | `hf`, `nemo` | | `lora_config.max_cpu_loras` | `Optional[int]` | `value` | | | | `lora_config.max_lora_rank` | `` | `value` | | | diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 0dac51684c91..1405c8f6aaca 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -909,11 +909,13 @@ def load( getattr(module, "_requires_standard_hf_loading", False) for module in model.modules()) # Pinned-host parameters must be materialized and filled in place. - # AUTO uses the HF mapper that preserves their stable addresses; - # AUTO may still resolve to MX, so check the resolved format too. - uses_standard_hf_loader = (load_format == LoadFormat.AUTO - and checkpoint_loader.checkpoint_format - != "MX") + # AUTO and LAZY_SAFETENSORS both go through the HF mapper that + # preserves their stable addresses -- the lazy format only changes + # how shards are opened. AUTO may still resolve to MX, so check the + # resolved format too. + uses_standard_hf_loader = ( + load_format in (LoadFormat.AUTO, LoadFormat.LAZY_SAFETENSORS) + and checkpoint_loader.checkpoint_format != "MX") if requires_standard_hf_loading and not uses_standard_hf_loader: raise ValueError( "Host-resident model weights currently require the standard " diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 04b714b8f4ab..84a03b5b8b11 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5804,7 +5804,7 @@ def validate_mla_skip_correction_config(self) -> 'TorchLlmArgs': description= "How to load the model weights. By default, detect the weight type from the model checkpoint.", telemetry=TelemetryField.categorical("auto", "dummy", "vision_only", - "gms")) + "gms", "lazy_safetensors")) enable_min_latency: bool = Field( default=False, diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index aef78b3a38ce..f3738d8e2c78 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1030,7 +1030,8 @@ "auto", "dummy", "vision_only", - "gms" + "gms", + "lazy_safetensors" ], "annotation": "Union[str, tensorrt_llm.llmapi.llm_args.LoadFormat]", "converter": "allowlist",