address feedbacks

Signed-off-by: Anthony Chang <[email protected]>

Author: rosenrodt

cpp/include/tensorrt_llm/common/utils.h

@@ -17,10 +17,8 @@
 #pragma once
 
 #include <algorithm>
-#include <csignal>
 #include <initializer_list>
 #include <string>
-#include <unistd.h>
 
 #ifndef _WIN32
 #include <pthread.h>

cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/KernelRunner.cpp

@@ -153,49 +153,69 @@ TrtllmGenBatchedGemmRunner::TrtllmGenBatchedGemmRunner(TrtllmGenBatchedGemmRunne
         if (!acceptIf(options.mDtypeA == mOptions.dtypeA,
                 fmtstr("dtypeA mismatch (kernel: %s, expected: %s)", tg::dtypeToString(options.mDtypeA).c_str(),
                     tg::dtypeToString(mOptions.dtypeA).c_str())))
+        {
             continue;
+        }
 
         if (!acceptIf(options.mDtypeB == mOptions.dtypeB,
                 fmtstr("dtypeB mismatch (kernel: %s, expected: %s)", tg::dtypeToString(options.mDtypeB).c_str(),
                     tg::dtypeToString(mOptions.dtypeB).c_str())))
+        {
             continue;
+        }
 
         if (!acceptIf(options.mDtypeC == mOptions.dtypeC,
                 fmtstr("dtypeC mismatch (kernel: %s, expected: %s)", tg::dtypeToString(options.mDtypeC).c_str(),
                     tg::dtypeToString(mOptions.dtypeC).c_str())))
+        {
             continue;
+        }
 
         if (!acceptIf(options.mUseDeepSeekFp8 == mOptions.deepSeekFp8,
                 fmtstr(
                     "deepSeekFp8 mismatch (kernel: %d, expected: %d)", options.mUseDeepSeekFp8, mOptions.deepSeekFp8)))
+        {
             continue;
+        }
 
         if (!acceptIf(options.mTransposeMmaOutput == mOptions.transposeMmaOutput,
                 fmtstr("transposeMmaOutput mismatch (kernel: %d, expected: %d)", options.mTransposeMmaOutput,
                     mOptions.transposeMmaOutput)))
+        {
             continue;
+        }
 
         if (!acceptIf((!doesRouteImplUseNoRoute(options.mRouteImpl)) == mOptions.routeAct,
                 fmtstr("routeAct mismatch (kernel: %d, expected: %d)", !doesRouteImplUseNoRoute(options.mRouteImpl),
                     mOptions.routeAct)))
+        {
             continue;
+        }
 
         if (!acceptIf(options.mFusedAct == mOptions.fusedAct,
                 fmtstr("fusedAct mismatch (kernel: %d, expected: %d)", options.mFusedAct, mOptions.fusedAct)))
+        {
             continue;
+        }
 
         if (!acceptIf(options.mIsStaticBatch == mOptions.staticBatch,
                 fmtstr(
                     "staticBatch mismatch (kernel: %d, expected: %d)", options.mIsStaticBatch, mOptions.staticBatch)))
+        {
             continue;
+        }
 
         if (!acceptIf(tileSize == mOptions.tileSize,
                 fmtstr("tileSize mismatch (kernel: %d, expected: %d)", tileSize, mOptions.tileSize)))
+        {
             continue;
+        }
 
         if (!acceptIf(isSMCompatible(gpuSM, configs[i].mSm),
                 fmtstr("SM not compatible (gpuSM: %d, kernelSM: %d)", gpuSM, static_cast<int>(configs[i].mSm))))
+        {
             continue;
+        }
 
         auto sm = configs[i].mSm;
         if (sm != SmVersion::Sm100f)
@@ -205,28 +225,36 @@ TrtllmGenBatchedGemmRunner::TrtllmGenBatchedGemmRunner(TrtllmGenBatchedGemmRunne
             {
                 if (!acceptIf(sm == SmVersion::Sm100a,
                         fmtstr("SM version 100 requires Sm100a (kernel has: %d)", static_cast<int>(sm))))
+                {
                     continue;
+                }
             }
             else if (smVersion == 103)
             {
                 if (!acceptIf(sm == SmVersion::Sm103a,
                         fmtstr("SM version 103 requires Sm103a (kernel has: %d)", static_cast<int>(sm))))
+                {
                     continue;
+                }
             }
         }
 
         if (options.mUseDeepSeekFp8)
         {
             if (!acceptIf(options.mUseShuffledMatrixA == false, "useShuffledMatrixA should be false for DeepSeek Fp8"))
+            {
                 continue;
+            }
         }
 
         if (options.mFusedAct)
         {
             if (!acceptIf(options.mActType == static_cast<batchedGemm::gemmGatedAct::ActType>(mOptions.actType),
                     fmtstr("actType mismatch (kernel: %d, expected: %d)", static_cast<int>(options.mActType),
                         static_cast<int>(mOptions.actType))))
+            {
                 continue;
+            }
         }
 
         // FIXME: Disables a few static scheduler kernels (schedS) that appears to have issues;
@@ -236,14 +264,18 @@ TrtllmGenBatchedGemmRunner::TrtllmGenBatchedGemmRunner(TrtllmGenBatchedGemmRunne
         if (!acceptIf(!(options.mTileScheduler == TileScheduler::Static && options.mUseTmaOobOpt == true
                           && options.mTileN == 64),
                 "Static scheduler with TmaOobOpt and TileN=64 (known issue)"))
+        {
             continue;
+        }
 
         if (mOptions.transposeMmaOutput)
         {
             if (!acceptIf(options.mEpilogueTileM == mOptions.epilogueTileM,
                     fmtstr("epilogueTileM mismatch (kernel: %d, expected: %d)", options.mEpilogueTileM,
                         mOptions.epilogueTileM)))
+            {
                 continue;
+            }
         }
 
         // Kernel passed all filters

cpp/tensorrt_llm/thop/fp8PerTensorScaleMoe.cpp

@@ -37,6 +37,9 @@ torch::Tensor fp8_per_tensor_scale_moe_runner(torch::optional<torch::Tensor> con
     torch::optional<torch::Tensor> const& topk_weights, torch::optional<torch::Tensor> const& topk_ids)
 {
     TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "Only SM100f is supported by FP8 block scale MOE");
+    TORCH_CHECK(tile_tokens_dim == 8 || tile_tokens_dim == 16 || tile_tokens_dim == 32 || tile_tokens_dim == 64
+            || tile_tokens_dim == 128 || tile_tokens_dim == 192 || tile_tokens_dim == 256,
+        "tile_tokens_dim must be 8, 16, 32, 64, 128, 256");
     if (topk_ids.has_value() && topk_weights.has_value())
     {
         TORCH_CHECK(topk_ids.value().scalar_type() == at::ScalarType::Int, "topk_ids must be int");

tests/unittest/_torch/modules/test_fused_moe.py

@@ -55,14 +55,20 @@
 
 @contextmanager
 def moe_trtllm_debug_msg(enable=False):
+    TLLM_BATCHED_GEMM_PRINT_NAME = os.environ.get(
+        "TLLM_BATCHED_GEMM_PRINT_NAME", "0")
+    TLLM_BATCHED_GEMM_PRINT_CONFIGS = os.environ.get(
+        "TLLM_BATCHED_GEMM_PRINT_CONFIGS", "0")
     if enable:
         os.environ["TLLM_BATCHED_GEMM_PRINT_NAME"] = "1"
         os.environ["TLLM_BATCHED_GEMM_PRINT_CONFIGS"] = "1"
     try:
         yield
     finally:
-        os.environ["TLLM_BATCHED_GEMM_PRINT_NAME"] = "0"
-        os.environ["TLLM_BATCHED_GEMM_PRINT_CONFIGS"] = "0"
+        os.environ[
+            "TLLM_BATCHED_GEMM_PRINT_NAME"] = TLLM_BATCHED_GEMM_PRINT_NAME
+        os.environ[
+            "TLLM_BATCHED_GEMM_PRINT_CONFIGS"] = TLLM_BATCHED_GEMM_PRINT_CONFIGS
 
 
 def round_up(x, alignment):