[None][chore] Update cutedsl to 4.8.0 dev

3rdparty/vendor_sources.lock.yaml

@@ -1,13 +1,13 @@
 schema_version: 1
 vendors:
   flashinfer-prims-ts:
-    url: https://github.com/yuxianq/flashinfer.git
+    url: https://github.com/rosong11/flashinfer.git
     branch: trtllm-prims-ts-dev
-    commit: e500966b575ab83db7c0e84e5a0f8fde6a4f3505
+    commit: bad2bdb15aac3553934e7a7a164dcb2ca4fda7f1
     source: flashinfer/attention/prims_ts
     destination: tensorrt_llm/_torch/attention/backends/prims_ts
     include:
     - '**/*.py'
     patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch
     patch_digest: sha256:0e2f58c6633f57fee03df42049bc78d4d038063b810ad3a2f0ba6d62f8183887
-    digest: sha256-tree-v1:e9af5482f6406af3128e711c3fb2d044359fb6e1d1bc9907743dc86d006d23ac
+    digest: sha256-tree-v1:5766418e4de80e50f6c64a9f3d37076e9b7e815118b5d8a7a22c27d321fd7576

ATTRIBUTIONS-Python.md

@@ -33239,7 +33239,7 @@ License: `NVIDIA Proprietary Software`
   - `Homepage`: https://developer.nvidia.com/cusparselt
 
 
-## nvidia-cutlass-dsl (4.7.0)
+## nvidia-cutlass-dsl (4.8.0.dev0)
 
 ### Licenses
 License: `None`

constraints.txt

@@ -1,13 +1,13 @@
 # setup.py appends these constraints to wheel install_requires, so every entry
 # must also be a valid TensorRT-LLM runtime dependency outside the base image.
 # Keep the base-image and wheel CuTe DSL stack aligned with requirements.txt.
-nvidia-cutlass-dsl==4.7.0
+nvidia-cutlass-dsl==4.8.0.dev0
 # Keep the CuTe DSL and FlashAttention 4 runtime stack aligned with requirements.txt.
 apache-tvm-ffi==0.1.13.post2
 # Newer published QuACK releases constrain CUTLASS DSL to 4.6.x.
 quack-kernels==0.5.0
 torch-c-dlpack-ext==0.1.3
-# CUTLASS DSL 4.7 libraries require protobuf 6.x.
+# CUTLASS DSL libraries require protobuf 6.x.
 protobuf>=6.30.2,<7
 # The `nvidia-cutlass-dsl` package does not pin numpy at all, which can be problematic in certain CI
 # stages.

docker/Dockerfile.multi

@@ -77,7 +77,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
         nvidia-cutlass-dsl-libs-core nvidia-cutlass-dsl-libs-cu12 nvidia-cutlass-dsl-libs-cu13 numpy wandb || true && \
     # Remove any leftover namespace dirs or dist-info that pip missed
     rm -rf $(python3 -c "import site; print(site.getsitepackages()[0])")/nvidia_cutlass_dsl* && \
-    pip3 install -r /tmp/constraints.txt "nvidia-cutlass-dsl[cu13]==4.7.0" && \
+    pip3 install -r /tmp/constraints.txt "nvidia-cutlass-dsl[cu13]==4.8.0.dev0" && \
     rm /tmp/constraints.txt
 
 # Install UCX, NIXL, etcd

requirements.txt

@@ -80,15 +80,16 @@ xdsl>=0.59.0 # Optional: required for MLIR-based elementwise fusion in AutoDeplo
 tiktoken
 blobfile
 openai-harmony==0.0.4
-nvidia-cutlass-dsl[cu13]==4.7.0; python_version >= "3.10" # required by PrimTS task-scheduled kernels
+# 4.8.0.dev0 is the first release exposing Rubin/SM107.
+nvidia-cutlass-dsl[cu13]==4.8.0.dev0; python_version >= "3.10" # required by PrimTS task-scheduled kernels and Rubin CuTe DSL kernels
 nvidia-matmul-heuristics==0.1.0.27; python_version >= "3.10" # analytic GEMM heuristics for CuTe DSL autotuner tactic pruning
