Cant-Be-Late Multi-Region Scheduling Problem
================================

Problem Setting
---------------

You are given a long-running compute job that must complete before a fixed hard deadline.
At each time step, you must choose which AWS region to run in and which type of cloud compute resource to use:

- **Spot instances**
  - Very cheap
  - May become unavailable at certain timesteps
  - Can be preempted at any time, the job will incur a **restart overhead**

- **On-demand instances**
  - Guaranteed available
  - Expensive
  - Never interrupted

- **Multi-region execution**
  - You may switch to another AWS region at any timestep
  - Switching regions forces a restart overhead (same as losing the work of the current timestep)
  - Spot availability differs per region based on real traces

Your strategy must decide at every timestep whether to use Spot, use On-Demand, or pause (NONE).

Your strategy can also switch to a different region at each step.

Restart overheads do not stack: launching a new instance while an old overhead is still pending will replace the previous remaining restart overhead with the new one.

Your goal is to **finish before the deadline** while **minimizing cost**.

The evaluation uses many real spot-availability traces.
---

API Specification
-----------------

Implement a `Solution` class that extends `MultiRegionStrategy`:

```python
import json
from argparse import Namespace

from sky_spot.strategies.multi_strategy import MultiRegionStrategy
from sky_spot.utils import ClusterType

class Solution(MultiRegionStrategy):
    """Your multi-region scheduling strategy."""

    NAME = "my_strategy"  # REQUIRED: unique identifier

    def solve(self, spec_path: str) -> "Solution":
        """
        Initialize the solution from spec_path config.

        The spec file contains:
        - deadline: deadline in hours
        - duration: task duration in hours
        - overhead: restart overhead in hours
        - trace_files: list of trace file paths (one per region)
        """
        with open(spec_path) as f:
            config = json.load(f)

        args = Namespace(
            deadline_hours=float(config["deadline"]),
            task_duration_hours=[float(config["duration"])],
            restart_overhead_hours=[float(config["overhead"])],
            inter_task_overhead=[0.0],
        )
        super().__init__(args)
        return self

    def _step(self, last_cluster_type: ClusterType, has_spot: bool) -> ClusterType:
        """
        Decide next action based on current state.

        Available attributes:
        - self.env.get_current_region(): Get current region index
        - self.env.get_num_regions(): Get total number of regions
        - self.env.switch_region(idx): Switch to region by index
        - self.env.elapsed_seconds: Current time elapsed
        - self.task_duration: Total task duration needed (seconds)
        - self.deadline: Deadline time (seconds)
        - self.restart_overhead: Restart overhead (seconds)
        - self.task_done_time: List of completed work segments
        - self.remaining_restart_overhead: Current pending overhead

        Returns: ClusterType.SPOT, ClusterType.ON_DEMAND, or ClusterType.NONE
        """
        # Your decision logic here
        if has_spot:
            return ClusterType.SPOT
        return ClusterType.ON_DEMAND
```

Parameters:
---------------
### ClusterType:
ClusterType has 3 members:

ClusterType.SPOT: Spot type cluster.

ClusterType.ON_DEMAND: On Demand type cluster.

ClusterType.None: None, no cluster.

#### You are given some fixed parameters:

env.gap_seconds: The size of each time step, in seconds.

task_duration: The total amount of work time required to finish the task (in seconds).

deadline: The task’s deadline (in seconds).

restart_overhead: The time overhead incurred when a job restarts.

You should implement the function to return the next cluster type to use as described above.

####  At each time step, you are given:

env.elapsed_seconds: Current time elapsed (in second).

env.cluster_type: The current cluster type running your task.

task_done_time: A list of completed work segments, where sum(self.task_done_time) = the amount of successful work time accumulated so far.

has_spot: A boolean indicating whether the Spot cluster is available in the current time step. If False, the strategy must not return ClusterType.SPOT (doing so will raise an error).

### You can use:

env.get_current_region(): Get your current region index (0-8).

env.switch_region(idx): Switch to region by index (no cost).

#### You should return:

ClusterType.SPOT: if you want to run the next time step on the Spot cluster.

ClusterType.ON_DEMAND: if you want to run the next time step on the On-Demand cluster.

ClusterType.NONE: if you choose not to run on any cluster during the next time step; this incurs no cost.

Scoring
-------
```
combined_score = -average_cost_across_all_scenarios
```

Negative cost: Lower cost = higher (less negative) score.

Notice that if you fail to finish the task before the deadline, you will receive a penalty score of -100000.

Evaluation Details
------------------
**Stage 1**: Quick check on 2-region scenario (must pass to proceed)
**Stage 2**: Full evaluation on 4 scenarios:
- 2 zones west (8 traces)
- 3 zones west (6 traces)
- 2 regions west-east2 (8 traces)
- 5 regions mixed (4 traces)

- Task duration: 24 hours
- Deadline: 36 hours (12-hour slack)
- Restart overhead: 0.05 hours (3 minutes)
- Price of on-demand is 3.06$/hr
- Price of Spot is 0.9701$/hr

- Notice your solution will be tested on real traces with low Spot availability.

Your program has a total time limit of 300 seconds. You may be evaluated for up to 36 × 60 × 60 = 129600 time steps. Please ensure that your code is efficient under python.

Implementation Notes
---------------------
**Required Elements (Missing these will cause evaluation failures):**
- `NAME` attribute must be defined on your Solution class
- `solve(self, spec_path)` method must initialize the strategy and return `self`
- `_step(self, last_cluster_type, has_spot)` method must return a ClusterType
- Ensure proper handling of ClusterType.NONE return values

Concrete Step Example:
----------------------
Here is a concrete example demonstrating our environment.
Assume we are:
```
Parameter                | Value
-------------------------|------------------------
env.gap_seconds          | 3600.0
env.elapsed_seconds      | 18000
task_done_time           | [3600, 3600, 2880, 3600, 3600]
has_spot                 | True
env.cluster_type         | ClusterType.SPOT
env.get_current_region() | 0
```
If we use env.switch_region(1), we will have:
```
Parameter                | Value
-------------------------|------------------------
env.gap_seconds          | 3600.0
env.elapsed_seconds      | 18000
task_done_time           | [3600, 3600, 2880, 3600, 3600]
has_spot                 | True
env.cluster_type         | ClusterType.SPOT
env.get_current_region() | 1
```
If our strategy returns ClusterType.ON_DEMAND, there will be a restart overhead:
```
Parameter                | Value
-------------------------|------------------------
env.gap_seconds          | 3600.0
env.elapsed_seconds      | 23400
task_done_time           | [3600, 3600, 2880, 3600, 3600, 2880]
has_spot                 | True
env.cluster_type         | ClusterType.ON_DEMAND
env.get_current_region() | 1
```
If our strategy returns ClusterType.SPOT:
```
Parameter                | Value
-------------------------|------------------------
env.gap_seconds          | 3600.0
env.elapsed_seconds      | 21600
task_done_time           | [3600, 3600, 2880, 3600, 3600, 3600]
has_spot                 | True
env.cluster_type         | ClusterType.SPOT
env.get_current_region() | 1
```


Cloudcast Broadcast Optimization Problem
========================================

Problem Setting
---------------
Design broadcast topology optimization for multi-cloud data distribution. Given a source node and multiple destination nodes across AWS, Azure, and GCP, find the optimal broadcast paths that minimize transfer cost while respecting bandwidth constraints.

The data is split into multiple partitions that can be transmitted independently. Different partitions can take different routes to the same destination, allowing for parallel transmission and load balancing across the network.

**Optimization Goal**: Minimize total cost ($)

$$
\text{Total Cost} = C_{\text{egress}} + C_{\text{instance}}
$$

Where:

$$
C_{\text{egress}} = \sum_{e \in E} \left( |P_e| \times s_{\text{partition}} \times c_e \right)
$$

$$
C_{\text{instance}} = |V| \times n_{\text{vm}} \times \frac{r_{\text{instance}}}{3600} \times t_{\text{transfer}}
$$

**Notation**:
- $E$: Set of all edges used in the broadcast topology (union of all partition paths)
- $P_e$: Set of partitions using edge $e$ (automatically computed by evaluator)
- $s_{\text{partition}}$: Size of each partition (GB) = $\frac{\text{data\_vol}}{\text{num\_partitions}}$
- $c_e$: Cost per GB for edge $e$ (\$/GB)
- $V$: Set of all nodes appearing in any partition path (automatically computed by evaluator)
- $n_{\text{vm}}$: Number of VMs per region (default: 2)
- $r_{\text{instance}}$: Instance hourly rate (\$/hour) = \$0.54
- $t_{\text{transfer}}$: Total transfer time (seconds) = $\max_{d \in D} \max_{p \in [0, n_p)} \max_{e \in \text{path}(d,p)} \frac{|P_e| \times s_{\text{partition}} \times 8}{f_e}$
  - $D$: Set of destination nodes
  - $n_p$: Number of partitions
  - $f_e$: Actual throughput (flow) on edge $e$ after bandwidth constraint enforcement (Gbps)

API Specification
-----------------
Implement a `Solution` class that returns a search algorithm:

```python
class Solution:
    def solve(self, spec_path: str = None) -> dict:
        """
        Returns a dict with either:
        - {"code": "python_code_string"}
        - {"program_path": "path/to/algorithm.py"}
        """
        # Your implementation
        pass
```

Your algorithm code must implement:

```python
import networkx as nx

def search_algorithm(src: str, dsts: list[str], G: nx.DiGraph, num_partitions: int) -> BroadCastTopology:
    """
    Design routing paths for broadcasting data partitions to multiple destinations.

    Args:
        src: Source node (e.g., "aws:ap-northeast-1")
        dsts: List of destination nodes (e.g., ["aws:us-east-1", "gcp:us-central1"])
        G: NetworkX DiGraph with edge attributes:
           - "cost": float ($/GB) - egress cost for transferring data
           - "throughput": float (Gbps) - maximum bandwidth capacity
        num_partitions: Number of data partitions to broadcast

    Returns:
        BroadCastTopology object with routing paths for all (destination, partition) pairs
    """
    pass

class BroadCastTopology:
    def __init__(self, src: str, dsts: list[str], num_partitions: int):
        self.src = src
        self.dsts = dsts
        self.num_partitions = int(num_partitions)
        # Structure: {dst: {partition_id: [edges]}}
        # Each edge is [src_node, dst_node, edge_data_dict]
        self.paths = {dst: {str(i): None for i in range(self.num_partitions)} for dst in dsts}

    def append_dst_partition_path(self, dst: str, partition: int, path: list):
        """
        Append an edge to the path for a specific destination-partition pair.

        Args:
            dst: Destination node
            partition: Partition ID (0 to num_partitions-1)
            path: Edge represented as [src_node, dst_node, edge_data_dict]
                  where edge_data_dict = G[src_node][dst_node]
        """
        partition = str(partition)
        if self.paths[dst][partition] is None:
            self.paths[dst][partition] = []
        self.paths[dst][partition].append(path)

    def set_dst_partition_paths(self, dst: str, partition: int, paths: list[list]):
        """
        Set the complete path (list of edges) for a destination-partition pair.

        Args:
            dst: Destination node
            partition: Partition ID
            paths: List of edges, each edge is [src_node, dst_node, edge_data_dict]
        """
        partition = str(partition)
        self.paths[dst][partition] = paths

    def set_num_partitions(self, num_partitions: int):
        """Update number of partitions"""
        self.num_partitions = num_partitions
```

Bandwidth Constraints
---------------------
Each cloud provider has ingress/egress limits (Gbps) per region:
- AWS: 10 Gbps ingress, 5 Gbps egress
- GCP: 16 Gbps ingress, 7 Gbps egress
- Azure: 16 Gbps ingress, 16 Gbps egress

