Agentic PyTorch Inference Optimization & Kernel Benchmark

An agent-driven performance engineering framework designed to search, benchmark, and optimize Hugging Face Causal Language Models (e.g., gpt2, distilgpt2) on NVIDIA GPUs. The framework optimizes runtime configurations across torch.compile modes, numerical precisions, and execution parameters while evaluating compiled outputs against a hand-written Triton CUDA kernel.

---

Architecture Overview

The system operates as a closed-loop performance engineering pipeline. An LLM-driven or heuristic agent proposes inference configurations from a parameter grid, which are then evaluated inside an isolated PyTorch benchmarking harness. Performance and correctness metrics are recorded, enabling the agent to select the optimal runtime configuration under explicit resource constraints.

+------------------------------------------------------------------------------------+
|                                      main.py                                       |
|                            (CLI Driver & Pipeline Flow)                            |
+-----------------------------------------+------------------------------------------+
                                          |
               1. Propose Config          v          2. Benchmark Trial
    +-------------------------------------+--------------------------------------+
    | ExperimentAgent (agent.py)          |  run_benchmark (benchmark.py)        |
    | - LLM Router (AsyncOpenAI / HF)     |  - AutoModelForCausalLM Execution    |
    | - Deterministic Heuristic Fallback  |  - DType Casting (FP32 / FP16 / BF16)|
    | - Trial History & State Tracking    |  - torch.compile Dynamic Wrapping    |
    +-------------------------------------+  - CUDA Event Timing & VRAM Tracking |
                                          |  - Cached FP32 Reference Validation  |
                                          +--------------------------------------+
                                                             |
                                  3. Record Result & Select Best Pick
                                                             v
    +--------------------------------------------------------------------------------+
    | Results & Kernel Evaluation                                                    |
    | - Formatted Results Table & Constraint Filter (select_best)                    |
    | - Standalone Kernel Check: Eager vs. torch.compile vs. Triton (kernels.py)     |
    +--------------------------------------------------------------------------------+

---

Component Breakdown

• Data Schema Layer (config.py): Uses Pydantic models (ExperimentConfig, BenchmarkResult) to enforce strict type validation for candidate parameters (batch_size, seq_len, dtype, compile_mode) and measured execution metrics.
• Agentic Search Engine (agent.py): Manages the trial candidate grid. Uses an LLM agent via Hugging Face's router API (or a deterministic heuristic fallback) to select optimal trial steps based on past performance history and target constraints.
• Benchmarking Harness (benchmark.py): Loads the model, handles precision casting, wraps execution in torch.compile modes, and executes timing runs. It measures timing using torch.cuda.Event, tracks peak VRAM via torch.cuda.max_memory_allocated, and calculates top-1 token logit agreement against an FP32 reference model.
• Custom Triton CUDA Kernel (kernels.py): Contains a hand-written, numerically stable row-wise Softmax kernel written in Triton (triton_softmax). It runs a standalone comparative benchmark between PyTorch eager, torch.compile, and Triton on attention-score tensors of shape (batch, heads, seq, seq).
• Pipeline Driver (main.py): Entry point orchestrating the agent sweep, formatting ASCII metric tables, outputting the optimal configuration recommendation, and triggering the kernel-level softmax evaluation.

---

Key Features

• Automated Hyperparameter Sweep: Explores combinations of batch_size (1–32), seq_len (128–512), dtype (FP16/BF16), and torch.compile modes (eager, default, max-autotune).
• Hand-Written Triton Kernel: Custom JIT-compiled Triton kernel targeting the row-wise softmax operation core to self-attention computation.
• Hardware-Accurate CUDA Profiling: Employs GPU-native CUDA events and VRAM memory stats rather than CPU-side clock approximations.
• Output Correctness Verification: Validates model predictions by comparing top-1 token logit agreement (target ≥95%) against an FP32 eager baseline.

---

Installation & Setup

Requirements

• Python 3.10+
• NVIDIA GPU with working CUDA drivers

Quickstart

# Clone repository and install dependencies
pip install -r requirements.txt

# Run offline heuristic mode (skips LLM API calls)
python main.py --model distilgpt2 --budget 4 --no-llm

# Run optimization with specific memory or latency constraints
python main.py --model gpt2 --budget 8 \
    --constraint "Maximize throughput subject to peak memory under 4000 MB."