[None][feat] Update CuTeDSL MegaMoE kernels

tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/__init__.py

@@ -2,9 +2,8 @@
 # SPDX-License-Identifier: Apache-2.0
 """CuteDSL MegaMoE NVFP4 kernel package.
 
-Hosts the ported MegaMoE fused dispatch + FC1 + activation + FC2 + combine
-CuteDSL kernel (flattened from the upstream ``moe_nvfp4_swapab/`` + ``src/``
-split). The package is loaded
+Hosts the MegaMoE fused dispatch + FC1 + activation + FC2 + combine CuteDSL
+kernel. The package is loaded
 lazily by :mod:`tensorrt_llm._torch.modules.fused_moe.mega_moe.mega_moe_cute_dsl`
 through :func:`import_kernel` so environments without a CUDA 13 Cutlass DSL
 runtime can still import the backend file for capability probing.
@@ -39,7 +38,6 @@
     "from_blocked",
     "import_kernel",
     "import_sym_buffer_host",
-    "import_topk_reduce",
     "stack_byte_reinterpretable_tensors",
     "to_blocked",
 ]
@@ -72,20 +70,3 @@ def import_sym_buffer_host():
     # SymBufferHost lives at module scope as a factory; the upstream API
     # constructs the per-world-size variant inside sym_buffer.py.
     return sym_buffer
-
-
-def import_topk_reduce():
-    """Lazily import the standalone CuteDSL top-k reduce kernel API.
-
-    Returns ``(compile_topk_reduce, launch_compiled_topk_reduce)`` from
-    :mod:`.topk_reduce` (mirrors :func:`import_kernel`). The reduce kernel
-    is only needed by the opt-in transformers graph
-    (``apply_topk_in_fc1=False``); the deepgemm-default route reduces on
-    the host via ``combine_output.sum(dim=1)`` and never imports it. Like
-    ``import_kernel`` this stays lazy so non-SM100 / no-cutlass-dsl
-    environments can import the backend for capability probing without
-    pulling the heavyweight CuteDSL symbols.
-    """
-    from .topk_reduce import compile_topk_reduce, launch_compiled_topk_reduce
-
-    return compile_topk_reduce, launch_compiled_topk_reduce

tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py

@@ -1,5 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
 """Codegen-time finite mapping contracts for RMEM tensor handoff."""
 
 from __future__ import annotations
@@ -78,9 +80,8 @@ def linearize(self, coord: Sequence[int]) -> int:
         """Convert a coordinate tuple into a CuTe-style linear index."""
         coord_tuple = tuple(coord)
         if len(coord_tuple) != self.rank:
-            raise ContractError(
-                f"Coordinate rank mismatch for {self.names!r}: {len(coord_tuple)} != {self.rank}"
-            )
+            raise ContractError(f"Coordinate rank mismatch for {self.names!r}: "
+                                f"{len(coord_tuple)} != {self.rank}")
 
         linear = 0
         stride = 1
@@ -162,8 +163,8 @@ def normalize(self, *, domain: Space, codomain: Space) -> tuple[int, ...]:
         """Validate and return the canonical table for the given spaces."""
         if len(self.table) != domain.size:
             raise ContractError(
-                f"TableMapping length must equal domain size {domain.size}, got {len(self.table)}"
-            )
+                f"TableMapping length must equal domain size {domain.size}, "
+                f"got {len(self.table)}")
         for idx, value in enumerate(self.table):
             if value < 0 or value >= codomain.size:
                 raise ContractError(
@@ -395,3 +396,27 @@ class TensorWithContract:
 
     tensor: Any
     contract: Contract
+
+
+def eval_function_mapping(contract: Contract, **domain_coord) -> dict:
+    """Evaluate a FunctionMapping contract at runtime."""
+    if not isinstance(contract.mapping, FunctionMapping):
+        raise TypeError("runtime contract eval requires a FunctionMapping")
+
+    result = contract.mapping.function(**domain_coord)
+    if isinstance(result, dict):
+        return result
+    if isinstance(result, (tuple, list)):
+        if len(result) != contract.codomain.rank:
+            raise ValueError(
+                "FunctionMapping result rank does not match codomain rank: "
+                f"{len(result)} vs {contract.codomain.rank}")
+        return {
+            name: result[i]
+            for i, name in enumerate(contract.codomain.names)
+        }
+    if contract.codomain.rank == 1:
+        return {contract.codomain.names[0]: result}
+    raise TypeError(
+        "FunctionMapping runtime eval must return dict/tuple/list, or scalar "
+        "for rank-1 codomain")

tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/custom_ext.py

@@ -1,5 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
 """Sched extension for fused fc1+fc2 work-tile enrichment and GMEM slicing."""
 
 from typing import List, Optional, Tuple, Union
@@ -37,10 +39,9 @@ def __init__(
         cumulative_data_physical_row: Int32,
         cumulative_sf_physical_row: Int32,
         cumulative_token_block_count: Int32,
-        valid_tokens_in_tile: Int32,
+        valid_tokens_in_cta_tile: Int32,
         phase_and_peek: Int32,
     ):
-        # Slot 3 reuses base k_tile_cnt storage.
         super().__init__(
             expert_idx,
             tile_m_idx,
@@ -50,7 +51,7 @@ def __init__(
         self.cumulative_data_physical_row = self.k_tile_cnt
         self.cumulative_sf_physical_row = cumulative_sf_physical_row
         self.cumulative_token_block_count = cumulative_token_block_count
-        self.valid_tokens_in_tile = valid_tokens_in_tile
+        self.valid_tokens_in_cta_tile = valid_tokens_in_cta_tile
         # Slot 7 is the packed (BlockPhase | (peek_ready << 16)) field.
         # The ``.phase`` and ``.peek_ready`` properties below unpack it;
         # consumers call them directly so the codebase reads as if the
@@ -81,14 +82,14 @@ def __extract_mlir_values__(self) -> List[ir.Value]:
         values = super().__extract_mlir_values__()
         values.extend(extract_mlir_values(self.cumulative_sf_physical_row))
         values.extend(extract_mlir_values(self.cumulative_token_block_count))
-        values.extend(extract_mlir_values(self.valid_tokens_in_tile))
+        values.extend(extract_mlir_values(self.valid_tokens_in_cta_tile))
         values.extend(extract_mlir_values(self.phase_and_peek))
         return values
 
     def __new_from_mlir_values__(
             self, values: List[ir.Value]) -> "SwapABSwigluFp4Fc12WorkTileInfo":
         assert len(values) == 8
-        return SwapABSwigluFp4Fc12WorkTileInfo(
+        return type(self)(
             expert_idx=new_from_mlir_values(self.expert_idx, [values[0]]),
             tile_m_idx=new_from_mlir_values(self.tile_m_idx, [values[1]]),
             tile_n_idx=new_from_mlir_values(self.tile_n_idx, [values[2]]),
@@ -98,8 +99,8 @@ def __new_from_mlir_values__(
                 self.cumulative_sf_physical_row, [values[4]]),
             cumulative_token_block_count=new_from_mlir_values(
                 self.cumulative_token_block_count, [values[5]]),
-            valid_tokens_in_tile=new_from_mlir_values(self.valid_tokens_in_tile,
-                                                      [values[6]]),
+            valid_tokens_in_cta_tile=new_from_mlir_values(
+                self.valid_tokens_in_cta_tile, [values[6]]),
             phase_and_peek=new_from_mlir_values(self.phase_and_peek,
                                                 [values[7]]),
         )
@@ -112,7 +113,7 @@ def to_rmem(self) -> cute.Tensor:
         rmem[3] = self.k_tile_cnt  # = cumulative_data_physical_row
         rmem[4] = self.cumulative_sf_physical_row
         rmem[5] = self.cumulative_token_block_count
-        rmem[6] = self.valid_tokens_in_tile
+        rmem[6] = self.valid_tokens_in_cta_tile
         rmem[7] = self.phase_and_peek
         return rmem
 
@@ -125,7 +126,7 @@ def from_rmem(cls, rmem: cute.Tensor) -> "SwapABSwigluFp4Fc12WorkTileInfo":
             cumulative_data_physical_row=rmem[3],  # type: ignore[arg-type]
             cumulative_sf_physical_row=rmem[4],  # type: ignore[arg-type]
             cumulative_token_block_count=rmem[5],  # type: ignore[arg-type]
-            valid_tokens_in_tile=rmem[6],  # type: ignore[arg-type]
+            valid_tokens_in_cta_tile=rmem[6],  # type: ignore[arg-type]
             phase_and_peek=rmem[7],  # type: ignore[arg-type]
         )
 
@@ -160,7 +161,7 @@ def __init__(
         # shows enough arrivals.  Mirrors ``fc1_done_counter_ptr`` for the
         # fc1->fc2 link: this side is "fc1 input ready", that side is
         # "fc1 output done".  The threshold per-tile is the tile's
-        # ``valid_tokens_in_tile`` (dispatch does not pull padding
+        # ``valid_tokens_in_cta_tile`` (dispatch does not pull padding
         # tokens), read straight off the base work tile -- no separate
         # threshold field needed.  ``None`` in the lean fc1+fc2 path keeps
         # ``enrich_work_tile_info`` to its existing fc2-only peek shape and
@@ -236,7 +237,7 @@ def enrich_work_tile_info(
           ``cumulative_token_block_count + tile_n_idx`` against
           ``self.fc2_spin_threshold`` (work-tile-invariant const).
         - fc1 tiles peek the dispatch->fc1 ``fc1_ready_counter`` at the
-          same slot index but with ``valid_tokens_in_tile`` as threshold
+          same slot index (``tile_n_idx``) but with ``valid_tokens_in_cta_tile`` as threshold
           (per-tile dynamic).  This branch only emits when
           ``self.fc1_ready_counter_ptr is not None`` (MegaMoE mode).
         """
