fix: compile fresh each call — cached compile produces wrong TMA descriptors

The CuTeDSL kernel's TMA descriptors are bound to the
compilation-time tensor addresses. Caching the compiled kernel
and reusing it with different tensor allocations produces wrong
memory access patterns (cosine 0.5 instead of 0.99).

Fresh compilation is proven correct (cosine 0.989). We can
optimize later with proper TMA descriptor reinitialization.

cutedsl/bridge.py

@@ -318,19 +318,29 @@ def run_nvfp4_grouped_gemm(
     
     2Dx3D: A(tokens, K) x B(experts, K, N) -> C(tokens, N)
     
-    Compiles with real tensors (not dummy shapes) because the CuTeDSL
-    kernel's TMA descriptors are sized from compilation-time tensor shapes.
-    Compilation is cached per (num_experts, K, N) and happens once.
+    Compiles with real tensors each call. The CuTeDSL kernel's TMA
+    descriptors are bound to the compilation-time tensor addresses,
+    so caching across different tensor allocations produces wrong results.
     
     The forward call (after compilation) is cudagraph-safe.
     """
     num_experts = mat_b.shape[0]
-    n_dim = mat_b.shape[2]  # N dimension
+    n_dim = mat_b.shape[2]
     tokens_sum = mat_a.shape[0]
     device = mat_a.device
 
     out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=device)
 
+    kernel = ScaledGroupedGemmKernel(
+        scenario="2Dx3D",
+        sf_vec_size=SF_VEC_SIZE,
+        accumulate_on_output=False,
+        separate_tensormap_init=True,
+        consistent_token_padding=False,
+        mma_tiler_mnk=(*mma_tiler_mn, 256),
+        cluster_shape_mnk=(*cluster_shape_mn, 1),
+    )
+
     # Convert to CuTe tensors with dynamic layout
     def to_cute(t):
         ct = cutlass_torch.from_dlpack(t)
@@ -343,60 +353,23 @@ def run_nvfp4_grouped_gemm(
     c_c = to_cute(out)
     offs_c = to_cute(expert_offsets)
 
-    workspace_size = 0
+    workspace_size = kernel.get_workspace_size(num_experts)
+    workspace = torch.full((workspace_size,), 255, dtype=torch.uint8, device=device)
+    ws_c = to_cute(workspace)
+
     gsa_c = to_cute(global_scale_a) if global_scale_a is not None else None
     gsb_c = to_cute(global_scale_b) if global_scale_b is not None else None
 
     import cuda.bindings.driver as cuda
+    cluster_size = cluster_shape_mn[0] * cluster_shape_mn[1]
+    max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size)
     stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
 
-    # Check cache — compile with real tensors on first use per (experts, K, N)
-    K_packed = mat_a.shape[1]
-    N_packed = mat_b.shape[2]
-    cache_key = (num_experts, str(device), mma_tiler_mn, cluster_shape_mn, K_packed, N_packed)
-    use_cache = True  # TEMP: set False to always recompile (debug)
-    
-    if not (use_cache and cache_key in _compiled_kernel_cache):
-        kernel = ScaledGroupedGemmKernel(
-            scenario="2Dx3D",
-            sf_vec_size=SF_VEC_SIZE,
-            accumulate_on_output=False,
-            separate_tensormap_init=True,
-            consistent_token_padding=False,
-            mma_tiler_mnk=(*mma_tiler_mn, 256),
-            cluster_shape_mnk=(*cluster_shape_mn, 1),
-        )
-
-        cluster_size = cluster_shape_mn[0] * cluster_shape_mn[1]
-        max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size)
-
-        workspace_size = kernel.get_workspace_size(num_experts)
-        workspace = torch.full((workspace_size,), 255, dtype=torch.uint8, device=device)
-        ws_c = to_cute(workspace)
-
-        compiled = cute.compile(
-            kernel, a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c,
-            max_active_clusters, stream,
-            global_scale_a=gsa_c, global_scale_b=gsb_c,
-        )
-
-        # Warm up with real data
-        compiled(
-            a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c,
-            stream,
-            global_scale_a=gsa_c, global_scale_b=gsb_c,
-        )
-        torch.cuda.synchronize()
-
-        _compiled_kernel_cache[cache_key] = (compiled, kernel, max_active_clusters)
-
-    compiled, kernel, max_active_clusters = _compiled_kernel_cache[cache_key]
-
-    # Allocate workspace if not already done during compilation
-    if workspace_size == 0:
-        workspace_size = kernel.get_workspace_size(num_experts)
-    workspace = torch.full((workspace_size,), 255, dtype=torch.uint8, device=device)
-    ws_c = to_cute(workspace)
+    compiled = cute.compile(
+        kernel, a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c,
+        max_active_clusters, stream,
+        global_scale_a=gsa_c, global_scale_b=gsb_c,
+    )
 
     compiled(
         a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c,