These limits are multiplied by the number of VMs per region.

When multiple edges share a node and exceed its limits:
- Flow is **equally distributed** among incoming/outgoing edges (each edge gets $\frac{\text{limit}}{n_{\text{edges}}}$)
- Transfer time increases as actual throughput decreases
- Example: If a node has 3 outgoing edges and 5 Gbps egress limit, each edge gets min(original_flow, 5/3 Gbps)

**Strategy tip**: Different partitions can use different paths to the same destination, potentially avoiding bottlenecks by distributing load across the network.

Scoring (0-100)
---------------
```python
score = 1.0 / (1.0 + total_cost) * 100
```

Lower total cost → higher score

Example: Basic Implementation
------------------------------
```python
def search_algorithm(src, dsts, G, num_partitions):
    bc_topology = BroadCastTopology(src, dsts, num_partitions)

    for dst in dsts:
        path = nx.dijkstra_path(G, src, dst, weight="cost")
        for i in range(len(path) - 1):
            for partition_id in range(num_partitions):
                bc_topology.append_dst_partition_path(dst, partition_id,
                    [path[i], path[i + 1], G[path[i]][path[i + 1]]])

    return bc_topology
```

Evaluation Details
------------------
- **Test configurations**: 5 network scenarios
  - intra-AWS: Broadcasting within AWS regions
  - intra-Azure: Broadcasting within Azure regions
  - intra-GCP: Broadcasting within GCP regions
  - inter-AGZ: Broadcasting across AWS, GCP, Azure
  - inter-GAZ2: Another multi-cloud scenario
- **Network scale**: ~20-50 regions per provider
- **Default setup**: 2 VMs per region
- **Data volume**: Varies by configuration (e.g., 300 GB)
- **Partitions**: Varies by configuration (e.g., 10 partitions)
- **Instance cost**: $0.54/hour (based on m5.8xlarge spot instances)

Input Format
------------
The `spec_path` parameter is a string containing the file path to the specification JSON file.

**spec_path file format:**
```json
{
    "config_files": ["examples/config/intra_aws.json", ...],
    "num_vms": 2
}
```

Each config file contains:
```json
{
    "source_node": "aws:ap-northeast-1",
    "dest_nodes": ["aws:us-east-1", "aws:eu-west-1", ...],
    "data_vol": 300,
    "num_partitions": 10,
    "ingress_limit": {"aws": 10, "gcp": 16, "azure": 16},
    "egress_limit": {"aws": 5, "gcp": 7, "azure": 16}
}
```

Requirements and Constraints
-----------------------------
- All partitions (0 to num_partitions-1) must have valid paths to each destination
- Paths must start from the source node and end at the specified destination
- Self-loops are not allowed
- Different partitions can use different routes to the same destination
- Multiple destinations can share intermediate nodes (tree topology)
- The BroadCastTopology class is provided in the evaluation environment


Cross Entropy Optimization Problem
====================================

Problem Setting
---------------
Design and optimize high-performance Triton kernels for Cross Entropy loss computation on GPU. This problem focuses on implementing efficient cross entropy loss kernels using Triton's JIT compilation system.

The challenge involves optimizing:
- **Loss computation**: Efficient computation of negative log-likelihood loss
- **Memory access patterns**: Efficient loading and storing of logits and targets
- **Numerical stability**: Handling log-sum-exp operations with proper numerical stability
- **Block tiling**: Optimal block sizes for GPU execution across different batch sizes
- **Performance benchmarking**: Achieving speedup over baseline PyTorch implementations

Target
------
- **Primary**: Maximize geometric mean speedup over baseline (higher is better)
- **Secondary**: Ensure correctness across diverse batch sizes and vocabulary sizes
- **Tertiary**: Minimize kernel launch overhead and memory usage

API Specification
-----------------
Implement a `Solution` class that returns a Triton kernel implementation:

```python
class Solution:
    def solve(self, spec_path: str = None) -> dict:
        """
        Returns a dict with either:
        - {"code": "python_code_string"}
        - {"program_path": "path/to/kernel.py"}
        """
        # Your implementation
        pass
```

Your kernel implementation must provide:

```python
import torch
import triton
import triton.language as tl

def cross_entropy(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
    """
    Cross entropy loss computation.

    Args:
        logits: Input tensor of shape (M, N) - logits for M samples and N classes
        targets: Input tensor of shape (M,) - target class indices (int64)

    Returns:
        Output tensor of shape (M,) - negative log-likelihood loss for each sample
    """
    pass
```

API Usage Notes
---------------
- The evaluator looks for a `cross_entropy` function in the module namespace
- Function must handle tensor strides and memory layouts correctly
- Must use Triton JIT compilation for kernel definition
- Should leverage Triton's autotuning features for optimization
- Kernel must handle variable batch sizes and vocabulary sizes efficiently
- Output must be float32 tensor of shape (M,)

Scoring (0-100)
---------------
Performance is measured against GPU baseline implementations:

```
geometric_mean_gpu_time = geometric_mean(gpu_baseline_times)
geometric_mean_answer_time = geometric_mean(answer_times)

# Linear interpolation: 0 points = 1x GPU baseline, 100 points = 3x GPU baseline
target_time_0 = geometric_mean_gpu_time  # 0 points (1x GPU baseline)
target_time_100 = geometric_mean_gpu_time / 3.0  # 100 points (3x speedup over GPU)
score = 100 * (target_time_0 - geometric_mean_answer_time) / (target_time_0 - target_time_100)
```

- 0 points = 1x GPU baseline performance
- 100 points = 3x speedup over GPU baseline
- Score is linearly interpolated between these two points

Note: Correctness is verified against GPU baseline, and scoring spans from 1x GPU baseline (0 points) to 3x GPU baseline (100 points).

Evaluation Details
------------------
- Tested on multiple batch sizes: M ∈ {256, 512, 1024} (default)
- Fixed vocabulary size: N=8192 (configurable via metadata)
- Can also test custom shapes specified in metadata
- Correctness verified with tolerance: rtol=1e-3, atol=5e-4
- Performance measured using median execution time
- Requires CUDA backend and GPU support
- All tests must pass for any score > 0

dependencies:
  uv_project: resources
tag: hpc
runtime:
  environment: "Triton 3.2.0 with CUDA 12.2 (triton-tlx image)"
  docker:
    image: andylizf/triton-tlx:tlx-nv-cu122
    gpu: true


Decoding Attention Optimization Problem
========================================

Problem Setting
---------------
Design and optimize high-performance Triton kernels for Decoding Attention computation on GPU. This problem focuses on implementing efficient attention kernels for decoder-only transformer models using Triton's JIT compilation system.

The challenge involves optimizing:
- **Attention computation**: Efficient computation of scaled dot-product attention
- **Memory access patterns**: Efficient loading and storing of Q, K, V tensors
- **Numerical stability**: Handling softmax operations with proper numerical stability
- **Block tiling**: Optimal block sizes for GPU execution across different sequence lengths
- **Performance benchmarking**: Achieving speedup over baseline PyTorch implementations

Target
------
- **Primary**: Maximize geometric mean speedup over baseline (higher is better)
- **Secondary**: Ensure correctness across diverse sequence lengths and attention heads
- **Tertiary**: Minimize kernel launch overhead and memory usage

API Specification
-----------------
Implement a `Solution` class that returns a Triton kernel implementation:

```python
class Solution:
    def solve(self, spec_path: str = None) -> dict:
        """
        Returns a dict with either:
        - {"code": "python_code_string"}
        - {"program_path": "path/to/kernel.py"}
        """
        # Your implementation
        pass
```

Your kernel implementation must provide:

```python
import torch
import triton
import triton.language as tl

def decoding_attn(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor) -> torch.Tensor:
    """
    Decoding attention computation.

    Args:
        Q: Input tensor of shape (Z, H, M, Dq) - query tensor (float16)
        K: Input tensor of shape (Z, H, N, Dq) - key tensor (float16)
        V: Input tensor of shape (Z, H, N, Dv) - value tensor (float16)

    Returns:
        Output tensor of shape (Z, H, M, Dv) - attention output (float16)
    """
    pass
```

API Usage Notes
---------------
- The evaluator looks for a `decoding_attn` function in the module namespace
- Function must handle tensor strides and memory layouts correctly
- Must use Triton JIT compilation for kernel definition
- Should leverage Triton's autotuning features for optimization
- Kernel must handle variable sequence lengths efficiently
- Output must be float16 tensor of shape (Z, H, M, Dv)

Scoring (0-100)
---------------
Performance is measured against GPU baseline implementations:

```
geometric_mean_gpu_time = geometric_mean(gpu_baseline_times)
geometric_mean_answer_time = geometric_mean(answer_times)

# Linear interpolation: 0 points = 1x GPU baseline, 100 points = 3x GPU baseline
target_time_0 = geometric_mean_gpu_time  # 0 points (1x GPU baseline)
target_time_100 = geometric_mean_gpu_time / 3.0  # 100 points (3x speedup over GPU)
score = 100 * (target_time_0 - geometric_mean_answer_time) / (target_time_0 - target_time_100)
```

- 0 points = 1x GPU baseline performance
- 100 points = 3x speedup over GPU baseline
- Score is linearly interpolated between these two points

Note: Correctness is verified against GPU baseline, and scoring spans from 1x GPU baseline (0 points) to 3x GPU baseline (100 points).

Evaluation Details
------------------
- Tested on multiple sequence lengths: N ∈ {1024, 2048, 4096, 8192} (default)
- Fixed dimensions: Z=1, H=8, M=1, Dq=64, Dv=64 (configurable via metadata)
- Can also test custom shapes specified in metadata
- Correctness verified with tolerance: rtol=1e-2, atol=5e-3
- Performance measured using median execution time
- Requires CUDA backend and GPU support
- All tests must pass for any score > 0


Flash Attention Optimization Problem
=====================================

Problem Setting
---------------
Design and optimize high-performance Triton kernels for Flash Attention computation on GPU. This problem focuses on implementing efficient attention kernels with causal masking support using Triton's JIT compilation system.

The challenge involves optimizing:
- **Attention computation**: Efficient computation of scaled dot-product attention
- **Causal masking**: Handling causal attention masks efficiently
- **Memory access patterns**: Efficient loading and storing of Q, K, V tensors
- **Numerical stability**: Handling softmax operations with proper numerical stability using streaming softmax
- **Block tiling**: Optimal block sizes for GPU execution across different sequence lengths
- **Performance benchmarking**: Achieving speedup over baseline PyTorch implementations

Target
------
- **Primary**: Maximize geometric mean speedup over baseline (higher is better)
- **Secondary**: Ensure correctness across diverse sequence lengths and attention heads
- **Tertiary**: Minimize kernel launch overhead and memory usage

API Specification
-----------------
Implement a `Solution` class that returns a Triton kernel implementation:

```python
class Solution:
    def solve(self, spec_path: str = None) -> dict:
        """
        Returns a dict with either:
        - {"code": "python_code_string"}
        - {"program_path": "path/to/kernel.py"}
        """
        # Your implementation
        pass
```

Your kernel implementation must provide:

```python
import torch
import triton
import triton.language as tl

def flash_attn(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, causal: bool = True) -> torch.Tensor:
    """
    Flash attention computation with optional causal masking.

    Args:
        Q: Input tensor of shape (Z, H, M, Dq) - query tensor (float16)
        K: Input tensor of shape (Z, H, N, Dq) - key tensor (float16)
        V: Input tensor of shape (Z, H, N, Dv) - value tensor (float16)
        causal: Whether to apply causal masking (default True)

    Returns:
        Output tensor of shape (Z, H, M, Dv) - attention output (float16)
    """
    # Your implementation
    pass
```

