System Design 02 — NPU Graph Optimization Pipeline

=====================================================

Full design document for an automated pipeline that takes a PyTorch FX graph

and produces an optimized, deployed artifact for Qualcomm NPU (HTP).

Problem Statement

Build an automated graph optimization and deployment pipeline that:

• Accepts any PyTorch model (FP32 or FP16)
• Applies graph-level optimizations (op fusion, layout transforms, quantization)
• Compiles to Qualcomm HTP (Hexagon Tensor Processor) binary
• Validates accuracy against FP32 reference
• Emits a production-ready .dlc or .qnncache artifact with benchmark report

Scale: 50+ model architectures, 3 hardware targets (Snapdragon 8 Gen 3, 888, 7 Gen 3), weekly pipeline runs.

Requirements

Functional

• Automated end-to-end: PyTorch → ONNX → QNN/SNPE → Validated artifact
• Per-layer accuracy analysis (cosine similarity, SNR)
• Mixed-precision support (per-layer FP16/INT8/INT4 assignment)
• AI Hub integration for cloud on-device profiling
• Artifact registry with model metadata, benchmark results, and reproducibility info

Non-Functional

• Pipeline SLA: < 4 hours per model
• Accuracy gate: CosSim ≥ 0.99 per layer, OR FP16 fallback for failing layers
• Latency gate: measured HTP latency within 10% of theoretical estimate
• Reproducibility: given same inputs + hardware, same output

Architecture Overview

┌──────────────────────────────────────────────────────────────────────────────┐
│                        NPU Graph Optimization Pipeline                       │
└──────────────────────────────────────────────────────────────────────────────┘

[1. Ingestion]
  PyTorch Model + Calibration Dataset (500 samples) + Config YAML
      │
      ▼
[2. FX Graph Analysis]
  - torch.export() → ExportedProgram
  - Op coverage check: which ops HTP supports natively
  - Identify unsupported ops → fallback candidates
  - Build layer dependency graph
      │
      ▼
[3. Graph Optimization Passes]
  Pass 1: Op Fusion (Conv+BN+ReLU, Linear+GeLU, QKV projection merge)
  Pass 2: Layout Transform (NCHW → NHWC for HTP efficiency)
  Pass 3: Dead Code Elimination
  Pass 4: Constant Folding (pre-compute static ops)
      │
      ▼
[4. Quantization]
  Mode A (PTQ): Collect activation histograms → calibrate scales → INT8/INT4
  Mode B (Mixed Precision): Run sensitivity analysis → assign precision per layer
  Mode C (QAT): Load pre-trained QAT checkpoint → extract fake-quant scales
      │
      ▼
[5. ONNX Export & Verification]
  torch.onnx.export(model, example_inputs, "model.onnx", opset=17)
  onnxruntime validation: max output diff < 1e-4 vs PyTorch
  onnx-simplifier + onnxoptimizer
      │
      ▼
[6. Compilation]
  Target: SNPE → snpe-onnx-to-dlc → model.dlc
  Target: QNN  → qnn-onnx-converter → qnn-context-binary-generator → model.bin
      │
      ▼
[7. On-Device Validation (AI Hub)]
  qai_hub.submit_compile_job(source_model=model.onnx, device="Snapdragon 8 Gen 3")
  qai_hub.submit_profile_job(model=compiled_model, device=..., input_specs=...)
  Metrics: latency (median, p95), memory peak, per-layer time
      │
      ▼
[8. Accuracy Gate]
  For each layer: compute cosine_similarity(fp32_out, int8_out)
  If CosSim < 0.99: mark layer as "sensitive" → re-quantize at FP16
  If overall task accuracy drop > threshold: FAIL pipeline, alert
      │
      ▼
[9. Artifact Registration]
  Store: model.dlc + metadata.json + benchmark_results.json + accuracy_report.json
  Metadata: model_sha, calibration_sha, hardware_target, precision_map, pipeline_version

Key Components

2. FX Graph Analysis

from torch.export import export

def analyze_graph(model, example_inputs):
    exported = export(model, example_inputs)
    fx_graph = exported.graph_module
    
    # Check op support
    unsupported = []
    for node in fx_graph.graph.nodes:
        if node.op == "call_function":
            if node.target not in HTP_SUPPORTED_OPS:
                unsupported.append((node.name, node.target))
    
    return {
        "n_ops": len(list(fx_graph.graph.nodes)),
        "unsupported_ops": unsupported,
        "coverage": 1 - len(unsupported) / max(1, n_ops),
    }

HTP unsupported ops (common): torch.nn.functional.scaled_dot_product_attention (before SDPA optimization), custom Python autograd functions, dynamic shape operators without static shapes.

