examples/llm-api/extra-llm-api-config.yml (1)

4-6: Normalize YAML indentation under speculative_config.

Indentation under speculative_config is inconsistent with the rest of the file. Use 2 spaces to avoid surprises with strict YAML parsers and to match the repo style.

 speculative_config:
-    decoding_type: MTP
-    num_nextn_predict_layers: 3
+  decoding_type: MTP
+  num_nextn_predict_layers: 3

tensorrt_llm/_torch/modules/embedding.py (1)

1-1: Add NVIDIA copyright/SPDX header.

Per coding guidelines, prepend the NVIDIA copyright header to all Python sources.

+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import math

tensorrt_llm/_torch/models/modeling_utils.py (1)

1-1: Add NVIDIA copyright/SPDX header.

Please add the required header at the top of this file.

+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import contextlib

tensorrt_llm/_torch/distributed/ops.py (1)

1-1: Add NVIDIA copyright/SPDX header.

Please add the required header to this Python source.

+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import logging

tensorrt_llm/_torch/models/modeling_deepseekv3.py (2)

164-170: Pre‑LM allgather: add tp_size guard and drop redundant local import

• The local import of allgather (Line 168) is redundant; this module already imports it at the top.
• Guarding on mapping.tp_size > 1 avoids unnecessary collective when TP is 1.

Apply this diff:

-        # Add pre-lm gather logic
-        if (self.model_config.mapping.enable_attention_dp and 
-            getattr(self.model_config.mapping, 'enable_lm_tp_in_adp', False)):
-            # ADP + LM TP mode: perform All-Gather before LM_head
-            from ..distributed import allgather
-            hidden_states = allgather(hidden_states, self.model_config.mapping, dim=-1)
+        # Add pre-LM gather logic
+        mapping = self.model_config.mapping
+        if (mapping.enable_attention_dp
+                and getattr(mapping, 'enable_lm_tp_in_adp', False)
+                and mapping.tp_size > 1):
+            # ADP + LM TP mode: gather hidden across the last dim (TP shards)
+            hidden_states = allgather(hidden_states, mapping, dim=-1)

---

1-20: Header/license consistency with repo guideline

Repo guideline requires an NVIDIA copyright header for Python sources. This file only carries the upstream MIT header. If the repo policy is to prepend NVIDIA’s header while preserving upstream attribution, please add it at the top.

tensorrt_llm/_torch/speculative/mtp.py (2)

1042-1057: Avoid casting token indices to float32 when packing

Casting indices to float can risk precision if vocab grows large. While 32‑bit floats are fine for today’s vocab sizes, consider keeping indices as int32 and only casting values, or use two allgathers (values and indices) as in the fix above to avoid packing altogether.

---

1060-1071: Minor: clarify assumptions in gathered layout

The even/odd slicing depends on the [idx,val] packing; add a brief comment documenting the layout and why float types are used to avoid surprises during maintenance. No functional issues otherwise.

tensorrt_llm/llmapi/llm_args.py (1)

226-227: Flag plumbing looks good; add validation to prevent silent misconfigurations

Propagation of enable_lm_tp_in_adp through _ParallelConfig and Mapping is correct. To prevent confusion, add a validator or a warning when:

• enable_lm_tp_in_adp is True but enable_attention_dp is False (flag has no effect),
• or tp_size == 1 (no LM TP to enable).

This keeps configs self‑consistent and avoids enabling a knob that does nothing.

I can add a small model_validator to emit a warning in these cases if you’d like.

Also applies to: 290-291, 1198-1202, 1450-1451

cpp/tensorrt_llm/thop/reducescatterOp.cpp (1)

43-45: reduceOp plumbed correctly; consider input validation and docs updates

• Passing ncclRedOp_t through the constructor and using it in ncclReduceScatter/ncclReduce looks correct. Nice.
• Please guard reduceOp against invalid enum values (e.g., out‑of‑range ints) and document the accepted values at the Python wrapper layer. A quick check with TLLM_CHECK can prevent undefined behavior.
• Minor: for the ragged path (sizes.has_value()), consider using nullptr for recvbuff on non‑root ranks in ncclReduce to make intent explicit.

