From ad1c7ebf6fb89fe5814b85bda250cbd1428ad051 Mon Sep 17 00:00:00 2001 From: farazkh80 <58580514+farazkh80@users.noreply.github.com> Date: Sat, 5 Sep 2026 13:59:06 +0000 Subject: [PATCH 1/3] [None][feat] Add SM107 CuTe DSL BF16 dense GEMM/BMM custom ops and dispatch The SM107 BF16 persistent dense GEMM kernels landed in #18369 but nothing called them: no custom op wrapped them and the BF16 dispatch sites still routed SM107 to the Blackwell op. Add the custom-op layer and route SM107 to it. - cute_dsl_custom_ops.py: `trtllm::cute_dsl_bf16_gemm_rubin` and `trtllm::cute_dsl_bf16_bmm_rubin` with `CuteDSLBf16RubinGemmRunner` / `CuteDSLBf16RubinBmmRunner`. The runners subclass the Blackwell runners for the shared TunableRunner plumbing but override tactic enumeration and launch in full (preferred-cluster kernel variant, SM107 tactic pruning, direct split-K on the GEMM), so the Blackwell classes are untouched. Both ops raise unless get_sm_version() == 107 and the CuTe DSL package ships the SM107 helpers (IS_CUTLASS_DSL_RUBIN_AVAILABLE). - linear.py (UnquantizedLinearMethod.apply), attention/mla.py (_bmm_bf16_out), modeling_deepseekv3.py (DeepseekV3Gate): pick the `*_rubin` op when get_sm_version() == 107, otherwise the existing `*_blackwell` op. SM100/SM103 call sites are unchanged. - tests: SM107-gated correctness tests (op path, base and preferred-cluster tactics, split-K 2/4/8 in bf16 and fp32 output, strided BMM views) plus dispatch tests that run on every architecture and check the SM107 ops reject other SMs, offer no autotuner tactics, and register fakes. Co-authored-by: Peace He <103117813+peaceh-nv@users.noreply.github.com> Co-authored-by: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com> Signed-off-by: farazkh80 <58580514+farazkh80@users.noreply.github.com> --- tensorrt_llm/_torch/attention/mla.py | 9 +- .../_torch/custom_ops/cute_dsl_custom_ops.py | 761 ++++++++++++++++++ .../_torch/models/modeling_deepseekv3.py | 6 +- tensorrt_llm/_torch/modules/linear.py | 14 +- .../test_cute_dsl_bf16_dense_rubin.py | 249 ++++++ 5 files changed, 1030 insertions(+), 9 deletions(-) create mode 100644 tests/unittest/_torch/thop/parallel/test_cute_dsl_bf16_dense_rubin.py diff --git a/tensorrt_llm/_torch/attention/mla.py b/tensorrt_llm/_torch/attention/mla.py index f7d9e2cc509a..49937775feec 100644 --- a/tensorrt_llm/_torch/attention/mla.py +++ b/tensorrt_llm/_torch/attention/mla.py @@ -1402,9 +1402,14 @@ def forward_context( ) def _bmm_bf16_out(self, a, b_no_transpose, b_transposed, output): - """BMM with optional CuTe DSL bf16 acceleration on Blackwell.""" + """BMM with optional CuTe DSL bf16 acceleration on Blackwell/SM107.""" if self.use_cute_dsl_bf16_bmm and is_sm_100f(): - torch.ops.trtllm.cute_dsl_bf16_bmm_blackwell(a, b_no_transpose, output) + bf16_bmm_op = ( + torch.ops.trtllm.cute_dsl_bf16_bmm_rubin + if get_sm_version() == 107 + else torch.ops.trtllm.cute_dsl_bf16_bmm_blackwell + ) + bf16_bmm_op(a, b_no_transpose, output) else: torch.ops.trtllm.bmm_out(a, b_transposed, output) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 405da76ebebe..a3a54970b388 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -9279,6 +9279,767 @@ def _( assert output.shape == ( m, n), "CuTe DSL bf16 gemm output shape is incorrect" + # ====================================================================== + # BF16 Dense Persistent GEMM / BMM (CuTe DSL) for SM107 + # ====================================================================== + # The SM107 runners subclass the Blackwell runners above only for the + # shared __init__/TunableRunner plumbing; tactic enumeration and launch + # are overridden in full so the Blackwell classes stay untouched. Every + # SM107 op raises unless get_sm_version() == 107 and the CuTe DSL package + # provides the SM107 helpers (IS_CUTLASS_DSL_RUBIN_AVAILABLE). + + if IS_CUTLASS_DSL_RUBIN_AVAILABLE: + from ..cute_dsl_kernels.rubin.dense_bf16_gemm_persistent import \ + PersistentDenseGemmKernel as Sm107Bf16PersistentDenseGemmKernel + from ..cute_dsl_kernels.rubin.dense_bf16_gemm_persistent import \ + PersistentDenseGemmKernelPreferredCluster as \ + Sm107Bf16PersistentDenseGemmKernelPreferredCluster + else: + Sm107Bf16PersistentDenseGemmKernel = None + Sm107Bf16PersistentDenseGemmKernelPreferredCluster = None + + def _is_sm107_cute_dsl_available() -> bool: + return get_sm_version() == 107 and IS_CUTLASS_DSL_RUBIN_AVAILABLE + + def _sm107_bf16_kernel_class(kernel_variant: str): + if not IS_CUTLASS_DSL_RUBIN_AVAILABLE: + raise RuntimeError( + "CuTe DSL SM107 BF16 GEMM requires a CuTe DSL package with " + "SM107 support.") + if kernel_variant == "preferred_cluster": + return Sm107Bf16PersistentDenseGemmKernelPreferredCluster + return Sm107Bf16PersistentDenseGemmKernel + + _SM107_BF16_MMA_TILER_MN_CANDIDATES = [(64, 128), (128, 128), (256, 128), + (256, 256)] + _SM107_BF16_CLUSTER_SHAPE_MN_CANDIDATES = [ + (1, 1), + (1, 2), + (1, 4), + (2, 1), + (2, 2), + (2, 4), + (4, 1), + (4, 2), + (4, 4), + ] + _SM107_BF16_PREFERRED_CLUSTER_SHAPE_MN = (4, 2) + _SM107_BF16_FALLBACK_CLUSTER_SHAPE_MN = (2, 1) + # Cap on A/B pipeline stages; 0 keeps the kernel's auto-computed depth. + _SM107_BF16_MAX_NUM_AB_STAGE_CANDIDATES = [0] + + def _sm107_bf16_valid_tactics( + m: int, + n: int, + k: int, + batch_size: int, + c_dtype_cutlass, + mma_tiler_mn_candi: List[Tuple[int, int]], + split_k_candi: Optional[List[int]], + ) -> List[Tuple]: + """Enumerate SM107 BF16 tactics shared by the GEMM and BMM runners. + + Base tactics are ``("base", use_2cta_instrs, mma_tiler_mn, + cluster_shape_mn, max_num_ab_stage[, split_k_slices])``; the trailing + split-K field is present only when ``split_k_candi`` is given (GEMM). + Preferred-cluster tactics are ``("preferred_cluster", use_2cta_instrs, + mma_tiler_mn, preferred_cluster_shape_mn, fallback_cluster_shape_mn, + max_num_ab_stage)``. + """ + a_major = "k" + b_major = "k" + c_major = "n" + use_2cta_instrs_candi = [False, True] + + valid_tactics = [] + base_kernel_class = _sm107_bf16_kernel_class("base") + for use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, max_num_ab_stage in itertools.product( + use_2cta_instrs_candi, mma_tiler_mn_candi, + _SM107_BF16_CLUSTER_SHAPE_MN_CANDIDATES, + _SM107_BF16_MAX_NUM_AB_STAGE_CANDIDATES): + # CTA_N=256 with cluster_n=2 is an illegal memory access on SM107. + if mma_tiler_mn[1] == 256 and cluster_shape_mn[1] == 2: + continue + # An M-cluster wider than the available M-tiles leaves phantom + # CTAs (e.g. cluster_m=4 on the M=1 decode BMM) -> illegal access. + if not _bf16_cluster_m_fits(m, use_2cta_instrs, mma_tiler_mn, + cluster_shape_mn): + continue + if not base_kernel_class.can_implement( + cutlass.BFloat16, cutlass.Float32, c_dtype_cutlass, + use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, m, n, k, + batch_size, a_major, b_major, c_major): + continue + if split_k_candi is None: + valid_tactics.append(("base", use_2cta_instrs, mma_tiler_mn, + cluster_shape_mn, max_num_ab_stage)) + else: + for split_k_slices in split_k_candi: + valid_tactics.append( + ("base", use_2cta_instrs, mma_tiler_mn, + cluster_shape_mn, max_num_ab_stage, split_k_slices)) + + preferred_cluster_shape_mn = _SM107_BF16_PREFERRED_CLUSTER_SHAPE_MN + fallback_cluster_shape_mn = _SM107_BF16_FALLBACK_CLUSTER_SHAPE_MN + preferred_kernel_class = _sm107_bf16_kernel_class("preferred_cluster") + for use_2cta_instrs, mma_tiler_mn, max_num_ab_stage in itertools.product( + use_2cta_instrs_candi, mma_tiler_mn_candi, + _SM107_BF16_MAX_NUM_AB_STAGE_CANDIDATES): + if not _bf16_cluster_m_fits(m, use_2cta_instrs, mma_tiler_mn, + preferred_cluster_shape_mn): + continue + if (_bf16_preferred_cluster_has_launchable_grid( + m, n, batch_size, use_2cta_instrs, mma_tiler_mn, + preferred_cluster_shape_mn, fallback_cluster_shape_mn) + and preferred_kernel_class.can_implement( + cutlass.BFloat16, cutlass.Float32, c_dtype_cutlass, + use_2cta_instrs, mma_tiler_mn, + fallback_cluster_shape_mn, m, n, k, batch_size, a_major, + b_major, c_major) + and preferred_kernel_class.can_implement( + cutlass.BFloat16, cutlass.Float32, c_dtype_cutlass, + use_2cta_instrs, mma_tiler_mn, + preferred_cluster_shape_mn, m, n, k, batch_size, + a_major, b_major, c_major)): + valid_tactics.append( + ("preferred_cluster", use_2cta_instrs, mma_tiler_mn, + preferred_cluster_shape_mn, fallback_cluster_shape_mn, + max_num_ab_stage)) + return valid_tactics + + def _parse_sm107_bf16_tactic(tactic): + """Decode a tactic into (kernel_variant, use_2cta_instrs, mma_tiler_mn, + preferred_cluster_shape_mn, cluster_shape_mn, max_num_ab_stage, + split_k_slices).""" + if (isinstance(tactic, tuple) and len(tactic) > 0 + and isinstance(tactic[0], str)): + kernel_variant = tactic[0] + if kernel_variant == "preferred_cluster": + _, use_2cta_instrs, mma_tiler_mn, preferred_cluster_shape_mn, cluster_shape_mn, max_num_ab_stage = tactic + return (kernel_variant, use_2cta_instrs, mma_tiler_mn, + preferred_cluster_shape_mn, cluster_shape_mn, + max_num_ab_stage, 1) + if len(tactic) == 6: + _, use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, max_num_ab_stage, split_k_slices = tactic + else: + _, use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, max_num_ab_stage = tactic + split_k_slices = 1 + return (kernel_variant, use_2cta_instrs, mma_tiler_mn, None, + cluster_shape_mn, max_num_ab_stage, split_k_slices) + if isinstance(tactic, tuple): + use_2cta_instrs, mma_tiler_mn, cluster_shape_mn = tactic + else: + use_2cta_instrs = False + mma_tiler_mn = (128, 128) + cluster_shape_mn = (1, 1) + return ("base", use_2cta_instrs, mma_tiler_mn, None, cluster_shape_mn, + 0, 1) + + def _sm107_bf16_gemm_kernel(kernel_variant: str, + use_2cta_instrs: bool, + mma_tiler_mn, + preferred_cluster_shape_mn, + cluster_shape_mn, + max_num_ab_stage: int, + split_k_slices: int = 1): + kernel_class = _sm107_bf16_kernel_class(kernel_variant) + if kernel_variant == "preferred_cluster": + return kernel_class( + acc_dtype=cutlass.Float32, + use_2cta_instrs=use_2cta_instrs, + mma_tiler_mn=mma_tiler_mn, + preferred_cluster_shape_mn=preferred_cluster_shape_mn, + fallback_cluster_shape_mn=cluster_shape_mn, + max_num_ab_stage=max_num_ab_stage, + ) + return kernel_class( + acc_dtype=cutlass.Float32, + use_2cta_instrs=use_2cta_instrs, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + max_num_ab_stage=max_num_ab_stage, + split_k_slices=split_k_slices, + ) + + def _sm107_bf16_occupancy(cluster_shape_mn, preferred_cluster_shape_mn): + # Persistent grid occupancy is baked into the compiled artifact. + max_active_clusters = get_max_activate_clusters(cluster_shape_mn[0] * + cluster_shape_mn[1]) + max_active_preferred_clusters = None + if preferred_cluster_shape_mn is not None: + max_active_preferred_clusters = get_max_activate_clusters( + preferred_cluster_shape_mn[0] * preferred_cluster_shape_mn[1]) + return max_active_clusters, max_active_preferred_clusters + + class CuteDSLBf16RubinBmmRunner(CuteDSLBf16BlackwellBmmRunner): + """SM107 BF16 batched GEMM: [B, M, K] @ [B, N, K] -> [B, M, N].""" + kernel_cache = dict() + + # The output's M dim (inputs[2] dim 1) must track input0's bucketed + # M: without the constraint, profiling-cache keys embed the raw + # tuning-time M and every other runtime M misses the cache and falls + # back to the default tactic. + tuning_config = TuningConfig( + dynamic_tensor_specs=(DynamicTensorSpec( + 0, 1, get_last_power_of_2_num_tokens_buckets, + last_positive_power_of_2), ), + constraint_specs=(ConstraintSpec( + 2, 1, lambda input_shapes: input_shapes[0][1]), ), + ) + + def get_valid_tactics( + self, + inputs: List[torch.Tensor], + profile: OptimizationProfile, + **kwargs, + ) -> List[Tuple]: + if not _is_sm107_cute_dsl_available(): + logger.debug( + f"CuteDSL: SM version {get_sm_version()} is not supported. " + "CuteDSL SM107 BF16 BMM skipping all tactics.") + return [] + batch_size, m, k = inputs[0].shape + n = inputs[1].shape[1] + return _sm107_bf16_valid_tactics( + m, n, k, batch_size, cutlass.BFloat16, + _SM107_BF16_MMA_TILER_MN_CANDIDATES, None) + + def forward( + self, + inputs: List[torch.Tensor], + tactic, + ) -> None: + (kernel_variant, use_2cta_instrs, mma_tiler_mn, + preferred_cluster_shape_mn, cluster_shape_mn, max_num_ab_stage, + _) = _parse_sm107_bf16_tactic(tactic) + + a_tensor, b_tensor, c_tensor = inputs + batch_size, m, k = a_tensor.shape + n = b_tensor.shape[1] + + # C is passed as [M, N, B]; from_dlpack captures the real strides so + # non-contiguous views are written by TMA without a copy. + c_tmp = c_tensor.permute(1, 2, 0) + c_layout_key = tuple(c_tmp.stride()) + + # A/B strides let the kernel consume non-contiguous views (e.g. + # [M,B,K].transpose(0,1), broadcast batch via expand, or sliced + # weight views) without a .contiguous() copy. CuTe tensors are + # (M, K, B) / (N, K, B): M/N stride = stride(1), B stride = + # stride(0), K must be innermost. The kernel is compiled K-major, + # so a transposed [B,N,K] view of a [B,K,N] tensor cannot be + # expressed and would silently compute the wrong product. + if a_tensor.shape[2] > 1 and a_tensor.stride(2) != 1: + raise ValueError( + "cute_dsl_bf16_bmm_rubin requires A with K innermost " + f"(stride 1 in the last dim); got a.stride()=" + f"{tuple(a_tensor.stride())} for a.shape=" + f"{tuple(a_tensor.shape)}.") + if b_tensor.shape[2] > 1 and b_tensor.stride(2) != 1: + raise ValueError( + "cute_dsl_bf16_bmm_rubin requires B with K innermost " + f"(stride 1 in the last dim); got b.stride()=" + f"{tuple(b_tensor.stride())} for b.shape=" + f"{tuple(b_tensor.shape)}.") + 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(cutlass.BFloat16, + a_tensor.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + b_ptr = make_ptr(cutlass.BFloat16, + b_tensor.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + c_cute_tensor = cute.runtime.from_dlpack( + c_tmp).mark_layout_dynamic(leading_dim=1) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + max_active_clusters, max_active_preferred_clusters = _sm107_bf16_occupancy( + cluster_shape_mn, preferred_cluster_shape_mn) + cache_key = ( + kernel_variant, + use_2cta_instrs, + mma_tiler_mn, + preferred_cluster_shape_mn, + cluster_shape_mn, + max_num_ab_stage, + self.use_tvm_ffi, + c_layout_key, + max_active_clusters, + max_active_preferred_clusters, + ) + if cache_key not in self.__class__.kernel_cache: + if self.use_tvm_ffi: + a_ptr = make_ptr(cutlass.BFloat16, + a_tensor.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + b_ptr = make_ptr(cutlass.BFloat16, + b_tensor.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + c_cute_tensor = cute.runtime.from_dlpack( + c_tmp).mark_layout_dynamic(leading_dim=1) + stream = cute.runtime.make_fake_stream( + use_tvm_ffi_env_stream=True) + + gemm = _sm107_bf16_gemm_kernel(kernel_variant, use_2cta_instrs, + mma_tiler_mn, + preferred_cluster_shape_mn, + cluster_shape_mn, + max_num_ab_stage) + compile_args = [ + m, n, k, batch_size, a_ptr, b_ptr, c_cute_tensor, + a_stride_m, a_stride_batch, b_stride_n, b_stride_batch + ] + if kernel_variant == "preferred_cluster": + compile_args.extend( + [max_active_preferred_clusters, max_active_clusters]) + else: + compile_args.append(max_active_clusters) + compiled_gemm = cute.compile( + gemm.wrapper_strided, + *compile_args, + stream=stream, + options="--opt-level 2 --enable-tvm-ffi" + if self.use_tvm_ffi else "--opt-level 2", + ) + self.__class__.kernel_cache[cache_key] = compiled_gemm + else: + compiled_gemm = self.__class__.kernel_cache[cache_key] + + if self.use_tvm_ffi: + compiled_gemm( + m, + n, + k, + batch_size, + a_tensor.data_ptr(), + b_tensor.data_ptr(), + c_tmp, + a_stride_m, + a_stride_batch, + b_stride_n, + b_stride_batch, + ) + else: + compiled_gemm( + m, + n, + k, + batch_size, + a_ptr, + b_ptr, + c_cute_tensor, + a_stride_m, + a_stride_batch, + b_stride_n, + b_stride_batch, + stream=stream, + ) + + @torch.library.custom_op("trtllm::cute_dsl_bf16_bmm_rubin", + mutates_args=("output", ), + device_types="cuda") + def cute_dsl_bf16_bmm_rubin( + input: torch.Tensor, + weight: torch.Tensor, + output: torch.Tensor, + use_tvm_ffi: bool = True, + ) -> None: + """SM107 counterpart of ``cute_dsl_bf16_bmm_blackwell``. + + - input: [B, M, K], weight: [B, N, K], output: [B, M, N] (bf16). + """ + if not _is_sm107_cute_dsl_available(): + raise ValueError( + f"CuteDSL: SM version {get_sm_version()} is not supported. " + "CuteDSL BF16 BMM SM107 requires SM107 and a CuTe DSL package " + "with SM107 support.") + + tuner = AutoTuner.get() + runner = CuteDSLBf16RubinBmmRunner(use_tvm_ffi=use_tvm_ffi) + inputs = [input, weight, output] + _, best_tactic = tuner.choose_one( + "trtllm::cute_dsl_bf16_bmm_rubin::gemm", + [runner], + runner.__class__.tuning_config, + inputs, + ) + runner(inputs, tactic=best_tactic) + + @torch.library.register_fake("trtllm::cute_dsl_bf16_bmm_rubin") + def _( + mat_a: torch.Tensor, + mat_b: torch.Tensor, + output: torch.Tensor, + use_tvm_ffi: bool = True, + ) -> None: + batch_size, m, k = mat_a.shape[0], mat_a.shape[1], mat_a.shape[2] + n = mat_b.shape[1] + assert output.dtype == torch.bfloat16, "CuTe DSL bf16 bmm output dtype must be bf16" + assert output.shape == ( + batch_size, m, n), "CuTe DSL bf16 bmm output shape is incorrect" + + class CuteDSLBf16RubinGemmRunner(CuteDSLBf16BlackwellGemmRunner): + """SM107 BF16 GEMM for Linear layers: [M, K] @ [N, K]^T -> [M, N]. + + Adds direct split-K on top of the Blackwell runner: the kernel expands + the persistent scheduler's L dimension by ``split_k_slices`` and every + CTA reduce-adds its partial into the pre-zeroed output through TMA, so + no FP32 workspace or reduction launch is needed. + """ + kernel_cache = dict() + # Split-K artifacts are shape-independent (tactic plus occupancy) and + # cached separately from the split=1 path. + split_k_gemm_cache = dict() + + # See CuteDSLBf16RubinBmmRunner.tuning_config. + tuning_config = TuningConfig( + dynamic_tensor_specs=(DynamicTensorSpec( + 0, 0, get_last_power_of_2_num_tokens_buckets, + last_positive_power_of_2), ), + constraint_specs=(ConstraintSpec( + 2, 0, lambda input_shapes: input_shapes[0][0]), ), + ) + + def __init__( + self, + use_tvm_ffi: bool = True, + output_dtype: Optional[torch.dtype] = None, + ): + super().__init__(use_tvm_ffi=use_tvm_ffi) + self.output_dtype = output_dtype + + def unique_id(self): + return (self.use_tvm_ffi, self.output_dtype) + + def get_valid_tactics( + self, + inputs: List[torch.Tensor], + profile: OptimizationProfile, + **kwargs, + ) -> List[Tuple]: + if not _is_sm107_cute_dsl_available(): + logger.debug( + f"CuteDSL: SM version {get_sm_version()} is not supported. " + "CuteDSL SM107 BF16 GEMM skipping all tactics.") + return [] + + m, k = inputs[0].shape + n = inputs[1].shape[0] + c_dtype_cutlass = _TORCH_TO_CUTLASS_DTYPE[inputs[2].dtype] + + mma_tiler_mn_candi = list(_SM107_BF16_MMA_TILER_MN_CANDIDATES) + # Skinny-M / small-N shapes under-use a 128x128 MMA tile (half the + # M rows wasted, coarse N tiling, few CTAs); offer smaller tiles so + # the autotuner can pick something closer to a 64x32 CTA tile. + if m <= 128 or n <= 512: + mma_tiler_mn_candi += [(64, 32), (64, 64), (128, 32), (128, 64)] + + # split>1 is only offered for large-K, few-output-tile shapes where + # the extra K-parallelism outweighs the atomic-reduce cost. Dynamic + # autotuning may bucket input M while leaving C at the full M; the + # split-K path assumes the two match, so skip those buckets. + split_k_candi = [1] + if k >= 4096 and n <= 512 and inputs[2].shape[0] == m: + split_k_candi = [1, 2, 4, 8] + + return _sm107_bf16_valid_tactics(m, n, k, 1, c_dtype_cutlass, + mma_tiler_mn_candi, split_k_candi) + + def forward( + self, + inputs: List[torch.Tensor], + tactic, + ) -> None: + (kernel_variant, use_2cta_instrs, mma_tiler_mn, + preferred_cluster_shape_mn, cluster_shape_mn, max_num_ab_stage, + split_k_slices) = _parse_sm107_bf16_tactic(tactic) + + a_tensor, b_tensor, c_tensor = inputs + m, k = a_tensor.shape + n = b_tensor.shape[0] + batch_size = 1 + if split_k_slices > 1 and tuple(c_tensor.shape) != (m, n): + raise RuntimeError( + "BF16 split-K GEMM requires an [M, N] output, got " + f"output.shape={tuple(c_tensor.shape)}, expected={(m, n)}.") + + a_tensor = a_tensor.contiguous() + b_tensor = b_tensor.contiguous() + c_needs_copy = not c_tensor.is_contiguous() + c_buf = torch.empty_like(c_tensor) if c_needs_copy else c_tensor + + a_batched = a_tensor.unsqueeze(0) # [1, M, K] + b_batched = b_tensor.unsqueeze(0) # [1, N, K] + c_tmp = c_buf.unsqueeze(-1) # [M, N, 1] + c_layout_key = tuple(c_tmp.stride()) + c_dtype_cutlass = _TORCH_TO_CUTLASS_DTYPE[c_tensor.dtype] + + if split_k_slices > 1: + self._forward_split_k( + a_batched, + b_batched, + c_buf, + m, + n, + k, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + max_num_ab_stage, + split_k_slices, + c_dtype_cutlass, + ) + if c_needs_copy: + c_tensor.copy_(c_buf) + return + + if not self.use_tvm_ffi: + a_ptr = make_ptr(cutlass.BFloat16, + a_batched.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + b_ptr = make_ptr(cutlass.BFloat16, + b_batched.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + c_cute_tensor = cute.runtime.from_dlpack( + c_tmp).mark_layout_dynamic(leading_dim=1) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + max_active_clusters, max_active_preferred_clusters = _sm107_bf16_occupancy( + cluster_shape_mn, preferred_cluster_shape_mn) + cache_key = ( + kernel_variant, + use_2cta_instrs, + mma_tiler_mn, + preferred_cluster_shape_mn, + cluster_shape_mn, + max_num_ab_stage, + self.use_tvm_ffi, + c_dtype_cutlass, + c_layout_key, + max_active_clusters, + max_active_preferred_clusters, + ) + if cache_key not in self.__class__.kernel_cache: + if self.use_tvm_ffi: + a_ptr = make_ptr(cutlass.BFloat16, + a_batched.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + b_ptr = make_ptr(cutlass.BFloat16, + b_batched.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + c_cute_tensor = cute.runtime.from_dlpack( + c_tmp).mark_layout_dynamic(leading_dim=1) + stream = cute.runtime.make_fake_stream( + use_tvm_ffi_env_stream=True) + + gemm = _sm107_bf16_gemm_kernel(kernel_variant, use_2cta_instrs, + mma_tiler_mn, + preferred_cluster_shape_mn, + cluster_shape_mn, + max_num_ab_stage) + compile_args = [ + m, n, k, batch_size, a_ptr, b_ptr, c_cute_tensor + ] + if kernel_variant == "preferred_cluster": + compile_args.extend( + [max_active_preferred_clusters, max_active_clusters]) + else: + compile_args.append(max_active_clusters) + compiled_gemm = cute.compile( + gemm.wrapper, + *compile_args, + stream=stream, + options="--opt-level 2 --enable-tvm-ffi" + if self.use_tvm_ffi else "--opt-level 2", + ) + self.__class__.kernel_cache[cache_key] = compiled_gemm + else: + compiled_gemm = self.__class__.kernel_cache[cache_key] + + if self.use_tvm_ffi: + compiled_gemm( + m, + n, + k, + batch_size, + a_batched.data_ptr(), + b_batched.data_ptr(), + c_tmp, + ) + else: + compiled_gemm( + m, + n, + k, + batch_size, + a_ptr, + b_ptr, + c_cute_tensor, + stream=stream, + ) + + if c_needs_copy: + c_tensor.copy_(c_buf) + + def _forward_split_k( + self, + a_batched: torch.Tensor, + b_batched: torch.Tensor, + c_buf: torch.Tensor, + m: int, + n: int, + k: int, + use_2cta_instrs: bool, + mma_tiler_mn, + cluster_shape_mn, + max_num_ab_stage: int, + split_k_slices: int, + c_dtype_cutlass, + ) -> None: + """Direct split-K: each K-slice reduce-adds into the zeroed C.""" + batch_size = 1 + c_tmp = c_buf.unsqueeze(-1) + c_layout_key = tuple(c_tmp.stride()) + max_active_clusters = get_max_activate_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1]) + gemm_key = ( + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + max_num_ab_stage, + split_k_slices, + self.use_tvm_ffi, + c_dtype_cutlass, + c_layout_key, + max_active_clusters, + ) + + a_ptr = make_ptr(cutlass.BFloat16, + a_batched.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + b_ptr = make_ptr(cutlass.BFloat16, + b_batched.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + if self.use_tvm_ffi: + stream = cute.runtime.make_fake_stream( + use_tvm_ffi_env_stream=True) + else: + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + def _c_cute_view(): + return cute.runtime.from_dlpack( + c_tmp, assumed_align=16).mark_layout_dynamic(leading_dim=1) + + compiled_gemm = self.__class__.split_k_gemm_cache.get(gemm_key) + if compiled_gemm is None: + gemm = _sm107_bf16_gemm_kernel("base", + use_2cta_instrs, + mma_tiler_mn, + None, + cluster_shape_mn, + max_num_ab_stage, + split_k_slices=split_k_slices) + compiled_gemm = cute.compile( + gemm.wrapper, + m, + n, + k, + batch_size, + a_ptr, + b_ptr, + _c_cute_view(), + max_active_clusters, + stream=stream, + options="--opt-level 2 --enable-tvm-ffi" + if self.use_tvm_ffi else "--opt-level 2", + ) + self.__class__.split_k_gemm_cache[gemm_key] = compiled_gemm + + # TMA ADD accumulates into C. Keep the zeroing on the launch stream + # and immediately adjacent to the GEMM so CUDA graph replay is safe. + c_buf.zero_() + if self.use_tvm_ffi: + compiled_gemm( + m, + n, + k, + batch_size, + a_batched.data_ptr(), + b_batched.data_ptr(), + c_tmp, + ) + else: + compiled_gemm( + m, + n, + k, + batch_size, + a_ptr, + b_ptr, + _c_cute_view(), + stream=stream, + ) + + @torch.library.custom_op("trtllm::cute_dsl_bf16_gemm_rubin", + mutates_args=("output", ), + device_types="cuda") + def cute_dsl_bf16_gemm_rubin( + input: torch.Tensor, + weight: torch.Tensor, + output: torch.Tensor, + use_tvm_ffi: bool = True, + ) -> None: + """SM107 counterpart of ``cute_dsl_bf16_gemm_blackwell``. + + Computes output = input @ weight^T + - input: [M, K], weight: [N, K], output: [M, N] (bf16 or fp32). + """ + if not _is_sm107_cute_dsl_available(): + raise ValueError( + f"CuteDSL: SM version {get_sm_version()} is not supported. " + "CuteDSL BF16 GEMM SM107 requires SM107 and a CuTe DSL package " + "with SM107 support.") + + tuner = AutoTuner.get() + runner = CuteDSLBf16RubinGemmRunner(use_tvm_ffi=use_tvm_ffi, + output_dtype=output.dtype) + inputs = [input, weight, output] + _, best_tactic = tuner.choose_one( + "trtllm::cute_dsl_bf16_gemm_rubin::gemm", + [runner], + runner.__class__.tuning_config, + inputs, + ) + runner(inputs, tactic=best_tactic) + + @torch.library.register_fake("trtllm::cute_dsl_bf16_gemm_rubin") + def _( + mat_a: torch.Tensor, + mat_b: torch.Tensor, + output: torch.Tensor, + use_tvm_ffi: bool = True, + ) -> None: + m, k = mat_a.shape[0], mat_a.shape[1] + n = mat_b.shape[0] + assert output.dtype in (torch.bfloat16, torch.float32), \ + "CuTe DSL bf16 gemm output dtype must be bf16 or fp32" + assert output.shape == ( + m, n), "CuTe DSL bf16 gemm output shape is incorrect" + # ------------------------------------------------------------------ # # CuTE DSL FP4 Paged MQA Logits (Blackwell SM100) # # ------------------------------------------------------------------ # diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index 9dc29b58d744..f23a7e1b3c41 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -896,8 +896,10 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: n, dtype=torch.float32, device=hidden_states.device) - torch.ops.trtllm.cute_dsl_bf16_gemm_blackwell( - input_2d.contiguous(), self.weight, output) + bf16_gemm_op = (torch.ops.trtllm.cute_dsl_bf16_gemm_rubin + if get_sm_version() == 107 else + torch.ops.trtllm.cute_dsl_bf16_gemm_blackwell) + bf16_gemm_op(input_2d.contiguous(), self.weight, output) logits = output.view(*hidden_states.shape[:-1], n) else: logits = torch.ops.trtllm.dsv3_router_gemm_op( diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 1286c619f524..87b69e8d9656 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -532,15 +532,16 @@ def get_tp_alignment(self, class UnquantizedLinearMethod(LinearMethodBase): """Linear method for unquantized (BF16 / FP16 / FP32) weights. - BF16 GEMM dispatch (priority order, Blackwell SM100/SM103) - ---------------------------------------------------------- + BF16 GEMM dispatch (priority order, Blackwell SM100/SM103 and SM107) + ------------------------------------------------------------------- 1. **low-m GEMM** (``TRTLLM_LOW_M_GEMM_BACKEND=auto``, M ≤ 32) CuTe-DSL low-m GEMM kernel for small-M decode batches on Blackwell. Orthogonal to ``use_cute_dsl_bf16_gemm`` — must be enabled independently via the env var. 2. **persistent GEMM** (``Linear(use_cute_dsl_bf16_gemm=True)``) - ``trtllm::cute_dsl_bf16_gemm_blackwell`` persistent CuTe-DSL kernel. + ``trtllm::cute_dsl_bf16_gemm_blackwell`` persistent CuTe-DSL kernel; + ``trtllm::cute_dsl_bf16_gemm_rubin`` on SM107. 3. **cublas_mm** (``Linear(use_custom_cublas_mm=True)``) ``trtllm::cublas_mm``; use when TP AllReduce fuse via NCCL @@ -586,7 +587,7 @@ def apply(self, module: Linear, input: torch.Tensor, output = apply_low_m_gemm(module, input, module.weight, bias) if output is not None: return output - # CuTe DSL BF16 GEMM path for Blackwell + # CuTe DSL BF16 GEMM path for Blackwell / SM107 if (module.use_cute_dsl_bf16_gemm and is_sm_100f() and module.weight.dtype == torch.bfloat16): # input: [*, K], weight: [N, K], output: [*, N] @@ -597,7 +598,10 @@ def apply(self, module: Linear, input: torch.Tensor, n, dtype=torch.bfloat16, device=input.device) - torch.ops.trtllm.cute_dsl_bf16_gemm_blackwell( + bf16_gemm_op = (torch.ops.trtllm.cute_dsl_bf16_gemm_rubin + if get_sm_version() == 107 else + torch.ops.trtllm.cute_dsl_bf16_gemm_blackwell) + bf16_gemm_op( input_2d.contiguous(), module.weight, output, diff --git a/tests/unittest/_torch/thop/parallel/test_cute_dsl_bf16_dense_rubin.py b/tests/unittest/_torch/thop/parallel/test_cute_dsl_bf16_dense_rubin.py new file mode 100644 index 000000000000..2cc3b4380406 --- /dev/null +++ b/tests/unittest/_torch/thop/parallel/test_cute_dsl_bf16_dense_rubin.py @@ -0,0 +1,249 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for the SM107 CuTe DSL BF16 dense GEMM/BMM custom ops. + +The dispatch tests run on every architecture with CuTe DSL installed and check +that ``trtllm::cute_dsl_bf16_{gemm,bmm}_rubin`` are unreachable off SM107. The +correctness tests need SM107 and a CuTe DSL package with the SM107 helpers. +""" + +import pytest +import torch +from torch._subclasses.fake_tensor import FakeTensorMode +from utils.util import getSMVersion + +import tensorrt_llm._torch.custom_ops # noqa: F401 +from tensorrt_llm._torch.autotuner import AutoTuner +from tensorrt_llm._torch.cute_dsl_utils import ( + IS_CUTLASS_DSL_AVAILABLE, + IS_CUTLASS_DSL_RUBIN_AVAILABLE, +) + +skip_no_cute_dsl = pytest.mark.skipif( + not IS_CUTLASS_DSL_AVAILABLE, reason="cutlass-dsl is not available" +) +skip_not_sm107 = pytest.mark.skipif( + getSMVersion() != 107 or not IS_CUTLASS_DSL_RUBIN_AVAILABLE, + reason="requires SM107 and a CuTe DSL package with SM107 support", +) +skip_on_sm107 = pytest.mark.skipif(getSMVersion() == 107, reason="checks the non-SM107 path") + + +def _runner_class(name): + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops + + return getattr(cute_dsl_custom_ops, name) + + +def _fresh_runner(name, **kwargs): + runner_class = _runner_class(name) + AutoTuner.get().clear_cache() + runner_class.kernel_cache.clear() + if hasattr(runner_class, "split_k_gemm_cache"): + runner_class.split_k_gemm_cache.clear() + return runner_class(**kwargs) + + +def _select_tactic(tactics, kernel_variant, split_k_slices=None): + candidates = [t for t in tactics if t[0] == kernel_variant and t[1] is False] + if split_k_slices is not None: + candidates = [t for t in candidates if len(t) == 6 and t[-1] == split_k_slices] + assert candidates, f"no {kernel_variant} tactic (split_k={split_k_slices})" + return candidates[0] + + +def _gemm_operands(m, n, k, c_dtype=torch.bfloat16): + act = torch.randn(m, k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + output = torch.empty(m, n, dtype=c_dtype, device="cuda") + return act, weight, output + + +def _bmm_operands(b, m, n, k): + act = torch.randn(b, m, k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(b, n, k, dtype=torch.bfloat16, device="cuda") + output = torch.empty(b, m, n, dtype=torch.bfloat16, device="cuda") + return act, weight, output + + +# -------------------------------------------------------------------------- +# Dispatch: the SM107 ops must be a no-op everywhere else. +# -------------------------------------------------------------------------- + + +@skip_no_cute_dsl +@skip_on_sm107 +def test_sm107_bf16_gemm_rejects_other_archs(): + act, weight, output = _gemm_operands(64, 128, 256) + with pytest.raises(ValueError, match="SM107"): + torch.ops.trtllm.cute_dsl_bf16_gemm_rubin(act, weight, output) + + +@skip_no_cute_dsl +@skip_on_sm107 +def test_sm107_bf16_bmm_rejects_other_archs(): + act, weight, output = _bmm_operands(2, 64, 128, 256) + with pytest.raises(ValueError, match="SM107"): + torch.ops.trtllm.cute_dsl_bf16_bmm_rubin(act, weight, output) + + +@skip_no_cute_dsl +@skip_on_sm107 +def test_sm107_bf16_runners_offer_no_tactics_off_sm107(): + gemm_runner = _fresh_runner("CuteDSLBf16RubinGemmRunner") + assert gemm_runner.get_valid_tactics(list(_gemm_operands(64, 128, 256)), None) == [] + bmm_runner = _fresh_runner("CuteDSLBf16RubinBmmRunner") + assert bmm_runner.get_valid_tactics(list(_bmm_operands(2, 64, 128, 256)), None) == [] + + +@skip_no_cute_dsl +def test_sm107_bf16_ops_fake_registration(): + with FakeTensorMode(): + act = torch.empty(8, 32, dtype=torch.bfloat16, device="cuda") + weight = torch.empty(16, 32, dtype=torch.bfloat16, device="cuda") + for c_dtype in (torch.bfloat16, torch.float32): + output = torch.empty(8, 16, dtype=c_dtype, device="cuda") + torch.ops.trtllm.cute_dsl_bf16_gemm_rubin(act, weight, output) + with pytest.raises(AssertionError): + torch.ops.trtllm.cute_dsl_bf16_gemm_rubin( + act, weight, torch.empty(8, 17, dtype=torch.bfloat16, device="cuda") + ) + + act = torch.empty(2, 8, 32, dtype=torch.bfloat16, device="cuda") + weight = torch.empty(2, 16, 32, dtype=torch.bfloat16, device="cuda") + output = torch.empty(2, 8, 16, dtype=torch.bfloat16, device="cuda") + torch.ops.trtllm.cute_dsl_bf16_bmm_rubin(act, weight, output) + with pytest.raises(AssertionError): + torch.ops.trtllm.cute_dsl_bf16_bmm_rubin( + act, weight, torch.empty(2, 8, 16, dtype=torch.float32, device="cuda") + ) + + +# -------------------------------------------------------------------------- +# SM107 correctness. +# -------------------------------------------------------------------------- + + +@skip_not_sm107 +@pytest.mark.parametrize("c_dtype", [torch.bfloat16, torch.float32]) +def test_cute_dsl_bf16_gemm_rubin_op(c_dtype): + torch.manual_seed(0) + AutoTuner.get().clear_cache() + act, weight, output = _gemm_operands(256, 1024, 2048, c_dtype) + torch.ops.trtllm.cute_dsl_bf16_gemm_rubin(act, weight, output) + torch.cuda.synchronize() + expected = act.float() @ weight.t().float() + torch.testing.assert_close(output.float(), expected, rtol=1e-2, atol=1.0) + + +@skip_not_sm107 +@pytest.mark.parametrize("kernel_variant", ["base", "preferred_cluster"]) +def test_cute_dsl_bf16_gemm_rubin_tactics(kernel_variant): + torch.manual_seed(1) + runner = _fresh_runner("CuteDSLBf16RubinGemmRunner", output_dtype=torch.bfloat16) + act, weight, output = _gemm_operands(1024, 2048, 1024) + tactics = runner.get_valid_tactics([act, weight, output], None) + tactic = _select_tactic(tactics, kernel_variant) + runner([act, weight, output], tactic=tactic) + torch.cuda.synchronize() + expected = act.float() @ weight.t().float() + torch.testing.assert_close(output.float(), expected, rtol=1e-2, atol=1.0) + + +@skip_not_sm107 +@pytest.mark.parametrize("split_k_slices", [2, 4, 8]) +@pytest.mark.parametrize("c_dtype", [torch.bfloat16, torch.float32]) +def test_cute_dsl_bf16_split_k_gemm_rubin(split_k_slices, c_dtype): + """Split-K matches the dense reference for large-K, small-N shapes.""" + torch.manual_seed(2026) + runner = _fresh_runner("CuteDSLBf16RubinGemmRunner", output_dtype=c_dtype) + + # Large K and small N so get_valid_tactics offers split>1 candidates. + act, weight, output = _gemm_operands(64, 256, 7168, c_dtype) + tactics = runner.get_valid_tactics([act, weight, output], None) + tactic = _select_tactic(tactics, "base", split_k_slices=split_k_slices) + + # Direct split-K rounds each partial to the output dtype before the + # atomic TMA ADD, so BF16 output needs tolerance for both rounding and + # arrival-order changes. + rtol, atol = (2e-2, 2.5) if c_dtype == torch.bfloat16 else (1e-2, 1.0) + + expected = act.float() @ weight.t().float() + runner([act, weight, output], tactic=tactic) + torch.cuda.synchronize() + torch.testing.assert_close(output.float(), expected, rtol=rtol, atol=atol) + + # A second launch must not accumulate on the previous output; poisoning C + # also catches a missing zero inside CUDA-graph replay and normal dispatch. + output.fill_(float("nan")) + runner([act, weight, output], tactic=tactic) + torch.cuda.synchronize() + torch.testing.assert_close(output.float(), expected, rtol=rtol, atol=atol) + + if c_dtype == torch.float32 and split_k_slices in (2, 4): + split1_output = torch.empty_like(output) + split1_tactic = _select_tactic(tactics, "base", split_k_slices=1) + runner([act, weight, split1_output], tactic=split1_tactic) + torch.cuda.synchronize() + # Both write FP32; direct split-K only changes the reduction order. + torch.testing.assert_close(output, split1_output, rtol=1e-3, atol=1e-2) + + +@skip_not_sm107 +def test_cute_dsl_bf16_bmm_rubin_op(): + torch.manual_seed(3) + AutoTuner.get().clear_cache() + act, weight, output = _bmm_operands(4, 256, 512, 1024) + torch.ops.trtllm.cute_dsl_bf16_bmm_rubin(act, weight, output) + torch.cuda.synchronize() + expected = torch.bmm(act.float(), weight.transpose(1, 2).float()) + torch.testing.assert_close(output.float(), expected, rtol=1e-2, atol=1.0) + + +@skip_not_sm107 +@pytest.mark.parametrize("kernel_variant", ["base", "preferred_cluster"]) +def test_cute_dsl_bf16_bmm_rubin_tactics(kernel_variant): + torch.manual_seed(4) + runner = _fresh_runner("CuteDSLBf16RubinBmmRunner") + act, weight, output = _bmm_operands(2, 1024, 1024, 512) + tactics = runner.get_valid_tactics([act, weight, output], None) + tactic = _select_tactic(tactics, kernel_variant) + runner([act, weight, output], tactic=tactic) + torch.cuda.synchronize() + expected = torch.bmm(act.float(), weight.transpose(1, 2).float()) + torch.testing.assert_close(output.float(), expected, rtol=1e-2, atol=1.0) + + +@skip_not_sm107 +def test_cute_dsl_bf16_bmm_rubin_strided_views(): + """Non-contiguous A/B views with K innermost are consumed without a copy; + a K-non-innermost view is rejected instead of computing the wrong product.""" + torch.manual_seed(5) + runner = _fresh_runner("CuteDSLBf16RubinBmmRunner") + b, m, n, k = 4, 128, 256, 512 + # [M, B, K] storage viewed as [B, M, K] and a broadcast batch for B. + act_storage = torch.randn(m, b, k, dtype=torch.bfloat16, device="cuda") + act = act_storage.transpose(0, 1) + weight = torch.randn(1, n, k, dtype=torch.bfloat16, device="cuda").expand(b, n, k) + output = torch.empty(b, m, n, dtype=torch.bfloat16, device="cuda") + tactics = runner.get_valid_tactics([act, weight, output], None) + runner([act, weight, output], tactic=_select_tactic(tactics, "base")) + torch.cuda.synchronize() + expected = torch.bmm(act.float(), weight.transpose(1, 2).float()) + torch.testing.assert_close(output.float(), expected, rtol=1e-2, atol=1.0) + + bad_weight = torch.randn(b, k, n, dtype=torch.bfloat16, device="cuda").transpose(1, 2) + with pytest.raises(ValueError, match="K innermost"): + runner([act, bad_weight, output], tactic=_select_tactic(tactics, "base")) From 8ef4fb83e7738600daaf4328494aba1ef40c607b Mon Sep 17 00:00:00 2001 From: farazkh80 <58580514+farazkh80@users.noreply.github.com> Date: Sat, 5 Sep 2026 19:16:55 +0000 Subject: [PATCH 2/3] [None][feat] Register the SM107 BF16 CuTe DSL ops as in-place for torch.compile Both ops mutate their output tensor; list them in the optional in-place op table so graph compilation tracks the mutation when the ops are registered. Signed-off-by: farazkh80 <58580514+farazkh80@users.noreply.github.com> --- tensorrt_llm/_torch/compilation/utils.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tensorrt_llm/_torch/compilation/utils.py b/tensorrt_llm/_torch/compilation/utils.py index a03e84014109..01ddc54f02b9 100644 --- a/tensorrt_llm/_torch/compilation/utils.py +++ b/tensorrt_llm/_torch/compilation/utils.py @@ -217,6 +217,12 @@ def inplace_info(): "fp8_block_scaling_bmm_out": { 1: "out" }, + "cute_dsl_bf16_bmm_rubin": { + 1: "output" + }, + "cute_dsl_bf16_gemm_rubin": { + 1: "output" + }, "gate_forward": { 1: "out_weights", 2: "out_indices" From a8121fbdeec67fb605a0665892eb6909c33a93f2 Mon Sep 17 00:00:00 2001 From: farazkh80 <58580514+farazkh80@users.noreply.github.com> Date: Wed, 9 Sep 2026 15:53:27 +0000 Subject: [PATCH 3/3] [None][fix] Apply the SM107 CTA_N=256/cluster_n=2 exclusion to the preferred-cluster tactic loop The base tactic loop already skips CTA_N=256 with cluster_n=2 (illegal memory access on SM107), but the preferred-cluster loop's fixed (4, 2) shape let (256, 256) mma tilers pass both can_implement checks and reach profiling. Apply the same guard. Found by CodeRabbit on #18761. Signed-off-by: farazkh80 <58580514+farazkh80@users.noreply.github.com> --- tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index a3a54970b388..45595de15037 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -9385,6 +9385,9 @@ def _sm107_bf16_valid_tactics( for use_2cta_instrs, mma_tiler_mn, max_num_ab_stage in itertools.product( use_2cta_instrs_candi, mma_tiler_mn_candi, _SM107_BF16_MAX_NUM_AB_STAGE_CANDIDATES): + # CTA_N=256 with cluster_n=2 is an illegal memory access on SM107. + if mma_tiler_mn[1] == 256 and preferred_cluster_shape_mn[1] == 2: + continue if not _bf16_cluster_m_fits(m, use_2cta_instrs, mma_tiler_mn, preferred_cluster_shape_mn): continue