[pull] main from NVIDIA:main

tensorrt_llm/_torch/models/modeling_bart.py

@@ -445,7 +445,6 @@ def __init__(self, model_config: ModelConfig[BartConfig]):
             dtype=config.torch_dtype,
             mapping=model_config.mapping,
             tensor_parallel_mode=TensorParallelMode.COLUMN,
-            gather_output=True,
         )
         self.embed_scale = (
             math.sqrt(config.d_model) if getattr(config, "scale_embedding", False) else 1.0
@@ -577,6 +576,14 @@ def load_weights(self, weights: Dict, **kwargs):
             weights, config, dtype=self.model_config.torch_dtype
         )
 
+        # __init__ aliases lm_head.weight to shared_embedding.weight when
+        # tie_word_embeddings=True, so checkpoints that omit lm_head.weight are
+        # handled correctly (lm_head picks up the loaded embedding automatically).
+        # When lm_head.weight is present in the checkpoint, break the alias so
+        # lm_head gets its own independent weight loaded from the checkpoint.
+        if "lm_head.weight" in weights:
+            self.lm_head.weight = nn.Parameter(torch.empty_like(self.lm_head.weight))
+
         for name, module in self.named_modules():
             if len(list(module.parameters(recurse=False))) == 0:
                 continue

tensorrt_llm/_torch/models/modeling_t5.py

@@ -755,7 +755,6 @@ def __init__(self, model_config: ModelConfig[T5Config]):
             dtype=config.torch_dtype,
             mapping=model_config.mapping,
             tensor_parallel_mode=TensorParallelMode.COLUMN,
-            gather_output=True,
         )
 
         self.encoder = T5Encoder(model_config)
@@ -908,6 +907,11 @@ def load_weights(self, weights: Dict, **kwargs):
         config = self.model_config.pretrained_config
         tllm_weights = _convert_hf_t5_weights(weights, config, dtype=self.model_config.torch_dtype)
 
+        # __init__ aliases lm_head.weight to shared_embedding.weight when
+        # tie_word_embeddings=True, so checkpoints that omit lm_head.weight are
+        # handled correctly (lm_head picks up the loaded embedding automatically).
+        # When lm_head.weight is present in the checkpoint, break the alias so
+        # lm_head gets its own independent weight loaded from the checkpoint.
         if "lm_head.weight" in weights:
             self.lm_head.weight = nn.Parameter(torch.empty_like(self.lm_head.weight))
 

tensorrt_llm/_torch/modules/cross_attention.py

@@ -24,6 +24,8 @@
 import torch
 from torch import nn
 
+from tensorrt_llm.mapping import Mapping
+
 from ..attention_backend import AttentionMetadata
 from ..attention_backend.interface import AttentionBackend, PredefinedAttentionMask
 from ..attention_backend.utils import create_attention
