fix combine kernel bug for full cudagraph

Author: liangel-02

hopper/flash_fwd_combine_launch_template.h

@@ -4,6 +4,8 @@
 
 #pragma once
 
+#include <cstdio>
+
 #include "cute/tensor.hpp"
 
 #include "cutlass/cutlass.h"
@@ -41,7 +43,23 @@ void run_flash_fwd_combine(Flash_fwd_params &params, cudaStream_t stream, bool e
 
     typename CombineKernel::Params kernel_params = CombineKernel::to_underlying_arguments(args);
     int num_blocks_k = cute::ceil_div(params.dv, kBlockK);
-    int num_blocks_m = cute::ceil_div(params.seqlen_q * params.h, kBlockM);
+    // [FIX] For varlen, size the combine grid from the actual total query extent
+    // (params.total_q) instead of the per-request max seqlen_q, which under a
+    // uniform-decode-captured CUDA graph is baked stale (=1), leaving the extra
+    // query rows of a mixed/prefill replay uncombined -> uninitialized/NaN.
+    int num_blocks_m = cute::ceil_div(
+        (Varlen ? params.total_q : params.seqlen_q) * params.h, kBlockM);
+    // [BUILD-VERIFY] one-time host print so we can confirm the rebuilt .so is loaded.
+    static bool _fa3_combine_fix_printed = false;
+    if (!_fa3_combine_fix_printed) {
+        _fa3_combine_fix_printed = true;
+        fprintf(stderr,
+            "[FA3-COMBINE-FIX ACTIVE] Varlen=%d seqlen_q=%d total_q=%d h=%d "
+            "kBlockM=%d -> num_blocks_m=%d\n",
+            (int)Varlen, (int)params.seqlen_q, (int)params.total_q,
+            (int)params.h, (int)kBlockM, num_blocks_m);
+        fflush(stderr);
+    }
     dim3 grid_m(num_blocks_m, num_blocks_k, params.b);
     auto kernel = cutlass::device_kernel<CombineKernel>;
     int smem_size = CombineKernel::SharedStorageSize;