debug: sync + printf before mega_moe kernel launch

2026-05-13 12:15:49 +00:00
parent ad335c38fb
commit c08a28888d
1 changed files with 6 additions and 0 deletions
--- a/csrc/jit_kernels/impls/sm100_fp8_nvfp4_mega_moe.hpp
+++ b/csrc/jit_kernels/impls/sm100_fp8_nvfp4_mega_moe.hpp
@@ -89,6 +89,12 @@ static void __instantiate_kernel() {{
    }

    static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
+        // Debug: sync before launch to flush TMA creation errors
+        if (get_env<int>("DG_JIT_DEBUG")) {
+            DG_CUDA_CHECK(cudaDeviceSynchronize());
+            printf("[MEGA_MOE_LAUNCH_DEBUG] About to launch kernel, y=%p, num_tokens=%d\n", args.y, args.num_tokens);
+            fflush(stdout);
+        }
        DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
            args.y,
            args.cumulative_local_expert_recv_stats,