[TRTLLM-16749][feat] Add MiniMax-H3 text encoder TP support

docs/source/models/supported-models.md

@@ -280,15 +280,15 @@ For full documentation, see the [Visual Generation](./visual-generation.md) page
 
 ### Feature Matrix
 
-| Model | FP8 blockwise | NVFP4 | TeaCache | CFG Parallelism | Ulysses Parallelism | Parallel VAE | CUDA Graph | torch.compile | trtllm-serve | Attention2D | Ring Attention | Tensor Parallelism |
-|---|---|---|---|---|---|---|---|---|---|--|--|--|
-| **FLUX.1** | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes |
-| **FLUX.2** | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes |
-| **Wan 2.1** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes |
-| **Wan 2.2** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes |
-| **LTX-2** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No |
-| **MiniMax-H3** | Yes | Yes | No | No | Yes | Yes | No | Yes | Yes | No | No | Yes |
-| **Qwen-Image** | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No |
-| **Qwen-Image-Layered** | Yes | No | No | No | No | No | Yes | Yes | Yes | No | No | No |
-| **Qwen-Image-Edit-2511** | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Yes | No | No | No |
-| **Cosmos3** | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes |
+| Model | FP8 blockwise | NVFP4 | TeaCache | CFG Parallelism | Ulysses Parallelism | Parallel VAE | CUDA Graph | torch.compile | trtllm-serve | Attention2D | Ring Attention | Tensor Parallelism | Text Encoder TP |
+|---|---|---|---|---|---|---|---|---|---|--|--|--|---|
+| **FLUX.1** | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No |
+| **FLUX.2** | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No |
+| **Wan 2.1** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No |
+| **Wan 2.2** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No |
+| **LTX-2** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No |
+| **MiniMax-H3** | Yes | Yes | No | No | Yes | Yes | No | Yes | Yes | No | No | Yes | Yes |
+| **Qwen-Image** | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No | No |
+| **Qwen-Image-Layered** | Yes | No | No | No | No | No | Yes | Yes | Yes | No | No | No | No |
+| **Qwen-Image-Edit-2511** | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Yes | No | No | No | No |
+| **Cosmos3** | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | No |

docs/source/models/visual-generation.md

@@ -56,22 +56,22 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models
 
 ### Feature Matrix
 
