diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 705d8af9091d..8d029fdedcf8 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/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 diff --git a/ATTRIBUTIONS-Python.md b/ATTRIBUTIONS-Python.md index f657a16809c7..28cd9ce24f91 100644 --- a/ATTRIBUTIONS-Python.md +++ b/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` diff --git a/constraints.txt b/constraints.txt index 4b920223cf9c..e1e5813855de 100644 --- a/constraints.txt +++ b/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. diff --git a/docker/Dockerfile.multi b/docker/Dockerfile.multi index 024b5bd37934..e49c76710816 100644 --- a/docker/Dockerfile.multi +++ b/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 diff --git a/requirements.txt b/requirements.txt index 1c47ff6183df..f965c4ac38a3 100644 --- a/requirements.txt +++ b/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 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_tasks.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_tasks.py index 4321a0d0d46e..a84b92985204 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_tasks.py +++ b/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,7 +774,6 @@ 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): @@ -783,7 +781,6 @@ def _run_packed_skip_iteration( tail_entries, work_tile, bookkeeping_domain, - context, ) @cute.jit @@ -791,14 +788,12 @@ 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 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/tasks.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/tasks.py index dba6349c05d1..cadd4fb66b5f 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/tasks.py +++ b/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,10 +188,7 @@ 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 @@ -204,13 +196,13 @@ def _run_task_body_persistent( 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) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/moe_as_dense_gemm/fc2.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/moe_as_dense_gemm/fc2.py index 669d070268d9..88f969258c86 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/moe_as_dense_gemm/fc2.py +++ b/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(), diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/dynamic_mainloop.py b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/dynamic_mainloop.py index 834c5a9cc576..d6b7a95f57da 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/dynamic_mainloop.py +++ b/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)