Input Specifications
--------------------
- **Q**: Query tensor of shape `(Z, H, M, Dq)` where:
  - `Z`: Batch size (typically 1)
  - `H`: Number of attention heads (typically 8)
  - `M`: Query sequence length (tested with 512, 1024, 2048)
  - `Dq`: Query/key feature dimension (typically 64)
- **K**: Key tensor of shape `(Z, H, N, Dq)` where `N` matches `M` for flash attention
- **V**: Value tensor of shape `(Z, H, N, Dv)` where:
  - `Dv`: Value feature dimension (typically 64)
- All inputs are `torch.float16` and on CUDA device
- `causal`: Boolean flag for causal masking (default True)

Output Specifications
--------------------
- Output tensor of shape `(Z, H, M, Dv)` matching the query batch/head dimensions
- Output dtype: `torch.float16`
- Output device: Same as input (CUDA)

Correctness Requirements
-------------------------
- Numerical correctness verified against PyTorch baseline implementation
- Relative tolerance: 1e-2, Absolute tolerance: 5e-3
- All test cases must pass for any score above 0
- Causal masking must be correctly implemented when `causal=True`

Scoring (0-100)
---------------
Performance is measured against GPU baseline implementations:

```
geometric_mean_gpu_time = geometric_mean(gpu_baseline_times)
geometric_mean_answer_time = geometric_mean(answer_times)

# Linear interpolation: 0 points = 1x GPU baseline, 100 points = 10x GPU baseline
target_time_0 = geometric_mean_gpu_time  # 0 points (1x GPU baseline)
target_time_100 = geometric_mean_gpu_time / 10.0  # 100 points (10x speedup over GPU)
score = 100 * (target_time_0 - geometric_mean_answer_time) / (target_time_0 - target_time_100)
```

- 0 points = 1x GPU baseline performance
- 100 points = 10x speedup over GPU baseline
- Score is linearly interpolated between these two points

Note: Correctness is verified against GPU baseline, and scoring spans from 1x GPU baseline (0 points) to 10x GPU baseline (100 points).

Evaluation Details
------------------
- Test cases: M = 512, 1024, 2048 (with N = M)
- Warmup phase: 10 iterations to stabilize GPU clocks and caches
- Random seed: Fixed seed (0) for reproducible data generation
- Strict correctness: Any test failure results in score of 0

Additional Notes
----------------
- The benchmark uses float32 for PyTorch baseline (for numerical stability) but float16 for answer evaluation
- Streaming softmax techniques are recommended for numerical stability
- Consider using block pointers (`tl.make_block_ptr`) for efficient memory access
- Causal masking requires careful attention to the masking pattern (lower triangular for causal attention)


Fused Linear Cross Entropy Optimization Problem
===============================================

Problem Setting
---------------
Design and optimize high-performance Triton kernels for Fused Linear Cross Entropy loss computation on GPU. This problem focuses on implementing efficient fused kernels that combine matrix multiplication (linear layer) with cross-entropy loss computation using Triton's JIT compilation system.

The challenge involves optimizing:
- **Fused computation**: Efficiently combining linear layer (X @ W + B) with cross-entropy loss
- **Memory access patterns**: Efficient loading and storing of X, W, B, and targets
- **Numerical stability**: Handling log-sum-exp operations with proper numerical stability
- **Two-pass algorithm**: Finding row-wise max in first pass, computing sumexp and target logit in second pass
- **Block tiling**: Optimal block sizes for GPU execution across different batch sizes
- **Performance benchmarking**: Achieving speedup over baseline PyTorch implementations

Target
------
- **Primary**: Maximize geometric mean speedup over baseline (higher is better)
- **Secondary**: Ensure correctness across diverse batch sizes and vocabulary sizes
- **Tertiary**: Minimize kernel launch overhead and memory usage

API Specification
-----------------
Implement a `Solution` class that returns a Triton kernel implementation:

```python
class Solution:
    def solve(self, spec_path: str = None) -> dict:
        """
        Returns a dict with either:
        - {"code": "python_code_string"}
        - {"program_path": "path/to/kernel.py"}
        """
        # Your implementation
        pass
```

Your kernel implementation must provide:

```python
import torch
import triton
import triton.language as tl

def fused_linear_ce(X: torch.Tensor, W: torch.Tensor, B: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
    """
    Fused linear layer with cross entropy loss computation.

    Args:
        X: Input tensor of shape (M, K) - input features (float16)
        W: Weight tensor of shape (K, N) - weight matrix (float16)
        B: Bias tensor of shape (N,) - bias vector (float32)
        targets: Target tensor of shape (M,) - target class indices (int64)

    Returns:
        Output tensor of shape (M,) - negative log-likelihood loss per sample (float32)
    """
    # Your implementation
    pass
```

Input Specifications
--------------------
- **X**: Input tensor of shape `(M, K)` where:
  - `M`: Batch size (tested with values from M_list)
  - `K`: Input feature dimension (typically 4096)
  - dtype: `torch.float16`
- **W**: Weight tensor of shape `(K, N)` where:
  - `N`: Number of classes / vocabulary size (typically 8192)
  - dtype: `torch.float16`
- **B**: Bias tensor of shape `(N,)`:
  - dtype: `torch.float32`
- **targets**: Target tensor of shape `(M,)`:
  - dtype: `torch.int64` (long)
- All inputs are on CUDA device

Output Specifications
--------------------
- Output tensor of shape `(M,)` matching the batch size
- Output dtype: `torch.float32`
- Output device: Same as input (CUDA)
- Each element is the negative log-likelihood loss for the corresponding sample

Correctness Requirements
-------------------------
- Numerical correctness verified against PyTorch baseline implementation
- Relative tolerance: 1e-2, Absolute tolerance: 0.5
- All test cases must pass for any score above 0
- The operation computes: logits = X @ W + B, then NLL = cross_entropy(logits, targets, reduction='none')

Scoring (0-100)
---------------
Performance is measured against CPU and GPU baseline implementations:

```
geometric_mean_cpu_time = geometric_mean(cpu_baseline_times)
geometric_mean_gpu_time = geometric_mean(gpu_baseline_times)
geometric_mean_answer_time = geometric_mean(answer_times)

# Linear interpolation: 0 points = 3x CPU baseline, 100 points = 7x GPU baseline
target_time_0 = geometric_mean_cpu_time / 3.0  # 0 points (3x speedup over CPU)
target_time_100 = geometric_mean_gpu_time / 7.0  # 100 points (7x speedup over GPU)
score = 100 * (target_time_0 - geometric_mean_answer_time) / (target_time_0 - target_time_100)
```

- 0 points = 3x speedup over CPU baseline
- 100 points = 7x speedup over GPU baseline
- Score is linearly interpolated between these two points

Note: Correctness is verified against GPU baseline. Scoring spans from 3x CPU baseline (0 points) to 7x GPU baseline (100 points).

Evaluation Details
------------------
- Test cases: M values from M_list (typically [128, 256, 512])
- N: Vocabulary size (typically 8192)
- K: Input feature dimension (typically 4096)
- Warmup phase: 10 iterations to stabilize GPU clocks and caches
- Random seed: Fixed seed (0) for reproducible data generation
- Strict correctness: Any test failure results in score of 0

Additional Notes
----------------
- The benchmark uses float32 for bias (for numerical stability)
- A two-pass algorithm is recommended:
  1. First pass: Compute logits and find row-wise maximum
  2. Second pass: Compute sumexp with fixed row_max and gather target logits
- Consider using block tiling for efficient matrix multiplication
- Numerical stability is crucial: use row_max for stable log-sum-exp computation


Fused Linear Jensen-Shannon Divergence Optimization Problem
==========================================================

Problem Setting
---------------
Design and optimize high-performance Triton kernels for Fused Linear Jensen-Shannon Divergence (JSD) computation on GPU. This problem focuses on implementing efficient fused kernels that combine two linear layers with JSD computation using Triton's JIT compilation system.

The challenge involves optimizing:
- **Fused computation**: Efficiently combining two linear layers (X @ W1 + B1, X @ W2 + B2) with JSD computation
- **Memory access patterns**: Efficient loading and storing of X, W1, W2, B1, B2
- **Numerical stability**: Handling log-sum-exp operations and log computations with proper numerical stability
- **Two-pass algorithm**: Computing log-sum-exp for both branches in first pass, computing JSD in second pass
- **Block tiling**: Optimal block sizes for GPU execution across different batch sizes
- **Performance benchmarking**: Achieving speedup over baseline PyTorch implementations

Target
------
- **Primary**: Maximize geometric mean speedup over baseline (higher is better)
- **Secondary**: Ensure correctness across diverse batch sizes and vocabulary sizes
- **Tertiary**: Minimize kernel launch overhead and memory usage

API Specification
-----------------
Implement a `Solution` class that returns a Triton kernel implementation:

```python
class Solution:
    def solve(self, spec_path: str = None) -> dict:
        """
        Returns a dict with either:
        - {"code": "python_code_string"}
        - {"program_path": "path/to/kernel.py"}
        """
        # Your implementation
        pass
```

Your kernel implementation must provide:

```python
import torch
import triton
import triton.language as tl

def fused_linear_jsd(X: torch.Tensor, W1: torch.Tensor, B1: torch.Tensor, W2: torch.Tensor, B2: torch.Tensor) -> torch.Tensor:
    """
    Fused linear layers with Jensen-Shannon Divergence computation.

    Args:
        X: Input tensor of shape (M, K) - input features (float16)
        W1: Weight tensor of shape (K, N) - first weight matrix (float16)
        B1: Bias tensor of shape (N,) - first bias vector (float32)
        W2: Weight tensor of shape (K, N) - second weight matrix (float16)
        B2: Bias tensor of shape (N,) - second bias vector (float32)

    Returns:
        Output tensor of shape (M,) - Jensen-Shannon Divergence per sample (float32)
    """
    # Your implementation
    pass
```

Input Specifications
--------------------
- **X**: Input tensor of shape `(M, K)` where:
  - `M`: Batch size (tested with values from M_list)
  - `K`: Input feature dimension (typically 2048)
  - dtype: `torch.float16`
- **W1**: Weight tensor of shape `(K, N)`:
  - `N`: Number of classes / vocabulary size (typically 4096)
  - dtype: `torch.float16`
- **B1**: Bias tensor of shape `(N,)`:
  - dtype: `torch.float32`
- **W2**: Weight tensor of shape `(K, N)`:
  - dtype: `torch.float16`
- **B2**: Bias tensor of shape `(N,)`:
  - dtype: `torch.float32`
- All inputs are on CUDA device

Output Specifications
--------------------
- Output tensor of shape `(M,)` matching the batch size
- Output dtype: `torch.float32`
- Output device: Same as input (CUDA)
- Each element is the Jensen-Shannon Divergence for the corresponding sample

Correctness Requirements
-------------------------
- Numerical correctness verified against PyTorch baseline implementation
- Relative tolerance: 1e-2, Absolute tolerance: 0.5
- All test cases must pass for any score above 0
- The operation computes:
  - logits1 = X @ W1 + B1, logits2 = X @ W2 + B2
  - P = softmax(logits1), Q = softmax(logits2)
  - M = 0.5 * (P + Q)
  - JSD = 0.5 * (KL(P||M) + KL(Q||M))

Scoring (0-100)
---------------
Performance is measured against CPU and GPU baseline implementations:

```
geometric_mean_cpu_time = geometric_mean(cpu_baseline_times)
geometric_mean_gpu_time = geometric_mean(gpu_baseline_times)
geometric_mean_answer_time = geometric_mean(answer_times)

# Linear interpolation: 0 points = 3x CPU baseline, 100 points = 7x GPU baseline
target_time_0 = geometric_mean_cpu_time / 3.0  # 0 points (3x speedup over CPU)
target_time_100 = geometric_mean_gpu_time / 7.0  # 100 points (7x speedup over GPU)
score = 100 * (target_time_0 - geometric_mean_answer_time) / (target_time_0 - target_time_100)
```

- 0 points = 3x speedup over CPU baseline
- 100 points = 7x speedup over GPU baseline
- Score is linearly interpolated between these two points

Note: Correctness is verified against GPU baseline. Scoring spans from 3x CPU baseline (0 points) to 7x GPU baseline (100 points).

Evaluation Details
------------------
- Test cases: M values from M_list (typically [128, 256])
- N: Vocabulary size (typically 4096)
- K: Input feature dimension (typically 2048)
- Warmup phase: 10 iterations to stabilize GPU clocks and caches
- Random seed: Fixed seed (0) for reproducible data generation
- Strict correctness: Any test failure results in score of 0

Additional Notes
----------------
- The benchmark uses float32 for biases (for numerical stability)
- A two-pass algorithm is recommended:
  1. First pass: Compute log-sum-exp for both logits1 and logits2
  2. Second pass: Compute softmax probabilities and accumulate JSD
- Consider using block tiling for efficient matrix multiplication
- Numerical stability is crucial: use log-sum-exp for stable softmax computation
- JSD computation requires careful handling of logarithms to avoid numerical issues


GDPA Attention Optimization Problem
===================================

Problem Setting
---------------
Design and optimize high-performance Triton kernels for GDPA (Gated Dot-Product Attention) computation on GPU. This problem focuses on implementing efficient attention kernels with gated Q and K tensors using Triton's JIT compilation system.

The challenge involves optimizing:
- **Gated attention computation**: Efficient computation of scaled dot-product attention with gated Q and K tensors
- **Gating mechanism**: Applying sigmoid gates to Q and K tensors before attention computation
- **Memory access patterns**: Efficient loading and storing of Q, K, V, GQ, GK tensors
- **Numerical stability**: Handling softmax operations with proper numerical stability using streaming softmax
- **Block tiling**: Optimal block sizes for GPU execution across different sequence lengths
- **Performance benchmarking**: Achieving speedup over baseline PyTorch implementations

Target
------
- **Primary**: Maximize geometric mean speedup over baseline (higher is better)
- **Secondary**: Ensure correctness across diverse sequence lengths and attention heads
- **Tertiary**: Minimize kernel launch overhead and memory usage

API Specification
-----------------
Implement a `Solution` class that returns a Triton kernel implementation:

```python
class Solution:
    def solve(self, spec_path: str = None) -> dict:
        """
        Returns a dict with either:
        - {"code": "python_code_string"}
        - {"program_path": "path/to/kernel.py"}
        """
        # Your implementation
        pass
```

Your kernel implementation must provide:

```python
import torch
import triton
import triton.language as tl

def gdpa_attn(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, GQ: torch.Tensor, GK: torch.Tensor) -> torch.Tensor:
    """
    GDPA attention computation with gated Q and K tensors.

    Args:
        Q: Input tensor of shape (Z, H, M, Dq) - query tensor (float16)
        K: Input tensor of shape (Z, H, N, Dq) - key tensor (float16)
        V: Input tensor of shape (Z, H, N, Dv) - value tensor (float16)
        GQ: Input tensor of shape (Z, H, M, Dq) - query gate tensor (float16)
        GK: Input tensor of shape (Z, H, N, Dq) - key gate tensor (float16)

    Returns:
        Output tensor of shape (Z, H, M, Dv) - attention output (float16)
    """
    # Your implementation
    pass
```

Input Specifications
--------------------
- **Q**: Query tensor of shape `(Z, H, M, Dq)` where:
  - `Z`: Batch size (typically 1)
  - `H`: Number of attention heads (typically 8)
  - `M`: Query sequence length (tested with 512, 1024)
  - `Dq`: Query/key feature dimension (typically 64)
- **K**: Key tensor of shape `(Z, H, N, Dq)` where `N` matches `M` for GDPA attention
- **V**: Value tensor of shape `(Z, H, N, Dv)` where:
  - `Dv`: Value feature dimension (typically 64)
- **GQ**: Query gate tensor of shape `(Z, H, M, Dq)`
- **GK**: Key gate tensor of shape `(Z, H, N, Dq)`
- All inputs are `torch.float16` and on CUDA device

Output Specifications
--------------------
- Output tensor of shape `(Z, H, M, Dv)` matching the query batch/head dimensions
- Output dtype: `torch.float16`
- Output device: Same as input (CUDA)

Correctness Requirements
------------------------
- Numerical correctness verified against PyTorch baseline implementation
- Relative tolerance: 1e-2, Absolute tolerance: 5e-3
- All test cases must pass for any score above 0
- Gating must be correctly applied: Qg = Q * sigmoid(GQ), Kg = K * sigmoid(GK)

Scoring (0-100)
---------------
Performance is measured against GPU baseline implementations:

```
geometric_mean_gpu_time = geometric_mean(gpu_baseline_times)
geometric_mean_answer_time = geometric_mean(answer_times)

# Linear interpolation: 0 points = 1x GPU baseline, 100 points = 3x GPU baseline
target_time_0 = geometric_mean_gpu_time  # 0 points (1x GPU baseline)
target_time_100 = geometric_mean_gpu_time / 3.0  # 100 points (3x speedup over GPU)
score = 100 * (target_time_0 - geometric_mean_answer_time) / (target_time_0 - target_time_100)
```

- 0 points = 1x GPU baseline performance
- 100 points = 3x speedup over GPU baseline
- Score is linearly interpolated between these two points

Note: Correctness is verified against GPU baseline, and scoring spans from 1x GPU baseline (0 points) to 3x GPU baseline (100 points).

Evaluation Details
------------------
- Test cases: M = 512, 1024 (with N = M)
- Warmup phase: 10 iterations to stabilize GPU clocks and caches
- Random seed: Fixed seed (0) for reproducible data generation
- Strict correctness: Any test failure results in score of 0

Additional Notes
----------------
- The benchmark uses float16 for both baseline and answer evaluation
- Streaming softmax techniques are recommended for numerical stability
- Consider using block pointers (`tl.make_block_ptr`) for efficient memory access
- Gating mechanism: Qg = Q * sigmoid(GQ), Kg = K * sigmoid(GK)
- Scale factor: 1.0 / sqrt(Dq)


GEMM Optimization Problem
=========================

Problem Setting
---------------
Design and optimize high-performance Triton kernels for General Matrix-Matrix Multiplication (GEMM) on GPU. This problem focuses on implementing efficient matrix multiplication kernels using Triton's JIT compilation system.

The challenge involves optimizing:
- **Memory access patterns**: Efficient loading and storing of matrix data
- **Block tiling**: Optimal block sizes for GPU execution
- **Autotuning**: Leveraging Triton's autotuning capabilities
- **Activation functions**: Implementing GELU activation within the kernel
- **Performance benchmarking**: Achieving speedup over baseline implementations

Target
------
- **Primary**: Maximize geometric mean speedup over baseline (higher is better)
- **Secondary**: Ensure correctness across diverse matrix shapes
- **Tertiary**: Minimize kernel launch overhead and memory usage

API Specification
-----------------
Implement a `Solution` class that returns a Triton kernel implementation:

```python
class Solution:
    def solve(self, spec_path: str = None) -> dict:
        """
        Returns a dict with either:
        - {"code": "python_code_string"}
        - {"program_path": "path/to/kernel.py"}
        """
        # Your implementation
        pass
```

Your kernel implementation must provide:

```python
import torch
import triton
import triton.language as tl

def matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """
    Matrix multiplication with GELU activation.

    Args:
        a: Input tensor of shape (M, K)
        b: Input tensor of shape (K, N)

    Returns:
        Output tensor of shape (M, N) with GELU activation applied
    """
    pass
```

Required GELU Implementation:
```python
@triton.jit
def gelu(x):
    return x * 0.5 * (1.0 + tl.extra.cuda.libdevice.erf(x * 0.7071067811865476))
```

API Usage Notes
---------------
- The evaluator looks for a `matmul` function in the module namespace
- Function must handle tensor strides and memory layouts correctly
- Must use Triton JIT compilation for kernel definition
- Should leverage Triton's autotuning features for optimization
- Kernel must apply GELU activation to the result before returning

Scoring (0-100)
---------------
Performance is measured against baseline implementations:

```
geometric_mean_speedup = geometric_mean(answer_times / baseline_times)
raw_score = min(geometric_mean_speedup, 3.0)  # Cap at 3x speedup
score = (raw_score - 1.0) / 2.0 * 100  # Map 1x-3x to 0-100
```

- 0 points = No speedup (1x baseline performance)
- 50 points = 2x speedup over baseline
- 100 points = 3x+ speedup over baseline

Evaluation Details (annoying variant)
------------------------------------
- Shapes focus on intentionally awkward, non-friendly dimensions:
  - (1000, 1000, 1000)
  - (1537, 1537, 1025)
  - (3001, 4093, 997)
  - (6143, 2003, 3079)
  - (5000, 3000, 1234)
  - (777, 3333, 2049)
- Correctness verified with tolerance: rtol=1e-2, atol=5e-3
- Performance measured using median execution time
- Requires CUDA backend and GPU support

dependencies:
  uv_project: resources
datasets: []
tag: hpc
runtime:
  docker:
    image: andylizf/triton-tlx:tlx-nv-cu122
    gpu: true
  environment: "Triton 3.2.0 with CUDA 12.2 (triton-tlx image)"


GEMM Optimization Problem
=========================

Problem Setting
---------------
Design and optimize high-performance Triton kernels for General Matrix-Matrix Multiplication (GEMM) on GPU. This problem focuses on implementing efficient matrix multiplication kernels using Triton's JIT compilation system.

The challenge involves optimizing:
- **Memory access patterns**: Efficient loading and storing of matrix data
- **Block tiling**: Optimal block sizes for GPU execution
- **Autotuning**: Leveraging Triton's autotuning capabilities
- **Activation functions**: Implementing GELU activation within the kernel
- **Performance benchmarking**: Achieving speedup over baseline implementations

Target
------
- **Primary**: Maximize geometric mean speedup over baseline (higher is better)
- **Secondary**: Ensure correctness across diverse matrix shapes
- **Tertiary**: Minimize kernel launch overhead and memory usage

API Specification
-----------------
Implement a `Solution` class that returns a Triton kernel implementation:

```python
class Solution:
    def solve(self, spec_path: str = None) -> dict:
        """
        Returns a dict with either:
        - {"code": "python_code_string"}
        - {"program_path": "path/to/kernel.py"}
        """
        # Your implementation
        pass
```

Your kernel implementation must provide:

```python
import torch
import triton
import triton.language as tl

def matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """
    Matrix multiplication with GELU activation.

    Args:
        a: Input tensor of shape (M, K)
        b: Input tensor of shape (K, N)

    Returns:
        Output tensor of shape (M, N) with GELU activation applied
    """
    pass
```

Required GELU Implementation:
```python
@triton.jit
def gelu(x):
    return x * 0.5 * (1.0 + tl.extra.cuda.libdevice.erf(x * 0.7071067811865476))
```

API Usage Notes
---------------
- The evaluator looks for a `matmul` function in the module namespace
- Function must handle tensor strides and memory layouts correctly
- Must use Triton JIT compilation for kernel definition
- Should leverage Triton's autotuning features for optimization
- Kernel must apply GELU activation to the result before returning

