Skip to content

Commit aef7d47

Browse files
[TRTLLM-13141][feat] Add backend-agnostic SourceIdentity gate for weight sharing (#14878)
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
1 parent ccc0708 commit aef7d47

11 files changed

Lines changed: 1571 additions & 165 deletions

File tree

tensorrt_llm/_torch/memory/gpu_memory_backend.py

Lines changed: 94 additions & 73 deletions
Large diffs are not rendered by default.

tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py

Lines changed: 122 additions & 53 deletions
Large diffs are not rendered by default.

tensorrt_llm/_torch/pyexecutor/model_loader.py

Lines changed: 95 additions & 38 deletions
Large diffs are not rendered by default.
Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
"""Backend-agnostic weight-sharing utilities (MX, GMS, ...)."""
16+
17+
from tensorrt_llm._torch.weight_sharing.source_identity import (
18+
SOURCE_IDENTITY_FORMAT_VERSION,
19+
IdentityCheckDecision,
20+
IdentityCheckPolicy,
21+
IdentityMatchResult,
22+
SourceIdentity,
23+
SourceIdentityMismatchError,
24+
check_weight_sharing_compatibility,
25+
)
26+
27+
__all__ = [
28+
"SOURCE_IDENTITY_FORMAT_VERSION",
29+
"SourceIdentity",
30+
"IdentityMatchResult",
31+
"IdentityCheckPolicy",
32+
"IdentityCheckDecision",
33+
"SourceIdentityMismatchError",
34+
"check_weight_sharing_compatibility",
35+
]

0 commit comments

Comments
 (0)