@@ -247,8 +248,7 @@ def enrich_work_tile_info(
         if is_valid:
             # Same slot index for both phases -- fc1 release-add (dispatch
             # pull) and fc2 release-add (fc1 epi) target the per-task-tile
-            # counter slot indexed by ``cumulative_token_block_count +
-            # tile_n_idx``.
+            # counter slot indexed by ``cumulative_token_block_count + tile_n_idx``.
             counter_slot = (base_work.cumulative_token_block_count +
                             base_work.tile_n_idx)
             is_fc1 = base_work.phase == Int32(int(BlockPhase.Linear1))
@@ -257,14 +257,14 @@ def enrich_work_tile_info(
             # MegaMoE-only: fc1 phase peek on fc1_ready_counter.  Threshold
             # is dynamic (per-tile valid count) because dispatch does not
             # pull padding tokens, so the counter's terminal value matches
-            # the tile's valid_tokens_in_tile (cluster_tile_m for full
+            # the tile's valid_tokens_in_cta_tile (cluster_tile_m for full
             # tiles, less for an expert's last partial tile).
             if cutlass.const_expr(self.fc1_ready_counter_ptr is not None):
                 if is_fc1:
                     counter_ptr = self.fc1_ready_counter_ptr + counter_slot
                     peek_ready = spin_wait(
                         counter_ptr,
-                        lambda v: v >= base_work.valid_tokens_in_tile,
+                        lambda v: v >= base_work.valid_tokens_in_cta_tile,
                         peek_only=True,
                     )
                     peek_bit = Int32(0)
@@ -299,7 +299,7 @@ def enrich_work_tile_info(
             cumulative_data_physical_row=base_work.cumulative_data_physical_row,
             cumulative_sf_physical_row=base_work.cumulative_sf_physical_row,
             cumulative_token_block_count=base_work.cumulative_token_block_count,
-            valid_tokens_in_tile=base_work.valid_tokens_in_tile,
+            valid_tokens_in_cta_tile=base_work.valid_tokens_in_cta_tile,
             phase_and_peek=new_phase_and_peek,
         )