Scoring (0-100)
---------------
Performance is measured against baseline implementations:

```
geometric_mean_speedup = geometric_mean(answer_times / baseline_times)
raw_score = min(geometric_mean_speedup, 3.0)  # Cap at 3x speedup
score = (raw_score - 1.0) / 2.0 * 100  # Map 1x-3x to 0-100
```

- 0 points = No speedup (1x baseline performance)
- 50 points = 2x speedup over baseline
- 100 points = 3x+ speedup over baseline

Evaluation Details (K-skewed variant)
------------------------------------
- Shapes emphasize very small and very large K with varied (M, N):
  - Base (M, N): (1024,1024), (2048,512), (512,2048), (4096,256), (256,4096)
  - K values: small K = [32, 48, 64, 96, 128]; huge K = [3072, 4096, 6144, 8192]
- Correctness verified with tolerance: rtol=1e-2, atol=5e-3
- Performance measured using median execution time
- Requires CUDA backend and GPU support


GEMM Optimization Problem
=========================

Problem Setting
---------------
Design and optimize high-performance Triton kernels for General Matrix-Matrix Multiplication (GEMM) on GPU. This problem focuses on implementing efficient matrix multiplication kernels using Triton's JIT compilation system.

The challenge involves optimizing:
- **Memory access patterns**: Efficient loading and storing of matrix data
- **Block tiling**: Optimal block sizes for GPU execution
- **Autotuning**: Leveraging Triton's autotuning capabilities
- **Activation functions**: Implementing GELU activation within the kernel
- **Performance benchmarking**: Achieving speedup over baseline implementations

Target
------
- **Primary**: Maximize geometric mean speedup over baseline (higher is better)
- **Secondary**: Ensure correctness across diverse matrix shapes
- **Tertiary**: Minimize kernel launch overhead and memory usage

API Specification
-----------------
Implement a `Solution` class that returns a Triton kernel implementation:

```python
class Solution:
    def solve(self, spec_path: str = None) -> dict:
        """
        Returns a dict with either:
        - {"code": "python_code_string"}
        - {"program_path": "path/to/kernel.py"}
        """
        # Your implementation
        pass
```

Your kernel implementation must provide:

```python
import torch
import triton
import triton.language as tl

def matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """
    Matrix multiplication with GELU activation.

    Args:
        a: Input tensor of shape (M, K)
        b: Input tensor of shape (K, N)

    Returns:
        Output tensor of shape (M, N) with GELU activation applied
    """
    pass
```

Required GELU Implementation:
```python
@triton.jit
def gelu(x):
    return x * 0.5 * (1.0 + tl.extra.cuda.libdevice.erf(x * 0.7071067811865476))
```

API Usage Notes
---------------
- The evaluator looks for a `matmul` function in the module namespace
- Function must handle tensor strides and memory layouts correctly
- Must use Triton JIT compilation for kernel definition
- Should leverage Triton's autotuning features for optimization
- Kernel must apply GELU activation to the result before returning

Scoring (0-100)
---------------
Performance is measured against baseline implementations:

```
geometric_mean_speedup = geometric_mean(answer_times / baseline_times)
raw_score = min(geometric_mean_speedup, 3.0)  # Cap at 3x speedup
score = (raw_score - 1.0) / 2.0 * 100  # Map 1x-3x to 0-100
```

- 0 points = No speedup (1x baseline performance)
- 50 points = 2x speedup over baseline
- 100 points = 3x+ speedup over baseline

Evaluation Details (near-tile variant)
-------------------------------------
- Shapes clustered around tile boundaries (tile M,N=128, K=64), including +/-1 and +7:
  - M in {127,128,129,135, 255, 385, 633}
  - N in {127,128,129,135, 257, 383, 643}
  - K in {63,64,65,71, 129, 191, 325}
- Only positive dimensions up to 8192 are included; Cartesian product filtered to limits
- Correctness verified with tolerance: rtol=1e-2, atol=5e-3
- Performance measured using median execution time
- Requires CUDA backend and GPU support

Implementation Notes for Solution Authors
----------------------------------------
- Triton `tl.arange(0, BLOCK_*)` requires the range to be a power of two. Choose `BLOCK_M`, `BLOCK_N`, and especially `BLOCK_K` from powers of two (e.g., 32/64/128/256) to avoid compilation errors.
- Return tensor dtype must match input dtype (fp16/bf16/fp32). Accumulate in fp32 inside the kernel, but allocate the output with `dtype=a.dtype` to pass correctness checks.
- Provide a `Solution.solve()` that returns a static code string via `{ "code": python_source }`. Avoid reflection-based approaches (e.g., `inspect.getsource`) as modules are imported under different names during evaluation.
- Respect arbitrary input strides; compute element-wise strides and use masked loads/stores for tail tiles.
- Autotuning: include strides in the autotune key (e.g., `a_stride_am`, `a_stride_ak`, `b_stride_bk`, `b_stride_bn`) to ensure correct kernel specialization across layouts.
- Recommended tile sets to cover near-tile cases:
  - `BLOCK_M/N`: {64, 128, 256}
  - `BLOCK_K`: {32, 64, 128} (avoid non-powers like 80)


GEMM Optimization Problem
=========================

Problem Setting
---------------
Design and optimize high-performance Triton kernels for General Matrix-Matrix Multiplication (GEMM) on GPU. This problem focuses on implementing efficient matrix multiplication kernels using Triton's JIT compilation system.

The challenge involves optimizing:
- **Memory access patterns**: Efficient loading and storing of matrix data
- **Block tiling**: Optimal block sizes for GPU execution
- **Autotuning**: Leveraging Triton's autotuning capabilities
- **Activation functions**: Implementing GELU activation within the kernel
- **Performance benchmarking**: Achieving speedup over baseline implementations

Target
------
- **Primary**: Maximize geometric mean speedup over baseline (higher is better)
- **Secondary**: Ensure correctness across diverse matrix shapes
- **Tertiary**: Minimize kernel launch overhead and memory usage

API Specification
-----------------
Implement a `Solution` class that returns a Triton kernel implementation:

```python
class Solution:
    def solve(self, spec_path: str = None) -> dict:
        """
        Returns a dict with either:
        - {"code": "python_code_string"}
        - {"program_path": "path/to/kernel.py"}
        """
        # Your implementation
        pass
```

Your kernel implementation must provide:

```python
import torch
import triton
import triton.language as tl

def matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """
    Matrix multiplication with GELU activation.

    Args:
        a: Input tensor of shape (M, K)
        b: Input tensor of shape (K, N)

    Returns:
        Output tensor of shape (M, N) with GELU activation applied
    """
    pass
```

Required GELU Implementation:
```python
@triton.jit
def gelu(x):
    return x * 0.5 * (1.0 + tl.extra.cuda.libdevice.erf(x * 0.7071067811865476))
```

API Usage Notes
---------------
- The evaluator looks for a `matmul` function in the module namespace
- Function must handle tensor strides and memory layouts correctly
- Must use Triton JIT compilation for kernel definition
- Should leverage Triton's autotuning features for optimization
- Kernel must apply GELU activation to the result before returning

Scoring (0-100)
---------------
Performance is measured against baseline implementations:

```
geometric_mean_speedup = geometric_mean(answer_times / baseline_times)
raw_score = min(geometric_mean_speedup, 3.0)  # Cap at 3x speedup
score = (raw_score - 1.0) / 2.0 * 100  # Map 1x-3x to 0-100
```

- 0 points = No speedup (1x baseline performance)
- 50 points = 2x speedup over baseline
- 100 points = 3x+ speedup over baseline

Evaluation Details (rectangles variant)
--------------------------------------
- Tall/skinny and short/wide rectangles; K is balanced:
  - Tall/skinny: M in [1024, 2048, 4096, 8192], N in [64, 128, 192, 256], K in [512, 1024, 2048]
  - Short/wide:  M in [64, 128, 192, 256],      N in [1024, 2048, 4096, 8192], K in [512, 1024, 2048]
- Correctness verified with tolerance: rtol=1e-2, atol=5e-3
- Performance measured using median execution time
- Requires CUDA backend and GPU support


GEMM Optimization Problem
=========================

Problem Setting
---------------
Design and optimize high-performance Triton kernels for General Matrix-Matrix Multiplication (GEMM) on GPU. This problem focuses on implementing efficient matrix multiplication kernels using Triton's JIT compilation system.

The challenge involves optimizing:
- **Memory access patterns**: Efficient loading and storing of matrix data
- **Block tiling**: Optimal block sizes for GPU execution
- **Autotuning**: Leveraging Triton's autotuning capabilities
- **Activation functions**: Implementing GELU activation within the kernel
- **Performance benchmarking**: Achieving speedup over baseline implementations

Target
------
- **Primary**: Maximize geometric mean speedup over baseline (higher is better)
- **Secondary**: Ensure correctness across diverse matrix shapes
- **Tertiary**: Minimize kernel launch overhead and memory usage

API Specification
-----------------
Implement a `Solution` class that returns a Triton kernel implementation:

```python
class Solution:
    def solve(self, spec_path: str = None) -> dict:
        """
        Returns a dict with either:
        - {"code": "python_code_string"}
        - {"program_path": "path/to/kernel.py"}
        """
        # Your implementation
        pass
```

Your kernel implementation must provide:

```python
import torch
import triton
import triton.language as tl

def matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """
    Matrix multiplication with GELU activation.

    Args:
        a: Input tensor of shape (M, K)
        b: Input tensor of shape (K, N)

    Returns:
        Output tensor of shape (M, N) with GELU activation applied
    """
    pass
```

Required GELU Implementation:
```python
@triton.jit
def gelu(x):
    return x * 0.5 * (1.0 + tl.extra.cuda.libdevice.erf(x * 0.7071067811865476))
```

API Usage Notes
---------------
- The evaluator looks for a `matmul` function in the module namespace
- Function must handle tensor strides and memory layouts correctly
- Must use Triton JIT compilation for kernel definition
- Should leverage Triton's autotuning features for optimization
- Kernel must apply GELU activation to the result before returning

Scoring (0-100)
---------------
Performance is measured against baseline implementations:

```
geometric_mean_speedup = geometric_mean(answer_times / baseline_times)
raw_score = min(geometric_mean_speedup, 3.0)  # Cap at 3x speedup
score = (raw_score - 1.0) / 2.0 * 100  # Map 1x-3x to 0-100
```

- 0 points = No speedup (1x baseline performance)
- 50 points = 2x speedup over baseline
- 100 points = 3x+ speedup over baseline

Evaluation Details (squares variant)
-----------------------------------
- Only square shapes with equal M=N=K from 512 to 8192, step 1024:
  - Shapes: (s, s, s) for s ∈ {512, 1536, 2560, 3584, 4608, 5632, 6656, 7680, 8192}
- Correctness verified with tolerance: rtol=1e-2, atol=5e-3
- Performance measured using median execution time
- Requires CUDA backend and GPU support


GEMM Optimization Problem
=========================

Problem Setting
---------------
Design and optimize high-performance Triton kernels for General Matrix-Matrix Multiplication (GEMM) on GPU. This problem focuses on implementing efficient matrix multiplication kernels using Triton's JIT compilation system.

The challenge involves optimizing:
- **Memory access patterns**: Efficient loading and storing of matrix data
- **Block tiling**: Optimal block sizes for GPU execution
- **Autotuning**: Leveraging Triton's autotuning capabilities
- **Activation functions**: Implementing GELU activation within the kernel
- **Performance benchmarking**: Achieving speedup over baseline implementations