Example guard near the cast site:

int64_t ro = reduceOp;
TLLM_CHECK_WITH_INFO(ro == ncclSum || ro == ncclProd || ro == ncclMax || ro == ncclMin,
    "Unsupported reduceOp %ld", static_cast<long>(ro));
ncclRedOp_t ncclReduceOp = static_cast<ncclRedOp_t>(ro);

Also applies to: 96-96, 107-107, 126-126, 134-144, 153-163, 175-176

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

Learnt from: timlee0212
PR: NVIDIA/TensorRT-LLM#6886
File: tensorrt_llm/_torch/models/modeling_deepseekv3.py:0-0
Timestamp: 2025-08-14T06:36:40.701Z
Learning: In DeepSeek V3 model (tensorrt_llm/_torch/models/modeling_deepseekv3.py), the disagreement between AllReduce.__init__ guard and _compute_mlp_tp_size logic for MNNVL usage is expected by design. The AllReduce component and MLP TP-size computation intentionally use different criteria for MNNVL availability decisions.

Learnt from: timlee0212
PR: NVIDIA/TensorRT-LLM#6886
File: tensorrt_llm/_torch/models/modeling_deepseekv3.py:0-0
Timestamp: 2025-08-14T06:36:40.701Z
Learning: In DeepSeek V3 model (tensorrt_llm/_torch/models/modeling_deepseekv3.py), the disagreement between AllReduce.__init__ guard and _compute_mlp_tp_size logic for MNNVL usage is expected by design. The AllReduce component and MLP TP-size computation intentionally use different criteria for MNNVL availability decisions.

examples/llm-api/extra-llm-api-config.yml (1)

5-6: All good – num_nextn_predict_layers is the correct key

Search across the codebase confirms that the configuration field is consistently defined and referenced as num_nextn_predict_layers, with no occurrences of variants like num_next_n_predict_layers or num_next_predict_layers. This key is used in:

• Parsing and argument definitions (tensorrt_llm/llmapi/llm_args.py)
• MTP decoding implementations and utilities
• Test suites (unit and integration tests)
• Examples and documentation

No action is required.

tensorrt_llm/_torch/modules/embedding.py (1)

40-43: Gating in LMHead init is correct — it disables any embedding parallelism when attention‐based DP is on (unless enable_lm_tp_in_adp=True), so by default both Embedding and LMHead fall back to tensor_parallel_mode=None and remain in sync. I’ve checked the major call sites—none explicitly pass a non-None tensor_parallel_mode—so tied embeddings will not trigger the TP mode/size assertions in DecoderModelForCausalLM.

• Double-check the two constructors in
tensorrt_llm/_torch/models/modeling_utils.py (lines 354–358)
to ensure that, if you do set enable_lm_tp_in_adp=True, you propagate the same tensor_parallel_mode into both the vocab embedding and the LM head.

With that confirmed, no code changes are needed here.

tensorrt_llm/_torch/models/modeling_utils.py (1)

355-383: Verified ADP+LM TP gating and gather_output behavior

• The enable_attention_dp and enable_lm_tp_in_adp flags in tensorrt_llm/mapping.py default to False, preserving the legacy ADP-only code path.
• In torch/models/modeling_utils.py, LMHead is instantiated without tensor parallelism when ADP is enabled but LM TP in ADP is disabled, and with gather_output=True only in the TP-aware branch.
• In torch/models/modeling_deepseekv3.py, the code performs an all-gather before the LM head when both ADP and LM TP in ADP are on, and dynamically toggles lm_head.gather_output around the forward call to ensure consistency.

A regex search confirmed that all gather_output assignments throughout the repo align with these gating conditions. Approving these changes.

tensorrt_llm/mapping.py (1)

294-305: Add enable_attention_dp and enable_lm_tp_in_adp to Mapping’s equality, hash, and serialization

To prevent cache collisions and stale workspaces (since Mapping instances are used as dict-keys in allreduce workspaces, KV-cache pools, etc.), please extend all three methods consistently:

• In Mapping.__eq__ (currently around lines 294–305, and similarly at 308–324): append

    and self.enable_attention_dp == other.enable_attention_dp
    and self.enable_lm_tp_in_adp == other.enable_lm_tp_in_adp

to the final return.

• In Mapping.__hash__ (around lines 300–315): include

    self.enable_attention_dp,
    self.enable_lm_tp_in_adp,

in the tuple passed to hash((…)).

• In Mapping.to_dict (around lines 500–515): add

    'enable_attention_dp': self.enable_attention_dp,
    'enable_lm_tp_in_adp': self.enable_lm_tp_in_adp,

to the returned dict.

Also, verify that Mapping.from_dict (if implemented) reads these two keys back—otherwise serialized configs won’t round-trip.

After applying these changes, please re-run your cache-collision check to ensure no unintended merges remain, e.g.:

rg -nP --type=py -C2 '\bMapping\b' tensorrt_llm | sed -n '1,200p'

…and confirm that every codepath which uses a Mapping as a dict key now distinguishes different DP settings.

tensorrt_llm/_torch/distributed/ops.py (2)

12-14: Re-exporting ReduceOp is fine.

This provides a convenient import path for call sites without pulling torch.distributed directly.

---

232-239: No existing call sites pass ints or ReduceOp
Ran rg -nP --type=py -C2 'reducescatter\([^)]*$' and inspected all callers of reducescatter (in tensorrt_llm/_torch/speculative/mtp.py, .../modules/fused_moe/interface.py, .../fused_moe_vanilla.py, and .../fused_moe_wide_ep.py). None of them supply a non-string reduce_op argument, so they’ll continue to use the default "sum". No call-site updates are required.

tensorrt_llm/llmapi/llm_args.py (1)

1251-1255: Mark the new public knob as prototype, per prior guidance

Previous discussion recommended marking brand‑new, dangling knobs as prototype until stabilized. Please set status="prototype" here to match that guidance.

Apply this diff:

-    enable_lm_tp_in_adp: bool = Field(
-        default=False,
-        description="Enable lm tp in attention dp.",
-        status="beta")
+    enable_lm_tp_in_adp: bool = Field(
+        default=False,
+        description="Enable LM tensor-parallel in Attention DP mode.",
+        status="prototype")

tensorrt_llm/_torch/models/modeling_deepseekv3.py (2)

163-169: Remove redundant local import; it's already imported at the file header

allgather is already imported at the top; keep imports centralized.

-            from ..distributed import allgather
             hidden_states = allgather(hidden_states, self.model_config.mapping, dim=-1)

---

170-177: Make gather_output toggle exception-safe and restore prior value

Current code may leave lm_head.gather_output in the wrong state on exceptions or nested calls. Compute the condition once, save the prior value, and restore it in finally.

Apply this diff:

-        # Temporarily disable gather_output when not in ADP mode or (in ADP mode and LM TP is enabled)
-        if (not self.model_config.mapping.enable_attention_dp) or (self.model_config.mapping.enable_attention_dp and
-                getattr(self.model_config.mapping, 'enable_lm_tp_in_adp', False)):
-            lm_head.gather_output = False
-        logits = lm_head(hidden_states)
-        if (not self.model_config.mapping.enable_attention_dp) or (self.model_config.mapping.enable_attention_dp and
-                getattr(self.model_config.mapping, 'enable_lm_tp_in_adp', False)):
-            lm_head.gather_output = True
+        # Temporarily disable gather_output when not in ADP, or when ADP+LM TP is enabled
+        mapping = self.model_config.mapping
+        disable_gather = (not mapping.enable_attention_dp) or getattr(mapping, 'enable_lm_tp_in_adp', False)
+        prev_gather = lm_head.gather_output
+        try:
+            if disable_gather:
+                lm_head.gather_output = False
+            logits = lm_head(hidden_states)
+        finally:
+            if disable_gather:
+                lm_head.gather_output = prev_gather

tensorrt_llm/llmapi/llm_args.py (1)

1503-1504: Good: flag is propagated into _ParallelConfig

Consider warning if enable_lm_tp_in_adp is set while enable_attention_dp is False (no effect). Optional.