-quack-kernels==0.5.0; python_version >= "3.10" # required for MiniMax-M3 MSA and FlashAttention 4; only published release compatible with CUTLASS DSL 4.7
+quack-kernels==0.5.0; python_version >= "3.10" # required for MiniMax-M3 MSA and FlashAttention 4; only published release compatible with CUTLASS DSL >= 4.7
 jinja2 # required for MinimaxM3 MSA
 plotly
 numexpr
 partial_json_parser
 mcp<2.0.0 # 2.0.0 removed mcp.server.fastmcp; scaffolding still uses the 1.x APIs
-apache-tvm-ffi==0.1.13.post2 # required by CUTLASS DSL 4.7's TVM-FFI provider
+apache-tvm-ffi==0.1.13.post2 # required by the CUTLASS DSL TVM-FFI provider
 torch-c-dlpack-ext==0.1.3 # used for reduce nvidia-cutlass-dsl host overhead, optional package for improved torch tensor calling perf
 flash-attn-4==4.0.0b19
 mistral-common>=1.10.0

tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_tasks.py

@@ -760,7 +760,6 @@ def _refresh_packed_q_bounds(self, work_tile: Any) -> None:
     def _run_packed_skip_iteration(
         self,
         work_tile: Any,
-        context: ResourceContext | None = None,
     ) -> None:
         """Advance one inactive tile through WorkQueue bookkeeping only."""
         # Packed schedules place every data-path entry inside ``skippable()``;
