[None][feat] Add SM107 CuTe DSL BF16 dense GEMM/BMM custom ops and dispatch

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)
 

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"