[None][feat] Add Rubin SM107 CuTe DSL foundation and BF16 kernels

cpp/tensorrt_llm/thop/cuteDslMoeUtilsOp.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2025-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -66,6 +66,8 @@ std::vector<torch::Tensor> moe_topk_sort_impl(torch::optional<torch::Tensor> con
     auto total_num_padded_tokens = torch::empty({1}, torch::dtype(torch::kInt32).device(torch::kCUDA));
     auto expanded_idx_to_permuted_idx
         = torch::empty({num_tokens, top_k}, torch::dtype(torch::kInt32).device(torch::kCUDA));
+    // The routing kernel writes only valid entries. Consumers use num_non_exiting_tiles and tile_idx_to_mn_limit to
+    // guard the padding region, so these tensors do not require initialization.
     auto permuted_idx_to_expanded_idx
         = torch::empty({max_num_padded_tokens}, torch::dtype(torch::kInt32).device(torch::kCUDA));
     auto num_tokens_per_expert = torch::empty({num_experts}, torch::dtype(torch::kInt32).device(torch::kCUDA));

tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py

@@ -51,7 +51,7 @@
 
 from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs, AttentionInputType
 from tensorrt_llm._torch.attention_backend.sparse.skip_softmax import SkipSoftmaxParams
-from tensorrt_llm._utils import get_sm_version, is_sm_100f, torch_dtype_to_binding
+from tensorrt_llm._utils import get_sm_version, torch_dtype_to_binding
 from tensorrt_llm.bindings import DataType
 from tensorrt_llm.bindings.internal import thop
 from tensorrt_llm.functional import AttentionMaskType
@@ -505,7 +505,7 @@ def is_available(cls, attn: "TrtllmAttention") -> bool:
             return False
 
         sm = get_sm_version()
-        if not is_sm_100f(sm):
+        if sm not in (100, 103):
             logger.debug(
                 f"FlashInfer TRTLLM-Gen FMHA is unavailable: requires SM100 or SM103, got SM{sm}."
             )

tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py

@@ -8490,6 +8490,8 @@ def forward(
             #   B stride  = a_tensor.stride(0)
             a_stride_m = a_tensor.stride(1)
             a_stride_batch = a_tensor.stride(0)
+            b_stride_n = b_tensor.stride(1)
+            b_stride_batch = b_tensor.stride(0)
 
             if not self.use_tvm_ffi:
                 a_ptr = make_ptr(
@@ -8555,6 +8557,8 @@ def forward(
                     c_cute_tensor,
                     a_stride_m,
                     a_stride_batch,
+                    b_stride_n,
+                    b_stride_batch,
                     max_active_clusters=max_active_clusters,
                     stream=stream,
                     options="--opt-level 2 --enable-tvm-ffi"
@@ -8576,6 +8580,8 @@ def forward(
                     c_tmp,
                     a_stride_m,
                     a_stride_batch,
+                    b_stride_n,
+                    b_stride_batch,
                 )
             else:
                 compiled_gemm(
@@ -8588,6 +8594,8 @@ def forward(
                     c_cute_tensor,
                     a_stride_m,
                     a_stride_batch,
+                    b_stride_n,
+                    b_stride_batch,
                     stream=stream,
                 )
 

tensorrt_llm/_torch/cute_dsl_kernels/blackwell/custom_pipeline.py

@@ -60,6 +60,12 @@
 
 _make_sync_object = _Sm100PipelineFactory._make_sync_object
 
+try:
+    from cutlass.pipeline.sm100 import PipelineTmaUmma as _Sm100PipelineTmaUmma
+    _sm100_make_sync = _Sm100PipelineTmaUmma._make_sync_object
+except (ImportError, AttributeError):
+    _sm100_make_sync = PipelineAsync._make_sync_object
+
 
 def pipeline_init_wait(cta_layout_vmnk: Optional[cute.Layout] = None):
     """Initializes the mbarrier and synchronizes the threadblock or cluster.
@@ -188,9 +194,9 @@ def create(
         producer = (producer_type, producer_group)
         consumer = (consumer_type, consumer_group)
 
-        sync_object_full = _make_sync_object(barrier_storage.align(min_align=8),
-                                             num_stages, producer, tx_count)
-        sync_object_empty = _make_sync_object(
+        sync_object_full = _sm100_make_sync(barrier_storage.align(min_align=8),
+                                            num_stages, producer, tx_count)
+        sync_object_empty = _sm100_make_sync(
             barrier_storage.align(min_align=8) + num_stages, num_stages,
             consumer)
 
@@ -343,9 +349,9 @@ def create(
         producer = (producer_type, producer_group)
         consumer = (consumer_type, consumer_group)
 
-        sync_object_full = _make_sync_object(barrier_storage.align(min_align=8),
-                                             num_stages, producer)
-        sync_object_empty = _make_sync_object(
+        sync_object_full = _sm100_make_sync(barrier_storage.align(min_align=8),
+                                            num_stages, producer)
+        sync_object_empty = _sm100_make_sync(
             barrier_storage.align(min_align=8) + num_stages, num_stages,
             consumer)
 
@@ -496,12 +502,12 @@ def create(
         producer = (producer_type, producer_group)
         consumer = (consumer_type, consumer_group)
 
-        sync_object_full = _make_sync_object(
+        sync_object_full = _sm100_make_sync(
             barrier_storage.align(min_align=8),
             num_stages,
             producer,
         )
-        sync_object_empty = _make_sync_object(
+        sync_object_empty = _sm100_make_sync(
             barrier_storage.align(min_align=8) + num_stages, num_stages,
             consumer)
 

tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_act_fusion.py

@@ -1450,7 +1450,6 @@ def kernel(
             c_producer_group = pipeline.CooperativeGroup(
                 pipeline.Agent.Thread,
                 32 * len(self.epilog_warp_id),
-                32 * len(self.epilog_warp_id),
             )
             c_pipeline = pipeline.PipelineTmaStore.create(
                 num_stages=self.num_c_stage,

tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py

@@ -43,7 +43,7 @@
 
 # This file is copied and modified from cutlass example https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py
 
-from typing import Optional, Tuple, Type, Union
+from typing import Literal, Optional, Tuple, Type, Union
 
 import cuda.bindings.driver as cuda
 import cutlass
@@ -53,7 +53,8 @@
 import cutlass.utils.blackwell_helpers as sm100_utils
 import cutlass.utils.blockscaled_layout as blockscaled_utils
 from cutlass._mlir.dialects import llvm
-from cutlass.cute.nvgpu import cpasync, tcgen05
+from cutlass.cute.nvgpu import OperandMajorMode, cpasync, tcgen05
+from cutlass.cute.runtime import make_ptr
 from cutlass.cutlass_dsl import dsl_user_op
 
 from .custom_pipeline import PipelineTmaUmma, PipelineUmmaAsync
@@ -197,13 +198,23 @@ def __init__(
         self.indexer_extra_warp_id = tuple(
             range(self.tma_warp_id + 1, self.tma_warp_id + 1 +
                   indexer_transform_warps - 4)) if self.indexer_q_fusion else ()
-        self.threads_per_cta = 32 * len(
+        self.threads_per_warp = 32
+        self.threads_per_cta = self.threads_per_warp * len(
             (self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id,
              *self.indexer_extra_warp_id))
         # Set barrier id for cta sync, epilogue sync and tmem ptr sync
         self.cta_sync_bar_id = 0
         self.epilog_sync_bar_id = 1
         self.tmem_ptr_sync_bar_id = 2
+        self.epilog_sync_barrier = pipeline.NamedBarrier(
+            barrier_id=self.epilog_sync_bar_id,
+            num_threads=self.threads_per_warp * len(self.epilog_warp_id),
+        )
+        self.tmem_alloc_barrier = pipeline.NamedBarrier(
+            barrier_id=self.tmem_ptr_sync_bar_id,
+            num_threads=self.threads_per_warp * len(
+                (self.mma_warp_id, *self.epilog_warp_id)),
+        )
         self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
         SM100_TMEM_CAPACITY_COLUMNS = 512
         self.num_tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS
@@ -1415,7 +1426,6 @@ def kernel(
             c_producer_group = pipeline.CooperativeGroup(
                 pipeline.Agent.Thread,
                 32 * len(self.epilog_warp_id),
-                32 * len(self.epilog_warp_id),
             )
             c_pipeline = pipeline.PipelineTmaStore.create(
                 num_stages=self.num_c_stage,
@@ -2088,15 +2098,28 @@ def _compute_grid(
         cluster_shape_mnl = (*cluster_shape_mn, 1)
 
         tile_sched_params = utils.PersistentTileSchedulerParams(
-            num_ctas_mnl,
-            cluster_shape_mnl,
-            swizzle_size=swizzle_size,
-            raster_along_m=raster_along_m)
+            num_ctas_mnl, cluster_shape_mnl, swizzle_size, raster_along_m)
         grid = utils.StaticPersistentTileScheduler.get_grid_shape(
             tile_sched_params, max_active_clusters)
 
         return tile_sched_params, grid
 
+    @staticmethod
+    def needs_unpack_tma(
+        a_dtype: Type[cutlass.Numeric],
+        b_dtype: Type[cutlass.Numeric],
+    ) -> bool:
+        """Decide whether TMA must use the UNPACK_U8 variant for operands.
+
+        Unpack is required when operand widths differ or either operand is
+        6-bit. Otherwise TMA can use the natural packed format.
+        """
+        if a_dtype.width != b_dtype.width:
+            return True
+        if a_dtype.width == 6 or b_dtype.width == 6:
+            return True
+        return False
+
     @staticmethod
     def is_valid_dtypes_and_scale_factor_vec_size(
         ab_dtype: Type[cutlass.Numeric],
@@ -2520,3 +2543,47 @@ def cvt_sf_MKL_to_M32x4xrm_K4xrk_L(
     for i in cutlass.range(cute.size(sf_ref_tensor)):
         mkl_coord = sf_ref_tensor.layout.get_hier_coord(i)
         sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord]
+
+
+def scaled_mm(
+    gemm_obj: Sm100BlockScaledPersistentDenseGemmKernel,
+    a_dtype: Type[cutlass.Numeric],
+    b_dtype: Type[cutlass.Numeric],
+    c_dtype: Type[cutlass.Numeric],
+    sf_dtype: Type[cutlass.Numeric],
+    a_major: Literal["m", "k"],
+    b_major: Literal["n", "k"],
+    c_major: Literal["m", "n"],
+    max_active_clusters: cutlass.Constexpr,
+    stream: cuda.CUstream,
+    epilogue_op: cutlass.Constexpr = lambda x: x,
+    options: str = "",
+):
+    """Compile the persistent dense blockscaled GEMM operation."""
+    a_ptr = make_ptr(a_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
+    b_ptr = make_ptr(b_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
+    c_ptr = make_ptr(c_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
+    sfa_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32)
+    sfb_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32)
+
+    a_major_mode = (OperandMajorMode.K
+                    if a_major == "k" else OperandMajorMode.MN)
+    b_major_mode = (OperandMajorMode.K
+                    if b_major == "k" else OperandMajorMode.MN)
+    c_layout = (utils.LayoutEnum.ROW_MAJOR
+                if c_major == "n" else utils.LayoutEnum.COL_MAJOR)
+    return cute.compile(
+        gemm_obj,
+        a_ptr,
+        b_ptr,
+        sfa_ptr,
+        sfb_ptr,
+        c_ptr,
+        (a_major_mode, b_major_mode, c_layout),
+        (cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(0),
+         cutlass.Int32(0)),
+        max_active_clusters,
+        stream,
+        epilogue_op,
+        options=options,
+    )