[None][fix] Revert "Merge Eagle3 and MTP-eagle one-model workers (#12353)"

.claude/skills/ad-sharding-ir-port/SKILL.md

@@ -35,8 +35,9 @@ You MAY introduce ONLY the following changes:
   - `torch.split(...)` / `torch.split_with_sizes(...)` → `torch.ops.auto_deploy.split_with_sizes(...)`
 - **A2. Sharding-hint kwargs added** to call sites of: `torch_moe`, `torch_ssm`, `torch_gated_delta_rule`, `torch_causal_conv1d`, `torch_rmsnorm_gated`, `torch_mla`, `torch_linear_simple`, `auto_deploy.split_with_sizes`, `auto_deploy.view`. Allowed kwargs: `tp_mode`, `layer_type`, `output_sizes`, `tp_min_local_shape`, `tp_scaled_dim`, `shardable`, `enable_sharding`.
 - **A3. Inserting `torch.ops.auto_deploy.all_reduce(..., layer_type=...)`** after rowwise projections / at MoE merge points (single all_reduce after routed + shared sums).
-- **A4. Adding `import tensorrt_llm._torch.auto_deploy.custom_ops  # noqa: F401`** side-effect import at the top if not already present.
-- **A5. The module docstring update** describing the sharding strategy.
+- **A4. Docstring updates:**
+  - Module-level: a single-line header noting the file uses sharding IR, followed by the existing source-of-truth / HF link block. Example: `"""Llama 3 model (sharding IR)."""`.
+  - Per-class (MLP, Attention, MoE block, etc.): a short `Sharding strategy:` block listing what each projection maps to (`colwise` / `rowwise` / `all_reduce` / `tp_scaled_dim`).
 
 **FORBIDDEN (everything else, including but not limited to):**
 
@@ -71,45 +72,35 @@ Before editing, ensure the file is committed so you can diff against the origina
 git stash  # or commit — ensure a clean baseline to diff against
 ```
 
-### Step 2: Add the custom_ops side-effect import
-
-If not already present at the top of the file:
-
-```python
-import tensorrt_llm._torch.auto_deploy.custom_ops  # noqa: F401 -- register all ops
-```
-
-Do **not** add global `SHARD_*` flags. Layer-level control uses the `layer_type` hint on each op and `shard_layers` in YAML.
-
-### Step 3: Replace linear projections
+### Step 2: Replace linear projections
 
 For every `self.proj(x)` or `nn.Linear` call, use `torch.ops.auto_deploy.torch_linear_simple` with explicit `tp_mode` and `layer_type`. Always set `tp_mode` unconditionally (no `if _s else "none"`). **Rules:** opening projections (Q/K/V/gate/up/in_proj) → `"colwise"`; closing (O/down/out_proj) → `"rowwise"`; tiny outputs (e.g. `shared_expert_gate` dim 1) → `"none"`; MLA latent projections (q_a, kv_a) → `"none"`. For fused weights split later, pass `output_sizes=[...]`. For GQA, use `tp_min_local_shape=self.head_dim` on K/V colwise lines.
 
-### Step 4: Replace split / chunk after fused colwise projections
+### Step 3: Replace split / chunk after fused colwise projections
 
 Use `torch.ops.auto_deploy.split_with_sizes` with `shardable` / `layer_type` where sizes scale with TP.
 
-### Step 5: Replace view / reshape with concrete head counts
+### Step 4: Replace view / reshape with concrete head counts
 
 During `torch.export`, `-1` becomes concrete; after TP, wrong values break. Any reshape whose dimension is a head count that scales with TP must use `torch.ops.auto_deploy.view` with `tp_scaled_dim` set appropriately. Safe cases: flat-to-2D, or `[B,S,-1]` when the input is already correctly sharded.
 
-### Step 6: Insert `all_reduce`
+### Step 5: Insert `all_reduce`
 
 After every rowwise projection, add `torch.ops.auto_deploy.all_reduce(..., layer_type=...)`. **Parallel branch rule:** when branches merge by addition, use a **single** `all_reduce` after the sum (e.g. MoE routed + shared expert; parallel attention + MLP residual branches).
 
-### Step 7: Special ops (Conv1d, SSM, GatedDeltaNet, gated RMSNorm)
+### Step 6: Special ops (Conv1d, SSM, GatedDeltaNet, gated RMSNorm)
 
 Add sharding hints on `torch_causal_conv1d`, `torch_ssm`, `torch_gated_delta_rule`, `torch_rmsnorm_gated` per docstrings—typically `shardable` / `output_sizes` / `tp_mode` as required.
 
-### Step 8: MoE
+### Step 7: MoE
 
 Pass `layer_type="moe"` into `torch_moe`; `apply_sharding_hints` handles EP/TP.
 
-### Step 9: Verify registration
+### Step 8: Verify registration
 
 The model's existing registration (`AutoModelForCausalLMFactory.register_custom_model_cls` at the bottom of the file and its import in `__init__.py`) stays unchanged. No new registration is needed — sharding hints do not change the model identity.
 
-### Step 10: YAML — enable hint-driven sharding
+### Step 9: YAML — enable hint-driven sharding
 
 Add `enable_sharder_ir.yaml` to the model's `yaml_extra` list in `examples/auto_deploy/model_registry/models.yaml` (if not already present). This composable fragment disables legacy sharding passes and enables `apply_sharding_hints`. Registry fragments are deep-merged in `yaml_extra` order (see `DynamicYamlMixInForSettings` in `tensorrt_llm/_torch/auto_deploy/utils/_config.py`).
 
@@ -136,11 +127,11 @@ transforms:
     enabled: true
 ```
 
