wafer-core 0.1.27__tar.gz → 0.1.28__tar.gz

{wafer_core-0.1.27 → wafer_core-0.1.28}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wafer-core
-Version: 0.1.27
+Version: 0.1.28
 Summary: Core utilities and environments for Wafer GPU kernel optimization
 Requires-Python: >=3.10
 Requires-Dist: aiohttp>=3.9.0

{wafer_core-0.1.27 → wafer_core-0.1.28}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "wafer-core"
-version = "0.1.27"
+version = "0.1.28"
 description = "Core utilities and environments for Wafer GPU kernel optimization"
 requires-python = ">=3.10"
 dependencies = [

{wafer_core-0.1.27 → wafer_core-0.1.28}/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-0.1.27 → wafer_core-0.1.28}/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": {

wafer_core-0.1.28/wafer_core/lib/trace_compare/fusion_analyzer.py

@@ -0,0 +1,464 @@
+"""Fusion analysis for detecting kernel fusion differences between platforms.
+
+Detects fusion differences between AMD and NVIDIA by analyzing:
+1. Kernel types unique to each platform (one platform fuses them away)
+2. Sequence patterns (what comes before/after unique kernels)
+3. Count imbalances (one platform has significantly more calls)
+
+This is the pattern-based approach which is more reliable than alignment-based detection.
+"""
+
+from collections import Counter, defaultdict
+from dataclasses import dataclass, field
+from typing import Any
+
+from .classifier import classify, Op
+
+
+@dataclass
+class FusionPattern:
+    """A detected fusion pattern."""
+
+    layer: int
+    operation: str  # The fused operation name (e.g., "RMSNorm+GEMM")
+    fused_platform: str  # Platform that fuses (has fewer kernels)
+    fused_kernel: str  # The actual fused kernel name
+    unfused_kernels: list[str]  # List of separate kernels on the other platform
+    count: int
+    evidence: str
+
+
+@dataclass
+class FusionAnalysis:
+    """Complete fusion analysis result."""
+
+    patterns: list[FusionPattern] = field(default_factory=list)
+    summary: dict[str, Any] = field(default_factory=dict)
+
+
+def _classify_kernel(name: str, platform: str = "AMD") -> str:
+    """Classify a kernel name to its operation type."""
+    op, _pattern = classify(name, platform)
+    return op.value
+
+
+def _find_fusion_mappings(
+    trace1_kernels: list[dict],
+    trace2_kernels: list[dict],
+    trace1_name: str = "Trace1",
+    trace2_name: str = "Trace2",
+    trace1_platform: str = "AMD",
+) -> list[dict]:
+    """Find fusion mappings by analyzing kernel execution sequence patterns.
+
+    This function identifies when one platform runs multiple kernels separately
+    while the other platform fuses them into a single kernel.
+
+    Args:
+        trace1_kernels: List of kernel events from first trace
+        trace2_kernels: List of kernel events from second trace
+        trace1_name: Name of first platform (e.g., "AMD")
+        trace2_name: Name of second platform (e.g., "NVIDIA")
+        trace1_platform: Platform string for classification ("AMD" or "NVIDIA")
+
+    Returns:
+        List of mapping dictionaries with fusion details
+    """
+    mappings = []
+    trace2_platform = "NVIDIA" if trace1_platform == "AMD" else "AMD"
+
+    # Sort kernels by timestamp
+    trace1_sorted = sorted(trace1_kernels, key=lambda k: k.get("ts", 0))
+    trace2_sorted = sorted(trace2_kernels, key=lambda k: k.get("ts", 0))
+
+    # Classify all kernels
+    trace1_types = [_classify_kernel(k.get("name", ""), trace1_platform) for k in trace1_sorted]
+    trace2_types = [_classify_kernel(k.get("name", ""), trace2_platform) for k in trace2_sorted]
+
+    # Find kernel types unique to each trace
+    trace1_type_set = set(trace1_types)
+    trace2_type_set = set(trace2_types)
+
+    trace1_only = trace1_type_set - trace2_type_set
+    trace2_only = trace2_type_set - trace1_type_set
+
+    # For each unique type in trace1, find common sequence patterns
+    for unique_type in trace1_only:
+        # Skip "Other" since it's too generic
+        if unique_type == "Other":
+            continue
+
+        # Find all occurrences of this type
+        indices = [i for i, t in enumerate(trace1_types) if t == unique_type]
+
+        if len(indices) < 5:  # Need enough samples to be meaningful
+            continue
+
+        # Analyze what comes before each occurrence
+        before_types: dict[str, int] = defaultdict(int)
+
+        for idx in indices:
+            if idx > 0:
+                before_types[trace1_types[idx - 1]] += 1
+
+        # Find the most common pattern
+        if not before_types:
+            continue
+        most_common_before = max(before_types.items(), key=lambda x: x[1])
+
+        # If there's a strong pattern (>80% of occurrences)
+        if most_common_before[1] / len(indices) > 0.8:
+            fusion_candidate = most_common_before[0]
+
+            # Verify trace2 has this type
+            if fusion_candidate in trace2_type_set:
+                trace1_fusion_count = trace1_types.count(fusion_candidate)
+                trace2_fusion_count = trace2_types.count(fusion_candidate)
+
+                mappings.append({
+                    "fused_platform": trace2_name,
+                    "fused_kernel_type": fusion_candidate,
+                    "fused_count": trace2_fusion_count,
+                    "unfused_platform": trace1_name,
+                    "unfused_sequence": [fusion_candidate, unique_type],
+                    "unfused_count_per_type": {
+                        fusion_candidate: trace1_fusion_count,
+                        unique_type: len(indices),
+                    },
+                    "pattern_count": len(indices),
+                    "pattern_confidence": most_common_before[1] / len(indices),
+                    "evidence": f"{trace1_name} runs {fusion_candidate}+{unique_type} separately, {trace2_name} fuses into {fusion_candidate}",
+                })
+
+    # Also check trace2-only types
+    for unique_type in trace2_only:
+        if unique_type == "Other":
+            continue
+
+        indices = [i for i, t in enumerate(trace2_types) if t == unique_type]
+
+        if len(indices) < 5:
+            continue
+
+        before_types = defaultdict(int)
+
+        for idx in indices:
+            if idx > 0:
+                before_types[trace2_types[idx - 1]] += 1
+
+        if not before_types:
+            continue
+        most_common_before = max(before_types.items(), key=lambda x: x[1])
+
+        if most_common_before[1] / len(indices) > 0.8:
+            fusion_candidate = most_common_before[0]
+
+            if fusion_candidate in trace1_type_set:
+                trace1_fusion_count = trace1_types.count(fusion_candidate)
+                trace2_fusion_count = trace2_types.count(fusion_candidate)
+
+                mappings.append({
+                    "fused_platform": trace1_name,
+                    "fused_kernel_type": fusion_candidate,
+                    "fused_count": trace1_fusion_count,
+                    "unfused_platform": trace2_name,
+                    "unfused_sequence": [fusion_candidate, unique_type],
+                    "unfused_count_per_type": {
+                        fusion_candidate: trace2_fusion_count,
+                        unique_type: len(indices),
+                    },
+                    "pattern_count": len(indices),
+                    "pattern_confidence": most_common_before[1] / len(indices),
+                    "evidence": f"{trace2_name} runs {fusion_candidate}+{unique_type} separately, {trace1_name} fuses into {fusion_candidate}",
+                })
+
+    return mappings
+
+
+def _find_count_imbalance_fusions(
+    trace1_kernels: list[dict],
+    trace2_kernels: list[dict],
+    trace1_name: str = "Trace1",
+    trace2_name: str = "Trace2",
+    trace1_platform: str = "AMD",
+) -> list[dict]:
+    """Find fusions by looking for significant count imbalances.
+
+    When one platform has significantly more kernel calls of a type (>1.5x),
+    it suggests the other platform fuses those operations.
+    """
+    mappings = []
+    trace2_platform = "NVIDIA" if trace1_platform == "AMD" else "AMD"
+
+    # Classify all kernels
+    trace1_types = [_classify_kernel(k.get("name", ""), trace1_platform) for k in trace1_kernels]
+    trace2_types = [_classify_kernel(k.get("name", ""), trace2_platform) for k in trace2_kernels]
+
+    # Count by type
+    trace1_counts = Counter(trace1_types)
+    trace2_counts = Counter(trace2_types)
+
+    # Find common types with significant differences
+    common_types = set(trace1_counts.keys()) & set(trace2_counts.keys())
+
+    for ktype in common_types:
+        if ktype == "Other":
+            continue
+
+        trace1_count = trace1_counts[ktype]
+        trace2_count = trace2_counts[ktype]
+
+        # Skip trivial counts
+        if trace1_count + trace2_count < 50:
+            continue
+
+        # Check if there's a significant imbalance (>1.5x)
+        if trace1_count == 0 or trace2_count == 0:
+            continue
+
+        ratio = max(trace1_count, trace2_count) / min(trace1_count, trace2_count)
+
+        if ratio < 1.5:
+            continue
+
+        # Determine which platform has more (unfused) and which has fewer (fused)
+        if trace1_count > trace2_count:
+            unfused_platform = trace1_name
+            fused_platform = trace2_name
+            unfused_count = trace1_count
+            fused_count = trace2_count
+        else:
+            unfused_platform = trace2_name
+            fused_platform = trace1_name
+            unfused_count = trace2_count
+            fused_count = trace1_count
+
+        mappings.append({
+            "fused_platform": fused_platform,
+            "fused_kernel_type": ktype,
+            "fused_count": fused_count,
+            "unfused_platform": unfused_platform,
+            "unfused_sequence": [ktype],
+            "unfused_count_per_type": {ktype: unfused_count},
+            "pattern_count": unfused_count - fused_count,
+            "pattern_confidence": (unfused_count - fused_count) / unfused_count,
+            "evidence": f"{unfused_platform} calls {ktype} {ratio:.1f}x more ({unfused_count} vs {fused_count}), {fused_platform} likely fuses",
+        })
+
+    return mappings
+
+
+def _find_explicit_fused_operations(
+    trace1_kernels: list[dict],
+    trace2_kernels: list[dict],
+    trace1_name: str = "Trace1",
+    trace2_name: str = "Trace2",
+    trace1_platform: str = "AMD",
+) -> list[dict]:
+    """Find explicit fused operations (kernels with '+' in their classification).
+
+    These are operations like 'RMSNorm+GEMM' that are explicitly classified as fused.
+    """
+    mappings = []
+    trace2_platform = "NVIDIA" if trace1_platform == "AMD" else "AMD"
+
+    def get_fused_ops(kernels: list[dict], platform: str) -> dict[str, list[str]]:
+        """Get fused operations and their kernel names."""
+        fused_ops: dict[str, list[str]] = defaultdict(list)
+        for k in kernels:
+            name = k.get("name", "")
+            op, _pattern = classify(name, platform)
+            if "+" in op.value or op == Op.FUSED_UNKNOWN:
+                fused_ops[op.value].append(name)
+        return dict(fused_ops)
+
+    trace1_fused = get_fused_ops(trace1_kernels, trace1_platform)
+    trace2_fused = get_fused_ops(trace2_kernels, trace2_platform)
+
+    # Find fused operations unique to trace1
+    for fused_op, kernels in trace1_fused.items():
+        if fused_op not in trace2_fused:
+            # Parse components from the fused op name
+            if "+" in fused_op:
+                components = [c.strip() for c in fused_op.split("+")]
+            else:
+                components = [fused_op]
+
+            mappings.append({
+                "fused_platform": trace1_name,
+                "fused_kernel_type": fused_op,
+                "fused_count": len(kernels),
+                "unfused_platform": trace2_name,
+                "unfused_sequence": components,
+                "unfused_count_per_type": {c: 0 for c in components},  # Unknown
+                "pattern_count": len(kernels),
+                "pattern_confidence": 1.0,
+                "evidence": f"{trace1_name} fuses {' + '.join(components)} into {fused_op} ({len(kernels)} calls)",
+                "fused_kernel_names": kernels[:3],  # Sample of kernel names
+            })
+
+    # Find fused operations unique to trace2
+    for fused_op, kernels in trace2_fused.items():
+        if fused_op not in trace1_fused:
+            if "+" in fused_op:
+                components = [c.strip() for c in fused_op.split("+")]
+            else:
+                components = [fused_op]
+
+            mappings.append({
+                "fused_platform": trace2_name,
+                "fused_kernel_type": fused_op,
+                "fused_count": len(kernels),
+                "unfused_platform": trace1_name,
+                "unfused_sequence": components,
+                "unfused_count_per_type": {c: 0 for c in components},
+                "pattern_count": len(kernels),
+                "pattern_confidence": 1.0,
+                "evidence": f"{trace2_name} fuses {' + '.join(components)} into {fused_op} ({len(kernels)} calls)",
+                "fused_kernel_names": kernels[:3],
+            })
+
+    return mappings
+
+
+def detect_fusion_patterns(
+    amd_kernels: list[dict],
+    nvidia_kernels: list[dict],
+) -> FusionAnalysis:
+    """Detect fusion patterns using pattern-based analysis.
+
+    This is the main entry point for fusion detection. It combines:
+    1. Explicit fused operations (kernels classified with '+' in name)
+    2. Sequence pattern analysis (unique kernel types with consistent patterns)
+    3. Count imbalance analysis (one platform has significantly more calls)
+
+    Args:
+        amd_kernels: List of AMD kernel events
+        nvidia_kernels: List of NVIDIA kernel events
+
+    Returns:
+        FusionAnalysis with detected patterns
+    """
+    all_mappings: list[dict] = []
+
+    # 1. Find explicit fused operations (highest confidence)
+    explicit_fusions = _find_explicit_fused_operations(
+        amd_kernels, nvidia_kernels,
+        trace1_name="AMD", trace2_name="NVIDIA",
+        trace1_platform="AMD",
+    )
+    all_mappings.extend(explicit_fusions)
+
+    # 2. Find sequence-based fusions
+    sequence_fusions = _find_fusion_mappings(
+        amd_kernels, nvidia_kernels,
+        trace1_name="AMD", trace2_name="NVIDIA",
+        trace1_platform="AMD",
+    )
+    # Deduplicate: skip if same fused_kernel_type already found
+    existing_types = {m["fused_kernel_type"] for m in all_mappings}
+    for fusion in sequence_fusions:
+        if fusion["fused_kernel_type"] not in existing_types:
+            all_mappings.append(fusion)
+            existing_types.add(fusion["fused_kernel_type"])
+
+    # 3. Find count-imbalance fusions
+    count_fusions = _find_count_imbalance_fusions(
+        amd_kernels, nvidia_kernels,
+        trace1_name="AMD", trace2_name="NVIDIA",
+        trace1_platform="AMD",
+    )
+    # Deduplicate
+    for fusion in count_fusions:
+        if fusion["fused_kernel_type"] not in existing_types:
+            all_mappings.append(fusion)
+            existing_types.add(fusion["fused_kernel_type"])
+
+    # Convert to FusionPattern objects
+    patterns: list[FusionPattern] = []
+    for mapping in all_mappings:
+        fused_kernel_names = mapping.get("fused_kernel_names", [])
+        fused_kernel = fused_kernel_names[0] if fused_kernel_names else mapping["fused_kernel_type"]
+        
+        patterns.append(
+            FusionPattern(
+                layer=0,  # Pattern-based analysis doesn't track layers
+                operation=mapping["fused_kernel_type"],
+                fused_platform=mapping["fused_platform"],
+                fused_kernel=fused_kernel,
+                unfused_kernels=mapping["unfused_sequence"],
+                count=mapping["pattern_count"],
+                evidence=mapping["evidence"],
+            )
+        )
+
+    amd_fuses = sum(1 for p in patterns if p.fused_platform == "AMD")
+    nvidia_fuses = sum(1 for p in patterns if p.fused_platform == "NVIDIA")
+
+    return FusionAnalysis(
+        patterns=patterns,
+        summary={
+            "amd_fuses": amd_fuses,
+            "nvidia_fuses": nvidia_fuses,
+            "total_fusion_opportunities": len(patterns),
+        },
+    )
+
+
+def analyze_fusion_from_alignment(
+    layer_alignments: list[Any],
+    amd_kernels: list[dict] | None = None,
+    nvidia_kernels: list[dict] | None = None,
+) -> dict[str, Any]:
+    """Analyze fusion from kernel data (for API compatibility).
+
+    Args:
+        layer_alignments: List of aligned layers (unused - kept for API compatibility)
+        amd_kernels: Optional list of AMD kernel events for pattern-based analysis
+        nvidia_kernels: Optional list of NVIDIA kernel events for pattern-based analysis
+
+    Returns:
+        Dictionary with fusion analysis results
+    """
+    # If raw kernels provided, use pattern-based analysis (preferred)
+    if amd_kernels is not None and nvidia_kernels is not None:
+        fusion_analysis = detect_fusion_patterns(amd_kernels, nvidia_kernels)
+    else:
+        # Fallback: empty analysis if no kernel data
+        fusion_analysis = FusionAnalysis(
+            patterns=[],
+            summary={"amd_fuses": 0, "nvidia_fuses": 0, "total_fusion_opportunities": 0},
+        )
+
+    fusion_opportunities = []
+    fusion_mappings = []
+
+    for pattern in fusion_analysis.patterns:
+        unfused_platform = "NVIDIA" if pattern.fused_platform == "AMD" else "AMD"
+
+        fusion_opportunities.append({
+            "kernel_type": pattern.operation,
+            "layer": pattern.layer,
+            "fused_by": pattern.fused_platform,
+            "fused_kernel": pattern.fused_kernel,
+            "unfused_kernels": pattern.unfused_kernels,
+            "count": pattern.count,
+            "evidence": pattern.evidence,
+        })
+
+        fusion_mappings.append({
+            "fused_platform": pattern.fused_platform,
+            "fused_kernel_type": pattern.operation,
+            "fused_kernel_name": pattern.fused_kernel,
+            "unfused_platform": unfused_platform,
+            "unfused_sequence": pattern.unfused_kernels,
+            "pattern_count": pattern.count,
+            "evidence": pattern.evidence,
+            "layer": pattern.layer,
+        })
+
+    return {
+        "fusion_opportunities": fusion_opportunities,
+        "fusion_mappings": fusion_mappings,
+        "summary": fusion_analysis.summary,
+    }

