[None][refactor] Dispatch both Sage smoothing flags in one place

launchWithSmoothK took the V-smoothing flag from its caller, which
dispatched it per stage. launchWithSmoothKV dispatches both flags for a
stage and instantiates only the combinations that exist: K smoothing
needs V staging, and V smoothing only changes the statistics launch.

Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>

Author: heyuhhh

cpp/tensorrt_llm/common/sageQuant.cu

@@ -726,33 +726,39 @@ void invokeSageQuantQkvImpl(SageQuantParams const& params)
             TLLM_CHECK_WITH_INFO(status == cudaSuccess, "%s", cudaGetErrorString(status));
         };
 
-        auto launchWithSmoothK = [&](auto smoothVStatic, auto vStageStatic)
+        // K smoothing needs V staging, and V smoothing only changes what the VStage==1 launch
+        // collects, so only those combinations are instantiated.
+        auto launchWithSmoothKV = [&](auto vStageStatic)
         {
-            if (params.kSmooth)
+            constexpr int VStage_ = vStageStatic;
+            auto launchWithSmoothV = [&](auto smoothKStatic)
             {
-                launchKernel(cute::true_type{}, smoothVStatic, vStageStatic);
-            }
-            else
+                if constexpr (VStage_ == 1)
+                {
+                    if (params.ptrVMean != nullptr)
+                    {
+                        launchKernel(smoothKStatic, cute::true_type{}, vStageStatic);
+                        return;
+                    }
+                }
+                launchKernel(smoothKStatic, cute::false_type{}, vStageStatic);
+            };
+            if constexpr (VStage_ != 0)
             {
-                launchKernel(cute::false_type{}, smoothVStatic, vStageStatic);
+                if (params.kSmooth)
+                {
+                    launchWithSmoothV(cute::true_type{});
+                    return;
+                }
             }
+            launchWithSmoothV(cute::false_type{});
         };
 
-        // V smoothing only changes what the VStage==1 launch collects.
         switch (params.vStage)
         {
-        case 0: launchKernel(cute::false_type{}, cute::false_type{}, Int<0>{}); return;
-        case 1:
-            if (params.ptrVMean != nullptr)
-            {
-                launchWithSmoothK(cute::true_type{}, Int<1>{});
-            }
-            else
-            {
-                launchWithSmoothK(cute::false_type{}, Int<1>{});
-            }
-            return;
-        case 2: launchWithSmoothK(cute::false_type{}, Int<2>{}); return;
+        case 0: launchWithSmoothKV(Int<0>{}); return;
+        case 1: launchWithSmoothKV(Int<1>{}); return;
+        case 2: launchWithSmoothKV(Int<2>{}); return;
         default: TLLM_THROW("Unsupported SageQuantV stage %d", params.vStage);
         }
     };