4. Sensitivity Analysis for Mixed Precision

def compute_per_layer_sensitivity(model, val_data, baseline_acc, bits_map):
    """
    Returns: {layer_name: accuracy_drop_when_quantized}
    """
    sensitivity = {}
    for name, module in model.named_modules():
        if not isinstance(module, nn.Linear):
            continue
        
        # Quantize only this layer
        model_copy = copy.deepcopy(model)
        layer = get_nested_attr(model_copy, name)
        quantize_layer(layer, bits=bits_map.get(name, 8))
        
        acc = evaluate(model_copy, val_data)
        sensitivity[name] = baseline_acc - acc
    
    return sensitivity

Mixed precision assignment rule:

• Sensitivity > 1.0%: FP16
• Sensitivity 0.3-1.0%: INT8
• Sensitivity < 0.3%: INT4

This typically results in: 5-10% of layers at FP16, 20% at INT8, 70% at INT4.

6. ONNX Export Best Practices

Common export failures and fixes:

7. AI Hub Integration

import qai_hub as hub

def deploy_and_profile(onnx_path, device_name, input_specs):
    """Submit compile + profile jobs to Qualcomm AI Hub."""
    # Compile to on-device format
    compile_job = hub.submit_compile_job(
        model=hub.upload_model(onnx_path),
        device=hub.Device(device_name),
        input_specs=input_specs,
        options="--target_runtime qnn_lib_aarch64_android"
    )
    compiled_model = compile_job.get_target_model()
    
    # Profile latency
    profile_job = hub.submit_profile_job(
        model=compiled_model,
        device=hub.Device(device_name),
        input_specs=input_specs,
    )
    profile = profile_job.download_profile()
    
    return {
        "latency_ms": profile["execution_summary"]["estimated_inference_time"] / 1000,
        "peak_memory_mb": profile["execution_summary"]["inference_memory_peak_range"][0] / (1024**2),
        "n_ops": profile["execution_summary"]["total_ops"],
        "htp_op_percentage": profile["execution_summary"].get("htp_ops_percentage", 0),
    }

8. Accuracy Validation

Layer-by-layer cosine similarity:

def validate_layer_accuracy(model_fp32, model_int8, calibration_inputs):
    """
    Attach hooks to both models, compare intermediate activations.
    Flag layers with CosSim < 0.99 as problematic.
    """
    activations_fp32 = {}
    activations_int8 = {}
    
    def make_hook(name, store):
        def hook(module, input, output):
            store[name] = output.detach()
        return hook
    
    # Register hooks
    hooks = []
    for name, module in model_fp32.named_modules():
        if isinstance(module, (nn.Linear, nn.Conv2d)):
            hooks.append(module.register_forward_hook(make_hook(name, activations_fp32)))
    
    # Run both models
    with torch.no_grad():
        model_fp32(calibration_inputs)
        model_int8(calibration_inputs)
    
    # Compute CosSim
    results = {}
    for name in activations_fp32:
        if name in activations_int8:
            cos_sim = F.cosine_similarity(
                activations_fp32[name].flatten(),
                activations_int8[name].flatten(),
                dim=0
            ).item()
            results[name] = {
                "cos_sim": cos_sim,
                "status": "PASS" if cos_sim > 0.99 else "FAIL"
            }
    
    # Cleanup
    for h in hooks:
        h.remove()
    
    return results

Failure Modes & Mitigations

Performance Characteristics

Artifact Schema

{
  "model_id": "llama-7b-instruct-v3",
  "model_sha": "a3f9b2...",
  "pipeline_version": "2.4.1",
  "hardware_target": "Snapdragon 8 Gen 3",
  "quantization": {
    "strategy": "mixed_precision",
    "default_bits": 4,
    "fp16_layers": ["model.layers.0.self_attn.q_proj", "..."],
    "int8_layers": ["model.layers.1.self_attn.q_proj", "..."],
    "calibration_samples": 500,
    "calibration_sha": "f3e2a1..."
  },
  "accuracy": {
    "task": "MMLU-abstract_algebra",
    "fp32_baseline": 0.723,
    "quantized": 0.716,
    "delta": -0.007,
    "gate": "PASS"
  },
  "latency": {
    "median_ms": 9.4,
    "p95_ms": 11.2,
    "peak_memory_mb": 285,
    "htp_op_coverage": 0.94
  },
  "artifacts": {
    "dlc": "s3://models/llama-7b-v3-snap8gen3-int4.dlc",
    "profile_json": "s3://models/llama-7b-v3-profile.json"
  }
}