Spaces:
Running
Running
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the BSD-style license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| """ | |
| Data models for the Kernel Env Environment. | |
| The agent receives a PyTorch reference `Model` and submits source code for a | |
| `ModelNew` that computes the same thing with custom GPU kernels. | |
| """ | |
| from typing import Any, Dict, List, Literal, Optional | |
| from openenv.core.env_server.types import Action, Observation, State | |
| from pydantic import Field | |
| Stage = Literal["ready", "format", "compile", "runtime", "timeout", "correctness", "timed"] | |
| class KernelAction(Action): | |
| """Full Python source defining `ModelNew`.""" | |
| code: str = Field(..., description="Python source defining class ModelNew") | |
| language: Literal["triton", "cuda"] = Field( | |
| default="triton", description="Kernel language used inside ModelNew" | |
| ) | |
| class KernelObservation(Observation): | |
| """Result of evaluating a submission (or the task prompt on reset).""" | |
| task_id: str = Field(default="", description="KernelBench task id, e.g. level1/19_ReLU") | |
| problem: str = Field(default="", description="Task prompt; only set on reset") | |
| stage: Stage = Field( | |
| default="ready", | |
| description="Furthest stage reached: format < compile < runtime/timeout < correctness < timed", | |
| ) | |
| feedback: str = Field(default="", description="Truncated error / mismatch / timing summary") | |
| correct: Optional[bool] = Field(default=None, description="Outputs matched reference") | |
| max_abs_err: Optional[float] = Field(default=None, description="Worst abs error across trials") | |
| speedup: Optional[float] = Field(default=None, description="t_ref / t_new (eager PyTorch ref)") | |
| t_ref_ms: Optional[float] = Field(default=None, description="Reference runtime, median ms") | |
| t_new_ms: Optional[float] = Field(default=None, description="Candidate runtime, median ms") | |
| t_min_ms: Optional[float] = Field(default=None, description="Roofline floor: max(FLOPs/peak, bytes/bandwidth)") | |
| bound: Optional[Literal["memory", "compute"]] = Field(default=None, description="Which roof sets t_min_ms") | |
| arithmetic_intensity: Optional[float] = Field(default=None, description="Counted FLOPs / compulsory bytes") | |
| roofline_frac: Optional[float] = Field(default=None, description="t_min / t_new: share of hardware peak reached") | |
| ref_roofline_frac: Optional[float] = Field(default=None, description="t_min / t_ref for the PyTorch reference") | |
| turn: int = Field(default=0, description="Submissions made this episode") | |
| max_turns: int = Field(default=4, description="Submissions allowed this episode") | |
| class KernelState(State): | |
| """Episode state.""" | |
| task_id: str = Field(default="") | |
| level: int = Field(default=0) | |
| best_speedup: float = Field(default=0.0, description="Best speedup among correct submissions") | |
| history: List[Dict[str, Any]] = Field(default_factory=list, description="Per-turn stage/speedup") | |