Skip to content

Commit f1930d7

Browse files
committed
feat(minimax-h3): add native Windows RTX path
Build MiniMax H3 plans in isolated processes and stream them directly to disk so Windows builds and inference do not co-reside every model stage. Add portable shared-library loading, explicit Windows SDK inputs, file-backed TensorRT-RTX deserialization with plan integrity checks, staged denoiser/VAE residency, and public-safe documentation. Keep the standard TensorRT path unchanged and omit machine-identifying provenance.
1 parent b3e1277 commit f1930d7

49 files changed

Lines changed: 3435 additions & 490 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

CMakeLists.txt

Lines changed: 261 additions & 62 deletions
Large diffs are not rendered by default.

cmake/trtmc_pipeline_plugins.cmake

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,25 @@ if(TRTMC_MODEL_PROOF_MODEL)
5353
endif()
5454
endif()
5555

56+
if(TRTMC_RUNTIME_MODELS)
57+
set(_trtmc_selected_model_manifests)
58+
foreach(_trtmc_model IN LISTS TRTMC_RUNTIME_MODELS)
59+
if(NOT _trtmc_model MATCHES "^[A-Za-z0-9_.-]+$")
60+
message(FATAL_ERROR
61+
"TRTMC_RUNTIME_MODELS contains an unsafe model id: '${_trtmc_model}'")
62+
endif()
63+
set(_trtmc_selected_manifest
64+
"${PROJECT_SOURCE_DIR}/src/runtime/models/${_trtmc_model}/MODEL.toml")
65+
if(NOT EXISTS "${_trtmc_selected_manifest}")
66+
message(FATAL_ERROR
67+
"TRTMC_RUNTIME_MODELS requests unknown model '${_trtmc_model}'")
68+
endif()
69+
list(APPEND _trtmc_selected_model_manifests "${_trtmc_selected_manifest}")
70+
endforeach()
71+
list(REMOVE_DUPLICATES _trtmc_selected_model_manifests)
72+
set(TRTMC_RUNTIME_MODEL_MANIFESTS ${_trtmc_selected_model_manifests})
73+
endif()
74+
5675
set(TRTMC_RUNTIME_MODEL_IDS)
5776
foreach(_trtmc_model_manifest IN LISTS TRTMC_RUNTIME_MODEL_MANIFESTS)
5877
get_filename_component(_trtmc_model_dir "${_trtmc_model_manifest}" DIRECTORY)
@@ -71,7 +90,10 @@ foreach(_trtmc_model_manifest IN LISTS TRTMC_RUNTIME_MODEL_MANIFESTS)
7190

7291
_trtmc_model_manifest_string("${_trtmc_model_manifest_text}" "runtime_library"
7392
_trtmc_runtime_library)
74-
if(NOT _trtmc_runtime_library)
93+
if(WIN32)
94+
set(_trtmc_runtime_library
95+
"${CMAKE_SHARED_LIBRARY_PREFIX}trtmc_model_${_trtmc_model}${CMAKE_SHARED_LIBRARY_SUFFIX}")
96+
elseif(NOT _trtmc_runtime_library)
7597
set(_trtmc_runtime_library "libtrtmc_model_${_trtmc_model}.so")
7698
endif()
7799

include/trtmc/runtime/distributed_runtime.h

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,10 +22,9 @@ struct DistributedRuntimeGroup {
2222

2323
// Initialize an NCCL communicator for TensorRT 11.0+ distributed collective layers.
2424
//
25-
// This intentionally avoids compile-time MPI/NCCL dependencies: ranks are
26-
// discovered from common mpirun environment variables, and NCCL is loaded with
27-
// dlopen at runtime. Rank 0 writes the NCCL unique ID to a small rendezvous
28-
// file under /tmp unless TRTMC_NCCL_RENDEZVOUS points elsewhere.
25+
// Linux avoids compile-time MPI/NCCL dependencies: ranks are discovered from
26+
// common mpirun environment variables and NCCL is loaded at runtime. Native
27+
// Windows builds support only the single-device case and reject tp_size > 1.
2928
DistributedRuntimeGroup initialize_tensor_parallel_group(int tp_size);
3029

3130
} // namespace trtmc

include/trtmc/runtime/pipeline_plugin_loader.h

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,8 +37,9 @@ std::optional<std::string> legacy_runtime_strategy_alias_target(const std::strin
3737
const std::string& config_text);
3838

