config.py

"""
Typed configuration and result models for the benchmark.

Kept as pydantic models (same pattern as the original RAG project) so the
agent, benchmark harness, and reporting code all share one validated schema
instead of passing raw dicts around.
"""
from __future__ import annotations

from enum import Enum
from typing import Optional

import torch
from pydantic import BaseModel, Field


class CompileMode(str, Enum):
    EAGER = "eager"
    DEFAULT = "default"
    REDUCE_OVERHEAD = "reduce-overhead"
    MAX_AUTOTUNE = "max-autotune"


class DType(str, Enum):
    FP32 = "fp32"
    FP16 = "fp16"
    BF16 = "bf16"

    def to_torch(self) -> torch.dtype:
        return {
            DType.FP32: torch.float32,
            DType.FP16: torch.float16,
            DType.BF16: torch.bfloat16,
        }[self]


class ExperimentConfig(BaseModel):
    batch_size: int = Field(gt=0, le=256)
    seq_len: int = Field(gt=0, le=8192)
    dtype: DType
    compile_mode: CompileMode
    use_triton_softmax: bool = False

    def key(self) -> str:
        return (
            f"bs{self.batch_size}_sl{self.seq_len}_{self.dtype.value}"
            f"_{self.compile_mode.value}_triton{int(self.use_triton_softmax)}"
        )


class BenchmarkResult(BaseModel):
    config: ExperimentConfig
    mean_latency_ms: float
    std_latency_ms: float
    throughput_tokens_per_s: float
    peak_memory_mb: float
    compile_overhead_s: Optional[float] = None
    max_abs_error: Optional[float] = None
    top1_agreement: Optional[float] = None
    correctness_passed: Optional[bool] = None
    error: Optional[str] = None

    @property
    def ok(self) -> bool:
        return self.error is None