Spaces:
Running
Running
File size: 3,039 Bytes
e313085 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | # 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")
|