-| Model | FP8 blockwise | NVFP4 | TeaCache | Cache-DiT | CPU Offloading | CFG Parallelism | Ulysses Parallelism | Parallel VAE | CUDA Graph | torch.compile | trtllm-serve | Attention2D | Ring Attention | Tensor Parallelism\* | SOL-Attn |
-|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
-| **FLUX.1** | Yes | Yes | Yes | Yes | No | No | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No |
-| **FLUX.2** | Yes | Yes | Yes | Yes | No | No | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No |
-| **Wan 2.1** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes |
-| **Wan 2.1 VSA** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | Yes |
-| **Wan 2.2** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes |
-| **FastWan 2.2** | Yes | Yes | No | No | No | No | No | No | Yes | Yes | Yes | No | No | No | No |
-| **LTX-2** | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | Yes |
-| **MiniMax-H3** | Yes | Yes | No | No | No | No | Yes | Yes | No | Yes | Yes | No | No | Yes | No |
-| **Qwen-Image** | Yes | Yes | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No | No |
-| **Qwen-Image-Layered** | Yes | No | No | Yes | No | No | No | No | Yes | Yes | Yes | No | No | No | No |
-| **Qwen-Image-Edit-2511** | Yes | Yes | No | No | No | Yes | Yes | No | Yes | Yes | Yes | No | No | No | No |
-| **Cosmos3** | Yes | Yes | No | No | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | No |
-| **HunyuanVideo 1.5** | Yes | Yes | No | No | No | No | No | No | No | No | Yes | No | No | No | No |
-| **GlmImage** | Yes | Yes | No | No | No | No | No | No | No | No | Yes | No | No | No | No |
+| Model | FP8 blockwise | NVFP4 | TeaCache | Cache-DiT | CPU Offloading | CFG Parallelism | Ulysses Parallelism | Parallel VAE | CUDA Graph | torch.compile | trtllm-serve | Attention2D | Ring Attention | Tensor Parallelism\* | Text Encoder TP | SOL-Attn |
+|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
+| **FLUX.1** | Yes | Yes | Yes | Yes | No | No | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No | No |
+| **FLUX.2** | Yes | Yes | Yes | Yes | No | No | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No | No |
+| **Wan 2.1** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | Yes |
+| **Wan 2.1 VSA** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | No | Yes |
+| **Wan 2.2** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | Yes |
+| **FastWan 2.2** | Yes | Yes | No | No | No | No | No | No | Yes | Yes | Yes | No | No | No | No | No |
+| **LTX-2** | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes |
+| **MiniMax-H3** | Yes | Yes | No | No | No | No | Yes | Yes | No | Yes | Yes | No | No | Yes | Yes | No |
+| **Qwen-Image** | Yes | Yes | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No | No | No |
+| **Qwen-Image-Layered** | Yes | No | No | Yes | No | No | No | No | Yes | Yes | Yes | No | No | No | No | No |
+| **Qwen-Image-Edit-2511** | Yes | Yes | No | No | No | Yes | Yes | No | Yes | Yes | Yes | No | No | No | No | No |
+| **Cosmos3** | Yes | Yes | No | No | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | No | No |
+| **HunyuanVideo 1.5** | Yes | Yes | No | No | No | No | No | No | No | No | Yes | No | No | No | No | No |
+| **GlmImage** | Yes | Yes | No | No | No | No | No | No | No | No | Yes | No | No | No | No | No |
 
 \* Wan 2.1 and Wan 2.2 (not VSA) also support a token-sharded TP layout (prototype), enabled with `parallel_config.tp_layout: token_sharded`; see [Multi-GPU Parallelism](#multi-gpu-parallelism).
 
@@ -406,6 +406,10 @@ Configured under `VisualGenArgs.parallel_config`. Modes can be combined:
     - **Ring Attention** (`ring_size: N`): Shards the sequence axis across a 1D ring of `N` ranks, streaming K/V blocks (CP degree = `N`; total SP degree = `N · ulysses_size`; mutually exclusive with Attention2D).
 - **Tensor Parallelism** (`tp_size: N`): Splits attention heads and transformer MLPs across GPUs for faster compute and reduced memory usage.
     - **Token-sharded TP** (`tp_layout: token_sharded` in `parallel_config`, prototype): Keeps the residual stream sharded by tokens across the TP group, so each all-reduce becomes a reduce-scatter plus an all-gather and the norms, modulation and quantization between projections run on `1/tp_size` of the tokens. Requires `tp_size > 1`; composes with CFG parallelism, not with Ulysses, Ring Attention, Attention2D or VSA. Supported for Wan 2.1 and Wan 2.2.
+- **Text encoder tensor parallelism** (`text_encoder_tp_size: N`): MiniMax-H3
+  only. Shards the Qwen3-VL text encoder over the first N ranks for text-only
+  prompt encoding, then broadcasts prompt embeddings to the remaining VisualGen
+  ranks. This does not enable DiT tensor parallelism for MiniMax-H3.
 
 For multi-node execution, VisualGen relies on the external launcher to terminate
 the remaining ranks when any rank exits. `torchrun` provides this behavior. With

tensorrt_llm/_torch/attention/attention.py

@@ -33,6 +33,16 @@
 from .rotary_embedding import MRotaryEmbedding, RotaryEmbedding
 
 
+def _inherit_process_group_attrs(mapping: Mapping, source: Mapping) -> Mapping:
+    """Preserve out-of-band process group routing on derived TP mappings."""
+    tp_group_pg = getattr(source, "_tp_group_pg", None)
+    if tp_group_pg is not None:
+        mapping._tp_group_pg = tp_group_pg
+    if getattr(source, "_use_torch_distributed_allreduce", False):
+        mapping._use_torch_distributed_allreduce = True
+    return mapping
+
+
 def extract_extra_attrs(layer_idx: str, attn_type: str):
     assert attn_type in ["mla", "attn"], "Invalid attention type"
     extra_attrs = get_model_extra_attrs()
@@ -614,6 +624,7 @@ def __init__(
             gpus_per_node=self.mapping.gpus_per_node,
             enable_attention_dp=self.mapping.enable_attention_dp,
         )
+        mapping = _inherit_process_group_attrs(mapping, self.mapping)
         self.tp_size = tp_size
         self.cp_size = cp_size
         self.tp_rank = mapping.tp_rank
@@ -672,6 +683,7 @@ def __init__(
             gpus_per_node=self.mapping.gpus_per_node,
             enable_attention_dp=self.mapping.enable_attention_dp,
         )
+        mapping_o = _inherit_process_group_attrs(mapping_o, self.mapping)
         self.mapping_o = mapping_o
 
         self.o_proj = Linear(

tensorrt_llm/_torch/device_mesh.py

@@ -79,7 +79,9 @@ class DeviceMeshTopologyImpl(_MappingBaseForTypeCheck):
     # attribute loads, matching MpiTopology's precomputed rank lists.
     @torch.compiler.disable
     def _resolve_mesh_dim(self, name: str) -> tuple:
-        pg = self._get_mesh_dim_by_name(name).get_group()
+        pg = getattr(self, '_tp_group_pg', None) if name == 'tp' else None
+        if pg is None:
+            pg = self._get_mesh_dim_by_name(name).get_group()
         group_name = pg.group_name if hasattr(pg, 'group_name') else ''
         resolved = (pg, pg.rank(), self._get_group_ranks(pg), group_name)
         # Before torch.distributed is initialized (single-process runs) the

tensorrt_llm/_torch/distributed/ops.py

@@ -1310,6 +1310,25 @@ def forward(
         if all_reduce_params is None:
             all_reduce_params = AllReduceParams()
 
+        use_torch_distributed_allreduce = self._disable_mpi and (
+            getattr(self.mapping, "_use_torch_distributed_allreduce", False)
+            or getattr(self.mapping, "_tp_group_pg", None) is not None)
+
+        if use_torch_distributed_allreduce:
+            if all_reduce_params.fusion_op != AllReduceFusionOp.NONE or any(
+                    value is not None for value in (
+                        all_reduce_params.residual,
+                        all_reduce_params.norm_weight,
+                        all_reduce_params.scale,
+                        all_reduce_params.bias,
+                    )):
+                raise NotImplementedError(
+                    "Torch distributed all-reduce fallback supports only "
+                    "plain all-reduce.")
+            output = input.clone()
+            torch.distributed.all_reduce(output, group=self.mapping.tp_group_pg)
+            return output
+
         # Try Symmetric Memory AllReduce first if available
         # Note: Currently only supports NONE fusion op (plain allreduce)
         if self.symm_mem_allreduce and all_reduce_params.fusion_op == AllReduceFusionOp.NONE:

tensorrt_llm/_torch/models/modeling_qwen3.py

@@ -215,6 +215,7 @@ def __init__(self,
             mapping=config.mapping,
             tensor_parallel_mode=TensorParallelMode.COLUMN,
             gather_output=True,
+            allreduce_strategy=config.allreduce_strategy,
         )
         self.layers = nn.ModuleList([
             Qwen3DecoderLayer(

tensorrt_llm/_torch/modules/embedding.py

@@ -8,7 +8,7 @@
 import torch.nn.functional as F
 from torch.nn.parameter import Parameter
 
-from tensorrt_llm.functional import AllReduceParams
+from tensorrt_llm.functional import AllReduceParams, AllReduceStrategy
 from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import Mapping
 from tensorrt_llm.math_utils import ceil_div
@@ -62,6 +62,7 @@ def __init__(
         use_custom_cublas_mm: bool = False,
         quant_config=None,
         output_dtype: Optional[torch.dtype] = None,
+        allreduce_strategy: AllReduceStrategy = AllReduceStrategy.AUTO,
     ):
         local_in_features = embedding_dim
         local_out_features = num_embeddings
@@ -105,6 +106,7 @@ def __init__(
             reduce_output=reduce_output,
             use_custom_cublas_mm=use_custom_cublas_mm,
             quant_config=quant_config,
+            allreduce_strategy=allreduce_strategy,
         )
 
         if tensor_parallel_mode == TensorParallelMode.ROW:
@@ -327,6 +329,7 @@ def __init__(
         reduce_output: bool = True,
         enable_torch_compile_for_embedding: Optional[bool] = False,
         use_custom_cublas_mm: bool = False,
+        allreduce_strategy: AllReduceStrategy = AllReduceStrategy.AUTO,
     ):
         super().__init__(
             embedding_dim=embedding_dim,
@@ -337,6 +340,7 @@ def __init__(
             gather_output=gather_output,
             reduce_output=reduce_output,
             use_custom_cublas_mm=use_custom_cublas_mm,
+            allreduce_strategy=allreduce_strategy,
         )
 
         self.enable_torch_compile_for_embedding = enable_torch_compile_for_embedding

tensorrt_llm/_torch/visual_gen/checkpoints/weight_loader.py

@@ -72,13 +72,16 @@ def load_weights(
         Args:
             checkpoint_dir: Path to checkpoint (pipeline root or component dir)
             mapping: Distributed mapping (for future TP/PP support)
+            kwargs: Optional loading controls. ``prefetch=False`` disables
+                safetensors host page-cache prefetch before loading.
 
         Returns:
             - If single component: Dict mapping weight names to tensors
             - If multiple components: Dict mapping component names to weight dicts
               Example: {"transformer": {...}, "transformer_2": {...}}
         """
         checkpoint_path = Path(checkpoint_dir)
+        prefetch = kwargs.get("prefetch", True)
 
         # Standard and modular Diffusers manifests use the same component
         # subdirectory layout for weight-bearing models.
@@ -109,7 +112,7 @@ def load_weights(
             if not weight_files:
                 raise ValueError(f"No weight files found in {weight_dir}")
 
-            if all(wf.endswith(".safetensors") for wf in weight_files):
+            if prefetch and all(wf.endswith(".safetensors") for wf in weight_files):
                 prefetch_files_to_host_cache(
                     weight_files,
                     description="visual-gen checkpoint",