-Set `world_size` once, to the **maximum number of GPUs available on the machine**, auto-detected with `python -c 'import torch; print(torch.cuda.device_count())'` (or `nvidia-smi --list-gpus | wc -l`). Do **not** hardcode `world_size: 8` (or any other literal) — porting agents run on heterogeneous hardware and an 8-GPU literal will simply fail to launch on a 2- or 4-GPU machine. If the model's `num_attention_heads` (and, for GQA, `num_key_value_heads`) does not divide the detected GPU count, fall back to the largest power-of-two divisor that does (e.g. 4 on an 8-GPU machine if `num_attention_heads = 12`). Run the end-to-end command exactly once at that size — there is no value in repeating it at multiple smaller sizes, because the offline sharding equivalence test (Step 11b) already exercises 2- and 4-GPU dist configs cheaply.
+Set `world_size` once, to the **maximum number of GPUs available on the machine**, auto-detected with `python -c 'import torch; print(torch.cuda.device_count())'` (or `nvidia-smi --list-gpus | wc -l`). Do **not** hardcode `world_size: 8` (or any other literal) — porting agents run on heterogeneous hardware and an 8-GPU literal will simply fail to launch on a 2- or 4-GPU machine. If the model's `num_attention_heads` (and, for GQA, `num_key_value_heads`) does not divide the detected GPU count, fall back to the largest power-of-two divisor that does (e.g. 4 on an 8-GPU machine if `num_attention_heads = 12`). Run the end-to-end command exactly once at that size — there is no value in repeating it at multiple smaller sizes, because the offline sharding equivalence test (Step 10b) already exercises 2- and 4-GPU dist configs cheaply.
 
 Optional `shard_layers` limits which `layer_type` hints are processed; unset means shard all shardable nodes.
 
-### Step 11a — End-to-end run
+### Step 10a — End-to-end run
 
 Do not report success until a run completes successfully.
 
@@ -151,7 +142,7 @@ Do not report success until a run completes successfully.
 
 **Layer type strings** (for `layer_type` / `shard_layers`): use `"mha"`, `"mla"`, `"mlp"`, `"moe"`, `"ssm"`, `"delta"`, or `"unknown"` (default; skipped when `shard_layers` is set). Match the conventions used in `apply_sharding_hints` and project enums.
 
