wafer-core 0.1.27__py3-none-any.whl → 0.1.28__py3-none-any.whl

wafer_core/lib/trace_compare/aligner.py

@@ -214,21 +214,28 @@ def align_kernels_within_layer(
         # The platform that HAS the kernel IS fusing; the other runs components separately
         is_fused_op = "+" in op_str
         
+        # Operations that can't be "fused away" - absence means alignment issue, not fusion
+        non_fusable_ops = {
+            "Attention (Prefill)", "Attention (Decode)", "Dense GEMM", 
+            "KV Cache", "MoE GEMM", "MoE Routing"
+        }
+        is_non_fusable = op_str in non_fusable_ops
+        
         fusion_note = None
         if amd_count > 0 and nvidia_count == 0:
             if is_fused_op:
                 # AMD has a fused kernel like "RMSNorm+GEMM" → AMD IS fusing
                 fusion_note = f"AMD fuses {op_str} into {amd_kernel_name}"
-            else:
-                # AMD has a regular kernel that NVIDIA doesn't need → NVIDIA fuses it elsewhere
-                fusion_note = f"AMD runs {amd_kernel_name}, NVIDIA fuses into another kernel"
+            elif not is_non_fusable:
+                # Only mark as fusion for ops that can legitimately be fused
+                fusion_note = f"AMD runs {amd_kernel_name}, NVIDIA may fuse into another kernel"
         elif amd_count == 0 and nvidia_count > 0:
             if is_fused_op:
                 # NVIDIA has a fused kernel → NVIDIA IS fusing
                 fusion_note = f"NVIDIA fuses {op_str} into {nvidia_kernel_name}"
-            else:
-                # NVIDIA has a regular kernel that AMD doesn't need → AMD fuses it elsewhere
-                fusion_note = f"NVIDIA runs {nvidia_kernel_name}, AMD fuses into another kernel"
+            elif not is_non_fusable:
+                # Only mark as fusion for ops that can legitimately be fused
+                fusion_note = f"NVIDIA runs {nvidia_kernel_name}, AMD may fuse into another kernel"
         elif amd_count > nvidia_count * 1.5 and nvidia_count > 0:
             # AMD runs more kernels = NVIDIA is fusing some
             fusion_note = f"AMD runs {amd_kernel_name} {amd_count / nvidia_count:.1f}x more → NVIDIA fuses"

wafer_core/lib/trace_compare/analyzer.py

@@ -429,13 +429,22 @@ def analyze_traces_aligned(
             "kernel_pairs": kernel_pairs,
         })
     
-    fusion_result = analyze_fusion_from_alignment(alignment.layer_alignments)
-    same_kernel_result = analyze_same_kernels_from_alignment(alignment.layer_alignments)
-    
+    # Determine which trace is AMD vs NVIDIA for fusion analysis
     if trace1.platform == "AMD":
         amd_trace, nvidia_trace = trace1, trace2
+        fusion_amd_kernels = amd_kernels
+        fusion_nvidia_kernels = nvidia_kernels
     else:
         amd_trace, nvidia_trace = trace2, trace1
+        fusion_amd_kernels = nvidia_kernels
+        fusion_nvidia_kernels = amd_kernels
+
+    fusion_result = analyze_fusion_from_alignment(
+        alignment.layer_alignments,
+        amd_kernels=fusion_amd_kernels,
+        nvidia_kernels=fusion_nvidia_kernels,
+    )
+    same_kernel_result = analyze_same_kernels_from_alignment(alignment.layer_alignments)
     
     return {
         "metadata": {