Example validator to add (outside the selected lines):

@model_validator(mode="after")
def warn_lm_tp_in_adp_without_adp(self):
    if self.enable_lm_tp_in_adp and not self.enable_attention_dp:
        logger.warning("enable_lm_tp_in_adp has no effect unless enable_attention_dp is True.")
    return self

examples/wide_ep/slurm_scripts/submit_e2e.sh (1)

14-15: mtp_size default change

Bumping mtp_size to 3 aligns with the new MTP path. Consider exposing this as an env override similar to other parameters.

-mtp_size=3
+mtp_size="${MTP_SIZE:-3}"

cpp/tensorrt_llm/thop/reducescatterOp.cpp (1)

2-2: Nit: update copyright year

Header still says 1993-2024; repository guideline asks for current year.

- * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

Learnt from: timlee0212
PR: NVIDIA/TensorRT-LLM#6886
File: tensorrt_llm/_torch/models/modeling_deepseekv3.py:0-0
Timestamp: 2025-08-14T06:36:40.701Z
Learning: In DeepSeek V3 model (tensorrt_llm/_torch/models/modeling_deepseekv3.py), the disagreement between AllReduce.__init__ guard and _compute_mlp_tp_size logic for MNNVL usage is expected by design. The AllReduce component and MLP TP-size computation intentionally use different criteria for MNNVL availability decisions.

Learnt from: timlee0212
PR: NVIDIA/TensorRT-LLM#6886
File: tensorrt_llm/_torch/models/modeling_deepseekv3.py:0-0
Timestamp: 2025-08-14T06:36:40.701Z
Learning: In DeepSeek V3 model (tensorrt_llm/_torch/models/modeling_deepseekv3.py), the disagreement between AllReduce.__init__ guard and _compute_mlp_tp_size logic for MNNVL usage is expected by design. The AllReduce component and MLP TP-size computation intentionally use different criteria for MNNVL availability decisions.

tensorrt_llm/_torch/modules/logits_processor.py (2)

70-76: Avoid shadowed imports; import allgather once at module scope

The second allgather relies on a prior in-block import. Import from ..distributed.ops at module scope for clarity and safety.

-            logits = allgather(logits, self.model_config.mapping, dim=-1)
+            # logits are local-shard vocab; gather to full vocab
+            from ..distributed.ops import allgather  # consider hoisting to module imports
+            logits = allgather(logits, self.model_config.mapping, dim=-1)
             batch_size = logits.shape[0]
             local_batch_size = batch_size // self.model_config.mapping.tp_size