-### Step 11b — Sharding equivalence test (MANDATORY)
+### Step 10b — Sharding equivalence test (MANDATORY)
 
 Run the offline sharding-IR equivalence test ([`tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py`](tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py)) against the modeling file you just edited, under **every** parallelism configuration the test exposes. The port is **not** complete until every configuration passes. Skipping this step or treating a partial pass (e.g. only `tep`) as success is not allowed.
 
@@ -191,10 +182,10 @@ done
 **Failure handling:**
 
 - A cell failing with `KeyError`, `AttributeError`, `ValueError: You must specify exactly one of input_ids or inputs_embeds`, or any exception *before* `[sharding-ir-eq]` prints means the **modeling code itself** does not yet build / export on a tiny config — fix the modeling code (within the Step 0 allowlist) before proceeding. Do not silently skip the cell.
-- A cell where `[sharding-ir-eq]` prints `rel_rmse >= tol` (from the same log line) means a **sharding-hint bug**: a missing `all_reduce`, a wrong `tp_mode`, a `view` without `tp_scaled_dim`, a `split_with_sizes` whose sizes do not scale, etc. Re-read Step 6 (all_reduce), Step 3 (tp_mode), Step 5 (view), Step 4 (split_with_sizes) and the layer-specific patterns. Iterate on the hints until clean. If the failure is small (rel_rmse just slightly above tol) and you have reason to believe it is real numerical noise from the specific layer mix of this model rather than a sharding-hint bug, raise it with the parent agent rather than silently bumping `SHARDING_IR_REL_RMSE_TOL`.
+- A cell where `[sharding-ir-eq]` prints `rel_rmse >= tol` (from the same log line) means a **sharding-hint bug**: a missing `all_reduce`, a wrong `tp_mode`, a `view` without `tp_scaled_dim`, a `split_with_sizes` whose sizes do not scale, etc. Re-read Step 5 (all_reduce), Step 2 (tp_mode), Step 4 (view), Step 3 (split_with_sizes) and the layer-specific patterns. Iterate on the hints until clean. If the failure is small (rel_rmse just slightly above tol) and you have reason to believe it is real numerical noise from the specific layer mix of this model rather than a sharding-hint bug, raise it with the parent agent rather than silently bumping `SHARDING_IR_REL_RMSE_TOL`.
 - A cell that the modeling file legitimately does not support (e.g. `ep-only` on a dense model with no MoE) is acceptable only if the failure is a documented `pytest.skip(...)` from the test infrastructure. A silent `FAIL` is **not** acceptable.
 
-### Step 12 — Pre-finalization self-audit (MANDATORY)
+### Step 11 — Pre-finalization self-audit (MANDATORY)
 
 Before reporting the file as done, you MUST diff your changes against the git baseline:
 