@@ -775,30 +774,26 @@ def _run_packed_skip_iteration(
                     head_entries,
                     work_tile,
                     bookkeeping_domain,
-                    context,
                 )
         for is_skippable_tail, tail_entries in self._tail_exec_groups:
             if cutlass.const_expr(not is_skippable_tail):
                 self._run_tail_entry_group(
                     tail_entries,
                     work_tile,
                     bookkeeping_domain,
-                    context,
                 )
 
     @cute.jit
     def _run_task_body_impl(
         self,
         work_tile: cute.Coord,
         skip_work_tile: Any = None,
-        context: ResourceContext | None = None,
     ) -> None:
         """Run one ordinary task tile and synchronize attention-sink tails."""
         Task._run_task_body_impl(
             self,
             work_tile,
             skip_work_tile,
-            context=context,
         )
         if cutlass.const_expr(
             self.cfg is not None
@@ -826,10 +821,7 @@ def _run_task_body_impl(
                 prims.barrier_cta_sync(12, thread_count=16 * 32)
 
     @cute.jit
-    def _run_task_body_persistent(
-        self,
-        context: ResourceContext | None = None,
-    ) -> None:
+    def _run_task_body_persistent(self) -> None:
         """Drain inactive packed tiles before each unconditional active body."""
         use_packed_early_stop = (
             self.cfg is not None
@@ -838,14 +830,14 @@ def _run_task_body_persistent(
             and self._has_skip_if
         )
         if cutlass.const_expr(not use_packed_early_stop):
-            Task._run_task_body_persistent(self, context)
+            Task._run_task_body_persistent(self)
             return
 
         assert self.work_queue is not None
         work_tile = self.work_queue.initial_work_tile_info()
         self.work_queue._set_consumer_var_from_ts("work_tile", work_tile)
 
-        self._run_pre_work_loop_entries(work_tile, context)
+        self._run_pre_work_loop_entries(work_tile)
         work_tile = self.work_queue._get_consumer_var_from_ts("work_tile")
         for resource in self.dst_resources:
             if cutlass.const_expr(
@@ -859,7 +851,7 @@ def _run_task_body_persistent(
         # inner loop executes only the non-skippable WorkQueue tail, so no TMA,
         # descriptor, pipeline, task data, or sink barrier is issued.
         while work_tile.is_valid_tile and self._should_skip_work_tile(work_tile):
-            self._run_packed_skip_iteration(work_tile, context)
+            self._run_packed_skip_iteration(work_tile)
             work_tile = self.work_queue._get_consumer_var_from_ts("work_tile")
             self.dummy = cutlass.Boolean(True)
 
@@ -868,18 +860,18 @@ def _run_task_body_persistent(
             # The tile is known active here. Running the complete schedule
             # without a dynamic skip guard keeps HEAD-produced pipeline state
             # in scope for LOOP and TAIL.
-            Task._run_task_body_impl(self, work_tile, None, context=context)
+            Task._run_task_body_impl(self, work_tile, None)
             if cutlass.const_expr(self.cfg.use_attention_sinks):
                 prims.barrier_cta_sync(12, thread_count=16 * 32)
             work_tile = self.work_queue._get_consumer_var_from_ts("work_tile")
             self.dummy = cutlass.Boolean(True)
 
             while work_tile.is_valid_tile and self._should_skip_work_tile(work_tile):
-                self._run_packed_skip_iteration(work_tile, context)
+                self._run_packed_skip_iteration(work_tile)
                 work_tile = self.work_queue._get_consumer_var_from_ts("work_tile")
                 self.dummy = cutlass.Boolean(True)
 
-        self._run_post_work_loop_entries(work_tile, context)
+        self._run_post_work_loop_entries(work_tile)
         for resource in self.dst_resources:
             if cutlass.const_expr(
                 resource.pipeline_config is not None

tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/tasks.py

@@ -43,7 +43,6 @@
 import cutlass
 import cutlass.cute as cute
 
-from cutlass.experimental.task_scheduling.memory import ResourceContext
 from cutlass.experimental.task_scheduling.schedule_builder import (
     domain_loop,
     schedule,
@@ -150,18 +149,14 @@ class MlaTask(Task):
     """Task subclass that recomputes MLA k-domain per persistent work tile."""
 
     @cute.jit
-    def _run_one_mla_work_tile(
-        self,
-        work_tile,
-        context: ResourceContext | None = None,
-    ) -> None:
+    def _run_one_mla_work_tile(self, work_tile) -> None:
         """Run a persistent work tile using the cached split-KV domain."""
 
         # WorkQueue decomposes the MLA persistent tile and caches the K-domain.
         # Keep task bodies on that cached value so page-offset/TMA/MMA/softmax
         # paths do not each rebuild the same split-KV arithmetic.
         self.domain = work_tile.k_tile_count
-        self._run_task_body_impl(work_tile, context=context)
+        self._run_task_body_impl(work_tile)
 
     @cute.jit
     def _drain_mla_work_tile_tails(self) -> None:
@@ -193,24 +188,21 @@ def _drain_mla_work_tile_tails(self) -> None:
             self.work_queue.producer_tail()
 
     @cute.jit
-    def _run_task_body_persistent(
-        self,
-        context: ResourceContext | None = None,
-    ) -> None:
+    def _run_task_body_persistent(self) -> None:
         """Schedule one or more persistent work tiles for this task instance."""
 
         params = self.work_queue.tile_sched_params
 
         if cutlass.const_expr(not params.is_persistent):
             work_tile = self.work_queue._work_tile_from_block_idx(cute.arch.block_idx())
             self.work_queue._set_consumer_var_from_ts("work_tile", work_tile)
-            self._run_pre_work_loop_entries(work_tile, context)
+            self._run_pre_work_loop_entries(work_tile)
             # Runtime K/Q metadata can make a statically launched split empty.
             # Keep the CTA on the ordinary initialized-pipeline path, but skip
             # its captured HEAD/LOOP/TAIL data work when the domain is zero.
             if work_tile.k_tile_count > cutlass.Int32(0):
-                self._run_one_mla_work_tile(work_tile, context)
-            self._run_post_work_loop_entries(work_tile, context)
+                self._run_one_mla_work_tile(work_tile)
+            self._run_post_work_loop_entries(work_tile)
             self._drain_mla_work_tile_tails()
             return
 
@@ -224,7 +216,7 @@ def _run_task_body_persistent(
         work_tile = self.work_queue._work_tile_from_linear_idx(current_work_linear_idx)
         self.work_queue._set_consumer_var_from_ts("work_tile", work_tile)
 
-        self._run_pre_work_loop_entries(work_tile, context)
+        self._run_pre_work_loop_entries(work_tile)
         while current_work_linear_idx < num_blocks:
             work_tile.update_from(
                 self.work_queue._work_tile_from_linear_idx(current_work_linear_idx)
@@ -235,7 +227,7 @@ def _run_task_body_persistent(
             # splits empty. Skip them and continue grid-striding rather than
             # running a captured HEAD/TAIL sequence with domain zero.
             if work_tile.k_tile_count > cutlass.Int32(0):
-                self._run_one_mla_work_tile(work_tile, context)
+                self._run_one_mla_work_tile(work_tile)
 
             # Each warp branch advances from the same scalar tile id, keeping
             # the persistent loop state compact across task bodies.
@@ -247,7 +239,7 @@ def _run_task_body_persistent(
             self.work_queue._work_tile_from_linear_idx(current_work_linear_idx)
         )
         self.work_queue._set_consumer_var_from_ts("work_tile", work_tile)
-        self._run_post_work_loop_entries(work_tile, context)
+        self._run_post_work_loop_entries(work_tile)
         self._drain_mla_work_tile_tails()
 
 
@@ -264,11 +256,7 @@ def __init__(self, *args, **kwargs) -> None:
         self._fixed_loop_end = cutlass.Int32(self.domain_start)
 
     @cute.jit
-    def _run_one_mla_work_tile(
-        self,
-        work_tile,
-        context: ResourceContext | None = None,
-    ) -> None:
+    def _run_one_mla_work_tile(self, work_tile) -> None:
         """Run one tile by mapping its local offsets onto this fixed lane."""
 
         fixed_lane = cutlass.Int32(self.domain_start)
@@ -281,7 +269,7 @@ def _run_one_mla_work_tile(
         ) // cutlass.Int32(2)
         self._fixed_loop_end = fixed_lane + lane_iterations * cutlass.Int32(2)
         self.domain = self._fixed_loop_end
-        self._run_task_body_impl(work_tile, context=context)
+        self._run_task_body_impl(work_tile)
         self._cumulative_k_parity = (
             self._cumulative_k_parity + self._actual_domain
         ) % cutlass.Int32(2)
@@ -297,7 +285,6 @@ def _create_stage_info(
         label=None,
         schedule_stage=None,
         routing_slot=None,
-        context: ResourceContext | None = None,
     ) -> StageInfo:
         """Return base pipeline state with the work tile's actual K offset."""
 
@@ -311,7 +298,6 @@ def _create_stage_info(
             label,
             schedule_stage,
             routing_slot,
-            context=context,
         )
         actual_loop_offset = self._mapped_domain_start + (
             cutlass.Int32(base_info.loop_offset) - cutlass.Int32(self.domain_start)

tensorrt_llm/_torch/cute_dsl_kernels/blackwell/moe_as_dense_gemm/fc2.py

@@ -815,12 +815,10 @@ def kernel(
         alpha_scale_pipeline_producer_group = pipeline.CooperativeGroup(
             pipeline.Agent.Thread,
             32 * 1,  # alpha_scale_load_warp_id threads
-            32 * 1,
         )
         alpha_scale_pipeline_consumer_group = pipeline.CooperativeGroup(
             pipeline.Agent.Thread,
             32 * len(self.epilog_warp_id),  # epilogue warps
-            32 * len(self.epilog_warp_id),
         )
         alpha_scale_pipeline = pipeline.PipelineCpAsync.create(
             barrier_storage=storage.alpha_scale_load_mbar_ptr.data_ptr(),

tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/dynamic_mainloop.py

@@ -295,17 +295,20 @@ def issue_dynamic_block_scaled_mma_tile(
                     _nvvm_raw.Tcgen05MMAScaleVecSize.X4,
                 }
             else:
+                # CUTLASS DSL 4.8 renamed the tcgen05.mma.block_scale operands
+                # and split Tcgen05MMAScaleVecSize into a 1X/2X/4X selector and
+                # Tcgen05MMABlockScale.
                 nvvm_args = {
-                    "mma_kind":
+                    "kind":
                     _nvvm_raw.Tcgen05MMAKind.MXF4NVF4,
                     "cta_group":
                     _nvvm_raw.CTAGroupKind.CTA_2 if mma_tiler_mnk[0] == 256 else
                     _nvvm_raw.CTAGroupKind.CTA_1,
-                    "d":
+                    "matrix_d":
                     operand_d_ptr,
-                    "a":
+                    "matrix_a":
                     operand_a,
-                    "b":
+                    "matrix_b":
                     operand_b,
                     "idesc":
                     idesc.ir_value(),
@@ -315,7 +318,7 @@ def issue_dynamic_block_scaled_mma_tile(
                     operand_sfa_ptr,
                     "scale_b":
                     operand_sfb_ptr,
-                    "scale_vec_size":
-                    _nvvm_raw.Tcgen05MMAScaleVecSize.BLOCK16,
+                    "block_scale":
+                    _nvvm_raw.Tcgen05MMABlockScale.BLOCK16,
                 }
             nvvm.tcgen05_mma_block_scale(**nvvm_args)