3939
// Load the model plugin that owns strategy. Search paths are directories that
40-
// contain libtrtmc_model_<model>.so; TRTMC_MODEL_PLUGIN_DIR and build/install
41-
// defaults are consulted after these explicit paths. When
40+
// contain the platform-native trtmc_model_<model> shared library;
41+
// TRTMC_MODEL_PLUGIN_DIR and build/install defaults are consulted after these
42+
// explicit paths. When
4243
// TRTMC_MODEL_PLUGIN_STRICT=1, only explicit paths and TRTMC_MODEL_PLUGIN_DIR
4344
// are used, so an installed or stale build-tree DSO cannot satisfy a CI proof.
4445
void load_model_plugin_for_strategy(const std::string& strategy,

python/tensorrt_model_connect/bundle_writer.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -178,8 +178,9 @@ def _open_atomic_bundle_output(destination: Path):
178178
except FileExistsError:
179179
continue
180180
try:
181-
if destination_mode is not None:
182-
os.fchmod(descriptor, destination_mode)
181+
fchmod = getattr(os, "fchmod", None)
182+
if destination_mode is not None and fchmod is not None:
183+
fchmod(descriptor, destination_mode)
183184
return os.fdopen(descriptor, "wb"), temporary_path
184185
except Exception:
185186
os.close(descriptor)

python/tensorrt_model_connect/engine_builder.py

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1722,6 +1722,35 @@ def _build_diffusion_bundle(
17221722
if "_transformer_config" in weights:
17231723
config.raw["_transformer_config"] = weights["_transformer_config"]
17241724

1725+
# A model may require isolated component processes and file-backed bundle
1726+
# sections when built for RTX. Keep that policy model-owned; ordinary TRT
1727+
# diffusion builds continue through build_components() below.
1728+
build_staged_bundle = getattr(plugin, "build_staged_bundle", None)
1729+
if rtx and callable(build_staged_bundle):
1730+
if fp8_scales or save_fp8_scales:
1731+
raise ValueError("Staged TensorRT-RTX builds do not support FP8 calibration")
1732+
staged_t0 = time.monotonic()
1733+
build_staged_bundle(
1734+
str(model_dir_path),
1735+
output_path,
1736+
config,
1737+
weights,
1738+
precision=precision,
1739+
verbose=verbose,
1740+
parallel_config=parallel,
1741+
max_batch_size=max_batch_size,
1742+
)
1743+
staged_elapsed = time.monotonic() - staged_t0
1744+
_add_build_timing(build_timing, "trt_compile_diffusion_components_s", staged_elapsed)
1745+
build_timing["total_s"] = time.monotonic() - t0
1746+
_write_build_timing(build_timing)
1747+
print(
1748+
f"[trtmc build] Staged RTX bundle saved: {output_path} "
1749+
f"[{staged_elapsed:.1f}s]",
1750+
file=sys.stderr,
1751+
)
1752+
return
1753+
17251754
# Prefer a family-provided scale asset before running live calibration.
17261755
if fp8_scales == "auto":
17271756
precomputed_fn = getattr(plugin, "fp8_precomputed_scales", None)

python/tensorrt_model_connect/families/minimax_h3/adaln_builder.py

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212

1313
import gc
1414
import sys
15+
from pathlib import Path
1516

1617
from tensorrt_model_connect import trt_compat
1718

@@ -44,20 +45,23 @@ def checkpoint_keys(
4445
return tuple(names)
4546

4647

48+
@op.cleanup_failed_build
4749
def build_adaln_precompute_engine(
4850
weights: dict,
4951
profile: MiniMaxH3Config,
5052
*,
5153
verbose: bool = False,
5254
consume_weights: bool = False,
5355
workspace_bytes: int | None = None,
54-
) -> bytes:
56+
weight_streaming: bool = False,
57+
output_path: str | Path | None = None,
58+
) -> bytes | dict[str, int | str]:
5559
profile.validate()
5660
logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING)
5761
builder = trt.Builder(logger)
5862
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
5963
config = builder.create_builder_config()
60-
op.configure_builder(config)
64+
op.configure_builder(config, weight_streaming=weight_streaming)
6165
op.configure_workspace(
6266
config,
6367
workspace_bytes,
@@ -128,14 +132,21 @@ def build_adaln_precompute_engine(
128132
f"timesteps={profile.max_timestep_count}",
129133
file=sys.stderr,
130134
)
135+
plan = None
136+
record = None
131137
try:
132-
plan = builder.build_serialized_network(network, config)
138+
if output_path is None:
139+
plan = builder.build_serialized_network(network, config)
140+
else:
141+
record = trt_compat.build_serialized_network_to_file(
142+
builder, network, config, output_path
143+
)
133144
finally:
134145
op.release_weight_buffers(network)
135146
if consume_weights:
136147
weights.clear()
137-
if plan is None:
148+
if output_path is None and plan is None:
138149
raise RuntimeError("TensorRT failed to build MiniMax-H3 AdaLN precompute engine")
139150
del network, config, builder
140151
gc.collect()
141-
return bytes(plan)
152+
return record if record is not None else bytes(plan)

python/tensorrt_model_connect/families/minimax_h3/config.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,13 @@
1818
DENOISER_DEFAULT_WORKSPACE_BYTES = 96 << 30
1919
VAE_TILE_DECODER_DEFAULT_WORKSPACE_BYTES = 96 << 30
2020

21+
# The RTX path builds each plan in a fresh process, so one conservative
22+
# workspace and runtime budget cover every stage without coupling the public
23+
# artifact to a particular workstation identity.
24+
RTX_STAGED_WORKSPACE_BYTES = 16 << 30
25+
RTX_WEIGHT_STREAMING_BUDGET_BYTES = 32 << 30
26+
RTX_CUDA_MAJOR = 12
27+
2128
DEFAULT_WORKSPACE_LIMIT_BYTES = {
2229
"text_encoder.plan": TEXT_ENCODER_DEFAULT_WORKSPACE_BYTES,
2330
"adaln_precompute.plan": ADALN_PRECOMPUTE_DEFAULT_WORKSPACE_BYTES,

0 commit comments

Comments
 (0)