Target
------
- **Primary**: Maximize geometric mean speedup over baseline (higher is better)
- **Secondary**: Ensure correctness across diverse matrix shapes
- **Tertiary**: Minimize kernel launch overhead and memory usage

API Specification
-----------------
Implement a `Solution` class that returns a Triton kernel implementation:

```python
class Solution:
    def solve(self, spec_path: str = None) -> dict:
        """
        Returns a dict with either:
        - {"code": "python_code_string"}
        - {"program_path": "path/to/kernel.py"}
        """
        # Your implementation
        pass
```

Your kernel implementation must provide:

```python
import torch
import triton
import triton.language as tl

def matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """
    Matrix multiplication with GELU activation.

    Args:
        a: Input tensor of shape (M, K)
        b: Input tensor of shape (K, N)

    Returns:
        Output tensor of shape (M, N) with GELU activation applied
    """
    pass
```

Required GELU Implementation:
```python
@triton.jit
def gelu(x):
    return x * 0.5 * (1.0 + tl.extra.cuda.libdevice.erf(x * 0.7071067811865476))
```

API Usage Notes
---------------
- The evaluator looks for a `matmul` function in the module namespace
- Function must handle tensor strides and memory layouts correctly
- Must use Triton JIT compilation for kernel definition
- Should leverage Triton's autotuning features for optimization
- Kernel must apply GELU activation to the result before returning

Scoring (0-100)
---------------
Performance is measured against baseline implementations:

```
geometric_mean_speedup = geometric_mean(answer_times / baseline_times)
raw_score = min(geometric_mean_speedup, 3.0)  # Cap at 3x speedup
score = (raw_score - 1.0) / 2.0 * 100  # Map 1x-3x to 0-100
```

- 0 points = No speedup (1x baseline performance)
- 50 points = 2x speedup over baseline
- 100 points = 3x+ speedup over baseline

Evaluation Details (transformer-ish variant)
-------------------------------------------
- Transformer-like shapes targeting common attention/FFN dimensions:
  - (2048, 4096, 4096)
  - (4096, 4096, 4096)
  - (8192, 4096, 4096)
  - (8192, 8192, 4096)
  - (4096, 11008, 4096)
  - (4096, 4096, 11008)
- Correctness verified with tolerance: rtol=1e-2, atol=5e-3
- Performance measured using median execution time
- Requires CUDA backend and GPU support


Group GEMM Optimization Problem
================================

Problem Setting
---------------
Design and optimize high-performance Triton kernels for Batched Matrix-Matrix Multiplication (BMM) on GPU. This problem focuses on implementing efficient batched matrix multiplication kernels using Triton's JIT compilation system.

The challenge involves optimizing:
- **Batched operations**: Efficient handling of multiple matrix pairs in a single kernel launch
- **Memory access patterns**: Efficient loading and storing of batched matrix data
- **Block tiling**: Optimal block sizes for GPU execution across different batch sizes
- **Performance benchmarking**: Achieving speedup over baseline PyTorch implementations

Target
------
- **Primary**: Maximize geometric mean speedup over baseline (higher is better)
- **Secondary**: Ensure correctness across diverse batch sizes and matrix shapes
- **Tertiary**: Minimize kernel launch overhead and memory usage

API Specification
-----------------
Implement a `Solution` class that returns a Triton kernel implementation:

```python
class Solution:
    def solve(self, spec_path: str = None) -> dict:
        """
        Returns a dict with either:
        - {"code": "python_code_string"}
        - {"program_path": "path/to/kernel.py"}
        """
        # Your implementation
        pass
```

Your kernel implementation must provide:

```python
import torch
import triton
import triton.language as tl

def bmm(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
    """
    Batched matrix multiplication.

    Args:
        A: Input tensor of shape (B, M, K) - batch of M×K matrices
        B: Input tensor of shape (B, K, N) - batch of K×N matrices

    Returns:
        Output tensor of shape (B, M, N) - batch of M×N result matrices
    """
    pass
```

API Usage Notes
---------------
- The evaluator looks for a `bmm` function in the module namespace
- Function must handle tensor strides and memory layouts correctly
- Must use Triton JIT compilation for kernel definition
- Should leverage Triton's autotuning features for optimization
- Kernel must handle variable batch sizes efficiently

Common Pitfalls & Implementation Requirements
---------------------------------------------
**Triton Autotune Keys:**
- Autotune `key` parameter must only include actual kernel parameters (e.g., `["M", "N", "K"]`)
- Do NOT include non-kernel parameters in the autotune key (e.g., `dtype_a`, `dtype_b`)
- Example (correct): `@triton.autotune(configs=[...], key=["M", "N", "K"])`
- Example (incorrect): `@triton.autotune(configs=[...], key=["M", "N", "K", "dtype_a"])`

**Dtype Casting:**
- Use `tl.float16` directly for output dtype casting: `acc.to(tl.float16)` (correct)
- Do NOT use `tl.dtype_elementwise(C_ptr)` - this function doesn't exist in Triton 3.4.0 (incorrect)
- The problem requires float16 output, so always cast accumulator to `tl.float16`

**Kernel Parameters:**
- Only pass actual kernel parameters as arguments to the kernel
- Do NOT pass Python objects (like `dtype`) as keyword arguments unless they're defined as kernel parameters
- Example (correct): `_bmm_kernel[grid](A, B, C, Batches, M, N, K, ..., BLOCK_M, BLOCK_N, BLOCK_K)`

**Correctness Requirements:**
- All tests must pass (correctness check) for any score > 0
- Output dtype must be float16 (match baseline behavior)
- Output shape must be (B, M, N) where B is batch size

**Kernel Implementation Pattern:**
- Initialize pointers inside the K-loop for each iteration (computes pointers per K-slice)
- Use proper boundary masking: `k_mask = (k + offs_k) < K` or `k_idxs = k0 + offs_k` with `k_idxs < K`
- Load data and convert to float32 BEFORE accumulation: `a = tl.load(A_ptrs, mask=a_mask, other=0.0).to(tl.float32)`
- Accumulate in float32: `acc += tl.dot(a, b)` where `acc` is `dtype=tl.float32`
- Example K-loop pattern:
  ```python
  k0 = 0
  while k0 < K:
      k_idxs = k0 + offs_k
      A_ptrs = A_batch_ptr + (offs_m[:, None] * stride_am) + (k_idxs[None, :] * stride_ak)
      B_ptrs = B_batch_ptr + (k_idxs[:, None] * stride_bk) + (offs_n[None, :] * stride_bn)
      a_mask = (offs_m[:, None] < M) & (k_idxs[None, :] < K)
      b_mask = (offs_n[None, :] < N) & (k_idxs[:, None] < K)
      a = tl.load(A_ptrs, mask=a_mask, other=0.0).to(tl.float32)
      b = tl.load(B_ptrs, mask=b_mask, other=0.0).to(tl.float32)
      acc += tl.dot(a, b)
      k0 += BLOCK_K
  ```

**Solution.solve() Method:**
- Read file directly using `Path(__file__).read_text()` (correct)
- Do NOT use `inspect.getsource(sys.modules[__name__])` - fails when module is dynamically loaded (incorrect)
- Example (correct):
  ```python
  def solve(self, spec_path: Optional[str] = None) -> Dict[str, str]:
      from pathlib import Path
      current_file = Path(__file__).resolve()
      return {"code": current_file.read_text(encoding="utf-8")}
  ```

**Performance Optimization Tips:**
- Use FP32 accumulator for numerical stability: `acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)`
- Load data as float32: `a = tl.load(...).to(tl.float32)` - critical for correctness
- Cast to float16 only at the end: `tl.store(c_ptrs, acc.to(tl.float16), mask=c_mask)`
- Consider using autotune to find optimal block sizes and warp configurations
- Test with warmup phase to stabilize GPU clocks before benchmarking

Scoring (0-100)
---------------
Performance is measured against baseline implementations:

```
geometric_mean_speedup = geometric_mean(baseline_times / answer_times)
raw_score = min(geometric_mean_speedup, 5.0)  # Cap at 5x speedup
score = (raw_score - 1.0) / 4.0 * 100  # Map 1x-5x to 0-100
```

- 0 points = No speedup (1x baseline performance)
- 25 points = 2x speedup over baseline
- 50 points = 3x speedup over baseline
- 75 points = 4x speedup over baseline
- 100 points = 5x+ speedup over baseline

Evaluation Details
------------------
- Tested on multiple batch sizes: B ∈ {64, 256, 1024} (default)
- Fixed matrix dimensions: M=64, N=64, K=64 (configurable via metadata)
- Can also test custom shapes specified in metadata
- Correctness verified with tolerance: rtol=1e-2, atol=5e-3
- Performance measured using median execution time
- Requires CUDA backend and GPU support
- All tests must pass for any score > 0


ImageNet Pareto Optimization - 1M Parameter Variant
===================================================

Problem Setting
---------------
Train a neural network on a synthetic ImageNet-like dataset to maximize accuracy while staying within a parameter budget of 1,000,000 parameters.

Objective: Achieve the highest possible accuracy without exceeding the parameter constraint.

Target
------
**Primary**: Maximize test accuracy
**Secondary**: Maintain model efficiency (stay under parameter budget)

API Specification
----------------
Implement a `Solution` class:

```python
import torch
import torch.nn as nn

class Solution:
    def solve(self, train_loader, val_loader, metadata: dict = None) -> torch.nn.Module:
        """
        Train a model and return it.

        Args:
            train_loader: PyTorch DataLoader with training data
            val_loader: PyTorch DataLoader with validation data
            metadata: Dict with keys:
                - num_classes: int (128)
                - input_dim: int (384)
                - param_limit: int (1,000,000)
                - baseline_accuracy: float (0.8)
                - train_samples: int
                - val_samples: int
                - test_samples: int
                - device: str ("cpu")

        Returns:
            Trained torch.nn.Module ready for evaluation
        """
        # Your implementation
        pass
```

**Implementation Requirements**:
- Use `metadata["input_dim"]` and `metadata["num_classes"]` for model architecture
- Keep model parameters <= 1,000,000 (hard constraint - models exceeding this receive 0 score)
- Return a trained model ready for evaluation
- Ensure model works with the provided device

Parameter Constraint
--------------------
**HARD LIMIT: 1,000,000 trainable parameters**

- This is an absolute constraint enforced during evaluation
- Models exceeding 1,000,000 parameters will receive a score of 0.0
- The constraint cannot be waived under any circumstances
- You must design your architecture carefully to stay under this limit

Example: A model with 1,000,001 parameters → Score 0.0 (constraint violated)
Example: A model with 1,000,000 parameters → Score based on accuracy

Baseline Accuracy
-----------------
**Baseline Accuracy for this variant: 80%**

- This is the expected performance level for a simple model at this parameter budget
- Solutions must achieve accuracy **above** this baseline to receive a positive score
- Accuracy **below** baseline results in 0 points
- Accuracy improvements are scored linearly

Scoring Formula
---------------

The scoring is based purely on **linear accuracy scaling** from baseline to 100%:

```
If model exceeds parameter limit (1,000,000):
    Score = 0.0  (constraint violation)

Else:
    Score = (accuracy - 0.8) / (1.0 - 0.8) × 100.0

    Where:
    - accuracy = achieved test accuracy (0.0 to 1.0)
    - 0.8 = baseline accuracy for this variant
    - 1.0 = target (100% accuracy = 100 points)

    Score is clamped to [0, 100] range
```

**Linearly Scaled Scoring for 1M variant:**

