Reviewer comment

Signed-off-by: Dom Brown <[email protected]>

Author: DomBrown

cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/KernelRunner.cpp

@@ -52,7 +52,8 @@ TrtllmGenGemmRunner::TrtllmGenGemmRunner(tg::Dtype eltType, tg::Dtype outputType
     mGemmConfig = &configs[selectedIndex[0]];
 }
 
-size_t TrtllmGenGemmRunner::getWorkspaceSizeInBytes(int32_t m, int32_t n, int32_t k) const
+size_t TrtllmGenGemmRunner::getWorkspaceSizeInBytes(
+    int32_t m, int32_t n, int32_t k, tg::Dtype eltType, tg::Dtype outputType) const
 {
     gemm::GemmData gemmData;
     gemmData.mProblemDimensions.mM = m;

cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/KernelRunner.h

@@ -33,7 +33,8 @@ class TrtllmGenGemmRunner
 public:
     explicit TrtllmGenGemmRunner(tg::Dtype eltType, tg::Dtype outputType);
 
-    [[nodiscard]] size_t getWorkspaceSizeInBytes(int32_t m, int32_t n, int32_t k) const;
+    [[nodiscard]] size_t getWorkspaceSizeInBytes(
+        int32_t m, int32_t n, int32_t k, tg::Dtype eltType, tg::Dtype outputType) const;
 
     void run(int32_t m, int32_t n, int32_t k, void const* a, float const* aScale, void const* b, float const* bScale,
         void* c, float* cScale, void* workspace, CUstream stream, int device);

cpp/tensorrt_llm/thop/fp4GemmTrtllmGen.cpp

@@ -39,7 +39,7 @@ void runGemm(at::Tensor& out, at::Tensor const& mat1, at::Tensor const& mat2, at
 
     tensorrt_llm::kernels::TrtllmGenGemmRunner runner(eltType, out_dtype);
 
-    int64_t const numBytesWorkspace = runner.getWorkspaceSizeInBytes(m, n, k);
+    int64_t const numBytesWorkspace = runner.getWorkspaceSizeInBytes(m, n, k, eltType, out_dtype);
     at::Tensor workspace
         = at::detail::empty_cuda({numBytesWorkspace}, at::ScalarType::Char, torch::kCUDA, std::nullopt);