From f4e692d2191fa6385647f4c89a12454b47af7886 Mon Sep 17 00:00:00 2001 From: TensorRT LLM AI Agent <296075020+trtllm-agent@users.noreply.github.com> Date: Fri, 24 Jul 2026 00:33:48 +0800 Subject: [PATCH 1/4] [None][infra] Waive 4 failed cases for main in pre-merge 49550 (#16798) Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 80e697faf3be..79384bbde5ce 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -366,6 +366,10 @@ unittest/auto_deploy/multigpu/custom_ops SKIP (https://nvbugs/6403920) unittest/disaggregated/test_kv_transfer.py SKIP (https://nvbugs/6403793) unittest/disaggregated/test_kv_transfer.py::test_transfer_worker_v2[tp1_pp4_to_tp2_pp2] SKIP (https://nvbugs/6445316) unittest/disaggregated/test_kv_transfer.py::test_transfer_worker_v2[tp4_pp1_to_tp2_pp2] SKIP (https://nvbugs/6426834) +unittest/executor/test_proxy_fast_death.py SKIP (https://nvbugs/6499882) +unittest/executor/test_proxy_fast_death.py::test_shutdown_does_not_block_on_dead_engine SKIP (https://nvbugs/6499882) +unittest/executor/test_proxy_fast_death.py::test_shutdown_does_not_shut_down_external_session SKIP (https://nvbugs/6499882) +unittest/executor/test_proxy_fast_death.py::test_shutdown_keeps_blocking_semantics_when_engine_alive SKIP (https://nvbugs/6499882) unittest/executor/test_rpc.py::TestRpcCorrectness::test_incremental_task_async SKIP (https://nvbugs/5741476) unittest/executor/test_rpc_proxy.py SKIP (https://nvbugs/5605741) unittest/executor/test_rpc_worker.py SKIP (https://nvbugs/5605741) From 5a69240c9ed7e7513a2b7324993f0115cc3989fe Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Thu, 23 Jul 2026 10:42:45 -0700 Subject: [PATCH 2/4] [None][feat] Bind SourceIdentity to checkpoint artifacts (#16159) Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../_torch/pyexecutor/model_loader.py | 47 +++- .../_torch/weight_sharing/__init__.py | 6 + .../weight_sharing/artifact_identity.py | 221 ++++++++++++++++++ .../_torch/weight_sharing/source_identity.py | 60 +++-- .../_torch/executor/test_model_loader_gms.py | 52 ++++- .../_torch/executor/test_model_loader_mx.py | 74 +++++- .../mx/test_mx_checkpoint_loader.py | 23 +- .../weight_sharing/_source_identity_fakes.py | 36 ++- .../weight_sharing/test_artifact_identity.py | 159 +++++++++++++ .../test_gms_source_identity_gate.py | 8 + .../test_mx_source_identity_gate.py | 7 + .../weight_sharing/test_source_identity.py | 73 +++++- 12 files changed, 725 insertions(+), 41 deletions(-) create mode 100644 tensorrt_llm/_torch/weight_sharing/artifact_identity.py create mode 100644 tests/unittest/_torch/weight_sharing/test_artifact_identity.py diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 688debc80692..9f3218f2e192 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -13,10 +13,10 @@ from tensorrt_llm._torch.models.checkpoints.base_checkpoint_loader import ( AutoCheckpointMapper, BaseCheckpointLoader) from tensorrt_llm._torch.weight_sharing import ( - IdentityCheckPolicy, PostTransformFeature, PostTransformProfile, - PostTransformProfileRegistry, PostTransformQualificationDecision, - PostTransformTransferScope, SourceIdentity, - check_weight_sharing_compatibility) + ArtifactIdentity, IdentityCheckPolicy, PostTransformFeature, + PostTransformProfile, PostTransformProfileRegistry, + PostTransformQualificationDecision, PostTransformTransferScope, + SourceIdentity, check_weight_sharing_compatibility) from tensorrt_llm._utils import str_dtype_to_torch from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig, ExecutorMemoryType, @@ -432,6 +432,39 @@ def _needs_source_identity(checkpoint_loader: BaseCheckpointLoader, """ return load_format == LoadFormat.GMS or checkpoint_loader.checkpoint_format == "MX" + @staticmethod + def _build_source_identity( + config: ModelConfig, + model: DecoderModelForCausalLM, + *, + checkpoint_dir: str, + model_name: str, + fallback_on_artifact_error: bool, + ) -> Optional[SourceIdentity]: + """Build the local identity without weakening artifact validation. + + Artifact construction remains fail-closed. MX may convert an artifact + error into an unavailable local identity so its compatibility gate + falls back to disk; GMS propagates the error because it has no fallback. + """ + try: + artifact_identity = ArtifactIdentity.from_checkpoint(checkpoint_dir) + except (OSError, RuntimeError, ValueError) as error: + if not fallback_on_artifact_error: + raise + logger.warning( + "Unable to build checkpoint artifact identity for MX checkpoint " + f"{checkpoint_dir}; falling back to regular checkpoint loading: {error}" + ) + return None + + return SourceIdentity.from_model_config( + config, + model, + artifact_identity=artifact_identity, + model_name=model_name, + ) + def load( self, checkpoint_dir: str, @@ -475,12 +508,16 @@ def load( # ground truth; building it here (post-construction, # pre-weight-load) gives producer and consumer a common, # comparable lifecycle point. - self._source_identity = SourceIdentity.from_model_config( + self._source_identity = self._build_source_identity( config, model, + checkpoint_dir=checkpoint_dir, model_name=str( getattr(self.llm_args, "model", None) or checkpoint_dir), + fallback_on_artifact_error=( + load_format != LoadFormat.GMS + and checkpoint_loader.checkpoint_format == "MX"), ) memo: dict[torch.Tensor, torch.Tensor] = {} diff --git a/tensorrt_llm/_torch/weight_sharing/__init__.py b/tensorrt_llm/_torch/weight_sharing/__init__.py index 80b70c6c515d..38ec2471c41c 100644 --- a/tensorrt_llm/_torch/weight_sharing/__init__.py +++ b/tensorrt_llm/_torch/weight_sharing/__init__.py @@ -14,6 +14,10 @@ # limitations under the License. """Backend-agnostic weight-sharing utilities (MX, GMS, ...).""" +from tensorrt_llm._torch.weight_sharing.artifact_identity import ( + ARTIFACT_IDENTITY_FORMAT_VERSION, + ArtifactIdentity, +) from tensorrt_llm._torch.weight_sharing.post_transform_profiles import ( PostTransformFeature, PostTransformProfile, @@ -33,6 +37,8 @@ ) __all__ = [ + "ARTIFACT_IDENTITY_FORMAT_VERSION", + "ArtifactIdentity", "SOURCE_IDENTITY_FORMAT_VERSION", "PostTransformFeature", "PostTransformProfile", diff --git a/tensorrt_llm/_torch/weight_sharing/artifact_identity.py b/tensorrt_llm/_torch/weight_sharing/artifact_identity.py new file mode 100644 index 000000000000..d1b8908e241a --- /dev/null +++ b/tensorrt_llm/_torch/weight_sharing/artifact_identity.py @@ -0,0 +1,221 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""Immutable checkpoint identity for shared-weight compatibility checks.""" + +from __future__ import annotations + +import hashlib +import json +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +ARTIFACT_IDENTITY_FORMAT_VERSION = 1 + +_HF_SNAPSHOT_SCHEME = "hf_snapshot_revision" +_CHECKPOINT_MANIFEST_SCHEME = "checkpoint_manifest_sha256" +_SUPPORTED_SCHEMES = frozenset({_HF_SNAPSHOT_SCHEME, _CHECKPOINT_MANIFEST_SCHEME}) +_IGNORED_DIRECTORY_NAMES = frozenset({".cache", ".git", "__pycache__"}) +_IGNORED_FILE_NAMES = frozenset({".DS_Store"}) +_HASH_CHUNK_SIZE = 1024 * 1024 + + +def _canonical_hash(value: Any) -> str: + payload = json.dumps(value, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def _is_hex(value: str, lengths: tuple[int, ...]) -> bool: + return len(value) in lengths and all(char in "0123456789abcdef" for char in value) + + +def _hf_snapshot_descriptor(path: Path) -> tuple[str, str] | None: + """Return an immutable HF revision and repository-relative subpath.""" + parts = path.resolve().parts + for index, part in enumerate(parts[:-1]): + if part != "snapshots" or index == 0: + continue + if not parts[index - 1].startswith("models--"): + continue + + revision = parts[index + 1].lower() + if not _is_hex(revision, (40, 64)): + continue + subpath = "/".join(parts[index + 2 :]) + return revision, subpath + return None + + +def _raise_walk_error(error: OSError) -> None: + raise error + + +def _checkpoint_files(path: Path) -> tuple[Path, list[Path]]: + if path.is_file(): + return path.parent, [path] + + files = [] + for directory, directory_names, file_names in os.walk(path, onerror=_raise_walk_error): + retained_directories = [] + for directory_name in directory_names: + if directory_name in _IGNORED_DIRECTORY_NAMES: + continue + nested_directory = Path(directory) / directory_name + if nested_directory.is_symlink(): + raise ValueError( + "Checkpoint manifests do not support nested symlinked directories: " + f"{nested_directory}" + ) + retained_directories.append(directory_name) + directory_names[:] = retained_directories + for file_name in file_names: + if file_name in _IGNORED_FILE_NAMES: + continue + candidate = Path(directory) / file_name + if candidate.is_file(): + files.append(candidate) + files.sort(key=lambda candidate: candidate.relative_to(path).as_posix()) + if not files: + raise ValueError(f"Checkpoint path contains no files: {path}") + return path, files + + +def _sha256_file(path: Path) -> tuple[int, str]: + before = path.stat() + digest = hashlib.sha256() + with path.open("rb") as checkpoint_file: + for chunk in iter(lambda: checkpoint_file.read(_HASH_CHUNK_SIZE), b""): + digest.update(chunk) + after = path.stat() + if (before.st_size, before.st_mtime_ns) != (after.st_size, after.st_mtime_ns): + raise RuntimeError(f"Checkpoint file changed while being fingerprinted: {path}") + return after.st_size, digest.hexdigest() + + +def _checkpoint_manifest_digest(path: Path) -> str: + root, files = _checkpoint_files(path) + manifest = [] + for checkpoint_file in files: + size, digest = _sha256_file(checkpoint_file) + manifest.append( + { + "path": checkpoint_file.relative_to(root).as_posix(), + "size": size, + "sha256": digest, + } + ) + return _canonical_hash( + { + "format_version": ARTIFACT_IDENTITY_FORMAT_VERSION, + "files": manifest, + } + ) + + +@dataclass(frozen=True) +class ArtifactIdentity: + """Versioned identity of the immutable checkpoint artifact being loaded. + + `SourceIdentity` embeds this value as a global compatibility component. + Hugging Face cache snapshots use their immutable commit revision; local + checkpoints use a canonical manifest of relative paths, sizes, and file + content digests. Absolute paths are intentionally excluded. + """ + + format_version: int + scheme: str + digest: str + + def __post_init__(self) -> None: + if not isinstance(self.format_version, int) or isinstance(self.format_version, bool): + raise ValueError("ArtifactIdentity format version must be an integer") + if self.format_version != ARTIFACT_IDENTITY_FORMAT_VERSION: + raise ValueError(f"Unsupported ArtifactIdentity format version: {self.format_version}") + if not isinstance(self.scheme, str): + raise ValueError("ArtifactIdentity scheme must be a string") + if self.scheme not in _SUPPORTED_SCHEMES: + raise ValueError(f"Unsupported ArtifactIdentity scheme: {self.scheme}") + if not isinstance(self.digest, str): + raise ValueError("ArtifactIdentity digest must be a string") + + normalized_digest = self.digest.lower() + if not _is_hex(normalized_digest, (64,)): + raise ValueError("ArtifactIdentity digest must be a 64-character hex value") + object.__setattr__(self, "digest", normalized_digest) + + @classmethod + def from_checkpoint(cls, checkpoint_path: str | os.PathLike[str]) -> "ArtifactIdentity": + """Build an identity from an immutable snapshot or local checkpoint. + + Args: + checkpoint_path: A model checkpoint file or directory. + + Returns: + The path-independent checkpoint identity. + + Raises: + FileNotFoundError: If `checkpoint_path` does not exist. + ValueError: If a local checkpoint directory contains no files. + RuntimeError: If a local checkpoint changes while it is hashed. + + Note: + Local checkpoints have no authoritative immutable revision, so + their regular files are read in full to derive a content-bound + manifest. Hugging Face cache snapshots use the resolved immutable + revision without rereading model shards. + """ + path = Path(checkpoint_path).expanduser() + if not path.exists(): + raise FileNotFoundError(f"Checkpoint path does not exist: {path}") + + snapshot_descriptor = _hf_snapshot_descriptor(path) + if snapshot_descriptor is not None: + revision, subpath = snapshot_descriptor + digest = _canonical_hash( + { + "scheme": _HF_SNAPSHOT_SCHEME, + "revision": revision, + "subpath": subpath, + } + ) + return cls( + format_version=ARTIFACT_IDENTITY_FORMAT_VERSION, + scheme=_HF_SNAPSHOT_SCHEME, + digest=digest, + ) + + return cls( + format_version=ARTIFACT_IDENTITY_FORMAT_VERSION, + scheme=_CHECKPOINT_MANIFEST_SCHEME, + digest=_checkpoint_manifest_digest(path), + ) + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-serializable representation.""" + return { + "format_version": self.format_version, + "scheme": self.scheme, + "digest": self.digest, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "ArtifactIdentity": + """Reconstruct and validate a serialized artifact identity.""" + return cls( + format_version=data["format_version"], + scheme=data["scheme"], + digest=data["digest"], + ) diff --git a/tensorrt_llm/_torch/weight_sharing/source_identity.py b/tensorrt_llm/_torch/weight_sharing/source_identity.py index 18168703c868..c6b010634760 100644 --- a/tensorrt_llm/_torch/weight_sharing/source_identity.py +++ b/tensorrt_llm/_torch/weight_sharing/source_identity.py @@ -14,17 +14,19 @@ # limitations under the License. """Backend-agnostic source identity for weight-sharing receivers. -A :class:`SourceIdentity` is a serializable fingerprint of configuration choices -that affect how a model's weights are laid out in memory. It exists so that a -*receiver* of pre-laid-out weights (e.g. MX peer-to-peer transfer, or a GMS -read-only materialize) can verify that both the producer ("source") and the -consumer built identities and agree on every layout-affecting choice before the -receiver consumes shared weights. +A :class:`SourceIdentity` is a serializable fingerprint of an immutable +checkpoint artifact and the configuration choices that affect how its weights +are laid out in memory. It exists so that a *receiver* of pre-laid-out weights +(e.g. MX peer-to-peer transfer, or a GMS read-only materialize) can verify that +both the producer ("source") and the consumer built identities and agree on the +artifact and every layout-affecting choice before consuming shared weights. The identity is intentionally decoupled from any specific weight-sharing technology (neither MX nor GMS appears here). Both consume it identically: - local = SourceIdentity.from_model_config(model_config) + local = SourceIdentity.from_model_config( + model_config, checkpoint_dir="/path/to/checkpoint" + ) decision = check_weight_sharing_compatibility(local, source_identity, policy) if decision.should_share: ... # pull / materialize shared weights @@ -40,8 +42,9 @@ -------------- The fingerprint is split so comparison can be selective: -* **global fingerprint** -- rank-invariant model identity, quantization, - backend selection, fusion flags, and parallel *sizes* (TP/PP/EP/CP). +* **global fingerprint** -- immutable checkpoint artifact, rank-invariant + model identity, quantization, backend selection, fusion flags, and parallel + *sizes* (TP/PP/EP/CP). * **shard fingerprint** -- this rank's TP/PP/EP/CP *rank* slice plus the realized local parameter/buffer `(shape, dtype)` layout. Receiver rank `N` must align with the source rank that produced shard `N`. @@ -67,6 +70,7 @@ from enum import Enum from typing import TYPE_CHECKING, Any, List, Optional +from tensorrt_llm._torch.weight_sharing.artifact_identity import ArtifactIdentity from tensorrt_llm.logger import logger if TYPE_CHECKING: @@ -78,7 +82,7 @@ # Bump when the fingerprint projection changes in a way that makes previously # stored identities incomparable. Two identities with different format versions # never match. -SOURCE_IDENTITY_FORMAT_VERSION = 1 +SOURCE_IDENTITY_FORMAT_VERSION = 2 _PRETRAINED_METADATA_FIELDS = frozenset( { @@ -224,6 +228,7 @@ class SourceIdentity: format_version: int # --- global parts (must match across all ranks) --- + artifact_identity: ArtifactIdentity model_fingerprint: str quant_fingerprint: str backend_fingerprint: str @@ -246,6 +251,8 @@ def from_model_config( model_config: "ModelConfig", model: Optional["nn.Module"] = None, *, + checkpoint_dir: Optional[str] = None, + artifact_identity: Optional[ArtifactIdentity] = None, model_name: Optional[str] = None, ) -> "SourceIdentity": """Build an identity from a torch-backend :class:`ModelConfig`. @@ -258,6 +265,12 @@ def from_model_config( Producer and consumer must build the identity at the same lifecycle point (model construction, before weight load). When `None`, the shard fingerprint contains no tensor-layout data. + checkpoint_dir: Checkpoint file or directory used to derive the + nested artifact identity. Required unless `artifact_identity` + is supplied explicitly. + artifact_identity: Precomputed immutable checkpoint identity, + primarily for callers that resolve provenance outside this + method. Mutually exclusive with `checkpoint_dir`. model_name: Human-readable model identity used by discovery layers (e.g. the MX server's source catalog). Does not affect the compatibility fingerprints. @@ -265,7 +278,17 @@ def from_model_config( Returns: A fully populated :class:`SourceIdentity` for `model_config.mapping.rank`. + + Raises: + ValueError: If neither or both artifact identity inputs are given. """ + if checkpoint_dir is None and artifact_identity is None: + raise ValueError("Exactly one of checkpoint_dir or artifact_identity must be provided") + if checkpoint_dir is not None and artifact_identity is not None: + raise ValueError("Exactly one of checkpoint_dir or artifact_identity must be provided") + if artifact_identity is None: + artifact_identity = ArtifactIdentity.from_checkpoint(checkpoint_dir) + mapping = model_config.mapping rank = getattr(mapping, "rank", 0) @@ -275,6 +298,7 @@ def from_model_config( return cls( format_version=SOURCE_IDENTITY_FORMAT_VERSION, + artifact_identity=artifact_identity, model_fingerprint=cls._build_model_fingerprint(model_config), quant_fingerprint=cls._build_quant_fingerprint(model_config), backend_fingerprint=cls._build_backend_fingerprint(model_config), @@ -418,6 +442,7 @@ def global_fingerprint(self) -> str: return _canonical_hash( { "format_version": self.format_version, + "artifact": self.artifact_identity.to_dict(), "model": self.model_fingerprint, "quant": self.quant_fingerprint, "backend": self.backend_fingerprint, @@ -447,6 +472,8 @@ def matches( mismatched.append("format_version") if compare_global: + if self.artifact_identity != other.artifact_identity: + mismatched.append("artifact_identity") for name in ( "model_fingerprint", "quant_fingerprint", @@ -473,6 +500,7 @@ def to_dict(self) -> dict: """ return { "format_version": self.format_version, + "artifact_identity": self.artifact_identity.to_dict(), "model_fingerprint": self.model_fingerprint, "quant_fingerprint": self.quant_fingerprint, "backend_fingerprint": self.backend_fingerprint, @@ -496,8 +524,13 @@ def from_dict(cls, data: dict) -> "SourceIdentity": Returns: The reconstructed :class:`SourceIdentity`. """ + format_version = data["format_version"] + if format_version != SOURCE_IDENTITY_FORMAT_VERSION: + raise ValueError(f"Unsupported SourceIdentity format version: {format_version}") + return cls( - format_version=data["format_version"], + format_version=format_version, + artifact_identity=ArtifactIdentity.from_dict(data["artifact_identity"]), model_fingerprint=data["model_fingerprint"], quant_fingerprint=data["quant_fingerprint"], backend_fingerprint=data["backend_fingerprint"], @@ -548,7 +581,8 @@ def check_weight_sharing_compatibility( result = IdentityMatchResult(matched=False, mismatched_fields=missing_fields) message = ( "SourceIdentity unavailable for fields " - f"{missing_fields}; receiver cannot verify source weight layout." + f"{missing_fields}; receiver cannot verify the source checkpoint " + "artifact and weight layout." ) if policy is IdentityCheckPolicy.STRICT: raise SourceIdentityMismatchError(message) @@ -574,7 +608,7 @@ def check_weight_sharing_compatibility( message = ( "SourceIdentity mismatch on fields " f"{result.mismatched_fields}; receiver and source disagree on " - "weight layout." + "the checkpoint artifact or weight layout." ) if policy is IdentityCheckPolicy.STRICT: raise SourceIdentityMismatchError(message) diff --git a/tests/unittest/_torch/executor/test_model_loader_gms.py b/tests/unittest/_torch/executor/test_model_loader_gms.py index 34dd8223a6a5..b56935da6378 100644 --- a/tests/unittest/_torch/executor/test_model_loader_gms.py +++ b/tests/unittest/_torch/executor/test_model_loader_gms.py @@ -14,6 +14,9 @@ from tensorrt_llm._torch.pyexecutor import model_loader as model_loader_mod from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader from tensorrt_llm._torch.weight_sharing import ( + ARTIFACT_IDENTITY_FORMAT_VERSION, + SOURCE_IDENTITY_FORMAT_VERSION, + ArtifactIdentity, PostTransformProfile, PostTransformProfileRegistry, PostTransformTransferScope, @@ -21,7 +24,12 @@ from tensorrt_llm.llmapi.llm_args import LoadFormat _SOURCE_IDENTITY = model_loader_mod.SourceIdentity( - format_version=1, + format_version=SOURCE_IDENTITY_FORMAT_VERSION, + artifact_identity=ArtifactIdentity( + format_version=ARTIFACT_IDENTITY_FORMAT_VERSION, + scheme="checkpoint_manifest_sha256", + digest="0" * 64, + ), model_fingerprint="model", quant_fingerprint="quant", backend_fingerprint="backend", @@ -100,10 +108,25 @@ def _make_loader(monkeypatch, *, events, spec_config=None): monkeypatch.setattr(model_loader_mod, "MetaInitMode", lambda: nullcontext()) # These tests stub ModelConfig, while SourceIdentity has dedicated # coverage. Keep this file focused on ModelLoader GMS branch behavior. + + def _build_artifact_identity(_cls, checkpoint_dir): + assert checkpoint_dir == "/ckpt" + return _SOURCE_IDENTITY.artifact_identity + + monkeypatch.setattr( + model_loader_mod.ArtifactIdentity, + "from_checkpoint", + classmethod(_build_artifact_identity), + ) + + def _build_source_identity(_cls, *_args, **kwargs): + assert kwargs["artifact_identity"] is _SOURCE_IDENTITY.artifact_identity + return _SOURCE_IDENTITY + monkeypatch.setattr( model_loader_mod.SourceIdentity, "from_model_config", - classmethod(lambda cls, *_args, **_kwargs: _SOURCE_IDENTITY), + classmethod(_build_source_identity), ) monkeypatch.setattr( model_loader_mod.AutoModelForCausalLM, @@ -181,6 +204,31 @@ def _tiny_profile_registry() -> PostTransformProfileRegistry: ) +def test_gms_artifact_identity_failure_remains_fatal(monkeypatch): + loader = _make_loader(monkeypatch, events=[]) + artifact_error = ValueError( + "Checkpoint manifests do not support nested symlinked directories: /ckpt/shards" + ) + monkeypatch.setattr( + model_loader_mod.ArtifactIdentity, + "from_checkpoint", + MagicMock(side_effect=artifact_error), + ) + source_identity_factory = MagicMock() + monkeypatch.setattr( + model_loader_mod.SourceIdentity, + "from_model_config", + source_identity_factory, + ) + checkpoint_loader = MagicMock(name="checkpoint_loader") + checkpoint_loader.checkpoint_format = "MX" + + with pytest.raises(ValueError, match="nested symlinked directories"): + loader.load("/ckpt", checkpoint_loader) + + source_identity_factory.assert_not_called() + + @pytest.mark.parametrize( "is_rw, expected_events", [ diff --git a/tests/unittest/_torch/executor/test_model_loader_mx.py b/tests/unittest/_torch/executor/test_model_loader_mx.py index 82a92622fd30..ccdf960ea786 100644 --- a/tests/unittest/_torch/executor/test_model_loader_mx.py +++ b/tests/unittest/_torch/executor/test_model_loader_mx.py @@ -24,6 +24,9 @@ from tensorrt_llm._torch.pyexecutor import model_loader as model_loader_mod from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader from tensorrt_llm._torch.weight_sharing import ( + ARTIFACT_IDENTITY_FORMAT_VERSION, + SOURCE_IDENTITY_FORMAT_VERSION, + ArtifactIdentity, PostTransformFeature, PostTransformProfile, PostTransformProfileRegistry, @@ -33,7 +36,12 @@ from tensorrt_llm.llmapi.llm_args import LoadFormat _SOURCE_IDENTITY = model_loader_mod.SourceIdentity( - format_version=1, + format_version=SOURCE_IDENTITY_FORMAT_VERSION, + artifact_identity=ArtifactIdentity( + format_version=ARTIFACT_IDENTITY_FORMAT_VERSION, + scheme="checkpoint_manifest_sha256", + digest="0" * 64, + ), model_fingerprint="model", quant_fingerprint="quant", backend_fingerprint="backend", @@ -188,10 +196,25 @@ def _make_loader(monkeypatch, *, events, spec_config=None): monkeypatch.setattr(model_loader_mod, "MetaInitMode", lambda: nullcontext()) # These tests stub ModelConfig, while SourceIdentity has dedicated # coverage. Keep this file focused on ModelLoader MX branch behavior. + + def _build_artifact_identity(_cls, checkpoint_dir): + assert checkpoint_dir == "/ckpt" + return _SOURCE_IDENTITY.artifact_identity + + monkeypatch.setattr( + model_loader_mod.ArtifactIdentity, + "from_checkpoint", + classmethod(_build_artifact_identity), + ) + + def _build_source_identity(_cls, *_args, **kwargs): + assert kwargs["artifact_identity"] is _SOURCE_IDENTITY.artifact_identity + return _SOURCE_IDENTITY + monkeypatch.setattr( model_loader_mod.SourceIdentity, "from_model_config", - classmethod(lambda cls, *_args, **_kwargs: _SOURCE_IDENTITY), + classmethod(_build_source_identity), ) monkeypatch.setattr( model_loader_mod.AutoModelForCausalLM, @@ -484,6 +507,53 @@ def test_mx_fallback_runs_standard_weight_mapping(monkeypatch): ) +def test_mx_artifact_identity_failure_falls_back_to_disk(monkeypatch): + events = [] + loader = _make_loader(monkeypatch, events=events) + monkeypatch.setattr( + ModelLoader, + "_POST_TRANSFORM_PROFILE_REGISTRY", + _tiny_profile_registry(), + ) + artifact_error = ValueError( + "Checkpoint manifests do not support nested symlinked directories: /ckpt/shards" + ) + monkeypatch.setattr( + model_loader_mod.ArtifactIdentity, + "from_checkpoint", + MagicMock(side_effect=artifact_error), + ) + source_identity_factory = MagicMock() + monkeypatch.setattr( + model_loader_mod.SourceIdentity, + "from_model_config", + source_identity_factory, + ) + warning = MagicMock() + monkeypatch.setattr(model_loader_mod.logger, "warning", warning) + + checkpoint_loader = MagicMock(name="checkpoint_loader") + checkpoint_loader.checkpoint_format = "MX" + checkpoint_loader.is_weights_preloaded.return_value = False + checkpoint_loader.load_weights.return_value = {"weight": MagicMock()} + checkpoint_loader.get_initialized_weight_mapper.return_value = MagicMock() + + model, _ = loader.load("/ckpt", checkpoint_loader) + + assert loader._source_identity is None + source_identity_factory.assert_not_called() + warning.assert_called_once() + assert "falling back to regular checkpoint loading" in warning.call_args.args[0] + _args, kwargs = checkpoint_loader.load_weights.call_args + assert kwargs["source_identity"] is None + checkpoint_loader.post_load_publish.assert_called_once_with( + model, + checkpoint_dir="/ckpt", + weights_preloaded=False, + source_identity=None, + ) + + class _HookRecorder(nn.Module): def __init__( self, diff --git a/tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py b/tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py index c8e845c5d3b0..e80026d31114 100644 --- a/tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py +++ b/tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py @@ -42,14 +42,24 @@ _resolve_mx_model_name, _serialize_source_identity, ) -from tensorrt_llm._torch.weight_sharing import SourceIdentity +from tensorrt_llm._torch.weight_sharing import ( + ARTIFACT_IDENTITY_FORMAT_VERSION, + SOURCE_IDENTITY_FORMAT_VERSION, + ArtifactIdentity, + SourceIdentity, +) _MISSING = object() def _identity(rank: int = 0, suffix: str = "same") -> SourceIdentity: return SourceIdentity( - format_version=1, + format_version=SOURCE_IDENTITY_FORMAT_VERSION, + artifact_identity=ArtifactIdentity( + format_version=ARTIFACT_IDENTITY_FORMAT_VERSION, + scheme="checkpoint_manifest_sha256", + digest="0" * 64, + ), model_fingerprint=f"model-{suffix}", quant_fingerprint=f"quant-{suffix}", backend_fingerprint=f"backend-{suffix}", @@ -753,12 +763,9 @@ def _publish_side_effect(model, **_kwargs): def test_serialized_identity_ignores_local_checkpoint_path(self): donor_identity = _identity() - receiver_identity = SourceIdentity( - **{ - **donor_identity.to_dict(), - "model_name": "/tmp/no-shards/TinyLlama", - } - ) + receiver_payload = donor_identity.to_dict() + receiver_payload["model_name"] = "/tmp/no-shards/TinyLlama" + receiver_identity = SourceIdentity.from_dict(receiver_payload) assert donor_identity.model_name != receiver_identity.model_name assert _serialize_source_identity(donor_identity) == _serialize_source_identity( diff --git a/tests/unittest/_torch/weight_sharing/_source_identity_fakes.py b/tests/unittest/_torch/weight_sharing/_source_identity_fakes.py index c68b10506b7a..ff567b742cec 100644 --- a/tests/unittest/_torch/weight_sharing/_source_identity_fakes.py +++ b/tests/unittest/_torch/weight_sharing/_source_identity_fakes.py @@ -19,9 +19,14 @@ `test_source_identity.py` and `test_mx_source_identity_gate.py`. """ +import hashlib from typing import Optional, Sequence -from tensorrt_llm._torch.weight_sharing import SourceIdentity +from tensorrt_llm._torch.weight_sharing import ( + ARTIFACT_IDENTITY_FORMAT_VERSION, + ArtifactIdentity, + SourceIdentity, +) _UNSET = object() @@ -204,16 +209,37 @@ def named_buffers(self): return list(self._buffers.items()) -def identity_from(config: FakeModelConfig, *, model_name: Optional[str] = None) -> SourceIdentity: +def make_artifact_identity(key: str = "same") -> ArtifactIdentity: + """Build a deterministic local-checkpoint identity for tests.""" + return ArtifactIdentity( + format_version=ARTIFACT_IDENTITY_FORMAT_VERSION, + scheme="checkpoint_manifest_sha256", + digest=hashlib.sha256(key.encode("utf-8")).hexdigest(), + ) + + +def identity_from( + config: FakeModelConfig, + *, + model_name: Optional[str] = None, + artifact_key: str = "same", +) -> SourceIdentity: """Build a :class:`SourceIdentity` from a fake config and derived model.""" return SourceIdentity.from_model_config( - config, FakeModel(config.pretrained_config), model_name=model_name + config, + FakeModel(config.pretrained_config), + artifact_identity=make_artifact_identity(artifact_key), + model_name=model_name, ) def make_identity( - *, attn_backend: str = "TRTLLM", rank: int = 0, model_name: str = "m" + *, + attn_backend: str = "TRTLLM", + rank: int = 0, + model_name: str = "m", + artifact_key: str = "same", ) -> SourceIdentity: """Build a :class:`SourceIdentity` from a fake config for `rank`.""" cfg = FakeModelConfig(mapping=FakeMapping(rank=rank, tp_rank=rank), attn_backend=attn_backend) - return identity_from(cfg, model_name=model_name) + return identity_from(cfg, model_name=model_name, artifact_key=artifact_key) diff --git a/tests/unittest/_torch/weight_sharing/test_artifact_identity.py b/tests/unittest/_torch/weight_sharing/test_artifact_identity.py new file mode 100644 index 000000000000..30526ce77c66 --- /dev/null +++ b/tests/unittest/_torch/weight_sharing/test_artifact_identity.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""Tests for immutable checkpoint artifact identities.""" + +from pathlib import Path + +import pytest + +from tensorrt_llm._torch.weight_sharing import ArtifactIdentity + + +def _write_checkpoint(path: Path, weights: bytes = b"weights") -> None: + path.mkdir(parents=True) + (path / "config.json").write_text('{"architectures":["LlamaForCausalLM"]}') + (path / "model.safetensors").write_bytes(weights) + + +def test_local_checkpoint_identity_is_path_independent(tmp_path: Path) -> None: + left = tmp_path / "left" / "checkpoint" + right = tmp_path / "right" / "checkpoint" + _write_checkpoint(left) + _write_checkpoint(right) + + assert ArtifactIdentity.from_checkpoint(left) == ArtifactIdentity.from_checkpoint(right) + + +def test_local_checkpoint_identity_binds_file_contents(tmp_path: Path) -> None: + left = tmp_path / "left" + right = tmp_path / "right" + _write_checkpoint(left, weights=b"fine-tune-a") + _write_checkpoint(right, weights=b"fine-tune-b") + + assert ArtifactIdentity.from_checkpoint(left) != ArtifactIdentity.from_checkpoint(right) + + +def test_local_checkpoint_identity_ignores_cache_and_scm_metadata(tmp_path: Path) -> None: + left = tmp_path / "left" + right = tmp_path / "right" + _write_checkpoint(left) + _write_checkpoint(right) + (left / ".cache").mkdir() + (left / ".cache" / "download.lock").write_text("transient") + (right / ".git").mkdir() + (right / ".git" / "HEAD").write_text("ref: refs/heads/main") + + assert ArtifactIdentity.from_checkpoint(left) == ArtifactIdentity.from_checkpoint(right) + + +def test_local_checkpoint_identity_rejects_nested_directory_symlink(tmp_path: Path) -> None: + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + external_weights = tmp_path / "external-weights" + external_weights.mkdir() + (external_weights / "model.safetensors").write_bytes(b"weights") + (checkpoint / "weights").symlink_to(external_weights, target_is_directory=True) + + with pytest.raises(ValueError, match="nested symlinked directories"): + ArtifactIdentity.from_checkpoint(checkpoint) + + +def test_hf_snapshot_identity_binds_revision_across_cache_roots(tmp_path: Path) -> None: + revision = "a" * 40 + left = tmp_path / "cache-a" / "models--org--model" / "snapshots" / revision + right = tmp_path / "cache-b" / "models--org--model" / "snapshots" / revision + left.mkdir(parents=True) + right.mkdir(parents=True) + + left_identity = ArtifactIdentity.from_checkpoint(left) + right_identity = ArtifactIdentity.from_checkpoint(right) + assert left_identity == right_identity + assert left_identity.scheme == "hf_snapshot_revision" + + +def test_hf_snapshot_identity_binds_revision_and_subpath(tmp_path: Path) -> None: + snapshot = tmp_path / "models--org--model" / "snapshots" + revision_a = snapshot / ("a" * 40) + revision_b = snapshot / ("b" * 40) + (revision_a / "variant-a").mkdir(parents=True) + (revision_a / "variant-b").mkdir() + revision_b.mkdir(parents=True) + + root_identity = ArtifactIdentity.from_checkpoint(revision_a) + assert root_identity != ArtifactIdentity.from_checkpoint(revision_b) + assert root_identity != ArtifactIdentity.from_checkpoint(revision_a / "variant-a") + assert ArtifactIdentity.from_checkpoint( + revision_a / "variant-a" + ) != ArtifactIdentity.from_checkpoint(revision_a / "variant-b") + + +def test_serialization_roundtrip(tmp_path: Path) -> None: + checkpoint = tmp_path / "checkpoint" + _write_checkpoint(checkpoint) + identity = ArtifactIdentity.from_checkpoint(checkpoint) + + assert ArtifactIdentity.from_dict(identity.to_dict()) == identity + + +def test_rejects_unknown_format_version(tmp_path: Path) -> None: + checkpoint = tmp_path / "checkpoint" + _write_checkpoint(checkpoint) + payload = ArtifactIdentity.from_checkpoint(checkpoint).to_dict() + payload["format_version"] += 1 + + with pytest.raises(ValueError, match="Unsupported ArtifactIdentity format version"): + ArtifactIdentity.from_dict(payload) + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("scheme", "unknown", "Unsupported ArtifactIdentity scheme"), + ("scheme", [], "scheme must be a string"), + ("digest", "not-a-digest", "64-character hex value"), + ("digest", 1, "digest must be a string"), + ], +) +def test_rejects_invalid_serialized_fields( + tmp_path: Path, field: str, value: object, message: str +) -> None: + checkpoint = tmp_path / "checkpoint" + _write_checkpoint(checkpoint) + payload = ArtifactIdentity.from_checkpoint(checkpoint).to_dict() + payload[field] = value + + with pytest.raises(ValueError, match=message): + ArtifactIdentity.from_dict(payload) + + +@pytest.mark.parametrize("version", [True, "1"]) +def test_rejects_non_integer_format_version(tmp_path: Path, version: object) -> None: + checkpoint = tmp_path / "checkpoint" + _write_checkpoint(checkpoint) + payload = ArtifactIdentity.from_checkpoint(checkpoint).to_dict() + payload["format_version"] = version + + with pytest.raises(ValueError, match="format version must be an integer"): + ArtifactIdentity.from_dict(payload) + + +def test_rejects_missing_or_empty_checkpoint(tmp_path: Path) -> None: + with pytest.raises(FileNotFoundError): + ArtifactIdentity.from_checkpoint(tmp_path / "missing") + + empty = tmp_path / "empty" + empty.mkdir() + with pytest.raises(ValueError, match="contains no files"): + ArtifactIdentity.from_checkpoint(empty) diff --git a/tests/unittest/_torch/weight_sharing/test_gms_source_identity_gate.py b/tests/unittest/_torch/weight_sharing/test_gms_source_identity_gate.py index f6689bfd0caa..99f21b4bd9b9 100644 --- a/tests/unittest/_torch/weight_sharing/test_gms_source_identity_gate.py +++ b/tests/unittest/_torch/weight_sharing/test_gms_source_identity_gate.py @@ -79,6 +79,14 @@ def test_gate_raises_on_mismatch(): loader._check_gms_source_identity(_FakeGMSBackend(writer)) +def test_gate_raises_on_checkpoint_artifact_mismatch(): + local = _identity(artifact_key="fine-tune-a") + writer = _identity(artifact_key="fine-tune-b") + loader = _new_loader(local) + with pytest.raises(SourceIdentityMismatchError): + loader._check_gms_source_identity(_FakeGMSBackend(writer)) + + def test_gate_raises_when_writer_identity_unavailable(): # Publisher metadata not wired yet (get_source_identity returns None); # GMS has no disk fallback, so unverified sharing must raise. diff --git a/tests/unittest/_torch/weight_sharing/test_mx_source_identity_gate.py b/tests/unittest/_torch/weight_sharing/test_mx_source_identity_gate.py index 2a3d3af7d179..e2658ca86d21 100644 --- a/tests/unittest/_torch/weight_sharing/test_mx_source_identity_gate.py +++ b/tests/unittest/_torch/weight_sharing/test_mx_source_identity_gate.py @@ -51,6 +51,13 @@ def test_gate_falls_back_on_mismatch(): assert loader._source_metadata_identity_compatible(_build_mx_source_metadata(source)) is False +def test_gate_falls_back_on_checkpoint_artifact_mismatch(): + local = _identity(artifact_key="fine-tune-a") + source = _identity(artifact_key="fine-tune-b") + loader = _new_loader(local) + assert loader._source_metadata_identity_compatible(_build_mx_source_metadata(source)) is False + + def test_gate_falls_back_when_no_local_identity(): # MX must not consume shared weights unless the receiver identity exists. loader = _new_loader(None) diff --git a/tests/unittest/_torch/weight_sharing/test_source_identity.py b/tests/unittest/_torch/weight_sharing/test_source_identity.py index a02995c4f520..779dc8c355bd 100644 --- a/tests/unittest/_torch/weight_sharing/test_source_identity.py +++ b/tests/unittest/_torch/weight_sharing/test_source_identity.py @@ -19,6 +19,7 @@ """ import copy +from pathlib import Path import pytest from _source_identity_fakes import ( @@ -29,6 +30,7 @@ FakeQuantConfig, FakeQuantConfigWithPythonOnlyField, identity_from, + make_artifact_identity, ) from tensorrt_llm._torch.weight_sharing import ( @@ -48,6 +50,24 @@ def test_identical_configs_match(): assert bool(result) is True +def test_from_model_config_derives_artifact_identity(tmp_path: Path) -> None: + (tmp_path / "model.safetensors").write_bytes(b"checkpoint") + identity = SourceIdentity.from_model_config(FakeModelConfig(), checkpoint_dir=str(tmp_path)) + assert identity.artifact_identity.scheme == "checkpoint_manifest_sha256" + + +def test_from_model_config_requires_one_artifact_source() -> None: + config = FakeModelConfig() + with pytest.raises(ValueError, match="Exactly one"): + SourceIdentity.from_model_config(config) + with pytest.raises(ValueError, match="Exactly one"): + SourceIdentity.from_model_config( + config, + checkpoint_dir="/checkpoint", + artifact_identity=make_artifact_identity(), + ) + + def test_rank_defaults_from_mapping(): cfg = FakeModelConfig(mapping=FakeMapping(rank=3, tp_rank=3)) identity = identity_from(cfg) @@ -83,10 +103,14 @@ def test_param_dtype_override_flags_shard(): # the realized-layout fingerprint catches it. cfg = FakeModelConfig() a = SourceIdentity.from_model_config( - cfg, FakeModel(cfg.pretrained_config, dtype="torch.bfloat16") + cfg, + FakeModel(cfg.pretrained_config, dtype="torch.bfloat16"), + artifact_identity=make_artifact_identity(), ) b = SourceIdentity.from_model_config( - cfg, FakeModel(cfg.pretrained_config, dtype="torch.float16") + cfg, + FakeModel(cfg.pretrained_config, dtype="torch.float16"), + artifact_identity=make_artifact_identity(), ) result = a.matches(b) assert not result.matched @@ -109,11 +133,24 @@ def test_cross_architecture_same_shapes_flags_global(): assert "model_fingerprint" in result.mismatched_fields +def test_different_checkpoint_artifacts_flag_global(): + a = identity_from(FakeModelConfig(), artifact_key="fine-tune-a") + b = identity_from(FakeModelConfig(), artifact_key="fine-tune-b") + result = a.matches(b) + assert not result.matched + assert result.mismatched_fields == ["artifact_identity"] + assert a.global_fingerprint != b.global_fingerprint + + def test_no_model_degrades_to_architecture_only(): - # Without a module, the fingerprint still builds (architecture-only) and - # two identical configs still match. - a = SourceIdentity.from_model_config(FakeModelConfig(), None) - b = SourceIdentity.from_model_config(FakeModelConfig(), None) + # Without a module, the shard fingerprint has no realized tensor layout, + # while matching artifacts and configurations remain comparable. + a = SourceIdentity.from_model_config( + FakeModelConfig(), None, artifact_identity=make_artifact_identity() + ) + b = SourceIdentity.from_model_config( + FakeModelConfig(), None, artifact_identity=make_artifact_identity() + ) assert a.matches(b).matched @@ -168,6 +205,29 @@ def test_serialization_roundtrip(): restored = SourceIdentity.from_dict(a.to_dict()) assert restored == a assert a.matches(restored).matched + assert restored.artifact_identity == a.artifact_identity + + +def test_deserialization_rejects_missing_artifact_identity(): + payload = identity_from(FakeModelConfig()).to_dict() + payload.pop("artifact_identity") + with pytest.raises(KeyError): + SourceIdentity.from_dict(payload) + + +def test_deserialization_rejects_unknown_format_version(): + payload = identity_from(FakeModelConfig()).to_dict() + payload["format_version"] += 1 + with pytest.raises(ValueError, match="Unsupported SourceIdentity format version"): + SourceIdentity.from_dict(payload) + + +def test_deserialization_rejects_v1_identity_without_artifact_binding(): + payload = identity_from(FakeModelConfig()).to_dict() + payload["format_version"] = 1 + payload.pop("artifact_identity") + with pytest.raises(ValueError, match="Unsupported SourceIdentity format version"): + SourceIdentity.from_dict(payload) def test_check_warn_fallback_on_mismatch(): @@ -219,6 +279,7 @@ def test_format_version_mismatch_never_matches(): if hasattr(copy, "replace") else SourceIdentity( format_version=a.format_version + 1, + artifact_identity=a.artifact_identity, model_fingerprint=a.model_fingerprint, quant_fingerprint=a.quant_fingerprint, backend_fingerprint=a.backend_fingerprint, From 526c1e3229c0fd983a9bd08437e7fd14f965fe2e Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Thu, 23 Jul 2026 12:27:16 -0700 Subject: [PATCH 3/4] [https://nvbugs/6448152][perf] make C++ context-transfer consensus asynchronous (#16634) Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 22 + .../contextTransferCoordinator.h | 109 ++++ .../tensorrt_llm/runtime/utils/mpiTags.h | 8 +- .../tensorrt_llm/runtime/utils/mpiUtils.h | 13 +- cpp/tensorrt_llm/batch_manager/CMakeLists.txt | 3 +- .../batch_manager/cacheTransceiver.cpp | 126 ++++- .../contextTransferCoordinator.cpp | 475 ++++++++++++++++++ .../unit_tests/batch_manager/CMakeLists.txt | 1 + .../contextTransferCoordinatorTest.cpp | 143 ++++++ .../multi_gpu/cacheTransceiverTest.cpp | 92 ++++ .../unit_tests/multi_gpu/mpiUtilsTest.cpp | 27 +- 11 files changed, 997 insertions(+), 22 deletions(-) create mode 100644 cpp/include/tensorrt_llm/batch_manager/contextTransferCoordinator.h create mode 100644 cpp/tensorrt_llm/batch_manager/contextTransferCoordinator.cpp create mode 100644 cpp/tests/unit_tests/batch_manager/contextTransferCoordinatorTest.cpp diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 9d28fa26c4ee..735716e59e92 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -27,6 +27,7 @@ #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/runtime/utils/pgUtils.h" +#include #include #include #include @@ -55,6 +56,7 @@ class BaseKVCacheManager; class CacheSender; class CacheReceiver; +class ContextTransferCoordinator; class CacheTransceiverComm { @@ -149,6 +151,25 @@ class CacheTransceiverComm TLLM_THROW("Input arguments only supported in mpi"); } + [[nodiscard]] std::unique_ptr sendAsync( + void const* buffer, std::size_t size, mpi::MpiType dtype, int dest, mpi::MpiTag tag) const + { + TLLM_CHECK_WITH_INFO(isMpi(), "Point-to-point cache-transceiver status messages require MPI."); + return mMpiComm->sendAsync(buffer, size, dtype, dest, tag); + } + + [[nodiscard]] bool iprobe(int source, mpi::MpiTag tag, MPI_Status* status) const + { + TLLM_CHECK_WITH_INFO(isMpi(), "Point-to-point cache-transceiver status messages require MPI."); + return mMpiComm->iprobe(source, tag, status); + } + + void recv(void* buffer, std::size_t size, mpi::MpiType dtype, int source, mpi::MpiTag tag) const + { + TLLM_CHECK_WITH_INFO(isMpi(), "Point-to-point cache-transceiver status messages require MPI."); + static_cast(mMpiComm->recv(buffer, size, dtype, source, tag)); + } + CacheTransceiverComm split(int color, int key) { if (isMpi()) @@ -307,6 +328,7 @@ class CacheTransceiver : public BaseCacheTransceiver std::shared_ptr mGroupComm; std::shared_ptr mGroupTensorParaComm, mGroupPipeParaComm, mGroupDataComm, mGroupTPInDPComm; + std::unique_ptr mContextTransferCoordinator; executor::kv_cache::CommState const* mCommState; std::unique_ptr mCacheState; diff --git a/cpp/include/tensorrt_llm/batch_manager/contextTransferCoordinator.h b/cpp/include/tensorrt_llm/batch_manager/contextTransferCoordinator.h new file mode 100644 index 000000000000..9b156c8887a0 --- /dev/null +++ b/cpp/include/tensorrt_llm/batch_manager/contextTransferCoordinator.h @@ -0,0 +1,109 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * 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. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager +{ + +class CacheTransceiverComm; + +enum class ContextTransferVote : std::uint64_t +{ + kCompleted = 1, + kFailed = 2, +}; + +struct ContextTransferConsensusResult +{ + std::unordered_set completedRequestIds; + std::unordered_set failedRequestIds; + std::unordered_set timedOutRequestIds; +}; + +//! Accumulates one immutable terminal vote per participant and request. +class ContextTransferVoteReducer +{ +public: + explicit ContextTransferVoteReducer(int participantCount); + + void recordVote(int participantRank, std::uint64_t requestId, ContextTransferVote vote); + + //! Record a sticky nonterminal timeout proposal. The final outcome still waits for every terminal vote. + void recordTimeout(std::uint64_t requestId); + + [[nodiscard]] ContextTransferConsensusResult takeReady(); + + void clear() noexcept; + +private: + struct RequestVotes + { + explicit RequestVotes(int participantCount) + : votes(static_cast(participantCount), 0) + { + } + + std::vector votes; + int terminalCount{0}; + bool failed{false}; + bool timedOut{false}; + bool timeoutPending{false}; + }; + + int mParticipantCount; + std::unordered_map mRequestVotes; +}; + +//! Coordinates asynchronous context-transfer outcomes across a rank group. +//! +//! Timeout and terminal events share one ordered stream toward the coordinator. Timeout updates and final commits +//! share another ordered stream toward followers, so a request that times out can never commit success first. +class ContextTransferCoordinator +{ +public: + explicit ContextTransferCoordinator(std::shared_ptr comm); + ~ContextTransferCoordinator(); + + ContextTransferCoordinator(ContextTransferCoordinator const&) = delete; + ContextTransferCoordinator& operator=(ContextTransferCoordinator const&) = delete; + + //! Publish this rank's immutable local terminal outcome without waiting for peers. + void publishLocalOutcome(std::uint64_t requestId, bool failed); + + //! Publish an idempotent, sticky, nonterminal timeout proposal without waiting for peers. + void publishTimeout(std::uint64_t requestId); + + //! Make nonblocking protocol progress and return newly committed global outcomes. + [[nodiscard]] ContextTransferConsensusResult poll(); + + //! Exchange ordered close markers so no active MPI request outlives its backing buffer. Failure aborts closed. + void shutdown() noexcept; + +private: + class Impl; + std::unique_ptr mImpl; +}; + +} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h b/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h index 32c086c84ee9..49c50e49b2e2 100644 --- a/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h +++ b/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2021-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -71,7 +71,11 @@ enum class MpiTag : int // KvCacheEventManager kKvCacheEventSize = 1026, - kKvCacheEvent = 1027 + kKvCacheEvent = 1027, + + // Asynchronous context-transfer coordination. + kContextTransferEvent = 1028, + kContextTransferUpdate = 1029 }; } // namespace tensorrt_llm::mpi diff --git a/cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h b/cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h index 75ec7a534815..b474b4b12cbf 100644 --- a/cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h +++ b/cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2021-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -235,6 +235,17 @@ class MpiRequest #endif } + [[nodiscard]] bool isCompleted() + { +#if ENABLE_MULTI_DEVICE + int completed = 0; + TLLM_MPI_CHECK(MPI_Test(&mRequest, &completed, MPI_STATUS_IGNORE)); + return completed != 0; +#else + TLLM_THROW("Multi device support is disabled."); +#endif + } + void cancel() { #if ENABLE_MULTI_DEVICE diff --git a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt index d893d12fa3e3..c7fea62c023a 100644 --- a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt +++ b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & # AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); you may not @@ -30,6 +30,7 @@ set(SRCS capacityScheduler.cpp createNewDecoderRequests.cpp contextProgress.cpp + contextTransferCoordinator.cpp dataTransceiver.cpp decoderBuffers.cpp kvCacheManager.cpp diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index f730fb2aaf1e..69adf4246591 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -36,6 +36,7 @@ #include "tensorrt_llm/batch_manager/cacheFormatter.h" #include "tensorrt_llm/batch_manager/cacheTransceiver.h" #include "tensorrt_llm/batch_manager/contextProgress.h" +#include "tensorrt_llm/batch_manager/contextTransferCoordinator.h" #include "tensorrt_llm/batch_manager/dataTransceiver.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/batch_manager/kvCacheType.h" @@ -569,6 +570,34 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa mCacheSender = std::make_unique(mManager.get(), worldConfig.getRank(), makeCacheTransferLayer()); mCacheReceiver = std::make_unique(mManager.get(), worldConfig.getRank(), makeCacheTransferLayer()); + // Keep automatic enablement within the currently qualified C++ NIXL/UCX TP1/CP1 pipeline topology. + bool const coordinatorTopologyEligible = worldConfig.getPipelineParallelism() > 1 && useMPI() + && backendType.value() == executor::CacheTransceiverConfig::BackendType::NIXL + && common::getEnvNixlBackend() == "UCX" && worldConfig.getTensorParallelism() == 1 + && worldConfig.getContextParallelism() == 1 && !mCacheState->getParallelConfig().mEnableAttentionDP; + if (worldConfig.getPipelineParallelism() > 1 && useMPI()) + { + TLLM_CHECK(mGroupPipeParaComm != nullptr); + constexpr std::uint64_t kCoordinatorProtocolVersion = 1; + std::uint64_t const localVersion = coordinatorTopologyEligible ? kCoordinatorProtocolVersion : 0; + bool const cancellationEnabled = common::getEnvDisaggEnableInflightCancel(); + std::uint64_t const localProtocolMode = (localVersion << 1) | static_cast(cancellationEnabled); + std::vector protocolModes(static_cast(mGroupPipeParaComm->getSize())); + mGroupPipeParaComm->allgather(&localProtocolMode, protocolModes.data(), 1, mpi::MpiType::kUINT64); + TLLM_CHECK_WITH_INFO(std::all_of(protocolModes.begin(), protocolModes.end(), + [&](std::uint64_t const mode) { return mode == localProtocolMode; }), + "Context-transfer consensus protocol version or cancellation mode differs across PP ranks."); + if (localVersion != 0) + { + mContextTransferCoordinator = std::make_unique(mGroupPipeParaComm); + TLLM_LOG_INFO( + "Enable asynchronous context-transfer consensus version %llu for PP group of size %d; in-flight " + "cancellation=%s.", + static_cast(kCoordinatorProtocolVersion), mGroupPipeParaComm->getSize(), + cancellationEnabled ? "enabled" : "disabled"); + } + } + initializeCommState(); } @@ -578,6 +607,7 @@ CacheTransceiver::~CacheTransceiver() // plugin are still alive. The workers can access both during termination. mCacheSender.reset(); mCacheReceiver.reset(); + mContextTransferCoordinator.reset(); if (mWrapperLibHandle) { @@ -868,6 +898,26 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( toCompleteIdSet.insert(request->mRequestId); } + auto recordTimeout = [&](RequestIdType const requestId) + { + bool const inserted = mTimedOutSenderIds.insert(requestId).second; + if (inserted && inflightCancelEnabled && mContextTransferCoordinator) + { + mContextTransferCoordinator->publishTimeout(requestId); + } + return inserted; + }; + auto recordOutcome + = [&](RequestIdType const requestId, std::shared_ptr const& request, bool const failed) + { + recordLocalTransferOutcome(requestId, request, failed, mCompletedSenderRequestIds, mFailedSenderRequestIds, + mSenderRequestsAwaitingConsensus); + if (mContextTransferCoordinator) + { + mContextTransferCoordinator->publishLocalOutcome(requestId, failed); + } + }; + // Record local terminal outcomes for requests selected this round. The // request is reported only after all ranks in the sync group agree that the // request reached a terminal state. @@ -881,7 +931,7 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( auto const elapsedMs = getTransferElapsedMs(request, LlmRequest::getSteadyClockNow()); if (elapsedMs > kvTransferTimeoutMs.value()) { - if (mTimedOutSenderIds.insert(requestId).second) + if (recordTimeout(requestId)) { TLLM_LOG_WARNING( "Context KV cache transfer for request %ld exceeded configured timeout: " @@ -893,6 +943,8 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( } if (blockAll || (toCompleteIdSet.find(requestId) != toCompleteIdSet.end())) { + bool terminal = false; + bool failed = false; try { auto const status = blockAll ? std::future_status::ready : future.wait_for(futureWaitInterval); @@ -902,7 +954,7 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( if (kvTransferTimeoutMs.has_value()) { auto const elapsedMs = getTransferElapsedMs(request, request->getKvCacheTransferEnd()); - if (elapsedMs > kvTransferTimeoutMs.value() && mTimedOutSenderIds.insert(requestId).second) + if (elapsedMs > kvTransferTimeoutMs.value() && recordTimeout(requestId)) { TLLM_LOG_WARNING( "Context KV cache transfer for request %ld completed after its deadline: " @@ -911,9 +963,8 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( inflightCancelEnabled ? "failing request" : "observe-only"); } } - recordLocalTransferOutcome(requestId, request, /*failed=*/false, mCompletedSenderRequestIds, - mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); - it = mSenderFutures.erase(it); + failed = request->getState() == LlmRequestState::kDISAGG_TRANS_ERROR; + terminal = true; } else if (status == std::future_status::timeout) { @@ -927,25 +978,28 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( { TLLM_LOG_ERROR( "Future returned unexpected status for request %ld. Recording as failed.", requestId); - - recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedSenderRequestIds, - mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); - it = mSenderFutures.erase(it); + failed = true; + terminal = true; } } catch (std::exception const& e) { TLLM_LOG_ERROR("Error occurred during context transfer for request %ld: %s", requestId, e.what()); - recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedSenderRequestIds, - mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); - it = mSenderFutures.erase(it); + failed = true; + terminal = true; } catch (...) { TLLM_LOG_ERROR("Unknown error occurred during context transfer for request %ld", requestId); - recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedSenderRequestIds, - mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); + failed = true; + terminal = true; + } + if (terminal) + { + auto terminalRequest = request; it = mSenderFutures.erase(it); + // Publish outside the transfer-future try/catch. A protocol error must not rewrite an immutable vote. + recordOutcome(requestId, terminalRequest, failed); } } else @@ -955,12 +1009,51 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( } RequestStatuses requestsStatus{}; - auto const consensusOutcome = reduceTransferStates(syncComm, mGroupPipeParaComm, mCompletedSenderRequestIds, - mFailedSenderRequestIds, inflightCancelEnabled ? mTimedOutSenderIds : std::unordered_set{}); + TransferConsensusOutcome consensusOutcome; + if (mContextTransferCoordinator) + { + auto mergeCoordinatorOutcome = [&]() + { + auto coordinatorOutcome = mContextTransferCoordinator->poll(); + consensusOutcome.completedRequestIds.insert( + coordinatorOutcome.completedRequestIds.begin(), coordinatorOutcome.completedRequestIds.end()); + consensusOutcome.failedRequestIds.insert( + coordinatorOutcome.failedRequestIds.begin(), coordinatorOutcome.failedRequestIds.end()); + consensusOutcome.timedOutRequestIds.insert( + coordinatorOutcome.timedOutRequestIds.begin(), coordinatorOutcome.timedOutRequestIds.end()); + }; + do + { + mergeCoordinatorOutcome(); + if (blockAll + && consensusOutcome.completedRequestIds.size() + consensusOutcome.failedRequestIds.size() + < mSenderRequestsAwaitingConsensus.size()) + { + std::this_thread::yield(); + } + } while (blockAll + && consensusOutcome.completedRequestIds.size() + consensusOutcome.failedRequestIds.size() + < mSenderRequestsAwaitingConsensus.size()); + + if (inflightCancelEnabled) + { + // A timeout update is transmitted once, but cancellation may be declined transiently. Keep the globally + // observed timeout active so rank-local cancellation is retried on every poll until terminal commit. + consensusOutcome.timedOutRequestIds.insert(mTimedOutSenderIds.begin(), mTimedOutSenderIds.end()); + } + } + else + { + consensusOutcome = reduceTransferStates(syncComm, mGroupPipeParaComm, mCompletedSenderRequestIds, + mFailedSenderRequestIds, inflightCancelEnabled ? mTimedOutSenderIds : std::unordered_set{}); + } if (inflightCancelEnabled) { for (auto const requestId : consensusOutcome.timedOutRequestIds) { + // Persist the global timeout even if this rank has not registered its local future yet. The one-shot + // coordinator update must remain actionable when that future appears on a later scheduler poll. + mTimedOutSenderIds.insert(requestId); auto const futureIt = std::find_if(mSenderFutures.begin(), mSenderFutures.end(), [requestId](auto const& entry) { return entry.first->mRequestId == requestId; }); if (futureIt == mSenderFutures.end() @@ -969,7 +1062,6 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( { continue; } - mTimedOutSenderIds.insert(requestId); if (requestCancellationNoThrow( requestId, "Context", [&]() { return mCacheSender->cancelRequest(*futureIt->first); })) { diff --git a/cpp/tensorrt_llm/batch_manager/contextTransferCoordinator.cpp b/cpp/tensorrt_llm/batch_manager/contextTransferCoordinator.cpp new file mode 100644 index 000000000000..909b7858c574 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/contextTransferCoordinator.cpp @@ -0,0 +1,475 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * 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. + */ + +#include "tensorrt_llm/batch_manager/contextTransferCoordinator.h" + +#include "tensorrt_llm/batch_manager/cacheTransceiver.h" +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/logger.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager +{ + +ContextTransferVoteReducer::ContextTransferVoteReducer(int const participantCount) + : mParticipantCount(participantCount) +{ + TLLM_CHECK_WITH_INFO(participantCount > 0, "Context-transfer consensus requires at least one participant."); +} + +void ContextTransferVoteReducer::recordVote( + int const participantRank, std::uint64_t const requestId, ContextTransferVote const vote) +{ + TLLM_CHECK_WITH_INFO(participantRank >= 0 && participantRank < mParticipantCount, + "Context-transfer consensus participant rank is out of range."); + TLLM_CHECK_WITH_INFO(vote == ContextTransferVote::kCompleted || vote == ContextTransferVote::kFailed, + "Context-transfer consensus received an invalid vote."); + + auto [requestIt, inserted] = mRequestVotes.try_emplace(requestId, mParticipantCount); + static_cast(inserted); + auto& requestVotes = requestIt->second; + auto& recordedVote = requestVotes.votes.at(static_cast(participantRank)); + auto const packedVote = static_cast(vote); + if (recordedVote != 0) + { + TLLM_CHECK_WITH_INFO( + recordedVote == packedVote, "Context-transfer participant changed its terminal vote for a request."); + return; + } + + recordedVote = packedVote; + ++requestVotes.terminalCount; + requestVotes.failed = requestVotes.failed || vote == ContextTransferVote::kFailed; +} + +void ContextTransferVoteReducer::recordTimeout(std::uint64_t const requestId) +{ + auto [requestIt, inserted] = mRequestVotes.try_emplace(requestId, mParticipantCount); + static_cast(inserted); + auto& requestVotes = requestIt->second; + if (!requestVotes.timedOut) + { + requestVotes.timedOut = true; + requestVotes.timeoutPending = true; + } +} + +ContextTransferConsensusResult ContextTransferVoteReducer::takeReady() +{ + ContextTransferConsensusResult result; + for (auto requestIt = mRequestVotes.begin(); requestIt != mRequestVotes.end();) + { + auto& requestVotes = requestIt->second; + if (requestVotes.timeoutPending) + { + result.timedOutRequestIds.insert(requestIt->first); + requestVotes.timeoutPending = false; + } + if (requestVotes.terminalCount != mParticipantCount) + { + ++requestIt; + continue; + } + + auto& terminalRequestIds + = (requestVotes.failed || requestVotes.timedOut) ? result.failedRequestIds : result.completedRequestIds; + terminalRequestIds.insert(requestIt->first); + requestIt = mRequestVotes.erase(requestIt); + } + return result; +} + +void ContextTransferVoteReducer::clear() noexcept +{ + mRequestVotes.clear(); +} + +class ContextTransferCoordinator::Impl +{ +public: + explicit Impl(std::shared_ptr comm) + : mComm(std::move(comm)) + , mCoordinatorRank(0) + , mReducer(mComm ? mComm->getSize() : 1) + { + TLLM_CHECK_WITH_INFO(mComm != nullptr, "Context-transfer coordination requires a communicator."); + TLLM_CHECK_WITH_INFO(mComm->isMpi(), "Asynchronous context-transfer coordination requires MPI."); + TLLM_CHECK_WITH_INFO( + mComm->getSize() > 1, "Asynchronous context-transfer coordination requires multiple participants."); + mCoordinatorRank = mComm->getSize() - 1; + } + + void publishLocalOutcome(std::uint64_t const requestId, bool const failed) + { + TLLM_CHECK_WITH_INFO(!mShutdown, "Cannot publish a context-transfer vote after coordinator shutdown."); + auto const vote = failed ? ContextTransferVote::kFailed : ContextTransferVote::kCompleted; + auto const [voteIt, inserted] = mPublishedLocalVotes.emplace(requestId, vote); + TLLM_CHECK_WITH_INFO( + inserted || voteIt->second == vote, "This rank changed its terminal vote for a context transfer."); + if (!inserted) + { + return; + } + + try + { + if (isCoordinator()) + { + mReducer.recordVote(mComm->getRank(), requestId, vote); + } + else + { + queuePacket( + requestId, static_cast(vote), mCoordinatorRank, mpi::MpiTag::kContextTransferEvent); + } + } + catch (...) + { + mPublishedLocalVotes.erase(voteIt); + throw; + } + } + + void publishTimeout(std::uint64_t const requestId) + { + TLLM_CHECK_WITH_INFO(!mShutdown, "Cannot publish a context-transfer timeout after coordinator shutdown."); + TLLM_CHECK_WITH_INFO(mPublishedLocalVotes.find(requestId) == mPublishedLocalVotes.end(), + "A context-transfer timeout must be published before its immutable terminal vote."); + auto const [timeoutIt, inserted] = mPublishedTimeouts.insert(requestId); + if (!inserted) + { + return; + } + + try + { + if (isCoordinator()) + { + mReducer.recordTimeout(requestId); + } + else + { + queuePacket(requestId, kTimedOutSignal, mCoordinatorRank, mpi::MpiTag::kContextTransferEvent); + } + } + catch (...) + { + mPublishedTimeouts.erase(timeoutIt); + throw; + } + } + + [[nodiscard]] ContextTransferConsensusResult poll() + { + TLLM_CHECK_WITH_INFO(!mShutdown, "Cannot poll context-transfer coordination after shutdown."); + return progress(); + } + + void shutdown() noexcept + { + if (mShutdown) + { + return; + } + mShutdown = true; + + try + { + if (!isCoordinator()) + { + queuePacket(/*requestId=*/0, kCloseMarker, mCoordinatorRank, mpi::MpiTag::kContextTransferEvent); + } + + auto const deadline = std::chrono::steady_clock::now() + kShutdownTimeout; + while (!shutdownComplete()) + { + static_cast(progress()); + if (!shutdownComplete()) + { + if (std::chrono::steady_clock::now() >= deadline) + { + TLLM_LOG_ERROR( + "Timed out shutting down asynchronous context-transfer coordinator; rank=%d " + "peer_closes=%zu/%d ack=%d pending_sends=%zu. Aborting to avoid freeing active MPI " + "requests.", + mComm->getRank(), mClosedPeers.size(), mComm->getSize() - 1, mCloseAcknowledged, + mPendingSends.size()); + std::abort(); + } + std::this_thread::yield(); + } + } + } + catch (std::exception const& error) + { + TLLM_LOG_ERROR("Failed to shut down asynchronous context-transfer coordinator: %s", error.what()); + std::abort(); + } + catch (...) + { + TLLM_LOG_ERROR("Failed to shut down asynchronous context-transfer coordinator with an unknown error."); + std::abort(); + } + } + +private: + static constexpr std::size_t kPacketFieldCount = 2; + static constexpr std::uint64_t kTimedOutSignal = 3; + static constexpr std::uint64_t kCloseMarker = 4; + static constexpr auto kShutdownTimeout = std::chrono::seconds(30); + + struct PendingSend + { + std::array packet{}; + std::unique_ptr request; + }; + + [[nodiscard]] bool isCoordinator() const + { + return mComm->getRank() == mCoordinatorRank; + } + + void queuePacket(std::uint64_t const requestId, std::uint64_t const value, int const peer, mpi::MpiTag const tag) + { + mPendingSends.emplace_back(); + auto& pendingSend = mPendingSends.back(); + pendingSend.packet = {requestId, value}; + try + { + pendingSend.request = mComm->sendAsync( + pendingSend.packet.data(), pendingSend.packet.size(), mpi::MpiType::kUINT64, peer, tag); + } + catch (...) + { + mPendingSends.pop_back(); + throw; + } + } + + void queueUpdateForPeers(std::uint64_t const requestId, std::uint64_t const update) + { + for (int peer = 0; peer < mComm->getSize(); ++peer) + { + if (peer != mCoordinatorRank) + { + queuePacket(requestId, update, peer, mpi::MpiTag::kContextTransferUpdate); + } + } + } + + void reapCompletedSends() + { + for (auto sendIt = mPendingSends.begin(); sendIt != mPendingSends.end();) + { + TLLM_CHECK(sendIt->request); + if (sendIt->request->isCompleted()) + { + sendIt = mPendingSends.erase(sendIt); + } + else + { + ++sendIt; + } + } + } + + void drainVotes() + { + for (int peer = 0; peer < mComm->getSize(); ++peer) + { + if (peer == mCoordinatorRank) + { + continue; + } + + MPI_Status status{}; + while (mComm->iprobe(peer, mpi::MpiTag::kContextTransferEvent, &status)) + { + std::array packet{}; + mComm->recv( + packet.data(), packet.size(), mpi::MpiType::kUINT64, peer, mpi::MpiTag::kContextTransferEvent); + if (packet.back() == kCloseMarker) + { + TLLM_CHECK_WITH_INFO( + mClosedPeers.insert(peer).second, "Received a duplicate context-transfer close marker."); + continue; + } + TLLM_CHECK_WITH_INFO(mClosedPeers.find(peer) == mClosedPeers.end(), + "Received a context-transfer vote after its peer close marker."); + if (packet.back() == kTimedOutSignal) + { + mReducer.recordTimeout(packet.front()); + continue; + } + mReducer.recordVote(peer, packet.front(), static_cast(packet.back())); + } + } + } + + ContextTransferConsensusResult completeCoordinatorUpdates() + { + auto result = mReducer.takeReady(); + for (auto const requestId : result.timedOutRequestIds) + { + queueUpdateForPeers(requestId, kTimedOutSignal); + } + for (auto const requestId : result.failedRequestIds) + { + queueUpdateForPeers(requestId, static_cast(ContextTransferVote::kFailed)); + mPublishedLocalVotes.erase(requestId); + mPublishedTimeouts.erase(requestId); + } + for (auto const requestId : result.completedRequestIds) + { + queueUpdateForPeers(requestId, static_cast(ContextTransferVote::kCompleted)); + mPublishedLocalVotes.erase(requestId); + mPublishedTimeouts.erase(requestId); + } + return result; + } + + ContextTransferConsensusResult drainUpdates() + { + ContextTransferConsensusResult result; + MPI_Status status{}; + while (mComm->iprobe(mCoordinatorRank, mpi::MpiTag::kContextTransferUpdate, &status)) + { + std::array packet{}; + mComm->recv(packet.data(), packet.size(), mpi::MpiType::kUINT64, mCoordinatorRank, + mpi::MpiTag::kContextTransferUpdate); + if (packet.back() == kCloseMarker) + { + TLLM_CHECK_WITH_INFO(!mCloseAcknowledged, "Received a duplicate coordinator close marker."); + mCloseAcknowledged = true; + mPublishedLocalVotes.clear(); + mPublishedTimeouts.clear(); + continue; + } + + if (packet.back() == kTimedOutSignal) + { + result.timedOutRequestIds.insert(packet.front()); + continue; + } + + auto const outcome = static_cast(packet.back()); + TLLM_CHECK_WITH_INFO(outcome == ContextTransferVote::kCompleted || outcome == ContextTransferVote::kFailed, + "Received an invalid context-transfer commit outcome."); + auto const localVoteIt = mPublishedLocalVotes.find(packet.front()); + TLLM_CHECK_WITH_INFO(localVoteIt != mPublishedLocalVotes.end(), + "Received a context-transfer commit before publishing the local terminal vote."); + if (outcome == ContextTransferVote::kFailed) + { + result.failedRequestIds.insert(packet.front()); + } + else + { + result.completedRequestIds.insert(packet.front()); + } + mPublishedLocalVotes.erase(localVoteIt); + mPublishedTimeouts.erase(packet.front()); + } + return result; + } + + ContextTransferConsensusResult progress() + { + reapCompletedSends(); + if (isCoordinator()) + { + drainVotes(); + auto result = completeCoordinatorUpdates(); + if (mShutdown && !mCloseSent && mClosedPeers.size() == static_cast(mComm->getSize() - 1)) + { + // Shutdown is an explicit abort epoch. A peer close proves that no more votes will arrive from that + // peer, so incomplete requests cannot reach a global decision and are intentionally abandoned. + mReducer.clear(); + mPublishedLocalVotes.clear(); + mPublishedTimeouts.clear(); + for (int peer = 0; peer < mCoordinatorRank; ++peer) + { + queuePacket(/*requestId=*/0, kCloseMarker, peer, mpi::MpiTag::kContextTransferUpdate); + } + mCloseSent = true; + } + return result; + } + return drainUpdates(); + } + + [[nodiscard]] bool shutdownComplete() const + { + if (isCoordinator()) + { + return mCloseSent && mPendingSends.empty(); + } + return mCloseAcknowledged && mPendingSends.empty(); + } + + std::shared_ptr mComm; + int mCoordinatorRank; + ContextTransferVoteReducer mReducer; + std::unordered_map mPublishedLocalVotes; + std::unordered_set mPublishedTimeouts; + std::unordered_set mClosedPeers; + std::list mPendingSends; + bool mShutdown{false}; + bool mCloseSent{false}; + bool mCloseAcknowledged{false}; +}; + +ContextTransferCoordinator::ContextTransferCoordinator(std::shared_ptr comm) + : mImpl(std::make_unique(std::move(comm))) +{ +} + +ContextTransferCoordinator::~ContextTransferCoordinator() +{ + shutdown(); +} + +void ContextTransferCoordinator::publishLocalOutcome(std::uint64_t const requestId, bool const failed) +{ + mImpl->publishLocalOutcome(requestId, failed); +} + +void ContextTransferCoordinator::publishTimeout(std::uint64_t const requestId) +{ + mImpl->publishTimeout(requestId); +} + +ContextTransferConsensusResult ContextTransferCoordinator::poll() +{ + return mImpl->poll(); +} + +void ContextTransferCoordinator::shutdown() noexcept +{ + if (mImpl) + { + mImpl->shutdown(); + } +} + +} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index a8f6c7dd4e8f..75329d192a8c 100644 --- a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt +++ b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt @@ -20,6 +20,7 @@ add_gtest(cacheTransBufferTest cacheTransBufferTest.cpp) add_gtest(bufferIndexHolderTest bufferIndexHolderTest.cpp) add_gtest(capacitySchedulerTest capacitySchedulerTest.cpp) add_gtest(contextProgressTest contextProgressTest.cu) +add_gtest(contextTransferCoordinatorTest contextTransferCoordinatorTest.cpp) add_gtest(evictionPolicyTest evictionPolicyTest.cpp) add_gtest(kvCacheManagerTest kvCacheManagerTest.cpp) add_gtest(kvCacheManagerFabricMemoryTest kvCacheManagerFabricMemoryTest.cpp) diff --git a/cpp/tests/unit_tests/batch_manager/contextTransferCoordinatorTest.cpp b/cpp/tests/unit_tests/batch_manager/contextTransferCoordinatorTest.cpp new file mode 100644 index 000000000000..7b0f7aba3fc4 --- /dev/null +++ b/cpp/tests/unit_tests/batch_manager/contextTransferCoordinatorTest.cpp @@ -0,0 +1,143 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * 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. + */ + +#include "tensorrt_llm/batch_manager/contextTransferCoordinator.h" + +#include + +namespace tensorrt_llm::batch_manager +{ +namespace +{ + +TEST(ContextTransferVoteReducerTest, WaitsForEveryParticipant) +{ + ContextTransferVoteReducer reducer(4); + reducer.recordVote(0, 17, ContextTransferVote::kCompleted); + reducer.recordVote(1, 17, ContextTransferVote::kCompleted); + reducer.recordVote(2, 17, ContextTransferVote::kCompleted); + EXPECT_TRUE(reducer.takeReady().completedRequestIds.empty()); + + reducer.recordVote(3, 17, ContextTransferVote::kCompleted); + auto const result = reducer.takeReady(); + EXPECT_EQ(result.completedRequestIds, std::unordered_set{17}); + EXPECT_TRUE(result.failedRequestIds.empty()); + EXPECT_TRUE(result.timedOutRequestIds.empty()); +} + +TEST(ContextTransferVoteReducerTest, FailureWinsAfterEveryParticipantVotes) +{ + ContextTransferVoteReducer reducer(4); + reducer.recordVote(0, 23, ContextTransferVote::kCompleted); + reducer.recordVote(1, 23, ContextTransferVote::kFailed); + reducer.recordVote(2, 23, ContextTransferVote::kCompleted); + reducer.recordVote(3, 23, ContextTransferVote::kCompleted); + + auto const result = reducer.takeReady(); + EXPECT_TRUE(result.completedRequestIds.empty()); + EXPECT_EQ(result.failedRequestIds, std::unordered_set{23}); +} + +TEST(ContextTransferVoteReducerTest, RejectsChangedTerminalVote) +{ + ContextTransferVoteReducer reducer(2); + reducer.recordVote(0, 29, ContextTransferVote::kCompleted); + reducer.recordVote(0, 29, ContextTransferVote::kCompleted); + EXPECT_ANY_THROW(reducer.recordVote(0, 29, ContextTransferVote::kFailed)); +} + +TEST(ContextTransferVoteReducerTest, AccumulatesInterleavedRequestsIndependently) +{ + ContextTransferVoteReducer reducer(2); + reducer.recordVote(0, 31, ContextTransferVote::kCompleted); + reducer.recordVote(1, 37, ContextTransferVote::kFailed); + reducer.recordVote(1, 31, ContextTransferVote::kCompleted); + + auto result = reducer.takeReady(); + EXPECT_EQ(result.completedRequestIds, std::unordered_set{31}); + EXPECT_TRUE(result.failedRequestIds.empty()); + + reducer.recordVote(0, 37, ContextTransferVote::kCompleted); + result = reducer.takeReady(); + EXPECT_TRUE(result.completedRequestIds.empty()); + EXPECT_EQ(result.failedRequestIds, std::unordered_set{37}); +} + +TEST(ContextTransferVoteReducerTest, RejectsInvalidInput) +{ + EXPECT_ANY_THROW(ContextTransferVoteReducer(0)); + ContextTransferVoteReducer reducer(2); + EXPECT_ANY_THROW(reducer.recordVote(-1, 41, ContextTransferVote::kCompleted)); + EXPECT_ANY_THROW(reducer.recordVote(2, 41, ContextTransferVote::kCompleted)); + EXPECT_ANY_THROW(reducer.recordVote(0, 41, static_cast(99))); +} + +TEST(ContextTransferVoteReducerTest, TimeoutIsVisibleBeforeTerminalFailure) +{ + ContextTransferVoteReducer reducer(2); + reducer.recordTimeout(43); + + auto result = reducer.takeReady(); + EXPECT_EQ(result.timedOutRequestIds, std::unordered_set{43}); + EXPECT_TRUE(result.completedRequestIds.empty()); + EXPECT_TRUE(result.failedRequestIds.empty()); + + result = reducer.takeReady(); + EXPECT_TRUE(result.timedOutRequestIds.empty()); + EXPECT_TRUE(result.failedRequestIds.empty()); + + reducer.recordVote(0, 43, ContextTransferVote::kCompleted); + reducer.recordVote(1, 43, ContextTransferVote::kCompleted); + result = reducer.takeReady(); + EXPECT_EQ(result.failedRequestIds, std::unordered_set{43}); + EXPECT_TRUE(result.completedRequestIds.empty()); + EXPECT_TRUE(result.timedOutRequestIds.empty()); +} + +TEST(ContextTransferVoteReducerTest, DuplicateTimeoutIsIdempotent) +{ + ContextTransferVoteReducer reducer(2); + reducer.recordTimeout(47); + reducer.recordTimeout(47); + + auto result = reducer.takeReady(); + EXPECT_EQ(result.timedOutRequestIds, std::unordered_set{47}); + + reducer.recordTimeout(47); + result = reducer.takeReady(); + EXPECT_TRUE(result.timedOutRequestIds.empty()); +} + +TEST(ContextTransferVoteReducerTest, TimeoutAfterPartialTerminalVotesStillWins) +{ + ContextTransferVoteReducer reducer(3); + reducer.recordVote(2, 53, ContextTransferVote::kCompleted); + reducer.recordTimeout(53); + reducer.recordVote(0, 53, ContextTransferVote::kCompleted); + + auto result = reducer.takeReady(); + EXPECT_EQ(result.timedOutRequestIds, std::unordered_set{53}); + EXPECT_TRUE(result.failedRequestIds.empty()); + + reducer.recordVote(1, 53, ContextTransferVote::kCompleted); + result = reducer.takeReady(); + EXPECT_EQ(result.failedRequestIds, std::unordered_set{53}); + EXPECT_TRUE(result.completedRequestIds.empty()); +} + +} // namespace +} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp index e294cdff9d49..5d1a958d37ce 100644 --- a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp @@ -31,6 +31,7 @@ #include "tensorrt_llm/batch_manager/cacheFormatter.h" #include "tensorrt_llm/batch_manager/cacheTransceiver.h" +#include "tensorrt_llm/batch_manager/contextTransferCoordinator.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" @@ -43,6 +44,7 @@ #include "tensorrt_llm/runtime/common.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/testing/kvCacheManagerTestUtil.h" +#include #include #include #include @@ -54,6 +56,7 @@ #include #include #include +#include #include "gtest/gtest.h" #include @@ -90,6 +93,95 @@ T serializeDeserialize(T const& val) } // namespace +TEST(ContextTransferCoordinatorTest, CommitsStaggeredSuccessAndFailureWithoutCollectivePolling) +{ + auto& world = tensorrt_llm::mpi::MpiComm::world(); + if (world.getSize() < 2) + { + GTEST_SKIP() << "mpirun with at least two processes is required to run this test."; + } + + auto comm = std::make_shared(std::addressof(world)); + { + ContextTransferCoordinator coordinator(comm); + auto waitForOutcome = [&](std::uint64_t const requestId, bool const expectFailure) + { + bool observed = false; + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds(10); + while (!observed && std::chrono::steady_clock::now() < deadline) + { + auto result = coordinator.poll(); + observed = expectFailure ? result.failedRequestIds.count(requestId) != 0 + : result.completedRequestIds.count(requestId) != 0; + if (!observed) + { + std::this_thread::yield(); + } + } + EXPECT_TRUE(observed); + }; + auto waitForTimeout = [&](std::uint64_t const requestId) + { + bool observed = false; + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds(10); + while (!observed && std::chrono::steady_clock::now() < deadline) + { + auto const result = coordinator.poll(); + observed = result.timedOutRequestIds.count(requestId) != 0; + if (!observed) + { + std::this_thread::yield(); + } + } + EXPECT_TRUE(observed); + }; + + constexpr std::uint64_t kCompletedRequestId = 644815201; + world.barrier(); + if (world.getRank() == world.getSize() - 1) + { + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + coordinator.publishLocalOutcome(kCompletedRequestId, /*failed=*/false); + if (world.getRank() != world.getSize() - 1) + { + auto const earlyResult = coordinator.poll(); + EXPECT_TRUE(earlyResult.completedRequestIds.empty()); + EXPECT_TRUE(earlyResult.failedRequestIds.empty()); + } + waitForOutcome(kCompletedRequestId, /*expectFailure=*/false); + + constexpr std::uint64_t kFailedRequestId = 644815202; + world.barrier(); + coordinator.publishLocalOutcome(kFailedRequestId, /*failed=*/world.getRank() == 0); + waitForOutcome(kFailedRequestId, /*expectFailure=*/true); + + constexpr std::uint64_t kTimedOutRequestId = 644815203; + world.barrier(); + if (world.getRank() == 0) + { + coordinator.publishTimeout(kTimedOutRequestId); + coordinator.publishTimeout(kTimedOutRequestId); + } + waitForTimeout(kTimedOutRequestId); + coordinator.publishLocalOutcome(kTimedOutRequestId, /*failed=*/false); + waitForOutcome(kTimedOutRequestId, /*expectFailure=*/true); + + constexpr std::uint64_t kInvalidOrderingRequestId = 644815204; + world.barrier(); + coordinator.publishLocalOutcome(kInvalidOrderingRequestId, /*failed=*/false); + EXPECT_ANY_THROW(coordinator.publishTimeout(kInvalidOrderingRequestId)); + waitForOutcome(kInvalidOrderingRequestId, /*expectFailure=*/false); + + // Exercise asymmetric but orderly teardown after every decision has committed. + if (world.getRank() == 0) + { + std::this_thread::sleep_for(std::chrono::milliseconds(25)); + } + } + world.barrier(); +} + class RequestInfoTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro-type-member-init) { public: diff --git a/cpp/tests/unit_tests/multi_gpu/mpiUtilsTest.cpp b/cpp/tests/unit_tests/multi_gpu/mpiUtilsTest.cpp index c8a04f658409..941d7cb53ed6 100644 --- a/cpp/tests/unit_tests/multi_gpu/mpiUtilsTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/mpiUtilsTest.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -24,6 +24,8 @@ #endif // ENABLE_MULTI_DEVICE #include +#include +#include namespace mpi = tensorrt_llm::mpi; namespace tr = tensorrt_llm::runtime; @@ -46,6 +48,29 @@ TEST(MPIUtils, WorldRankAndSize) EXPECT_LE(rank, size); } +#if ENABLE_MULTI_DEVICE +TEST(MPIUtils, AsyncSendCanBePolled) +{ + auto& comm = mpi::MpiComm::world(); + auto const rank = comm.getRank(); + auto const size = comm.getSize(); + auto const destination = (rank + 1) % size; + auto const source = (rank + size - 1) % size; + std::uint64_t const sentValue = static_cast(rank); + std::uint64_t receivedValue = 0; + + auto request + = comm.sendAsync(&sentValue, 1, mpi::MpiType::kUINT64, destination, mpi::MpiTag::kContextTransferEvent); + static_cast(comm.recv(&receivedValue, 1, mpi::MpiType::kUINT64, source, mpi::MpiTag::kContextTransferEvent)); + while (!request->isCompleted()) + { + std::this_thread::yield(); + } + + EXPECT_EQ(receivedValue, static_cast(source)); +} +#endif // ENABLE_MULTI_DEVICE + template void testBroadcast() { From 80b4eb3675383ba2a7d3e85ffed398858a931116 Mon Sep 17 00:00:00 2001 From: Aswin Visva <31215515+aswinvisva@users.noreply.github.com> Date: Thu, 23 Jul 2026 12:49:42 -0700 Subject: [PATCH 4/4] [None][perf] Optimize fusedQKNormRope Kernel (#16633) Signed-off-by: Aswin Visva <31215515+aswinvisva@users.noreply.github.com> --- .../kernels/fusedQKNormRopeKernel.cu | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu index e06b0f200e4b..6e26dfce65bf 100644 --- a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu +++ b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -174,6 +174,17 @@ __global__ void fusedQKNormRopeKernel( // pos_id is selected per rotary half-dim (interleaved mRoPE); for plain RoPE // selectMRopePosId always returns position_ids[tokenIdx]. + // Hoist log2(base) and the loop-invariant constants out of the per-thread + // per-elem loop. powf(base, -2*hd/rd) == exp2f(-2*hd/rd * log2(base)); base + // and rotary_dim are kernel-uniform, so one MUFU.LG2 per warp instead of + // one per (thread, iter). Uses the fast __log2f intrinsic (a few ULPs of + // error, absorbed by bf16 downcast at store time). + float const neg2_log2base_over_rd = -2.0f * __log2f(base) / static_cast(rotary_dim); + // rotary_dim is even by contract; when it's also a power of 2 (always in + // practice — 64/128/256) '% rotary_dim' becomes '& (rotary_dim - 1)'. + // The bool is warp-uniform → predicated select, no branch divergence. + int const rd_mask = rotary_dim - 1; + bool const rd_is_pow2 = ((rotary_dim & rd_mask) == 0); // TODO: cos sin calculation could be halved. if constexpr (interleave) { @@ -191,7 +202,7 @@ __global__ void fusedQKNormRopeKernel( int dim_idx = laneId * numElemsPerThread + i; int half_dim = dim_idx / 2; - float freq = powf(base, -2.0f * half_dim / static_cast(rotary_dim)); + float freq = exp2f(static_cast(half_dim) * neg2_log2base_over_rd); if (factor != 1.0f) { @@ -232,9 +243,9 @@ __global__ void fusedQKNormRopeKernel( } int dim_idx = laneId * numElemsPerThread + i; - dim_idx = (dim_idx * 2) % rotary_dim; + dim_idx = rd_is_pow2 ? ((dim_idx * 2) & rd_mask) : ((dim_idx * 2) % rotary_dim); int half_dim = dim_idx / 2; - float freq = powf(base, -2.0f * half_dim / static_cast(rotary_dim)); + float freq = exp2f(static_cast(half_dim) * neg2_log2base_over_rd); if (factor != 1.0f) {