| Accuracy | Score | Notes |
|----------|-------|-------|
| 80.0% | 0 | At baseline (0 points) |
| 85.0% | ~25 | 5% above baseline |
| 90.0% | ~50 | 10% above baseline |
| 95.0% | ~75 | 15% above baseline |
| 100% | 100 | Perfect accuracy (max score) |

Evaluation Process
------------------
The evaluator follows these steps:

### 1. Build Synthetic Dataset
```python
# Generate synthetic ImageNet-like data
train_loader, val_loader, test_loader = make_dataloaders()
# Each sample: (384,) feature vector, label in [0, 127]
```

### 2. Call Solution
```python
from solution import Solution
solution = Solution()
model = solution.solve(train_loader, val_loader, metadata)
# metadata contains: num_classes, input_dim, param_limit, baseline_accuracy, device
```

### 3. Validate Model
```python
param_count = sum(p.numel() for p in model.parameters() if p.requires_grad)
if param_count > 1000000:
    score = 0.0  # Constraint violation
```

### 4. Evaluate Accuracy
```python
model.eval()
correct = 0
total = 0
for inputs, targets in test_loader:
    outputs = model(inputs)
    preds = outputs.argmax(dim=1)
    correct += (preds == targets).sum().item()
    total += targets.numel()
accuracy = correct / total
```

### 5. Calculate Score
```python
score = (accuracy - 0.8) / (1.0 - 0.8) * 100.0
score = max(0.0, min(100.0, score))
```

Evaluation Details
------------------
- 128 classes, 384-dimensional feature vectors
- Training: 2,048 samples (16 per class)
- Validation: 512 samples (4 per class)
- Test: 1,024 samples (8 per class)
- Data generated synthetically with controlled noise

Environment Details
-------------------
- **Device**: CPU only (`device="cpu"`)
- **Python Environment**:
  - Python 3
  - PyTorch 2.2-2.4
  - NumPy ≥1.24
  - tqdm ≥4.64
- **Timeout**: 1 hour (3600 seconds) for entire evaluation

Key Points
----------
1. **Parameter Constraint is Hard**: Models exceeding 1,000,000 parameters always score 0
2. **Baseline is Lower Bound**: Must achieve 80%+ accuracy to score points
3. **Linear Scoring**: Every accuracy improvement scales linearly to the score
4. **100% is Target**: Achieving 100% accuracy gives full 100 points
5. **Accuracy is Primary**: Focus on accuracy within the parameter budget

Example: Simple Baseline
-------------------------
```python
import torch
import torch.nn as nn

class Solution:
    def solve(self, train_loader, val_loader, metadata: dict = None):
        # Simple 2-layer MLP
        input_dim = metadata["input_dim"]      # 384
        num_classes = metadata["num_classes"]  # 128
        hidden_dim = 512

        model = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, num_classes)
        )

        # Parameter count: 384*512 + 512 + 512*128 + 128 = ~262,784

        # Simple training loop
        optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
        criterion = nn.CrossEntropyLoss()

        for epoch in range(50):
            model.train()
            for inputs, targets in train_loader:
                optimizer.zero_grad()
                outputs = model(inputs)
                loss = criterion(outputs, targets)
                loss.backward()
                optimizer.step()

        return model
```

**Note**: This baseline achieves ~80% accuracy with ~263K parameters. To reach higher accuracy within the 1M budget, consider deeper networks, residual connections, or better optimization.

Implementation Tips
-------------------
- Monitor parameter count: `sum(p.numel() for p in model.parameters() if p.requires_grad)`
- Gradually improve architecture while staying under budget
- Use techniques like batch normalization, dropout, or residual connections
- Higher capacity (more parameters) generally improves accuracy up to the limit

Baseline Performance
--------------------
- **Baseline Accuracy**: 80%
- **Baseline Parameters**: Approximately 1,000,000
- This represents a simple model at this parameter budget


ImageNet Pareto Optimization - 200K Parameter Variant
=====================================================

Problem Setting
---------------
Train a neural network on a synthetic ImageNet-like dataset to maximize accuracy while staying within a parameter budget of 200,000 parameters.

Objective: Achieve the highest possible accuracy without exceeding the parameter constraint.

Target
------
**Primary**: Maximize test accuracy
**Secondary**: Maintain model efficiency (stay under parameter budget)

API Specification
----------------
Implement a `Solution` class:

```python
import torch
import torch.nn as nn

class Solution:
    def solve(self, train_loader, val_loader, metadata: dict = None) -> torch.nn.Module:
        """
        Train a model and return it.

        Args:
            train_loader: PyTorch DataLoader with training data
            val_loader: PyTorch DataLoader with validation data
            metadata: Dict with keys:
                - num_classes: int (128)
                - input_dim: int (384)
                - param_limit: int (200,000)
                - baseline_accuracy: float (0.65)
                - train_samples: int
                - val_samples: int
                - test_samples: int
                - device: str ("cpu")

        Returns:
            Trained torch.nn.Module ready for evaluation
        """
        # Your implementation
        pass
```

**Implementation Requirements**:
- Use `metadata["input_dim"]` and `metadata["num_classes"]` for model architecture
- Keep model parameters <= 200,000 (hard constraint - models exceeding this receive 0 score)
- Return a trained model ready for evaluation
- Ensure model works with the provided device

Parameter Constraint
--------------------
**HARD LIMIT: 200,000 trainable parameters**

- This is an absolute constraint enforced during evaluation
- Models exceeding 200,000 parameters will receive a score of 0.0
- The constraint cannot be waived under any circumstances
- You must design your architecture carefully to stay under this limit

Example: A model with 200,001 parameters → Score 0.0 (constraint violated)
Example: A model with 200,000 parameters → Score based on accuracy

Baseline Accuracy
-----------------
**Baseline Accuracy for this variant: 65%**

- This is the expected performance level for a simple model at this parameter budget
- Solutions must achieve accuracy **above** this baseline to receive a positive score
- Accuracy **below** baseline results in 0 points
- Accuracy improvements are scored linearly

Scoring Formula
---------------

The scoring is based purely on **linear accuracy scaling** from baseline to 100%:

```
If model exceeds parameter limit (200,000):
    Score = 0.0  (constraint violation)

Else:
    Score = (accuracy - 0.65) / (1.0 - 0.65) × 100.0

    Where:
    - accuracy = achieved test accuracy (0.0 to 1.0)
    - 0.65 = baseline accuracy for this variant
    - 1.0 = target (100% accuracy = 100 points)

    Score is clamped to [0, 100] range
```

**Linearly Scaled Scoring for 200K variant:**

| Accuracy | Score | Notes |
|----------|-------|-------|
| 65.0% | 0 | At baseline (0 points) |
| 70.0% | ~14 | 5% above baseline |
| 75.0% | ~28 | 10% above baseline |
| 80.0% | ~42 | 15% above baseline |
| 100% | 100 | Perfect accuracy (max score) |

Evaluation Process
------------------
The evaluator follows these steps:

### 1. Build Synthetic Dataset
```python
# Generate synthetic ImageNet-like data
train_loader, val_loader, test_loader = make_dataloaders()
# Each sample: (384,) feature vector, label in [0, 127]
```

### 2. Call Solution
```python
from solution import Solution
solution = Solution()
model = solution.solve(train_loader, val_loader, metadata)
# metadata contains: num_classes, input_dim, param_limit, baseline_accuracy, device
```

### 3. Validate Model
```python
param_count = sum(p.numel() for p in model.parameters() if p.requires_grad)
if param_count > 200000:
    score = 0.0  # Constraint violation
```

### 4. Evaluate Accuracy
```python
model.eval()
correct = 0
total = 0
for inputs, targets in test_loader:
    outputs = model(inputs)
    preds = outputs.argmax(dim=1)
    correct += (preds == targets).sum().item()
    total += targets.numel()
accuracy = correct / total
```

### 5. Calculate Score
```python
score = (accuracy - 0.65) / (1.0 - 0.65) * 100.0
score = max(0.0, min(100.0, score))
```

Evaluation Details
------------------
- 128 classes, 384-dimensional feature vectors
- Training: 2,048 samples (16 per class)
- Validation: 512 samples (4 per class)
- Test: 1,024 samples (8 per class)
- Data generated synthetically with controlled noise

Environment Details
-------------------
- **Device**: CPU only (`device="cpu"`)
- **Python Environment**:
  - Python 3
  - PyTorch 2.2-2.4
  - NumPy ≥1.24
  - tqdm ≥4.64
- **Timeout**: 1 hour (3600 seconds) for entire evaluation

Key Points
----------
1. **Parameter Constraint is Hard**: Models exceeding 200,000 parameters always score 0
2. **Baseline is Lower Bound**: Must achieve 65%+ accuracy to score points
3. **Linear Scoring**: Every accuracy improvement scales linearly to the score
4. **100% is Target**: Achieving 100% accuracy gives full 100 points
5. **Accuracy is Primary**: Focus on accuracy within the parameter budget

Example: Simple Baseline
-------------------------
```python
import torch
import torch.nn as nn

class Solution:
    def solve(self, train_loader, val_loader, metadata: dict = None):
        # Simple 2-layer MLP
        input_dim = metadata["input_dim"]      # 384
        num_classes = metadata["num_classes"]  # 128
        hidden_dim = 256

        model = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, num_classes)
        )

        # Parameter count: 384*256 + 256 + 256*128 + 128 = ~131,456

        # Simple training loop
        optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
        criterion = nn.CrossEntropyLoss()

        for epoch in range(50):
            model.train()
            for inputs, targets in train_loader:
                optimizer.zero_grad()
                outputs = model(inputs)
                loss = criterion(outputs, targets)
                loss.backward()
                optimizer.step()

        return model
```

**Note**: This baseline achieves ~65% accuracy with ~131K parameters. To reach higher accuracy within the 200K budget, consider deeper networks or better optimization.

Implementation Tips
-------------------
- Monitor parameter count: `sum(p.numel() for p in model.parameters() if p.requires_grad)`
- Gradually improve architecture while staying under budget
- Use techniques like batch normalization, dropout, or residual connections
- Higher capacity (more parameters) generally improves accuracy up to the limit

Baseline Performance
--------------------
- **Baseline Accuracy**: 65%
- **Baseline Parameters**: Approximately 200,000
- This represents a simple model at this parameter budget


ImageNet Pareto Optimization - 2.5M Parameter Variant
=====================================================

Problem Setting
---------------
Train a neural network on a synthetic ImageNet-like dataset to maximize accuracy while staying within a parameter budget of 2,500,000 parameters.

Objective: Achieve the highest possible accuracy without exceeding the parameter constraint.

Target
------
**Primary**: Maximize test accuracy
**Secondary**: Maintain model efficiency (stay under parameter budget)

API Specification
----------------
Implement a `Solution` class:

```python
import torch
import torch.nn as nn

class Solution:
    def solve(self, train_loader, val_loader, metadata: dict = None) -> torch.nn.Module:
        """
        Train a model and return it.

        Args:
            train_loader: PyTorch DataLoader with training data
            val_loader: PyTorch DataLoader with validation data
            metadata: Dict with keys:
                - num_classes: int (128)
                - input_dim: int (384)
                - param_limit: int (2,500,000)
                - baseline_accuracy: float (0.85)
                - train_samples: int
                - val_samples: int
                - test_samples: int
                - device: str ("cpu")

        Returns:
            Trained torch.nn.Module ready for evaluation
        """
        # Your implementation
        pass
```

**Implementation Requirements**:
- Use `metadata["input_dim"]` and `metadata["num_classes"]` for model architecture
- Keep model parameters <= 2,500,000 (hard constraint - models exceeding this receive 0 score)
- Return a trained model ready for evaluation
- Ensure model works with the provided device

