Operator Fusion & the Compiler Stack

# ————————————————————————————————————————
# SYMPTOM: Model runs at 15% SM utilization; target 65%+
# GOAL: Classify regime and apply correct compiler fix
# ————————————————————————————————————————

# STEP 1: Identify regime with nvidia-smi
nvidia-smi dmon -s u -d 1  # Watch SM utilization every 1 second
# SM util <30% during inference → overhead-bound (CUDA Graphs fix)
# SM util 30-80% → memory-bandwidth-bound (fusion fix)
# SM util >80% → compute-bound (quantization / TC alignment fix)

# STEP 2: Count graph breaks before applying torch.compile
import torch._dynamo as dynamo
model = LLaMAModel(config).cuda().bfloat16()
explanation = dynamo.explain(model)(dummy_input)
print(f"Graphs: {explanation.graph_count}")
print(f"Graph breaks: {explanation.break_count}")
print(f"Break reasons: {explanation.break_reasons}")
# Target: <5 graph breaks. If >10: fix breaks before compiling
# Common reason: "data-dependent control flow" → restructure conditionals
# Common reason: "Unsupported: .item()" → remove .item() from forward()

# STEP 3: Benchmark properly (common mistake: no warmup)
import torch, time

def bench(model, inputs, n_warmup=10, n_bench=100, label=""):
    for _ in range(n_warmup):
        with torch.no_grad(): _ = model(*inputs)
    torch.cuda.synchronize()
    
    times = []
    for _ in range(n_bench):
        s = torch.cuda.Event(enable_timing=True)
        e = torch.cuda.Event(enable_timing=True)
        s.record()
        with torch.no_grad(): _ = model(*inputs)
        e.record()
        torch.cuda.synchronize()
        times.append(s.elapsed_time(e))
    
    times.sort()
    print(f"{label}: p50={times[n_bench//2]:.2f}ms p99={times[int(n_bench*.99)]:.2f}ms")

model_eager = LLaMAModel(config).cuda().bfloat16()
bench(model_eager, inputs, label="Eager")

model_compiled = torch.compile(model_eager, mode="reduce-overhead")
bench(model_compiled, inputs, label="Compiled")  # First call compiles
# IMPORTANT: Measure p50 AND p99. Compilers can improve mean but worsen p99.
# If compiled p99 > eager p99: check for recompilation on variable shapes

# STEP 4: Inspect generated Triton code
import torch._inductor.config as config
config.trace.enabled = True
config.debug.save_output_code = True
# Run one forward pass → generated kernels saved to /tmp/torchinductor_*/
# Look for: fused_ops_0.py, fused_ops_1.py etc.
# Verify fusion: each file should contain multiple operations

# STEP 5: Profile with torch.profiler to see kernel breakdown
from torch.profiler import profile, ProfilerActivity
with profile(activities=[ProfilerActivity.CUDA], record_shapes=True) as prof:
    with torch.no_grad():
        for _ in range(5):
            _ = model_compiled(*inputs)
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))
# Look for: many small kernels → fusion opportunity missed
# Look for: kernel named "fused_*" → torch.compile fusion working
# Look for: "cudaLaunchKernel" overhead → consider CUDA Graphs

# STEP 6: Verify correctness after compilation
with torch.no_grad():
    out_eager = model_eager(*inputs)
    out_compiled = model_compiled(*inputs)
    max_diff = (out_eager - out_compiled).abs().max().item()
    print(f"Max output difference: {max_diff:.6f}")
    assert max_diff < 1e-3, f"Compilation introduced numeric error: {max_diff}"
# Typical compiled vs eager diff: 1e-5 to 1e-4 (numerical rounding only)
# If diff > 1e-3: check for FP16/BF16 precision issues in fused kernel