{wafer_core-0.1.27 → wafer_core-0.1.28}/wafer_core/targets/__init__.py

@@ -1,5 +1,43 @@
-"""Re-export targets from utils for convenience."""
+"""Target system: specs (config) + targets (live resources) + reconciliation.
 
+New API (preferred):
+    from wafer_core.targets import Target, TargetSpec, ReconcileResult
+    from wafer_core.targets.providers import get_provider
+    from wafer_core.targets.reconcile import reconcile
+    from wafer_core.targets.spec_store import load_spec, list_spec_names
+    from wafer_core.targets.state_cache import get_binding_hints
+
+Legacy API (still works, will be deprecated):
+    from wafer_core.targets import RunPodTarget, runpod_ssh_context, ...
+"""
+
+# ── New types ────────────────────────────────────────────────────────────────
+from wafer_core.targets.digitalocean import (
+    DigitalOceanError,
+    DigitalOceanSSHInfo,
+    cleanup_all_droplets,
+    cleanup_digitalocean_target,
+    digitalocean_ssh_context,
+    get_droplet_state,
+    list_running_droplets,
+)
+from wafer_core.targets.runpod import (
+    RunPodError,
+    RunPodSSHInfo,
+    cleanup_all_pods,
+    cleanup_target,
+    get_pod_state,
+    list_running_pods,
+    runpod_ssh_context,
+)
+from wafer_core.targets.types import (
+    ReconcileResult,
+    Target,
+    TargetProvider,
+    TargetSpec,
+)
+
+# ── Legacy re-exports (unchanged, for backwards compatibility) ───────────────
 from wafer_core.utils.kernel_utils.targets import (
     BaremetalTarget,
     DigitalOceanTarget,
@@ -18,26 +56,14 @@ from wafer_core.utils.kernel_utils.targets import (
     select_target_for_operation,
     target_to_deployment_config,
 )
-from wafer_core.targets.runpod import (
-    RunPodError,
-    RunPodSSHInfo,
-    cleanup_all_pods,
-    cleanup_target,
-    get_pod_state,
-    list_running_pods,
-    runpod_ssh_context,
-)
-from wafer_core.targets.digitalocean import (
-    DigitalOceanError,
-    DigitalOceanSSHInfo,
-    cleanup_all_droplets,
-    cleanup_digitalocean_target,
-    digitalocean_ssh_context,
-    get_droplet_state,
-    list_running_droplets,
-)
 
 __all__ = [
+    # New API
+    "Target",
+    "TargetSpec",
+    "TargetProvider",
+    "ReconcileResult",
+    # Legacy: target config types
     "BaremetalTarget",
     "VMTarget",
     "ModalTarget",
@@ -54,7 +80,7 @@ __all__ = [
     "check_target_available",
     "find_free_gpu",
     "run_operation_on_target",
-    # RunPod provisioning
+    # Legacy: RunPod provisioning
     "RunPodError",
     "RunPodSSHInfo",
     "runpod_ssh_context",
@@ -62,7 +88,7 @@ __all__ = [
     "cleanup_all_pods",
     "list_running_pods",
     "get_pod_state",
-    # DigitalOcean provisioning
+    # Legacy: DigitalOcean provisioning
     "DigitalOceanError",
     "DigitalOceanSSHInfo",
     "digitalocean_ssh_context",