Parameter Constraint
--------------------
**HARD LIMIT: 2,500,000 trainable parameters**

- This is an absolute constraint enforced during evaluation
- Models exceeding 2,500,000 parameters will receive a score of 0.0
- The constraint cannot be waived under any circumstances
- You must design your architecture carefully to stay under this limit

Example: A model with 2,500,001 parameters → Score 0.0 (constraint violated)
Example: A model with 2,500,000 parameters → Score based on accuracy

Baseline Accuracy
-----------------
**Baseline Accuracy for this variant: 85%**

- This is the expected performance level for a simple model at this parameter budget
- Solutions must achieve accuracy **above** this baseline to receive a positive score
- Accuracy **below** baseline results in 0 points
- Accuracy improvements are scored linearly

Scoring Formula
---------------

The scoring is based purely on **linear accuracy scaling** from baseline to 100%:

```
If model exceeds parameter limit (2,500,000):
    Score = 0.0  (constraint violation)

Else:
    Score = (accuracy - 0.85) / (1.0 - 0.85) × 100.0

    Where:
    - accuracy = achieved test accuracy (0.0 to 1.0)
    - 0.85 = baseline accuracy for this variant
    - 1.0 = target (100% accuracy = 100 points)

    Score is clamped to [0, 100] range
```

**Linearly Scaled Scoring for 2.5M variant:**

| Accuracy | Score | Notes |
|----------|-------|-------|
| 85.0% | 0 | At baseline (0 points) |
| 90.0% | ~33 | 5% above baseline |
| 95.0% | ~66 | 10% above baseline |
| 100.0% | ~100 | 15% above baseline |
| 100% | 100 | Perfect accuracy (max score) |

Evaluation Process
------------------
The evaluator follows these steps:

### 1. Build Synthetic Dataset
```python
# Generate synthetic ImageNet-like data
train_loader, val_loader, test_loader = make_dataloaders()
# Each sample: (384,) feature vector, label in [0, 127]
```

### 2. Call Solution
```python
from solution import Solution
solution = Solution()
model = solution.solve(train_loader, val_loader, metadata)
# metadata contains: num_classes, input_dim, param_limit, baseline_accuracy, device
```

### 3. Validate Model
```python
param_count = sum(p.numel() for p in model.parameters() if p.requires_grad)
if param_count > 2500000:
    score = 0.0  # Constraint violation
```

### 4. Evaluate Accuracy
```python
model.eval()
correct = 0
total = 0
for inputs, targets in test_loader:
    outputs = model(inputs)
    preds = outputs.argmax(dim=1)
    correct += (preds == targets).sum().item()
    total += targets.numel()
accuracy = correct / total
```

### 5. Calculate Score
```python
score = (accuracy - 0.85) / (1.0 - 0.85) * 100.0
score = max(0.0, min(100.0, score))
```

Evaluation Details
------------------
- 128 classes, 384-dimensional feature vectors
- Training: 2,048 samples (16 per class)
- Validation: 512 samples (4 per class)
- Test: 1,024 samples (8 per class)
- Data generated synthetically with controlled noise

Environment Details
-------------------
- **Device**: CPU only (`device="cpu"`)
- **Python Environment**:
  - Python 3
  - PyTorch 2.2-2.4
  - NumPy ≥1.24
  - tqdm ≥4.64
- **Timeout**: 1 hour (3600 seconds) for entire evaluation

Key Points
----------
1. **Parameter Constraint is Hard**: Models exceeding 2,500,000 parameters always score 0
2. **Baseline is Lower Bound**: Must achieve 85%+ accuracy to score points
3. **Linear Scoring**: Every accuracy improvement scales linearly to the score
4. **100% is Target**: Achieving 100% accuracy gives full 100 points
5. **Accuracy is Primary**: Focus on accuracy within the parameter budget

Example: Simple Baseline
-------------------------
```python
import torch
import torch.nn as nn

class Solution:
    def solve(self, train_loader, val_loader, metadata: dict = None):
        # Simple 3-layer MLP
        input_dim = metadata["input_dim"]      # 384
        num_classes = metadata["num_classes"]  # 128
        hidden_dim = 1024

        model = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, num_classes)
        )

        # Parameter count: 384*1024 + 1024 + 1024*1024 + 1024 + 1024*128 + 128 = ~1,577,728

        # Simple training loop
        optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
        criterion = nn.CrossEntropyLoss()

        for epoch in range(50):
            model.train()
            for inputs, targets in train_loader:
                optimizer.zero_grad()
                outputs = model(inputs)
                loss = criterion(outputs, targets)
                loss.backward()
                optimizer.step()

        return model
```

**Note**: This baseline achieves ~85% accuracy with ~1.58M parameters. To reach higher accuracy within the 2.5M budget, consider deeper networks or better optimization.

Implementation Tips
-------------------
- Monitor parameter count: `sum(p.numel() for p in model.parameters() if p.requires_grad)`
- Gradually improve architecture while staying under budget
- Use techniques like batch normalization, dropout, or residual connections
- Higher capacity (more parameters) generally improves accuracy up to the limit

Baseline Performance
--------------------
- **Baseline Accuracy**: 85%
- **Baseline Parameters**: Approximately 2,500,000
- This represents a simple model at this parameter budget


ImageNet Pareto Optimization - 500K Parameter Variant
=====================================================

Problem Setting
---------------
Train a neural network on a synthetic ImageNet-like dataset to maximize accuracy while staying within a parameter budget of 500,000 parameters.

Objective: Achieve the highest possible accuracy without exceeding the parameter constraint.

Target
------
**Primary**: Maximize test accuracy
**Secondary**: Maintain model efficiency (stay under parameter budget)

API Specification
----------------
Implement a `Solution` class:

```python
import torch
import torch.nn as nn

class Solution:
    def solve(self, train_loader, val_loader, metadata: dict = None) -> torch.nn.Module:
        """
        Train a model and return it.

        Args:
            train_loader: PyTorch DataLoader with training data
            val_loader: PyTorch DataLoader with validation data
            metadata: Dict with keys:
                - num_classes: int (128)
                - input_dim: int (384)
                - param_limit: int (500,000)
                - baseline_accuracy: float (0.72)
                - train_samples: int
                - val_samples: int
                - test_samples: int
                - device: str ("cpu")

        Returns:
            Trained torch.nn.Module ready for evaluation
        """
        # Your implementation
        pass
```

**Implementation Requirements**:
- Use `metadata["input_dim"]` and `metadata["num_classes"]` for model architecture
- Keep model parameters <= 500,000 (hard constraint - models exceeding this receive 0 score)
- Return a trained model ready for evaluation
- Ensure model works with the provided device

Parameter Constraint
--------------------
**HARD LIMIT: 500,000 trainable parameters**

- This is an absolute constraint enforced during evaluation
- Models exceeding 500,000 parameters will receive a score of 0.0
- The constraint cannot be waived under any circumstances
- You must design your architecture carefully to stay under this limit

Example: A model with 500,001 parameters → Score 0.0 (constraint violated)
Example: A model with 500,000 parameters → Score based on accuracy

Baseline Accuracy
-----------------
**Baseline Accuracy for this variant: 72%**

- This is the expected performance level for a simple model at this parameter budget
- Solutions must achieve accuracy **above** this baseline to receive a positive score
- Accuracy **below** baseline results in 0 points
- Accuracy improvements are scored linearly

Scoring Formula
---------------

The scoring is based purely on **linear accuracy scaling** from baseline to 100%:

```
If model exceeds parameter limit (500,000):
    Score = 0.0  (constraint violation)

Else:
    Score = (accuracy - 0.72) / (1.0 - 0.72) × 100.0

    Where:
    - accuracy = achieved test accuracy (0.0 to 1.0)
    - 0.72 = baseline accuracy for this variant
    - 1.0 = target (100% accuracy = 100 points)

    Score is clamped to [0, 100] range
```

**Linearly Scaled Scoring for 500K variant:**

| Accuracy | Score | Notes |
|----------|-------|-------|
| 72.0% | 0 | At baseline (0 points) |
| 77.0% | ~17 | 5% above baseline |
| 82.0% | ~35 | 10% above baseline |
| 87.0% | ~53 | 15% above baseline |
| 100% | 100 | Perfect accuracy (max score) |

Evaluation Process
------------------
The evaluator follows these steps:

### 1. Build Synthetic Dataset
```python
# Generate synthetic ImageNet-like data
train_loader, val_loader, test_loader = make_dataloaders()
# Each sample: (384,) feature vector, label in [0, 127]
```

### 2. Call Solution
```python
from solution import Solution
solution = Solution()
model = solution.solve(train_loader, val_loader, metadata)
# metadata contains: num_classes, input_dim, param_limit, baseline_accuracy, device
```

### 3. Validate Model
```python
param_count = sum(p.numel() for p in model.parameters() if p.requires_grad)
if param_count > 500000:
    score = 0.0  # Constraint violation
```

### 4. Evaluate Accuracy
```python
model.eval()
correct = 0
total = 0
for inputs, targets in test_loader:
    outputs = model(inputs)
    preds = outputs.argmax(dim=1)
    correct += (preds == targets).sum().item()
    total += targets.numel()
accuracy = correct / total
```

### 5. Calculate Score
```python
score = (accuracy - 0.72) / (1.0 - 0.72) * 100.0
score = max(0.0, min(100.0, score))
```

Evaluation Details
------------------
- 128 classes, 384-dimensional feature vectors
- Training: 2,048 samples (16 per class)
- Validation: 512 samples (4 per class)
- Test: 1,024 samples (8 per class)
- Data generated synthetically with controlled noise

Environment Details
-------------------
- **Device**: CPU only (`device="cpu"`)
- **Python Environment**:
  - Python 3
  - PyTorch 2.2-2.4
  - NumPy ≥1.24
  - tqdm ≥4.64
- **Timeout**: 1 hour (3600 seconds) for entire evaluation

Key Points
----------
1. **Parameter Constraint is Hard**: Models exceeding 500,000 parameters always score 0
2. **Baseline is Lower Bound**: Must achieve 72%+ accuracy to score points
3. **Linear Scoring**: Every accuracy improvement scales linearly to the score
4. **100% is Target**: Achieving 100% accuracy gives full 100 points
5. **Accuracy is Primary**: Focus on accuracy within the parameter budget

Example: Simple Baseline
-------------------------
```python
import torch
import torch.nn as nn

class Solution:
    def solve(self, train_loader, val_loader, metadata: dict = None):
        # Simple 2-layer MLP
        input_dim = metadata["input_dim"]      # 384
        num_classes = metadata["num_classes"]  # 128
        hidden_dim = 384

        model = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, num_classes)
        )

        # Parameter count: 384*384 + 384 + 384*128 + 128 = ~196,992

        # Simple training loop
        optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
        criterion = nn.CrossEntropyLoss()

        for epoch in range(50):
            model.train()
            for inputs, targets in train_loader:
                optimizer.zero_grad()
                outputs = model(inputs)
                loss = criterion(outputs, targets)
                loss.backward()
                optimizer.step()

        return model
```

**Note**: This baseline achieves ~72% accuracy with ~197K parameters. To reach higher accuracy within the 500K budget, consider deeper networks or better optimization.

Implementation Tips
-------------------
- Monitor parameter count: `sum(p.numel() for p in model.parameters() if p.requires_grad)`
- Gradually improve architecture while staying under budget
- Use techniques like batch normalization, dropout, or residual connections
- Higher capacity (more parameters) generally improves accuracy up to the limit

Baseline Performance
--------------------
- **Baseline Accuracy**: 72%
- **Baseline Parameters**: Approximately 500,000
- This represents a simple model at this parameter budget