[TRTLLM-6898] [feat] Add CuTe DSL fp8 gemm, bmm and group gemm ops on Blackwell architecuture

cpp/tensorrt_llm/thop/moeUtilOp.cpp

@@ -128,6 +128,10 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Te
     int const num_experts_per_node = num_experts_on_rank;
     auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
     int64_t num_moe_inputs = static_cast<int64_t>(experts_per_token * num_rows);
+    TORCH_CHECK(num_moe_inputs <= std::numeric_limits<int32_t>::max(),
+        "num_moe_inputs exceeds int32 range (because we use int32 for expert_first_token_offset_tensor output). "
+        "num_moe_inputs = ",
+        num_moe_inputs);
 
     auto permuted_row_to_unpermuted_row_tensor
         = torch::empty({num_moe_inputs}, torch::dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false));
@@ -224,6 +228,8 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Te
             "Invalid dtype, only supports input tensor with float32, float16 and bfloat16 dtype");
         break;
     }
+    expert_first_token_offset_tensor = expert_first_token_offset_tensor.to(torch::kInt32);
+
     return std::make_tuple(permuted_row_to_unpermuted_row_tensor, permuted_token_selected_experts_tensor,
         permuted_data_tensor, expert_first_token_offset_tensor, permuted_token_final_scales_tensor,
         unpermuted_row_to_permuted_row_tensor);

examples/llm-api/quickstart_advanced.py

@@ -147,6 +147,14 @@ def add_llm_args(parser):
                         default=False,
                         action='store_true')
     parser.add_argument('--logprobs', default=False, action='store_true')
+
+    # cute dsl op configs
+    parser.add_argument('--use_cute_dsl_blockscaling_mm',
+                        default=False,
+                        action='store_true')
+    parser.add_argument('--use_cute_dsl_blockscaling_bmm',
+                        default=False,
+                        action='store_true')
     return parser
 
 
@@ -246,6 +254,8 @@ def setup_llm(args, **kwargs):
         trust_remote_code=args.trust_remote_code,
         gather_generation_logits=args.return_generation_logits,
         max_beam_width=args.max_beam_width,
+        use_cute_dsl_blockscaling_mm=args.use_cute_dsl_blockscaling_mm,
+        use_cute_dsl_blockscaling_bmm=args.use_cute_dsl_blockscaling_bmm,
         **kwargs)
 
     use_beam_search = args.max_beam_width > 1

requirements.txt

@@ -67,3 +67,4 @@ soundfile
 triton==3.3.1; platform_machine == "x86_64"
 tiktoken
 blobfile
+nvidia-cutlass-dsl