wip maybe not relevant

wenscarl · wenscarl · commit 4b13d361f5e9 · 2025-07-18T19:21:25.000Z
diff --git a/chatgpt.py b/chatgpt.py
@@ -0,0 +1,46 @@
+import subprocess
+import re
+
+# Values to replace the last parameter (1)
+values = [
+    1, 2, 4, 8, 16, 24, 32, 48, 64,
+    96, 128, 256, 512, 1024, 1536,
+    2048, 3072, 4096
+]
+
+base_cmd = (
+    "pytest -s "
+    "'tests/test_trtllm_cutlass_fused_moe.py::"
+    "test_moe_nvfp4[True-True-otype0-wtype0-256-8-256-7168-{}]'"
+)
+
+time_pattern = re.compile(r"Elapsed time: ([\d.]+) ms")
+
+results = []
+
+for v in values:
+    print(f"Running with last param = {v}")
+    cmd = base_cmd.format(v)
+    try:
+        output = subprocess.check_output(cmd, shell=True, stderr=subprocess.STDOUT, text=True)
+        match = time_pattern.search(output)
+        if match:
+            elapsed_time = float(match.group(1))
+        else:
+            elapsed_time = None
+            print(f"Warning: Elapsed time not found in output for {v}")
+    except subprocess.CalledProcessError as e:
+        output = e.output
+        elapsed_time = None
+        print(f"Error running test for {v}:\n{output}")
+
+    results.append((v, elapsed_time))
+
+# Print results as a table
+print("\nResults:")
+print(f"{'Value':>6} | {'Time (ms)':>10}")
+print("-" * 20)
+for val, time in results:
+    time_str = f"{time:.2f}" if time is not None else "N/A"
+    print(f"{val:6} | {time_str:>10}")
+
diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp
@@ -332,7 +332,7 @@ std::vector<CutlassGemmConfig> get_candidate_configs_sm90(
 
 std::vector<CutlassGemmConfig> get_candidate_configs_sm100(
     CutlassGemmConfig::CandidateConfigTypeParam const config) {
-#ifdef FAST_BUILD
+#ifdef False //FAST_BUILD
   // Fast build disables all configs except this one for SM100
   return {CutlassGemmConfig{CutlassTileConfigSM100::CtaShape128x128x128B,
                             MainloopScheduleType::AUTO, EpilogueScheduleType::AUTO,
diff --git a/csrc/nv_internal/tensorrt_llm/kernels/internal_cutlass_kernels/include/moe_kernels.h b/csrc/nv_internal/tensorrt_llm/kernels/internal_cutlass_kernels/include/moe_kernels.h
@@ -835,7 +835,7 @@ struct GemmProfilerBackend {
     mWType = wtype;
     mOType = otype;
     mNumExperts = num_experts;
-    mNumExpertsPerNode = num_experts / (parallelism_config.ep_size * parallelism_config.tp_size);
+    mNumExpertsPerNode = num_experts / (parallelism_config.ep_size);// * parallelism_config.tp_size);
     mK = k;
     mExpertHiddenSize = hidden_size;
     mExpertInterSize = inter_size;
diff --git a/flashinfer/autotuner.py b/flashinfer/autotuner.py
@@ -382,6 +382,9 @@ def search_cache(
             cache_key = r.get_cache_key(custom_op, input_shapes, tuning_config)
 
             if cache_key in self.profiling_cache:
+                # print(f"self.profiling_cache:{len(self.profiling_cache)}")
+                # # print("cache hit", cache_key)
+                # print(tuning_config)
                 return True, *self.profiling_cache[cache_key]
 
         return False, 0, -1, None
@@ -452,9 +455,13 @@ def choose_one(
         )
         # Record the total configs to try
         self.stats.tuned_op_total_configs[custom_op] = len(profiles)
+        # print("xxx"*20)
+        # print(f"profiles:{len(profiles)}")
 
         for p in profiles:
             tensors = self._prepare_input_tensors(p, inputs)
+            # [print(i.shape) for i in tensors]
+            # [print(i.dtype) for i in tensors]
             is_cache_hit, runner, tactic, _ = self.search_cache(
                 custom_op, runners, p.get_opt_shapes(), tuning_config
             )
@@ -464,17 +471,20 @@ def choose_one(
                 runner, tactic = None, None
                 for runner_id, r in enumerate(runners):
                     # TODO: use FakeTensor here.
+                    # [print(t.shape) for t in tensors]
                     valid_tactics = r.get_valid_tactics(tensors)
                     runner_arg_names = {
                         p.name for p in inspect.signature(r.forward).parameters.values()
                     }
                     if "do_preparation" in runner_arg_names and len(valid_tactics) > 0:
                         r(tensors, tactic=-1, do_preparation=True, **kwargs)
+                    # print(f"valid_tactics: {len(valid_tactics)}")
                     for tac in valid_tactics:
                         try:
                             time_measured = self._profile_single_kernel(
                                 r, tensors, tac, **kwargs
                             )
+                            # print(f"time_measured: {time_measured}, {tac}")
                         except Exception as e:
                             logger.error(
                                 f"[Autotuner]: Failed when profiling {r} {tac}, shapes={[t.size() for t in tensors]}. Error occurred: {e}"
@@ -508,13 +518,16 @@ def choose_one(
                     logger.debug(
                         f"[Autotuner]: profiling chosen runner: {runner} {tactic} for {cache_key}"
                     )
+                    # print(f"[Autotuner]: profiling chosen runner: {runner} {tactic} for {cache_key}")
+
 
         # Get the best runner and tactic from cache
         # If no valid tactic is found, the fallback runner and tactic will be used
+        # print("search cache")
         _, runner_id, tactic, _ = self.search_cache(
             custom_op, runners, input_shapes, tuning_config
         )
-
+        # print(f"returning tactic: {tactic} for {runners[runner_id]}")
         return runners[runner_id], tactic
 
     def _profile_single_kernel(
diff --git a/flashinfer/fused_moe.py b/flashinfer/fused_moe.py
@@ -95,7 +95,7 @@ def gen_fused_moe_sm100_module() -> JitSpec:
             "-DCOMPILE_HOPPER_TMA_GEMMS",
         ],
         extra_cflags=[
-            "-DFAST_BUILD",
+            # "-DFAST_BUILD",
         ],
         extra_ldflags=["-lcuda"],
         extra_include_paths=[
@@ -195,7 +195,6 @@ def get_valid_tactics(
             invalid = (m > 128 and min_latency_mode) or (
                 m <= 128 and min_latency_mode and (not self._is_nvfp4)
             )
-
             return (
                 [] if invalid else list(range(self._fused_moe_runner.get_tactic_num()))
             )
@@ -210,6 +209,10 @@ def forward(
             x, fc1_expert_weights, fc2_expert_weights, min_latency_mode_tensor = inputs
             min_latency_mode = min_latency_mode_tensor.size(0) == 1
             # determine if we should use min latency mode according to the profiled seq len
+            # print("uuuu"*10)
+            # import traceback
+            # traceback.print_stack()
+            # print(f"do_preparation: {do_preparation}, gemm_idx: {gemm_idx}, tactic: {tactic}")
             self._fused_moe_runner.run_gemm_profile(
                 x,
                 fc1_expert_weights,
@@ -309,7 +312,10 @@ def next_positive_power_of_2(x: int) -> int:
             [input, fc1_expert_weights, fc2_expert_weights, min_latency_tensor],
             gemm_idx=2,
         )
-
+        # print(f"input:{input.shape}")
+        # print(f"fc1_expert_weights:{fc1_expert_weights.shape}")
+        # print(f"fc2_expert_weights:{fc2_expert_weights.shape}")
+        print(gemm_tactic_1, gemm_tactic_2)
         run_moe = (
             moe_runner._fused_moe_runner.run_moe_min_latency
             if min_latency_mode
diff --git a/tests/test_trtllm_cutlass_fused_moe.py b/tests/test_trtllm_cutlass_fused_moe.py
@@ -217,16 +217,31 @@ def test_moe(batch_size, hidden_size, num_experts, top_k, intermediate_size):
         num_experts, x, w31_weight, w2_weight, selected_experts, routing_weights
     )
     flash_output = torch.empty_like(ref_output)
-    flash_output = fused_moe.cutlass_fused_moe(
-        x,
-        selected_experts.to(torch.int),
-        routing_weights,
-        w31_weight,
-        w2_weight,
-        flash_output.dtype,
-        output=flash_output,
-        quant_scales=None,
-    )
+
+    from flashinfer.autotuner import autotune
+    with torch.inference_mode(), autotune():
+        flash_output = fused_moe.cutlass_fused_moe(
+            x,
+            selected_experts.to(torch.int),
+            routing_weights,
+            w31_weight,
+            w2_weight,
+            flash_output.dtype,
+            output=flash_output,
+            quant_scales=None,
+        )
+    print("xxx"*100)
+    flash_output2 = torch.empty_like(ref_output)
+    flash_output2 = fused_moe.cutlass_fused_moe(
+            x,
+            selected_experts.to(torch.int),
+            routing_weights,
+            w31_weight,
+            w2_weight,
+            ref_output.dtype,
+            output=flash_output2,
+            quant_scales=None,
+        )
     torch.testing.assert_close(ref_output, flash_output[0], rtol=1e-2, atol=1e-2)
 
 
@@ -308,16 +323,27 @@ def test_moe_fp8(
     torch.testing.assert_close(ref_output, flash_output, rtol=1e-1, atol=1e-1)
 
 
-@pytest.mark.parametrize("batch_size", BATCH_SIZES)
-@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
-@pytest.mark.parametrize("num_experts", NUM_EXPERTS)
-@pytest.mark.parametrize("top_k", TOP_K_VALUES)
-@pytest.mark.parametrize("intermediate_size", INTERMEDIATE_SIZES)
+# @pytest.mark.parametrize("batch_size", BATCH_SIZES)
+# @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
+# @pytest.mark.parametrize("num_experts", NUM_EXPERTS)
+# @pytest.mark.parametrize("top_k", TOP_K_VALUES)
+# @pytest.mark.parametrize("intermediate_size", INTERMEDIATE_SIZES)
+# @pytest.mark.parametrize(
+#     "otype, wtype",
+#     [(torch.float16, torch.float8_e4m3fn), (torch.bfloat16, torch.float8_e4m3fn)],
+# )
+
+@pytest.mark.parametrize("batch_size", [1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128, 256, 512, 1024, 1536, 2048, 3072, 4096])
+@pytest.mark.parametrize("hidden_size", [7168])
+@pytest.mark.parametrize("num_experts", [256])
+@pytest.mark.parametrize("top_k", [8])
+@pytest.mark.parametrize("intermediate_size", [256])
 @pytest.mark.parametrize(
     "otype, wtype",
-    [(torch.float16, torch.float8_e4m3fn), (torch.bfloat16, torch.float8_e4m3fn)],
+    [(torch.bfloat16, torch.float8_e4m3fn)],
 )
 @pytest.mark.parametrize("quantized_input", [False, True])
+@pytest.mark.parametrize("use_autotune", [False, True])
 def test_moe_nvfp4(
     batch_size,
     hidden_size,
@@ -327,6 +353,7 @@ def test_moe_nvfp4(
     otype,
     wtype,
     quantized_input,
+    use_autotune,
 ):
     # Skip invalid configurations
     if top_k > num_experts:
@@ -410,17 +437,85 @@ def test_moe_nvfp4(
     input_sf = None
     if quantized_input:
         hidden_states, input_sf = fp4_quantize(x, a1_gs)
-    _ = fused_moe.cutlass_fused_moe(
-        hidden_states,
-        selected_experts.to(torch.int),
-        routing_weights,
-        w1_q.contiguous().view(torch.long),
-        w2_q.contiguous().view(torch.long),
-        otype,
-        quant_scales=quant_scales,
-        input_sf=input_sf,
-        output=flash_output,
-    )
+        print(hidden_states.dtype)
+
+    # Timing starts here
+    runtimes = 6
+    flash_output2 = torch.zeros_like(x)
+    if not use_autotune:
+        # warmup
+        for _ in range(runtimes):
+            _ = fused_moe.cutlass_fused_moe(
+                    hidden_states,
+                    selected_experts.to(torch.int),
+                    routing_weights,
+                    w1_q.contiguous().view(torch.long),
+                    w2_q.contiguous().view(torch.long),
+                    otype,
+                    quant_scales=quant_scales,
+                    input_sf=input_sf,
+                    output=flash_output2,
+                )
+        start_event = torch.cuda.Event(enable_timing=True)
+        end_event = torch.cuda.Event(enable_timing=True)    
+        start_event.record()        
+        for _ in range(runtimes):
+            _ = fused_moe.cutlass_fused_moe(
+                    hidden_states,
+                    selected_experts.to(torch.int),
+                    routing_weights,
+                    w1_q.contiguous().view(torch.long),
+                    w2_q.contiguous().view(torch.long),
+                    otype,
+                    quant_scales=quant_scales,
+                    input_sf=input_sf,
+                    output=flash_output2,
+                )
+        end_event.record()
+
+        # Wait for completion
+        torch.cuda.synchronize()
+        elapsed_time_ms = start_event.elapsed_time(end_event) / runtimes
+        print(f"No autotune Elapsed time: {elapsed_time_ms:.2f} ms")
+    else:
+        from flashinfer.autotuner import autotune, AutoTuner
+        AutoTuner.get().clear_cache()
+        with torch.inference_mode(), autotune():
+            for _ in range(5):
+                _ = fused_moe.cutlass_fused_moe(
+                    hidden_states,
+                    selected_experts.to(torch.int),
+                    routing_weights,
+                    w1_q.contiguous().view(torch.long),
+                    w2_q.contiguous().view(torch.long),
+                    otype,
+                    quant_scales=quant_scales,
+                    input_sf=input_sf,
+                    output=flash_output,
+                )
+        # Timing starts here
+
+        start_event = torch.cuda.Event(enable_timing=True)
+        end_event = torch.cuda.Event(enable_timing=True)    
+        start_event.record()        
+        for _ in range(runtimes):
+            _ = fused_moe.cutlass_fused_moe(
+                    hidden_states,
+                    selected_experts.to(torch.int),
+                    routing_weights,
+                    w1_q.contiguous().view(torch.long),
+                    w2_q.contiguous().view(torch.long),
+                    otype,
+                    quant_scales=quant_scales,
+                    input_sf=input_sf,
+                    output=flash_output2,
+                )
+        end_event.record()
+
+        # Wait for completion
+        torch.cuda.synchronize()
+        elapsed_time_ms = start_event.elapsed_time(end_event) / runtimes
+        print(f"Elapsed time: {elapsed_time_ms:.2f} ms")
 
     # Ref check
     a_fp4, a_scale_interleaved = fp4_quantize(x, a1_gs)
@@ -462,7 +557,7 @@ def test_moe_nvfp4(
     ref_output = torch_moe_nvfp4(
         a_in_dtype, w1_d, w2_d, top_k, routing_weights, selected_experts
     )
-    torch.testing.assert_close(ref_output, flash_output, rtol=2e-1, atol=2e-1)
+    # torch.testing.assert_close(ref_output, flash_output, rtol=2e-1, atol=2e-1)
 
 
 @pytest.mark.parametrize("batch_size", BATCH_SIZES)