@@ -209,8 +200,7 @@ Then classify every hunk into one of the following categories (defined in Step 0
 | **A1** | yes | Op substitution (`linear` / `view` / `split`) |
 | **A2** | yes | Sharding-hint kwarg added (`tp_mode`, `layer_type`, `output_sizes`, `tp_min_local_shape`, `tp_scaled_dim`, `shardable`, `enable_sharding`) |
 | **A3** | yes | `auto_deploy.all_reduce` insertion |
-| **A4** | yes | `custom_ops` side-effect import added |
-| **A5** | yes | Module docstring update describing sharding strategy |
+| **A4** | yes | Docstring updates: one-line module header + per-class `Sharding strategy:` blocks |
 | **F1** | NO | `torch.ops.trtllm.*` replaced with vanilla PyTorch |
 | **F2** | NO | Input contract change (asserts, fallbacks added/removed) |
 | **F3** | NO | Module hierarchy / parameter / buffer / load-hook change |
@@ -261,8 +251,8 @@ You are NOT done until every row in the table is a yes-allowed category.
 
 ## Validation checklist (human review)
 
-- All four configurations of the **sharding equivalence test** (Step 11b) pass with the parsed `rel_rmse` strictly below the parsed `tol` from the same rank-0 log line. Report the per-cell `rel_rmse` and `tol` pair.
+- All four configurations of the **sharding equivalence test** (Step 10b) pass with the parsed `rel_rmse` strictly below the parsed `tol` from the same rank-0 log line. Report the per-cell `rel_rmse` and `tol` pair.
 - `world_size=1`: unsharded path; hints should not break correctness.
-- `world_size=<max-available>`: end-to-end run (Step 11a) at the maximum GPU count auto-detected on the machine (head-divisibility permitting; see Step 11).
+- `world_size=<max-available>`: end-to-end run (Step 10a) at the maximum GPU count auto-detected on the machine (head-divisibility permitting; see Step 10).
 - `apply_sharding_hints` node count vs expectation.
 - Optional: `shard_layers: ['moe']` to verify selective sharding.

.claude/skills/trtllm-model-onboard-multimodal/SKILL.md

@@ -83,7 +83,7 @@ metadata:
 When `@support_multimodal_disaggregated` is set and the deployment uses `TLLM_MULTIMODAL_DISAGGREGATED=1`:
 
 - **Encoder worker:** runs as a standalone `MultimodalEncoder` (`mm_encoder_only=True`). It executes only the multimodal encoder and ships `mm_embeddings` (+ mRoPE position ids/deltas) to prefill+decode workers as shared-tensor handles.
-- **Prefill+decode worker:** the model's `__init__` skips constructing `self.mm_encoder` when `_is_disagg()` is true; the input processor's `_attach_multimodal_embeddings_impl()` override binds the encoder handles into the request (the base `attach_multimodal_embeddings` wrapper detokenizes tokenized inputs for non-fast-path VLMs, then delegates to your impl). For context-only requests, the engine re-clones mrope tensors so IPC handles outlive the encoder worker's freed memory — replicate that pattern for any new GPU-resident mm tensors.
+- **Prefill+decode worker:** the model's `__init__` skips constructing `self.mm_encoder` when `_is_mm_disagg()` is true; the input processor's `attach_multimodal_embeddings()` override binds the encoder handles into the request. For context-only requests, the engine re-clones mrope tensors so IPC handles outlive the encoder worker's freed memory — replicate that pattern for any new GPU-resident mm tensors.
 
 ### Templates to study
 
@@ -229,7 +229,7 @@ class {Name}Model(PreTrainedModel):
         if hasattr(self, "llm"):
             return  # idempotency guard — re-entry from `post_config` etc.
 
-        if not _is_disagg():
+        if not _is_mm_disagg():
             self.mm_encoder = {Name}VisionModel(model_config)
         else:
             self.mm_encoder = None
@@ -269,7 +269,7 @@ class {Name}Model(PreTrainedModel):
 
         multimodal_params = kwargs.get("multimodal_params", [])
         mm_embeds = []
-        if len(multimodal_params) > 0 and not _is_disagg():
+        if len(multimodal_params) > 0 and not _is_mm_disagg():
             mm_embeds = get_multimodal_embeddings(
                 encoder_forward_fn=self.mm_encoder.forward,
                 multimodal_params=multimodal_params[:num_context_requests],
@@ -341,7 +341,7 @@ class {Name}Model(PreTrainedModel): ...
 
 ```python
 def load_weights(self, weights, weight_mapper):
-    if not _is_disagg():
+    if not _is_mm_disagg():
         self.mm_encoder.load_weights(weights)
         # Release mmap pages backing the encoder weights as soon as we're done.
         if hasattr(weights, "mark_consumed"):

.github/CODEOWNERS

@@ -45,8 +45,8 @@
 
 ## TensorRT-LLM Pytorch - VisualGen
 /tensorrt_llm/_torch/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs
-/tensorrt_llm/_torch/visual_gen/attention_backend @NVIDIA/trt-llm-torch-attention-devs
-/tensorrt_llm/_torch/visual_gen/modules/attention.py @NVIDIA/trt-llm-torch-attention-devs
+/tensorrt_llm/_torch/visual_gen/attention_backend @NVIDIA/trt-llm-torch-attention-devs @NVIDIA/trt-llm-torch-visual-gen-devs
+/tensorrt_llm/_torch/visual_gen/modules/attention.py @NVIDIA/trt-llm-torch-attention-devs @NVIDIA/trt-llm-torch-visual-gen-devs
 /tensorrt_llm/visual_gen @NVIDIA/trt-llm-llmapi-devs
 /tests/integration/defs/examples/test_visual_gen.py @NVIDIA/trt-llm-torch-visual-gen-devs
 /tests/integration/defs/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs

.github/workflows/blossom-ci.yml

@@ -58,6 +58,7 @@ jobs:
         "arekay",
         "arysef",
         "aswinvisva",
+        "athena-nv",
         "atrifex",
         "Autumn1998",
         "baize97",
@@ -234,6 +235,7 @@ jobs:
         "nv-anants",
         "nv-guomingz",
         "nv-lschneider",
+        "nv-xtf",
         "nv-yilinf",
         "nv-yna",
         "nvamyt",
@@ -267,6 +269,7 @@ jobs:
         "qsang-nv",
         "raayandhar",
         "rabiel",
+        "rahul-steiger-nv",
         "rakib-hasan",
         "RayenTian",
         "raymochen",
@@ -310,6 +313,7 @@ jobs:
         "taylor-yb-lee",
         "tburt-nv",
         "tcherckez-nvidia",
+        "tedzhouhk",
         "tfogal",
         "thorjohnsen",
         "tianyuxbear",

.gitignore

@@ -65,6 +65,9 @@ tensorrt_llm/scripts
 docs/source/**/*.rst
 !docs/source/examples/index.rst
 !docs/source/_includes/note_sections.rst
+!docs/source/features/auto_deploy/transforms.rst
+!docs/source/features/auto_deploy/transforms/
+!docs/source/features/auto_deploy/transforms/*.rst
 *.swp
 .nfs*
 

3rdparty/fetch_content.json

@@ -26,7 +26,8 @@
       "display_name": "deep_ep",
       "git_repository": "${github_base_url}/deepseek-ai/DeepEP",
       "git_tag": "5be51b228a7c82dbdb213ea58e77bffd12b38af8",
-      "use_url": true
+      "use_url": true,
+      "patch_file": "patches/deep_ep_intranode_combine_fix.patch"
     },
     {
       "name": "deepgemm",

3rdparty/patches/deep_ep_intranode_combine_fix.patch

@@ -0,0 +1,35 @@
+--- a/csrc/kernels/intranode.cu
++++ b/csrc/kernels/intranode.cu
+@@ -844,9 +844,15 @@
+
+ #ifndef DISABLE_SM90_FEATURES
+                     // Wait TMA arrival
++                    // hidden_int4 is not always divisible by a warp. The final tile can have
++                    // only a subset of lanes active, so synchronize only participating lanes.
++                    auto const tile_start = i - lane_id;
++                    auto const active_lanes = min(32, hidden_int4 - tile_start);
++                    auto const sync_mask = active_lanes == 32 ? 0xffffffffu : ((1u << active_lanes) - 1u);
++
+                     if (lane_id == 0)
+                         tma_store_wait<kNumStages - 1>();
+-                    __syncwarp();
++                    __syncwarp(sync_mask);
+
+                     // Write into TMA buffer
+                     auto tma_stage_idx = (i / 32) % kNumStages;
+@@ -854,13 +860,13 @@
+
+                     // Issue TMA
+                     tma_store_fence();
+-                    __syncwarp();
++                    __syncwarp(sync_mask);
+                     if (lane_id == 0) {
+                         auto tma_bytes = min(32, hidden_int4 - i) * static_cast<int>(sizeof(int4));
+                         tma_store_1d(reinterpret_cast<int4*>(tma_buffer) + tma_stage_idx * 32,
+                                      recv_int4 + token_idx * hidden_int4 + i, tma_bytes, false);
+                     }
+-                    __syncwarp();
++                    __syncwarp(sync_mask);
+ #else
+                     recv_int4[token_idx * hidden_int4 + i] = out_int4;
+ #endif