-            logits = logits.view(self.model_config.mapping.tp_size,
+            logits = logits.reshape(self.model_config.mapping.tp_size,
                                 local_batch_size, -1)
             logits = logits[self.model_config.mapping.tp_rank][:token_count]

---

33-34: Use reshape for robustness

Switch view(...) to reshape(...) to be resilient to non-contiguous hidden_states.

-        token_count = hidden_states.view(-1, hidden_states.shape[-1]).shape[0]
+        token_count = hidden_states.reshape(-1, hidden_states.shape[-1]).shape[0]

tensorrt_llm/_torch/speculative/mtp.py (1)

861-867: Remove commented debug prints or guard behind a verbose flag

Keep the hot path clean; these prints clutter diffs and can hurt traceability if re-enabled accidentally.

tensorrt_llm/_torch/modules/embedding.py (1)

40-42: Gating looks good; confirm downstream expectations

Disabling LM TP under Attention DP unless enable_lm_tp_in_adp is set aligns with the new flag. Please confirm that callers relying on LM TP under ADP explicitly set the flag in Mapping/args, and add a brief docstring note where the flag is defined.

tensorrt_llm/_torch/models/modeling_utils.py (1)

355-385: Replace prints with logger and simplify LMHead gating to a single code path

• Avoid print in multi-rank runs; use logger.debug (rank-gated if noisy).
• You can rely on LMHead’s internal ADP gating and always construct it with mapping and TP settings to remove duplicated branches.

Apply:

-        if config.mapping.enable_attention_dp and not getattr(config.mapping, 'enable_lm_tp_in_adp', False):
-            print(f"In DecoderModelForCausalLM, creating LMHead without TP")
-            self.lm_head = LMHead(
-                vocab_size,
-                hidden_size,
-                dtype=config.pretrained_config.torch_dtype,
-            )
-        else:
-            print(f"In DecoderModelForCausalLM, creating LMHead with TP")
-            # TODO(zhenhuanc): Currently lm_head Linear will not accept QuantConfig
-            # will considering per layer QuantConfig in the future.
-            if (hasattr(config, 'lora_config')
-                    and config.lora_config is not None
-                    and len(config.lora_config.lora_dir) == 1):
-                # Only check for custom lm_head in HF LoRA, not NeMo
-                if config.lora_config.lora_ckpt_source == "hf":
-                    lora_loader = HfLoraLoader(config.lora_config.lora_dir)
-                    if lora_loader.lm_head is not None and lora_loader.vocab_size != 0:
-                        weight = lora_loader.lm_head
-                        self.has_custom_lm_head = True
-                        vocab_size = lora_loader.vocab_size
-
-            self.lm_head = LMHead(
-                vocab_size,
-                hidden_size,
-                dtype=config.pretrained_config.torch_dtype,
-                mapping=config.mapping,
-                tensor_parallel_mode=TensorParallelMode.COLUMN,
-                gather_output=True,
-            )
+        # Detect custom LoRA LM head (HF), may override vocab_size.
+        if (hasattr(config, 'lora_config')
+                and config.lora_config is not None
+                and len(config.lora_config.lora_dir) == 1
+                and config.lora_config.lora_ckpt_source == "hf"):
+            lora_loader = HfLoraLoader(config.lora_config.lora_dir)
+            if lora_loader.lm_head is not None and lora_loader.vocab_size != 0:
+                weight = lora_loader.lm_head
+                self.has_custom_lm_head = True
+                vocab_size = lora_loader.vocab_size
+
+        # Always construct with mapping; LMHead internally disables TP for ADP when enable_lm_tp_in_adp=False.
+        self.lm_head = LMHead(
+            vocab_size,
+            hidden_size,
+            dtype=config.pretrained_config.torch_dtype,
+            mapping=config.mapping,
+            tensor_parallel_mode=TensorParallelMode.COLUMN,
+            gather_output=True,
+        )
+        logger.debug(
+            "DecoderModelForCausalLM: LMHead constructed (tp_mode=%s, tp_size=%s, gather_output=%s)",
+            getattr(self.lm_head, "tp_mode", None),
+            getattr(self.lm_head, "tp_size", 1),
+            getattr(self.lm_head, "gather_output", False),
+        )

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

Learnt from: djns99
PR: NVIDIA/TensorRT-LLM#6728
File: cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp:966-966
Timestamp: 2025-08-08T04:10:19.038Z
Learning: TensorRT plugins currently don't support padding functionality, and TensorRT is not getting new features (in maintenance mode). This means that duplicating parameters like mExpertHiddenSize in function calls, even with TODO comments, can be acceptable as pragmatic solutions within these constraints.

Learnt from: thorjohnsen
PR: NVIDIA/TensorRT-LLM#6910
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-14T21:04:50.248Z
Learning: In KV cache onboarding logic during prefill in cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, when calculating which blocks fall within the attention window, use getTokensPerBlock() to advance token indices rather than block->getUniqueTokens().size(), because the calculation needs to consider the post-prefill state where blocks will be filled to capacity, not their current token count.

Learnt from: timlee0212
PR: NVIDIA/TensorRT-LLM#6886
File: tensorrt_llm/_torch/models/modeling_deepseekv3.py:0-0
Timestamp: 2025-08-14T06:36:40.701Z
Learning: In DeepSeek V3 model (tensorrt_llm/_torch/models/modeling_deepseekv3.py), the disagreement between AllReduce.__init__ guard and _compute_mlp_tp_size logic for MNNVL usage is expected by design. The AllReduce component and MLP TP-size computation intentionally use different criteria for MNNVL availability decisions.