diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp index 7bba57a03d5e..bcc0eb1165ac 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2020-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -17,6 +17,7 @@ #include "tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.h" #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/cudaBf16Wrapper.h" +#include "tensorrt_llm/common/cudaUtils.h" #ifdef __GNUC__ // Check if the compiler is GCC or Clang #pragma GCC diagnostic push @@ -31,6 +32,7 @@ #pragma GCC diagnostic pop #endif // __GNUC +#include #include #include #include @@ -573,7 +575,7 @@ std::vector get_candidate_configs_sm120(CutlassGemmConfig::Ca std::vector candidate_configs; if (config & CutlassGemmConfig::FP8FP4_MIXED) { - // Mixed FP8 x FP4: restrict to 128x128x128B only + // Mixed FP8 x FP4 only supports the 128x128x128B tile. candidate_configs.push_back(CutlassGemmConfig{CutlassTileConfigSM120::CtaShape128x128x128B, MainloopScheduleType::AUTO, EpilogueScheduleType::AUTO, ClusterShape::ClusterShape_1x1x1}); return candidate_configs; @@ -589,9 +591,34 @@ std::vector get_candidate_configs_sm120(CutlassGemmConfig::Ca MainloopScheduleType::AUTO, EpilogueScheduleType::AUTO, ClusterShape::ClusterShape_1x1x1}); candidate_configs.push_back(CutlassGemmConfig{CutlassTileConfigSM120::CtaShape256x128x64B, MainloopScheduleType::AUTO, EpilogueScheduleType::AUTO, ClusterShape::ClusterShape_1x1x1}); - return candidate_configs; } - TLLM_THROW("Not Implemented: SM120 group GEMM only supports mxfp8-mxfp4 mixed or nvfp4."); + else + { + TLLM_THROW("Not Implemented: SM120 group GEMM only supports mxfp8-mxfp4 mixed or nvfp4."); + } + // Filter configs by device shared memory. SM100 (B200) has 228 KiB, but + // consumer Blackwell (SM120 RTX PRO 6000, SM121 GB10 / DGX Spark) has only + // 99 KiB. On these constrained devices, keep only CtaShape128x128x64B which + // fits within 99 KiB including FINALIZE epilogue (~80 KiB total). + // CtaShape128x256x64B/256x128x64B overflow with FINALIZE (~100 KiB). + // CtaShape128x128x128B also exceeds 99 KiB at typical stage counts. + { + constexpr int kMinSmemForFullTileSet = 120 * 1024; + int device = 0; + tensorrt_llm::common::check_cuda_error(cudaGetDevice(&device)); + int maxSmem = 0; + tensorrt_llm::common::check_cuda_error( + cudaDeviceGetAttribute(&maxSmem, cudaDevAttrMaxSharedMemoryPerBlockOptin, device)); + + if (maxSmem < kMinSmemForFullTileSet) + { + auto const it = std::remove_if(candidate_configs.begin(), candidate_configs.end(), + [](CutlassGemmConfig const& config) + { return config.tile_config_sm120 != CutlassTileConfigSM120::CtaShape128x128x64B; }); + candidate_configs.erase(it, candidate_configs.end()); + } + } + return candidate_configs; } else { diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl index 0044528b4dff..9c5ebbdaa19d 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl @@ -15,6 +15,8 @@ */ #pragma once +#include "tensorrt_llm/common/cudaUtils.h" + #include "cutlass/array.h" #include "cutlass/numeric_conversion.h" @@ -677,6 +679,19 @@ using namespace cutlass::epilogue; "Workspace is size %zu but only %zu were allocated", calculated_ws_size, \ tma_ws_input.gemm_workspace_size); \ \ + /* Check if kernel SMEM fits on the active device before launch. */ \ + { \ + using GemmKernel_ = typename GemmGrouped::GemmKernel; \ + int smem_size = static_cast(sizeof(typename GemmKernel_::SharedStorage)); \ + int device_ = 0; \ + tensorrt_llm::common::check_cuda_error(cudaGetDevice(&device_)); \ + int maxSmem_ = 0; \ + tensorrt_llm::common::check_cuda_error( \ + cudaDeviceGetAttribute(&maxSmem_, cudaDevAttrMaxSharedMemoryPerBlockOptin, device_)); \ + TLLM_CHECK_WITH_INFO(smem_size <= maxSmem_, \ + "MoE grouped GEMM requires %d bytes shared memory but device supports %d", smem_size, maxSmem_); \ + } \ + \ auto can_implement = gemm.can_implement(args); \ TLLM_CHECK_WITH_INFO(can_implement == cutlass::Status::kSuccess, \ "Grouped GEMM kernel will fail for params. Error: " \ diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.inl b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.inl index f37920dcf73c..cdf5ea8dc3c9 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.inl +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.inl @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2020-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -19,6 +19,8 @@ #pragma GCC diagnostic ignored "-Wstrict-aliasing" #endif // __GNUC__ +#include "tensorrt_llm/common/cudaUtils.h" + #include "cutlass/epilogue/collective/default_epilogue.hpp" #include "cutlass/epilogue/thread/linear_combination.h" #include "cutlass/gemm/collective/collective_builder.hpp" @@ -273,6 +275,17 @@ void sm90_generic_mixed_moe_gemm_kernelLauncher(GroupedGemmInput(sizeof(typename GemmKernel::SharedStorage)); + int device = 0; + tensorrt_llm::common::check_cuda_error(cudaGetDevice(&device)); + int maxSmem = 0; + tensorrt_llm::common::check_cuda_error( + cudaDeviceGetAttribute(&maxSmem, cudaDevAttrMaxSharedMemoryPerBlockOptin, device)); + TLLM_CHECK_WITH_INFO(smem_size <= maxSmem, + "Mixed dtype WS grouped GEMM requires %d bytes shared memory but device supports %d", smem_size, maxSmem); + } + auto can_implement = gemm.can_implement(arguments); if (can_implement != cutlass::Status::kSuccess) {