@@ -79,16 +81,48 @@ def __init__(
 
         self.mapping = config.mapping
         tp_size = self.mapping.tp_size
+        pp_size = self.mapping.pp_size
+        cp_size = self.mapping.cp_size
+        dp_size = 1
+        # Unreachable for enc-dec models today (validate_encoder_decoder_tp_scope
+        # rejects attention DP); kept for parity with Attention.__init__.
         if self.mapping.enable_attention_dp:
+            dp_size = tp_size
             tp_size = 1
 
         assert self.num_heads % tp_size == 0
+        if self.num_key_value_heads % tp_size != 0:
+            raise ValueError(
+                "Cross-attention requires the encoder KV head count "
+                f"({self.num_key_value_heads}) to be divisible by tp_size ({tp_size}); "
+                "KV head duplication is not supported for cross-attention."
+            )
+
         self.num_heads = self.num_heads // tp_size
         self.num_key_value_heads = (self.num_key_value_heads + tp_size - 1) // tp_size
         self.q_size = self.num_heads * self.head_dim
         self.kv_size = self.num_key_value_heads * self.head_dim
 
-        mapping = config.mapping
+        mapping = Mapping(
+            world_size=dp_size * tp_size * pp_size * cp_size,
+            tp_size=tp_size,
+            pp_size=pp_size * dp_size,
+            cp_size=cp_size,
+            cp_config=self.mapping.cp_config,
+            rank=self.mapping.rank,
+            gpus_per_node=self.mapping.gpus_per_node,
+            enable_attention_dp=self.mapping.enable_attention_dp,
+        )
+
+        mapping_o = Mapping(
+            world_size=dp_size * tp_size * pp_size * cp_size,
+            tp_size=tp_size * cp_size,
+            pp_size=pp_size * dp_size,
+            cp_size=1,
+            rank=self.mapping.rank,
+            gpus_per_node=self.mapping.gpus_per_node,
+            enable_attention_dp=self.mapping.enable_attention_dp,
+        )
 
         self.q_proj = Linear(
             self.hidden_size,
@@ -99,6 +133,7 @@ def __init__(
             tensor_parallel_mode=TensorParallelMode.COLUMN,
             quant_config=config.get_quant_config(),
             skip_create_weights_in_init=config.skip_create_weights_in_init,
+            allreduce_strategy=config.allreduce_strategy,
         )
 
         self.k_proj = Linear(
@@ -110,6 +145,7 @@ def __init__(
             tensor_parallel_mode=TensorParallelMode.COLUMN,
             quant_config=config.get_quant_config(),
             skip_create_weights_in_init=config.skip_create_weights_in_init,
+            allreduce_strategy=config.allreduce_strategy,
         )
 
         self.v_proj = Linear(
@@ -121,18 +157,20 @@ def __init__(
             tensor_parallel_mode=TensorParallelMode.COLUMN,
             quant_config=config.get_quant_config(),
             skip_create_weights_in_init=config.skip_create_weights_in_init,
+            allreduce_strategy=config.allreduce_strategy,
         )
 
         self.o_proj = Linear(
             tp_size * self.q_size,
             self.hidden_size,
             bias=dense_bias,
             dtype=dtype,
-            mapping=mapping,
+            mapping=mapping_o,
             tensor_parallel_mode=TensorParallelMode.ROW,
             quant_config=config.get_quant_config(),
             skip_create_weights_in_init=config.skip_create_weights_in_init,
             reduce_output=True,
+            allreduce_strategy=config.allreduce_strategy,
         )
 
         # Cross-attention backend selection honors ``ModelConfig.attn_backend``

tensorrt_llm/_torch/pyexecutor/model_loader.py

@@ -122,6 +122,34 @@ def validate_encoder_decoder_kv_cache_config(model_config: ModelConfig,
         )
 
 
+def validate_encoder_decoder_tp_scope(model_config: ModelConfig) -> None:
+    """Validate initially supported encoder-decoder TP combinations."""
+    if not model_config.is_encoder_decoder:
+        return
+
+    mapping = model_config.mapping
+    if mapping.enable_attention_dp:
+        raise ValueError(
+            "Encoder-decoder models do not support attention DP yet. "
+            "Set enable_attention_dp=False.")
+
+    if mapping.cp_size > 1:
+        raise ValueError(
+            "Encoder-decoder models do not support context parallelism yet. "
+            "Set context_parallel_size=1.")
+
+    if mapping.pp_size > 1:
+        raise ValueError(
+            "Encoder-decoder models do not support pipeline parallelism yet. "
+            "Set pipeline_parallel_size=1.")
+
+    if mapping.tp_size > 1 and model_config.attn_backend != "TRTLLM":
+        raise ValueError(
+            "Encoder-decoder tensor parallelism currently supports "
+            f"attn_backend='TRTLLM' only, but got "
+            f"attn_backend='{model_config.attn_backend}'.")
+
+
 def initialize_dummy_weights(
     model: torch.nn.Module,
     low: float = -1e-3,
@@ -1158,6 +1186,7 @@ def _load_and_validate_config(
                 f"{type(config.pretrained_config).__name__}: {e}. "
                 f"AllReduce pre-allocation will be skipped.")
 
+        validate_encoder_decoder_tp_scope(config)
         validate_encoder_decoder_kv_cache_config(config,
                                                  self.llm_args.kv_cache_config)
         validate_and_set_kv_cache_quant(config,

tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py

@@ -75,13 +75,19 @@ def _test_case(
     feature_id: str,
     cuda_graph_batch_sizes: list[int] | None = None,
     kv_cache_dtype: str = "auto",
+    tensor_parallel_size: int = 1,
+    marks=None,
 ):
     expected_output_token_ids = (
         [_EXPECTED_GREEDY_OUTPUT_TOKEN_IDS]
         if num_beams == 1
         else _EXPECTED_BEAM_OUTPUT_TOKEN_IDS_BY_BEAMS[num_beams]
     )
 
+    param_kwargs = {"id": f"{feature_id}-{_MODEL_NAME}"}
+    if marks is not None:
+        param_kwargs["marks"] = marks
+
     return pytest.param(
         expected_output_token_ids,
         torch_dtype,
@@ -92,7 +98,8 @@ def _test_case(
         exact_match,
         cuda_graph_batch_sizes,
         kv_cache_dtype,
-        id=f"{feature_id}-{_MODEL_NAME}",
+        tensor_parallel_size,
+        **param_kwargs,
     )
 
 
@@ -161,6 +168,29 @@ def _test_case(
         exact_match=True,
         feature_id="bf16-kv-v2-cuda-graph-on-greedy",
     ),
+    # Tensor parallelism (TP=2) coverage
+    _test_case(
+        torch_dtype="bfloat16",
+        use_kv_cache_manager_v2=False,
+        enable_cuda_graph=False,
+        num_beams=1,
+        num_return_sequences=1,
+        exact_match=True,
+        tensor_parallel_size=2,
+        feature_id="bf16-kv-v1-cuda-graph-off-greedy-tp2",
+        marks=pytest.mark.skip_less_device(2),
+    ),
+    _test_case(
+        torch_dtype="bfloat16",
+        use_kv_cache_manager_v2=False,
+        enable_cuda_graph=True,
+        num_beams=1,
+        num_return_sequences=1,
+        exact_match=True,
+        tensor_parallel_size=2,
+        feature_id="bf16-kv-v1-cuda-graph-on-greedy-tp2",
+        marks=pytest.mark.skip_less_device(2),
+    ),
 ]
 
 
@@ -330,16 +360,18 @@ def _run_bart_pytorch_generate_encoder_decoder(
     exact_match: bool,
     cuda_graph_batch_sizes: list[int] | None,
     kv_cache_dtype: str = "auto",
+    tensor_parallel_size: int = 1,
 ) -> None:
-    monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1")
+    if tensor_parallel_size == 1:
+        monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1")
     monkeypatch.setenv("TRTLLM_SKIP_KV_CACHE_ESTIMATION", "1")
 
     model_path = _get_bart_model_path()
     tokenizer = AutoTokenizer.from_pretrained(model_path)
     case_id = (
         f"model={_MODEL_NAME}, dtype={torch_dtype}, kv_v2={use_kv_cache_manager_v2}, "
         f"cuda_graph={enable_cuda_graph}, beams={num_beams}, returns={num_return_sequences}, "
-        f"kv_dtype={kv_cache_dtype}"
+        f"kv_dtype={kv_cache_dtype}, tp={tensor_parallel_size}"
     )
     sampling_params = _sampling_params(num_beams, num_return_sequences)
 
@@ -353,6 +385,7 @@ def _run_bart_pytorch_generate_encoder_decoder(
         disable_overlap_scheduler=True,
         dtype=torch_dtype,
         enable_chunked_prefill=False,
+        tensor_parallel_size=tensor_parallel_size,
         kv_cache_config=KvCacheConfig(
             enable_block_reuse=False,
             max_tokens=_MAX_KV_TOKENS,
@@ -385,17 +418,22 @@ def _run_bart_pytorch_generate_encoder_decoder(
             exact_match,
             expected_output_token_ids_by_output,
         )
-        _assert_decoder_cuda_graph_state(
-            llm,
-            enable_cuda_graph,
-            cuda_graph_batch_sizes,
-        )
+        # CUDA graph state introspection reaches into the in-process engine,
+        # which is only available when the executor runs single-process (TP=1).
+        # For TP>1 the executor is a multi-process proxy without a local engine,
+        # so we rely on the generated-output assertions above for correctness.
+        if tensor_parallel_size == 1:
+            _assert_decoder_cuda_graph_state(
+                llm,
+                enable_cuda_graph,
+                cuda_graph_batch_sizes,
+            )
 
 
 @pytest.mark.parametrize(
     "expected_output_token_ids_by_output,torch_dtype,use_kv_cache_manager_v2,"
     "enable_cuda_graph,num_beams,num_return_sequences,exact_match,cuda_graph_batch_sizes,"
-    "kv_cache_dtype",
+    "kv_cache_dtype,tensor_parallel_size",
     _TEST_CASES,
 )
 def test_bart_pytorch_generate_encoder_decoder_end_to_end(
@@ -409,6 +447,7 @@ def test_bart_pytorch_generate_encoder_decoder_end_to_end(
     exact_match: bool,
     cuda_graph_batch_sizes: list[int] | None,
     kv_cache_dtype: str,
+    tensor_parallel_size: int,
 ) -> None:
     _run_bart_pytorch_generate_encoder_decoder(
         monkeypatch,
@@ -421,6 +460,7 @@ def test_bart_pytorch_generate_encoder_decoder_end_to_end(
         exact_match,
         cuda_graph_batch_sizes,
         kv_cache_dtype,
+        tensor_parallel_size,
     )