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

Consider a CSV file with $N$ rows and $M$ columns, where $M \leq 10$. We feed each row to an LLM inference engine (with a prefix KV cache) by concatenating all column values in that row. For the $i$-th row with entries $A[i,1], A[i,2], \ldots, A[i,M]$, we construct the input string:

```math
S_i = \text{Concat}(\text{string}(A[i,1]), \text{string}(A[i,2]), \ldots, \text{string}(A[i,M]))
````

When requesting $S_i$ for $i > 1$, the prefix KV-cache hit rate depends on the longest common prefix with any previously seen request:

```math
\text{hit\_rate}_i =
\frac{\max_{1 \le j < i} \text{LCP}(S_i, S_j)}{|S_i|}
```

where $LCP(S, T)$ is the length of the longest common prefix between strings $S$ and $T$.

You are allowed to reorder the CSV columns. Let $p$ be a permutation of $\{1, 2, ..., M\}$. The reordered string for row $i$ becomes:

```math
S'_i = \text{Concat}(\text{string}(A[i,p_1]), \text{string}(A[i,p_2]), \ldots, \text{string}(A[i,p_M]))
```

The goal is to choose a permutation $p$ that maximizes the overall KV-cache hit rate:

```math
\max_p\;
\frac{\sum_{i=2}^N \max_{1 \le j < i} \text{LCP}(S'_i, S'_j)}
     {\sum_{i=1}^N |S'_i|}
```

Target
---
Maximize prefix hit rate shown above (higher is better)

- **Hard Constraint**: Average runtime per dataset must be $\leq 10$ seconds (score = 0 if exceeded) and correctly handle column merge constraint.

**Column Merges**:
- Column merge specs are provided per dataset
- Columns in each merge group are concatenated into a single column
- The merged column replaces the original columns
- Merge operations are applied before column reordering

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

```python
import pandas as pd

class Solution:
    def solve(
        self,
        df: pd.DataFrame,
        early_stop: int = 100000,
        row_stop: int = 4,
        col_stop: int = 2,
        col_merge: list = None,
        one_way_dep: list = None,
        distinct_value_threshold: float = 0.7,
        parallel: bool = True,
    ) -> pd.DataFrame:
        """
        Reorder columns in the DataFrame to maximize prefix hit rate.

        Args:
            df: Input DataFrame to optimize
            early_stop: Early stopping parameter (default: 100000)
            row_stop: Row stopping parameter (default: 4)
            col_stop: Column stopping parameter (default: 2)
            col_merge: List of column groups to merge (columns in each group are merged into one)
            one_way_dep: List of one-way dependencies (not used in this variant)
            distinct_value_threshold: Threshold for distinct values (default: 0.7)
            parallel: Whether to use parallel processing (default: True)

        Returns:
            DataFrame with reordered columns (same rows, different column order)
        """
        # Your implementation
        pass
```

**Evaluation Process**:
1. Column merges are applied if specified
2. Your `solve()` method reorders the remaining columns
3. Rows are concatenated (no spaces) and prefix hit rate is calculated

Scoring (0-100)
---

baseline_hit_rate = Average prefix hit rate using original column order (0-point anchor)
avg_hit_rate = Your solution's average prefix hit rate across all datasets

For each dataset:
    dataset_score = ((hit_rate - baseline_hit_rate) / (1.0 - baseline_hit_rate)) × 100

final_score = Average of individual dataset scores

Score is clamped to [0, 100] range

**Runtime Constraint**:
- Average runtime per dataset must be ≤ 10 seconds
- If average runtime exceeds 10 seconds, score = 0.0

**Scoring Examples**:
- baseline_hit_rate = 0.0 (worst), avg_hit_rate = 1.0 (perfect) → Score = 100
- baseline_hit_rate = 0.5, avg_hit_rate = 0.5 → Score = 0
- baseline_hit_rate = 0.5, avg_hit_rate = 0.75 → Score = 50
- baseline_hit_rate = 0.5, avg_hit_rate = 1.0 → Score = 100

Implementation Notes
---
- Row values are concatenated without spaces: `"".join(row.values)`
- Column reordering should optimize for maximum prefix overlap in the concatenated string representation
- Consider column dependencies, distinct value distributions, and merge requirements when reordering
- Large datasets with $M > 10$ columns require efficient algorithms due to larger search space
- In our smaller dataset, $15k \leq N \leq 28k$ and $4 \leq M \leq 9$

**Example input**
please ignore the $> 10$ column number here
---
```csv
ID,LIMIT_BAL,SEX,EDUCATION,MARRIAGE,AGE,PAY_0,PAY_2,PAY_3,PAY_4,PAY_5,PAY_6,BILL_AMT1,BILL_AMT2,BILL_AMT3,BILL_AMT4,BILL_AMT5,BILL_AMT6,PAY_AMT1,PAY_AMT2,PAY_AMT3,PAY_AMT4,PAY_AMT5,PAY_AMT6,default payment next month
1,20000,2,2,1,24,2,2,-1,-1,-2,-2,3913,3102,689,0,0,0,0,689,0,0,0,0,1
2,120000,2,2,2,26,-1,2,0,0,0,2,2682,1725,2682,3272,3455,3261,0,1000,1000,1000,0,2000,1
3,90000,2,2,2,34,0,0,0,0,0,0,29239,14027,13559,14331,14948,15549,1518,1500,1000,1000,1000,5000,0
4,50000,2,2,1,37,0,0,0,0,0,0,46990,48233,49291,28314,28959,29547,2000,2019,1200,1100,1069,1000,0
5,50000,1,2,1,57,-1,0,-1,0,0,0,8617,5670,35835,20940,19146,19131,2000,36681,10000,9000,689,679,0
...
```

**Example output**
---
```
ID,LIMIT_BAL,SEX,EDUCATION,MARRIAGE,AGE,PAY_0,PAY_2,PAY_3,PAY_4,PAY_5,PAY_6,BILL_AMT1,BILL_AMT2,BILL_AMT3,BILL_AMT4,BILL_AMT5,BILL_AMT6,PAY_AMT1,PAY_AMT2,PAY_AMT3,PAY_AMT4,PAY_AMT5,PAY_AMT6,default payment next month
1,20000,2,2,1,24,2,2,-1,-1,-2,-2,3913,3102,689,0,0,0,0,689,0,0,0,0,1
2,120000,2,2,2,26,-1,2,0,0,0,2,2682,1725,2682,3272,3455,3261,0,1000,1000,1000,0,2000,1
3,90000,2,2,2,34,0,0,0,0,0,0,29239,14027,13559,14331,14948,15549,1518,1500,1000,1000,1000,5000,0
4,50000,2,2,1,37,0,0,0,0,0,0,46990,48233,49291,28314,28959,29547,2000,2019,1200,1100,1069,1000,0
5,50000,1,2,1,57,-1,0,-1,0,0,0,8617,5670,35835,20940,19146,19131,2000,36681,10000,9000,689,679,0
```
($p$ = $1, 2, \ldots M$)


Mamba2 Scan Optimization Problem
==================================

Problem Setting
---------------
Design and optimize high-performance Triton kernels for Mamba2 scan computation on GPU. This problem focuses on implementing efficient sequential scan operations using chunked parallelism with Triton's JIT compilation system.

The challenge involves optimizing:
- **Sequential scan computation**: Efficient computation of y_t = a_t * y_{t-1} + b_t * x_t
- **Chunked parallelism**: Processing sequences in chunks to enable parallelism while maintaining correctness
- **State management**: Efficiently managing and propagating state between chunks
- **Memory access patterns**: Efficient loading and storing of X, A, B tensors and state
- **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 feature dimensions
- **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 chunk_scan(X: torch.Tensor, A: torch.Tensor, B: torch.Tensor, chunk: int = 128, BD: int = 128) -> torch.Tensor:
    """
    Mamba2 chunked scan computation.

    Args:
        X: Input tensor of shape (L, D) - input sequence (float16)
        A: Input tensor of shape (L, D) - decay factors (float16)
        B: Input tensor of shape (L, D) - input weights (float16)
        chunk: Chunk size for parallel processing (default 128)
        BD: Block dimension for feature dimension tiling (default 128)

    Returns:
        Output tensor of shape (L, D) - scan output (float16)
    """
    # Your implementation
    pass
```

Input Specifications
--------------------
- **X**: Input tensor of shape `(L, D)` where:
  - `L`: Sequence length (tested with 2048, 4096)
  - `D`: Feature dimension (typically 512)
- **A**: Decay factor tensor of shape `(L, D)` (float16, typically |A| < 0.5)
- **B**: Input weight tensor of shape `(L, D)` (float16)
- All inputs are `torch.float16` and on CUDA device
- `chunk`: Chunk size for parallel processing (default 128)
- `BD`: Block dimension for feature dimension tiling (default 128)
- **Constraint**: L must be divisible by chunk

Output Specifications
--------------------
- Output tensor of shape `(L, D)` matching the input 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
- Sequential dependency must be correctly maintained: y_t = a_t * y_{t-1} + b_t * x_t

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 = 200x GPU baseline
target_time_0 = geometric_mean_gpu_time  # 0 points (1x GPU baseline)
target_time_100 = geometric_mean_gpu_time / 200.0  # 100 points (200x 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 = 200x 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 200x GPU baseline (100 points).

Evaluation Details
------------------
- Test cases: L = 2048, 4096 (with D = 512)
- 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
- Chunk size: 128, BD: 128

Additional Notes
----------------
- The benchmark uses float32 for PyTorch baseline (for numerical stability) but float16 for answer evaluation
- Sequential scan operation: y_t = a_t * y_{t-1} + b_t * x_t
- Chunked parallelism: Process sequence in chunks, maintaining state between chunks
- State propagation: State must be correctly propagated from one chunk to the next
- Consider using block tiling along the feature dimension (BD) for parallelism

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


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

Problem Setting
---------------
Design and optimize high-performance Triton kernels for Mixed GEMM (Linear + Bias + GELU) computation on GPU. This problem focuses on implementing efficient fused kernels that combine matrix multiplication, bias addition, and GELU activation using Triton's JIT compilation system.

The challenge involves optimizing:
- **Fused computation**: Efficiently combining linear layer (X @ W + B) with GELU activation
- **Memory access patterns**: Efficient loading and storing of X, W, B tensors
- **Mixed precision**: Handling float16 inputs/outputs with float32 bias and accumulation
- **GELU activation**: Implementing efficient GELU computation using CUDA libdevice functions
- **Block tiling**: Optimal block sizes for GPU execution across different matrix 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 matrix 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 linear_gelu(X: torch.Tensor, W: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
    """
    Linear layer with GELU activation 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)

    Returns:
        Output tensor of shape (M, N) - output with GELU activation (float16)
    """
    # Your implementation
    pass
```

Input Specifications
--------------------
- **X**: Input tensor of shape `(M, K)` where:
  - `M`: Batch size (tested with 512, 1024)
  - `K`: Input feature dimension (typically 4096)
  - dtype: `torch.float16`
- **W**: Weight tensor of shape `(K, N)` where:
  - `N`: Output feature dimension (typically 4096)
  - dtype: `torch.float16`
- **B**: Bias tensor of shape `(N,)` where:
  - dtype: `torch.float32`
- All inputs are on CUDA device

Output Specifications
--------------------
- Output tensor of shape `(M, N)` matching the input batch and output feature 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
- GELU activation must be correctly implemented

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 = 4096, K = 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 PyTorch baseline (for numerical stability) but float16 for answer evaluation
- GELU formula: gelu(x) = x * 0.5 * (1.0 + erf(x * 0.7071067811865476))
- Consider using CUDA libdevice erf function: `tl.extra.cuda.libdevice.erf`
- Accumulation should use float32 for numerical stability
- Bias addition should be done after matrix multiplication but before GELU

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


tag: security
{
  "dependencies": {
    "uv_project": "resources"
  },
  "datasets": [
    "arvo:47101"
  ],
  "tag": "security"
}


QKNorm Optimization Problem
============================

Problem Setting
---------------
Design and optimize high-performance implementations for Query-Key Normalization (QKNorm) on GPU. This problem focuses on implementing efficient normalization kernels that apply RMSNorm to query and key tensors.

This is a **memory-bound** (even **launch-bound**) **tiny operator**. Performance optimization requires careful attention to:

1. **Memory Efficiency**: Focus on **vectorized memory access patterns**. Minimize memory transactions and maximize memory bandwidth utilization.

2. **Operation Fusion**: **Avoid additional transpose/contiguous kernels**. Fuse operations to reduce kernel launch overhead and memory traffic.

3. **Non-Contiguous Input Handling**: **Be aware that inputs may be non-contiguous** due to weight-QKV fusion. Your implementation should efficiently handle non-contiguous memory layouts without triggering expensive memory copies.

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

API Specification
-----------------
Implement a `Solution` class that returns a qknorm 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 flashinfer

def qknorm(q: torch.Tensor, k: torch.Tensor, norm_weight: torch.Tensor):
    """
    Apply RMSNorm to query and key tensors.

    Args:
        q: Query tensor of arbitrary shape (will be reshaped to 2D)
        k: Key tensor of arbitrary shape (will be reshaped to 2D)
        norm_weight: Normalization weight tensor of shape (hidden_dim,)

    Returns:
        Tuple of (q_normalized, k_normalized) tensors
    """
    pass
```

Required Default Implementation:
```python
def default_qknorm(q: torch.Tensor, k: torch.Tensor, norm_weight: torch.Tensor):
    q_2d = q.contiguous().view(-1, q.shape[-1])
    k_2d = k.contiguous().view(-1, k.shape[-1])
    q_o = torch.empty_like(q_2d)
    k_o = torch.empty_like(k_2d)
    flashinfer.norm.rmsnorm(q_2d, norm_weight, out=q_o)
    flashinfer.norm.rmsnorm(k_2d, norm_weight, out=k_o)
    return q_o.view(q.shape), k_o.view(k.shape)
```

Baseline Implementation:
```python
def customized_qknorm(q: torch.Tensor, k: torch.Tensor, norm_weight: torch.Tensor):
    q_o = torch.empty(q.shape, device=q.device, dtype=q.dtype)
    k_o = torch.empty(k.shape, device=k.device, dtype=k.dtype)
    flashinfer.norm.rmsnorm(q, norm_weight, out=q_o)
    flashinfer.norm.rmsnorm(k, norm_weight, out=k_o)
    return q_o, k_o
```

API Usage Notes
---------------
- The evaluator looks for a `qknorm` function in the module namespace
- Function must handle tensor reshaping correctly (q and k may have arbitrary shapes)
- Must use flashinfer.norm.rmsnorm for normalization
- Function returns a tuple of (q_normalized, k_normalized) tensors
- **Important**: Inputs q and k may be **non-contiguous** due to weight-QKV fusion
- **Avoid**: Additional `.contiguous()` or `.transpose()` calls that trigger memory copies
- **Focus**: Vectorized memory access and operation fusion to minimize kernel launches

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

```
geometric_mean_speedup = geometric_mean(answer_times / baseline_times)

if speedup < 0.5 or correctness is wrong:
    score = 0
elif speedup >= 0.5 and speedup < 1.0:
    score = 50
elif speedup >= 1.0:
    score = 100
```

- 0 points = Speedup < 0.5x OR correctness fails
- 50 points = Speedup >= 0.5x and < 1.0x
- 100 points = Speedup >= 1.0x

Evaluation Details
------------------
- Shapes focus on diverse batch-sizes, head-dim, num-kv-heads, num-qo-heads, e.g.:
  - (16, 8, 32, 128)
  - (128, 32, 32, 64)
- Correctness verified with tolerance: rtol=1e-2, atol=5e-3
- Performance measured using median execution time
- Requires CUDA backend and GPU support

CUDA 12.2, Python 3.11, PyTorch 2.0+, flashinfer 0.5.0, Triton 3.0+


Ragged Attention Optimization Problem
======================================

Problem Setting
---------------
Design and optimize high-performance Triton kernels for ragged attention computation on GPU. This problem focuses on implementing efficient kernels that handle variable-length sequences using ragged attention, where each query row can attend to a different number of key/value rows.

The challenge involves optimizing:
- **Ragged attention**: Efficiently handling variable-length sequences where each row has different attention lengths
- **Memory access patterns**: Efficient loading and storing of Q, K, V tensors with ragged masking
- **Streaming softmax**: Computing softmax in a streaming fashion for numerical stability
- **Row-wise masking**: Correctly masking attention scores based on row_lens
- **Mixed precision**: Handling float16 inputs/outputs with float32 accumulation
- **Block tiling**: Optimal block sizes for GPU execution across different matrix 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 matrix sizes and ragged lengths
- **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 ragged_attn(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, row_lens: torch.Tensor) -> torch.Tensor:
    """
    Ragged attention computation.

    Args:
        Q: Query tensor of shape (M, D) - query features (float16)
        K: Key tensor of shape (N, D) - key features (float16)
        V: Value tensor of shape (N, Dv) - value features (float16)
        row_lens: Row lengths tensor of shape (M,) - number of valid K/V rows per Q row (int32 or int64)

    Returns:
        Output tensor of shape (M, Dv) - attention output (float16)

    Semantics:
        For each query row i (0 <= i < M), compute attention over the first row_lens[i] key/value rows.
        Specifically:
        - scores[i, j] = (Q[i] @ K[j].T) * scale, for j < row_lens[i], else -inf
        - P[i] = softmax(scores[i])
        - O[i] = P[i] @ V[:row_lens[i]]
    """
    pass
```

Scoring
-------
The scoring system evaluates your implementation based on geometric mean speedup over GPU baseline:

- **0 points**: 1x GPU baseline (same speed as PyTorch GPU baseline)
- **100 points**: 3x GPU baseline (3x speedup over PyTorch GPU baseline)
- **Linear interpolation**: Scores between 0-100 are linearly interpolated based on speedup

The evaluation uses the following test cases:
- M (number of query rows): [512, 1024]
- N (number of key/value rows): 1024
- D (model dimension): 64
- Dv (value dimension): 64
- row_lens: Random integers between [min_ratio*N, N] where min_ratio=0.25

Correctness is verified using:
- Relative tolerance: 1e-2
- Absolute tolerance: 5e-3

All tests must pass for a non-zero score. If any test fails correctness, the score is 0.

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

@triton.jit
def _ragged_kernel(Q, K, V, O, ROW_LENS, ...):
    # Your kernel implementation
    pass

def ragged_attn(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, row_lens: torch.Tensor) -> torch.Tensor:
    # Your kernel launch logic
    pass
```

Constraints
-----------
- All tensors must be CUDA tensors (float16 for Q, K, V; int32/int64 for row_lens)
- Output must be float16
- The implementation must handle variable row lengths correctly
- Accumulation should use float32 for numerical stability
- Must use streaming softmax for numerical stability

Tips
----
1. Use efficient block tiling (BM, BN, BD, BDV) for optimal performance
2. Implement streaming softmax to handle large attention matrices
3. Correctly mask attention scores based on row_lens
4. Load row_lens once per program and broadcast for masking
5. Use proper masking for boundary conditions

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
  resources:
    accelerators: L4:1


Symbolic Regression Benchmark - McCormick Dataset
=================================================

Problem Setting
---------------
Learn a closed-form symbolic expression `f(x1, x2)` that predicts the target `y`.

This dataset is derived from the McCormick function, a classic 2D optimization test function featuring a combination of trigonometric and polynomial terms. The function exhibits a smooth, wavy surface with a global minimum.

Input Format
------------
- Your `Solution.solve` receives:
  - `X`: numpy.ndarray of shape `(n, 2)` containing feature values
  - `y`: numpy.ndarray of shape `(n,)` containing target values
- Dataset columns: `x1, x2, y`

Output Specification
--------------------
Implement a `Solution` class in `solution.py`:

```python
import numpy as np

class Solution:
    def __init__(self, **kwargs):
        pass

    def solve(self, X: np.ndarray, y: np.ndarray) -> dict:
        """
        Args:
            X: Feature matrix of shape (n, 2)
            y: Target values of shape (n,)

        Returns:
            dict with keys:
              - "expression": str, a Python-evaluable expression using x1, x2
              - "predictions": list/array of length n (optional)
              - "details": dict with optional "complexity" int
        """
        # Example: fit a symbolic expression to the data
        expression = "x1 + x2"  # placeholder
        return {
            "expression": expression,
            "predictions": None,  # will be computed from expression if omitted
            "details": {}
        }
```

Expression Requirements:
- Must be a valid Python expression string
- Use variable names: `x1`, `x2`
- Allowed operators: `+`, `-`, `*`, `/`, `**`
- Allowed functions: `sin`, `cos`, `exp`, `log`
- Numeric constants are allowed

Dependencies (pinned versions)
------------------------------
```
pysr==0.19.0
numpy==1.26.4
pandas==2.2.2
sympy==1.13.3
```

Minimal Working Examples
------------------------

**Example 1: Using PySR (recommended)**
```python
import numpy as np
from pysr import PySRRegressor

class Solution:
    def __init__(self, **kwargs):
        pass

    def solve(self, X: np.ndarray, y: np.ndarray) -> dict:
        model = PySRRegressor(
            niterations=40,
            binary_operators=["+", "-", "*", "/"],
            unary_operators=["sin", "cos", "exp", "log"],
            populations=15,
            population_size=33,
            maxsize=25,
            verbosity=0,
            progress=False,
            random_state=42,
        )
        model.fit(X, y, variable_names=["x1", "x2"])

        # Get best expression as sympy, convert to string
        best_expr = model.sympy()
        expression = str(best_expr)

        # Predictions
        predictions = model.predict(X)

        return {
            "expression": expression,
            "predictions": predictions.tolist(),
            "details": {}
        }
```

**Example 2: Manual expression (simple baseline)**
```python
import numpy as np

class Solution:
    def __init__(self, **kwargs):
        pass

    def solve(self, X: np.ndarray, y: np.ndarray) -> dict:
        # Simple linear combination as baseline
        x1, x2 = X[:, 0], X[:, 1]

        # Fit coefficients via least squares
        A = np.column_stack([x1, x2, np.ones_like(x1)])
        coeffs, _, _, _ = np.linalg.lstsq(A, y, rcond=None)
        a, b, c = coeffs

        expression = f"{a:.6f}*x1 + {b:.6f}*x2 + {c:.6f}"
        predictions = a * x1 + b * x2 + c

        return {
            "expression": expression,
            "predictions": predictions.tolist(),
            "details": {}
        }
```

**Example 3: Using sympy for expression manipulation**
```python
import numpy as np
import sympy as sp
from pysr import PySRRegressor

class Solution:
    def __init__(self, **kwargs):
        pass

    def solve(self, X: np.ndarray, y: np.ndarray) -> dict:
        model = PySRRegressor(
            niterations=30,
            binary_operators=["+", "-", "*", "/"],
            unary_operators=["sin", "cos"],
            verbosity=0,
            progress=False,
        )
        model.fit(X, y, variable_names=["x1", "x2"])

        # Get sympy expression and simplify
        sympy_expr = model.sympy()
        simplified = sp.simplify(sympy_expr)

        # Convert to evaluable string
        expression = str(simplified)

        return {
            "expression": expression,
            "predictions": None,  # evaluator will compute from expression
            "details": {}
        }
```

PySR API Notes (v0.19.0)
------------------------
- `model.fit(X, y, variable_names=["x1", "x2"])` - use variable_names to match expected output
- `model.sympy()` - returns best expression as sympy object
- `model.predict(X)` - returns predictions array
- `model.equations_` - DataFrame of all discovered equations
- Common parameters:
  - `niterations`: number of evolution iterations (more = better but slower)
  - `populations`: number of parallel populations
  - `maxsize`: maximum expression complexity
  - `verbosity=0, progress=False`: suppress output

Expression Format Requirements
------------------------------
- Must be a valid Python expression string
- Use variable names: `x1`, `x2`
- Allowed operators: `+`, `-`, `*`, `/`, `**`
- Allowed functions: `sin`, `cos`, `exp`, `log` (NO `np.` prefix)
- Numeric constants are allowed
- The evaluator uses `sympy.sympify()` to parse your expression

Scoring
-------
```
MSE = (1/n) Σ (y_i - ŷ_i)²
Score = 100 × clamp((m_base - MSE) / (m_base - m_ref), 0, 1) × 0.99^max(C - C_ref, 0)
```

- `m_base`: linear regression baseline MSE
- `m_ref`, `C_ref`: reference solution MSE and complexity
- `C = 2 × (#binary ops) + (#unary ops)`
- Lower MSE and lower complexity yield higher scores

Environment
-----------
Run `set_up_env.sh` to install dependencies.


Symbolic Regression Benchmark - Mixed PolyExp 4D Dataset
=========================================================

Problem Setting
---------------
Learn a closed-form symbolic expression `f(x1, x2, x3, x4)` that predicts the target `y`.

This is a higher-dimensional dataset (4 input features) combining polynomial interactions with exponential decay. The function involves cross-terms between variables and Gaussian-like damping, making it more challenging than the 2D variants.

Input Format
------------
- Your `Solution.solve` receives:
  - `X`: numpy.ndarray of shape `(n, 4)` containing feature values
  - `y`: numpy.ndarray of shape `(n,)` containing target values
- Dataset columns: `x1, x2, x3, x4, y`

Output Specification
--------------------
Implement a `Solution` class in `solution.py`:

```python
import numpy as np

class Solution:
    def __init__(self, **kwargs):
        pass

    def solve(self, X: np.ndarray, y: np.ndarray) -> dict:
        """
        Args:
            X: Feature matrix of shape (n, 4)
            y: Target values of shape (n,)

        Returns:
            dict with keys:
              - "expression": str, a Python-evaluable expression using x1, x2, x3, x4
              - "predictions": list/array of length n (optional)
              - "details": dict with optional "complexity" int
        """
        # Example: fit a symbolic expression to the data
        expression = "x1 + x2 + x3 + x4"  # placeholder
        return {
            "expression": expression,
            "predictions": None,  # will be computed from expression if omitted
            "details": {}
        }
```

Expression Requirements:
- Must be a valid Python expression string
- Use variable names: `x1`, `x2`, `x3`, `x4`
- Allowed operators: `+`, `-`, `*`, `/`, `**`
- Allowed functions: `sin`, `cos`, `exp`, `log`
- Numeric constants are allowed

Dependencies (pinned versions)
------------------------------
```
pysr==0.19.0
numpy==1.26.4
pandas==2.2.2
sympy==1.13.3
```

Minimal Working Examples
------------------------

**Example 1: Using PySR (recommended)**
```python
import numpy as np
from pysr import PySRRegressor

class Solution:
    def __init__(self, **kwargs):
        pass

    def solve(self, X: np.ndarray, y: np.ndarray) -> dict:
        model = PySRRegressor(
            niterations=50,  # more iterations for 4D
            binary_operators=["+", "-", "*", "/"],
            unary_operators=["sin", "cos", "exp", "log"],
            populations=20,
            population_size=40,
            maxsize=30,  # larger for 4D complexity
            verbosity=0,
            progress=False,
            random_state=42,
        )
        model.fit(X, y, variable_names=["x1", "x2", "x3", "x4"])

        # Get best expression as sympy, convert to string
        best_expr = model.sympy()
        expression = str(best_expr)

        # Predictions
        predictions = model.predict(X)

        return {
            "expression": expression,
            "predictions": predictions.tolist(),
            "details": {}
        }
```

**Example 2: Manual expression (simple baseline)**
```python
import numpy as np

class Solution:
    def __init__(self, **kwargs):
        pass

    def solve(self, X: np.ndarray, y: np.ndarray) -> dict:
        # Simple linear combination as baseline
        x1, x2, x3, x4 = X[:, 0], X[:, 1], X[:, 2], X[:, 3]

        # Fit coefficients via least squares
        A = np.column_stack([x1, x2, x3, x4, np.ones_like(x1)])
        coeffs, _, _, _ = np.linalg.lstsq(A, y, rcond=None)
        a, b, c, d, e = coeffs

        expression = f"{a:.6f}*x1 + {b:.6f}*x2 + {c:.6f}*x3 + {d:.6f}*x4 + {e:.6f}"
        predictions = a * x1 + b * x2 + c * x3 + d * x4 + e

        return {
            "expression": expression,
            "predictions": predictions.tolist(),
            "details": {}
        }
```

PySR API Notes (v0.19.0)
------------------------
- `model.fit(X, y, variable_names=["x1", "x2", "x3", "x4"])` - use variable_names to match expected output
- `model.sympy()` - returns best expression as sympy object
- `model.predict(X)` - returns predictions array
- `model.equations_` - DataFrame of all discovered equations
- Common parameters:
  - `niterations`: number of evolution iterations (more = better but slower)
  - `populations`: number of parallel populations
  - `maxsize`: maximum expression complexity
  - `verbosity=0, progress=False`: suppress output

Expression Format Requirements
------------------------------
- Must be a valid Python expression string
- Use variable names: `x1`, `x2`, `x3`, `x4`
- Allowed operators: `+`, `-`, `*`, `/`, `**`
- Allowed functions: `sin`, `cos`, `exp`, `log` (NO `np.` prefix)
- Numeric constants are allowed
- The evaluator uses `sympy.sympify()` to parse your expression

Scoring
-------
```
MSE = (1/n) Σ (y_i - ŷ_i)²
Score = 100 × clamp((m_base - MSE) / (m_base - m_ref), 0, 1) × 0.99^max(C - C_ref, 0)
```

- `m_base`: linear regression baseline MSE
- `m_ref`, `C_ref`: reference solution MSE and complexity
- `C = 2 × (#binary ops) + (#unary ops)`
- Lower MSE and lower complexity yield higher scores

Environment
-----------
Run `set_up_env.sh` to install dependencies.


Symbolic Regression Benchmark - Peaks Dataset
==============================================

Problem Setting
---------------
Learn a closed-form symbolic expression `f(x1, x2)` that predicts the target `y`.

This dataset is based on a peaks-like function, characterized by exponential terms that create localized peaks and valleys across the 2D input space. The underlying function involves interactions between polynomial and exponential components.

Input Format
------------
- Your `Solution.solve` receives:
  - `X`: numpy.ndarray of shape `(n, 2)` containing feature values
  - `y`: numpy.ndarray of shape `(n,)` containing target values
- Dataset columns: `x1, x2, y`

Output Specification
--------------------
Implement a `Solution` class in `solution.py`:

```python
import numpy as np

class Solution:
    def __init__(self, **kwargs):
        pass

    def solve(self, X: np.ndarray, y: np.ndarray) -> dict:
        """
        Args:
            X: Feature matrix of shape (n, 2)
            y: Target values of shape (n,)

        Returns:
            dict with keys:
              - "expression": str, a Python-evaluable expression using x1, x2
              - "predictions": list/array of length n (optional)
              - "details": dict with optional "complexity" int
        """
        # Example: fit a symbolic expression to the data
        expression = "x1 + x2"  # placeholder
        return {
            "expression": expression,
            "predictions": None,  # will be computed from expression if omitted
            "details": {}
        }
```

Expression Requirements:
- Must be a valid Python expression string
- Use variable names: `x1`, `x2`
- Allowed operators: `+`, `-`, `*`, `/`, `**`
- Allowed functions: `sin`, `cos`, `exp`, `log`
- Numeric constants are allowed

Dependencies (pinned versions)
------------------------------
```
pysr==0.19.0
numpy==1.26.4
pandas==2.2.2
sympy==1.13.3
```

Minimal Working Examples
------------------------

**Example 1: Using PySR (recommended)**
```python
import numpy as np
from pysr import PySRRegressor

class Solution:
    def __init__(self, **kwargs):
        pass

    def solve(self, X: np.ndarray, y: np.ndarray) -> dict:
        model = PySRRegressor(
            niterations=40,
            binary_operators=["+", "-", "*", "/"],
            unary_operators=["sin", "cos", "exp", "log"],
            populations=15,
            population_size=33,
            maxsize=25,
            verbosity=0,
            progress=False,
            random_state=42,
        )
        model.fit(X, y, variable_names=["x1", "x2"])

        # Get best expression as sympy, convert to string
        best_expr = model.sympy()
        expression = str(best_expr)

        # Predictions
        predictions = model.predict(X)

        return {
            "expression": expression,
            "predictions": predictions.tolist(),
            "details": {}
        }
```

**Example 2: Manual expression (simple baseline)**
```python
import numpy as np

class Solution:
    def __init__(self, **kwargs):
        pass

    def solve(self, X: np.ndarray, y: np.ndarray) -> dict:
        # Simple linear combination as baseline
        x1, x2 = X[:, 0], X[:, 1]

        # Fit coefficients via least squares
        A = np.column_stack([x1, x2, np.ones_like(x1)])
        coeffs, _, _, _ = np.linalg.lstsq(A, y, rcond=None)
        a, b, c = coeffs

        expression = f"{a:.6f}*x1 + {b:.6f}*x2 + {c:.6f}"
        predictions = a * x1 + b * x2 + c

        return {
            "expression": expression,
            "predictions": predictions.tolist(),
            "details": {}
        }
```

PySR API Notes (v0.19.0)
------------------------
- `model.fit(X, y, variable_names=["x1", "x2"])` - use variable_names to match expected output
- `model.sympy()` - returns best expression as sympy object
- `model.predict(X)` - returns predictions array
- `model.equations_` - DataFrame of all discovered equations
- Common parameters:
  - `niterations`: number of evolution iterations (more = better but slower)
  - `populations`: number of parallel populations
  - `maxsize`: maximum expression complexity
  - `verbosity=0, progress=False`: suppress output

Expression Format Requirements
------------------------------
- Must be a valid Python expression string
- Use variable names: `x1`, `x2`
- Allowed operators: `+`, `-`, `*`, `/`, `**`
- Allowed functions: `sin`, `cos`, `exp`, `log` (NO `np.` prefix)
- Numeric constants are allowed
- The evaluator uses `sympy.sympify()` to parse your expression

Scoring
-------
```
MSE = (1/n) Σ (y_i - ŷ_i)²
Score = 100 × clamp((m_base - MSE) / (m_base - m_ref), 0, 1) × 0.99^max(C - C_ref, 0)
```

- `m_base`: linear regression baseline MSE
- `m_ref`, `C_ref`: reference solution MSE and complexity
- `C = 2 × (#binary ops) + (#unary ops)`
- Lower MSE and lower complexity yield higher scores

Environment
-----------
Run `set_up_env.sh` to install dependencies.


Symbolic Regression Benchmark - Ripple Dataset
===============================================

Problem Setting
---------------
Learn a closed-form symbolic expression `f(x1, x2)` that predicts the target `y`.

This dataset is generated from a ripple-like function that combines polynomial amplitude modulation with high-frequency trigonometric oscillations. The function creates concentric wave patterns with varying intensity across the domain.

Input Format
------------
- Your `Solution.solve` receives:
  - `X`: numpy.ndarray of shape `(n, 2)` containing feature values
  - `y`: numpy.ndarray of shape `(n,)` containing target values
- Dataset columns: `x1, x2, y`

Output Specification
--------------------
Implement a `Solution` class in `solution.py`:

```python
import numpy as np

class Solution:
    def __init__(self, **kwargs):
        pass

    def solve(self, X: np.ndarray, y: np.ndarray) -> dict:
        """
        Args:
            X: Feature matrix of shape (n, 2)
            y: Target values of shape (n,)

        Returns:
            dict with keys:
              - "expression": str, a Python-evaluable expression using x1, x2
              - "predictions": list/array of length n (optional)
              - "details": dict with optional "complexity" int
        """
        # Example: fit a symbolic expression to the data
        expression = "x1 + x2"  # placeholder
        return {
            "expression": expression,
            "predictions": None,  # will be computed from expression if omitted
            "details": {}
        }
```

Expression Requirements:
- Must be a valid Python expression string
- Use variable names: `x1`, `x2`
- Allowed operators: `+`, `-`, `*`, `/`, `**`
- Allowed functions: `sin`, `cos`, `exp`, `log`
- Numeric constants are allowed

Dependencies (pinned versions)
------------------------------
```
pysr==0.19.0
numpy==1.26.4
pandas==2.2.2
sympy==1.13.3
```

Minimal Working Examples
------------------------

**Example 1: Using PySR (recommended)**
```python
import numpy as np
from pysr import PySRRegressor

class Solution:
    def __init__(self, **kwargs):
        pass

    def solve(self, X: np.ndarray, y: np.ndarray) -> dict:
        model = PySRRegressor(
            niterations=40,
            binary_operators=["+", "-", "*", "/"],
            unary_operators=["sin", "cos", "exp", "log"],
            populations=15,
            population_size=33,
            maxsize=25,
            verbosity=0,
            progress=False,
            random_state=42,
        )
        model.fit(X, y, variable_names=["x1", "x2"])

        # Get best expression as sympy, convert to string
        best_expr = model.sympy()
        expression = str(best_expr)

        # Predictions
        predictions = model.predict(X)

        return {
            "expression": expression,
            "predictions": predictions.tolist(),
            "details": {}
        }
```

**Example 2: Manual expression (simple baseline)**
```python
import numpy as np

class Solution:
    def __init__(self, **kwargs):
        pass

    def solve(self, X: np.ndarray, y: np.ndarray) -> dict:
        # Simple linear combination as baseline
        x1, x2 = X[:, 0], X[:, 1]

        # Fit coefficients via least squares
        A = np.column_stack([x1, x2, np.ones_like(x1)])
        coeffs, _, _, _ = np.linalg.lstsq(A, y, rcond=None)
        a, b, c = coeffs

        expression = f"{a:.6f}*x1 + {b:.6f}*x2 + {c:.6f}"
        predictions = a * x1 + b * x2 + c

        return {
            "expression": expression,
            "predictions": predictions.tolist(),
            "details": {}
        }
```

PySR API Notes (v0.19.0)
------------------------
- `model.fit(X, y, variable_names=["x1", "x2"])` - use variable_names to match expected output
- `model.sympy()` - returns best expression as sympy object
- `model.predict(X)` - returns predictions array
- `model.equations_` - DataFrame of all discovered equations
- Common parameters:
  - `niterations`: number of evolution iterations (more = better but slower)
  - `populations`: number of parallel populations
  - `maxsize`: maximum expression complexity
  - `verbosity=0, progress=False`: suppress output

Expression Format Requirements
------------------------------
- Must be a valid Python expression string
- Use variable names: `x1`, `x2`
- Allowed operators: `+`, `-`, `*`, `/`, `**`
- Allowed functions: `sin`, `cos`, `exp`, `log` (NO `np.` prefix)
- Numeric constants are allowed
- The evaluator uses `sympy.sympify()` to parse your expression

Scoring
-------
```
MSE = (1/n) Σ (y_i - ŷ_i)²
Score = 100 × clamp((m_base - MSE) / (m_base - m_ref), 0, 1) × 0.99^max(C - C_ref, 0)
```

- `m_base`: linear regression baseline MSE
- `m_ref`, `C_ref`: reference solution MSE and complexity
- `C = 2 × (#binary ops) + (#unary ops)`
- Lower MSE and lower complexity yield higher scores

Environment
-----------
Run `set_up_env.sh` to install dependencies.


Symbolic Regression Benchmark - SinCos Dataset
===============================================

Problem Setting
---------------
Learn a closed-form symbolic expression `f(x1, x2)` that predicts the target `y`.

This dataset features a function built from basic trigonometric operations. The target exhibits periodic behavior in both input dimensions, representing a straightforward but fundamental pattern for symbolic regression.

Input Format
------------
- Your `Solution.solve` receives:
  - `X`: numpy.ndarray of shape `(n, 2)` containing feature values
  - `y`: numpy.ndarray of shape `(n,)` containing target values
- Dataset columns: `x1, x2, y`

Output Specification
--------------------
Implement a `Solution` class in `solution.py`:

```python
import numpy as np

class Solution:
    def __init__(self, **kwargs):
        pass

    def solve(self, X: np.ndarray, y: np.ndarray) -> dict:
        """
        Args:
            X: Feature matrix of shape (n, 2)
            y: Target values of shape (n,)

        Returns:
            dict with keys:
              - "expression": str, a Python-evaluable expression using x1, x2
              - "predictions": list/array of length n (optional)
              - "details": dict with optional "complexity" int
        """
        # Example: fit a symbolic expression to the data
        expression = "x1 + x2"  # placeholder
        return {
            "expression": expression,
            "predictions": None,  # will be computed from expression if omitted
            "details": {}
        }
```

Expression Requirements:
- Must be a valid Python expression string
- Use variable names: `x1`, `x2`
- Allowed operators: `+`, `-`, `*`, `/`, `**`
- Allowed functions: `sin`, `cos`, `exp`, `log`
- Numeric constants are allowed

Dependencies (pinned versions)
------------------------------
```
pysr==0.19.0
numpy==1.26.4
pandas==2.2.2
sympy==1.13.3
```

Minimal Working Examples
------------------------

**Example 1: Using PySR (recommended)**
```python
import numpy as np
from pysr import PySRRegressor

class Solution:
    def __init__(self, **kwargs):
        pass

    def solve(self, X: np.ndarray, y: np.ndarray) -> dict:
        model = PySRRegressor(
            niterations=40,
            binary_operators=["+", "-", "*", "/"],
            unary_operators=["sin", "cos", "exp", "log"],
            populations=15,
            population_size=33,
            maxsize=25,
            verbosity=0,
            progress=False,
            random_state=42,
        )
        model.fit(X, y, variable_names=["x1", "x2"])

        # Get best expression as sympy, convert to string
        best_expr = model.sympy()
        expression = str(best_expr)

        # Predictions
        predictions = model.predict(X)

        return {
            "expression": expression,
            "predictions": predictions.tolist(),
            "details": {}
        }
```

**Example 2: Manual expression (simple baseline)**
```python
import numpy as np

class Solution:
    def __init__(self, **kwargs):
        pass

    def solve(self, X: np.ndarray, y: np.ndarray) -> dict:
        # Simple linear combination as baseline
        x1, x2 = X[:, 0], X[:, 1]

        # Fit coefficients via least squares
        A = np.column_stack([x1, x2, np.ones_like(x1)])
        coeffs, _, _, _ = np.linalg.lstsq(A, y, rcond=None)
        a, b, c = coeffs

        expression = f"{a:.6f}*x1 + {b:.6f}*x2 + {c:.6f}"
        predictions = a * x1 + b * x2 + c

        return {
            "expression": expression,
            "predictions": predictions.tolist(),
            "details": {}
        }
```

PySR API Notes (v0.19.0)
------------------------
- `model.fit(X, y, variable_names=["x1", "x2"])` - use variable_names to match expected output
- `model.sympy()` - returns best expression as sympy object
- `model.predict(X)` - returns predictions array
- `model.equations_` - DataFrame of all discovered equations
- Common parameters:
  - `niterations`: number of evolution iterations (more = better but slower)
  - `populations`: number of parallel populations
  - `maxsize`: maximum expression complexity
  - `verbosity=0, progress=False`: suppress output

Expression Format Requirements
------------------------------
- Must be a valid Python expression string
- Use variable names: `x1`, `x2`
- Allowed operators: `+`, `-`, `*`, `/`, `**`
- Allowed functions: `sin`, `cos`, `exp`, `log` (NO `np.` prefix)
- Numeric constants are allowed
- The evaluator uses `sympy.sympify()` to parse your expression

Scoring
-------
```
MSE = (1/n) Σ (y_i - ŷ_i)²
Score = 100 × clamp((m_base - MSE) / (m_base - m_ref), 0, 1) × 0.99^max(C - C_ref, 0)
```

- `m_base`: linear regression baseline MSE
- `m_ref`, `C_ref`: reference solution MSE and complexity
- `C = 2 × (#binary ops) + (#unary ops)`
- Lower MSE and lower complexity yield higher scores

Environment
-----------
Run `set_up_env.sh` to install dependencies.


VDB Design Problem - Balanced Tier
===================================

Problem Setting
---------------
Design a Vector Database index optimized for **recall** subject to a **latency constraint**. This tier uses latency-gated scoring: solutions exceeding the latency threshold receive zero points, while solutions meeting the constraint are scored purely by recall@1.

**Optimization Goal**: Maximize recall@1 within latency constraint

$$
\text{score} = \begin{cases}
0 & \text{if } t_{\text{query}} > t_{\text{max}} \\
100 & \text{if } t_{\text{query}} \leq t_{\text{max}} \text{ and } r \geq r_{\text{baseline}} \\
100 \cdot \frac{r - r_{\text{min}}}{r_{\text{baseline}} - r_{\text{min}}} & \text{if } t_{\text{query}} \leq t_{\text{max}} \text{ and } r < r_{\text{baseline}}
\end{cases}
$$

Where:
- $r$: Your recall@1
- $t_{\text{query}}$: Your average query latency (ms)
- $r_{\text{baseline}} = 0.9914$ (baseline recall)
- $r_{\text{min}} = 0.6939$ (minimum acceptable recall, 70% of baseline)
- $t_{\text{max}} = 5.775\text{ms}$ (maximum allowed latency, 150% of baseline 3.85ms)

**Key Insight**: Latency is a hard constraint. Only recall determines your score within the constraint.

Baseline Performance
--------------------
- Recall@1: **0.9914** (99.14%)
- Avg query time: **3.85ms**
- Baseline score: **100** (recall equals baseline within latency constraint)

Scoring Examples
----------------
Assuming all solutions meet latency constraint ($t \leq 5.775\text{ms}$):

| Recall@1 | Latency | Score Calculation | Score |
|----------|---------|-------------------|-------|
| 0.9914 | 3.85ms | $r = r_{\text{baseline}}$ → max score | **100** |
| 0.9950 | 3.00ms | $r > r_{\text{baseline}}$ → max score | **100** |
| 0.9500 | 2.50ms | $\frac{0.95 - 0.6939}{0.9914 - 0.6939} = 0.860$ | **86.0** |
| 0.8500 | 4.00ms | $\frac{0.85 - 0.6939}{0.9914 - 0.6939} = 0.524$ | **52.4** |
| 0.6939 | 5.00ms | $r = r_{\text{min}}$ → minimum score | **0** |
| 0.9900 | **6.00ms** | $t > t_{\text{max}}$ → latency gate fails | **0** |

**Note**: Faster latency does NOT increase score - only recall matters if constraint is met.

API Specification
-----------------
Implement a class with the following interface:

```python
import numpy as np
from typing import Tuple

class YourIndexClass:
    def __init__(self, dim: int, **kwargs):
        """
        Initialize the index for vectors of dimension `dim`.

        Args:
            dim: Vector dimensionality (e.g., 128 for SIFT1M)
            **kwargs: Optional parameters (e.g., M, ef_construction for HNSW)

        Example:
            index = YourIndexClass(dim=128, M=16, ef_search=64)
        """
        pass

    def add(self, xb: np.ndarray) -> None:
        """
        Add vectors to the index.

        Args:
            xb: Base vectors, shape (N, dim), dtype float32

        Notes:
            - Can be called multiple times (cumulative)
            - Must handle large N (e.g., 1,000,000 vectors)

        Example:
            index.add(xb)  # xb.shape = (1000000, 128)
        """
        pass

    def search(self, xq: np.ndarray, k: int) -> Tuple[np.ndarray, np.ndarray]:
        """
        Search for k nearest neighbors of query vectors.

        Args:
            xq: Query vectors, shape (nq, dim), dtype float32
            k: Number of nearest neighbors to return

        Returns:
            (distances, indices):
                - distances: shape (nq, k), dtype float32, L2 distances
                - indices: shape (nq, k), dtype int64, indices into base vectors

        Notes:
            - Must return exactly k neighbors per query
            - Indices should refer to positions in the vectors passed to add()
            - Lower distance = more similar

        Example:
            D, I = index.search(xq, k=1)  # xq.shape = (10000, 128)
            # D.shape = (10000, 1), I.shape = (10000, 1)
        """
        pass
```

**Implementation Requirements**:
- Class can have any name (evaluator auto-discovers classes with `add` and `search` methods)
- Must handle SIFT1M dataset: 1M base vectors, 10K queries, 128 dimensions
- Your `search` must return tuple `(distances, indices)` with shapes `(nq, k)`
- Distances should be L2 (Euclidean) or L2-squared
- No need to handle dataset loading - evaluator provides numpy arrays

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

### 1. Load Dataset
```python
from faiss.contrib.datasets import DatasetSIFT1M
ds = DatasetSIFT1M()
xb = ds.get_database()        # (1000000, 128) float32
xq = ds.get_queries()         # (10000, 128) float32
gt = ds.get_groundtruth()     # (10000, 100) int64 - ground truth indices
```

### 2. Build Index
```python
from solution import YourIndexClass  # Auto-discovered
d = xb.shape[1]                       # 128 for SIFT1M
index = YourIndexClass(d)             # Pass dimension as first argument
index.add(xb)                         # Add all 1M base vectors
```

### 3. Measure Performance (Batch Queries)
```python
import time
t0 = time.time()
D, I = index.search(xq, k=1)          # Search all 10K queries at once
t1 = time.time()

# Calculate metrics
recall_at_1 = (I[:, :1] == gt[:, :1]).sum() / len(xq)
avg_query_time_ms = (t1 - t0) * 1000.0 / len(xq)
```

**Important**: `avg_query_time_ms` from **batch queries** is used for scoring. Batch queries benefit from CPU cache and vectorization, typically faster than single queries.

### 4. Calculate Score
```python
if avg_query_time_ms > 5.775:
    score = 0.0
elif recall_at_1 >= 0.9914:
    score = 100.0
else:
    recall_range = 0.9914 - 0.6939
    recall_proportion = (recall_at_1 - 0.6939) / recall_range
    score = max(0.0, min(100.0, 100.0 * recall_proportion))
```

Dataset Details
---------------
- **Name**: SIFT1M
- **Base vectors**: 1,000,000 vectors of dimension 128
- **Query vectors**: 10,000 vectors
- **Ground truth**: Precomputed nearest neighbors (k=1)
- **Metric**: L2 (Euclidean distance)
- **Vector type**: float32

Runtime Platform
----------------
- **Infrastructure**: Evaluations run on SkyPilot-managed cloud instances (AWS, GCP, or Azure)
- **Compute**: CPU-only instances (no GPU required)
- **Environment**: Docker containerized execution with Python 3, NumPy ≥1.24, FAISS-CPU ≥1.7.4

Constraints
-----------
- **Timeout**: 1 hour for entire evaluation (index construction + queries)
- **Memory**: Use reasonable memory (index should fit in RAM)
- **Latency constraint**: avg_query_time_ms ≤ 5.775ms
- **Recall range**: 0.6939 ≤ recall@1 ≤ 1.0

Strategy Tips
-------------
1. **Focus on recall**: Latency only needs to meet threshold, doesn't improve score beyond that
2. **Batch optimization is key**: Your `search` should handle batch queries efficiently
3. **Parameter tuning**: Small changes (e.g., HNSW's M, ef_search) significantly affect recall
4. **Don't over-optimize latency**: Meeting 5.775ms is enough; focus energy on recall

Example: Simple Baseline
-------------------------
```python
import numpy as np

class SimpleIndex:
    def __init__(self, dim: int, **kwargs):
        self.dim = dim
        self.xb = None

    def add(self, xb: np.ndarray) -> None:
        if self.xb is None:
            self.xb = xb.copy()
        else:
            self.xb = np.vstack([self.xb, xb])

    def search(self, xq: np.ndarray, k: int) -> tuple:
        # Compute all pairwise L2 distances
        # xq: (nq, dim), xb: (N, dim)
        # distances: (nq, N)
        distances = np.sqrt(((xq[:, np.newaxis, :] - self.xb[np.newaxis, :, :]) ** 2).sum(axis=2))

        # Get k nearest neighbors
        indices = np.argpartition(distances, k-1, axis=1)[:, :k]
        sorted_indices = np.argsort(distances[np.arange(len(xq))[:, None], indices], axis=1)
        final_indices = indices[np.arange(len(xq))[:, None], sorted_indices]
        final_distances = distances[np.arange(len(xq))[:, None], final_indices]

        return final_distances, final_indices
```

**Note**: This baseline achieves perfect recall (100%) but is too slow for large datasets. Use approximate methods like HNSW, IVF, or LSH for better speed-recall tradeoffs.

Debugging Tips
--------------
- **Test locally**: Use a subset of data (e.g., 10K vectors) for faster iteration
- **Verify shapes**: Ensure `search` returns `(nq, k)` shaped arrays
- **Check recall calculation**: `(I[:, :1] == gt[:, :1]).sum() / len(xq)`
- **Profile latency**: Measure batch vs single query performance separately
- **Validate before submit**: Run full 1M dataset locally if possible


VDB Design Problem - High Recall Tier
======================================

Problem Setting
---------------
Design a Vector Database index optimized for **recall** subject to a **relaxed latency constraint**. This tier uses latency-gated scoring: solutions exceeding the latency threshold receive zero points, while solutions meeting the constraint are scored purely by recall@1.

**Optimization Goal**: Maximize recall@1 within latency constraint

$$
\text{score} = \begin{cases}
0 & \text{if } t_{\text{query}} > t_{\text{max}} \\
100 & \text{if } t_{\text{query}} \leq t_{\text{max}} \text{ and } r \geq r_{\text{baseline}} \\
100 \cdot \frac{r - r_{\text{min}}}{r_{\text{baseline}} - r_{\text{min}}} & \text{if } t_{\text{query}} \leq t_{\text{max}} \text{ and } r < r_{\text{baseline}}
\end{cases}
$$

Where:
- $r$: Your recall@1
- $t_{\text{query}}$: Your average query latency (ms)
- $r_{\text{baseline}} = 0.9914$ (baseline recall)
- $r_{\text{min}} = 0.9409$ (minimum acceptable recall, 95% of baseline)
- $t_{\text{max}} = 7.7\text{ms}$ (maximum allowed latency, 200% of baseline 3.85ms)

**Key Insight**: This tier provides 2× latency budget compared to balanced tier, allowing more thorough search for higher recall.

Baseline Performance
--------------------
- Recall@1: **0.9914** (99.14%)
- Avg query time: **3.85ms**
- Baseline score: **100** (recall equals baseline within latency constraint)

Scoring Examples
----------------
Assuming all solutions meet latency constraint ($t \leq 7.7\text{ms}$):

| Recall@1 | Latency | Score Calculation | Score |
|----------|---------|-------------------|-------|
| 0.9914 | 3.85ms | $r = r_{\text{baseline}}$ → max score | **100** |
| 0.9950 | 5.00ms | $r > r_{\text{baseline}}$ → max score | **100** |
| 0.9700 | 6.00ms | $\frac{0.97 - 0.9409}{0.9914 - 0.9409} = 0.576$ | **57.6** |
| 0.9500 | 4.00ms | $\frac{0.95 - 0.9409}{0.9914 - 0.9409} = 0.180$ | **18.0** |
| 0.9409 | 7.00ms | $r = r_{\text{min}}$ → minimum score | **0** |
| 0.9914 | **8.00ms** | $t > t_{\text{max}}$ → latency gate fails | **0** |

**Note**: The relaxed latency constraint (7.7ms vs 5.775ms in balanced) allows more aggressive search strategies for higher recall.

API Specification
-----------------
Implement a class with the following interface:

```python
import numpy as np
from typing import Tuple

class YourIndexClass:
    def __init__(self, dim: int, **kwargs):
        """
        Initialize the index for vectors of dimension `dim`.

        Args:
            dim: Vector dimensionality (e.g., 128 for SIFT1M)
            **kwargs: Optional parameters (e.g., M, ef_construction for HNSW)

        Example:
            index = YourIndexClass(dim=128, M=64, ef_search=800)
        """
        pass

    def add(self, xb: np.ndarray) -> None:
        """
        Add vectors to the index.

        Args:
            xb: Base vectors, shape (N, dim), dtype float32

        Notes:
            - Can be called multiple times (cumulative)
            - Must handle large N (e.g., 1,000,000 vectors)

        Example:
            index.add(xb)  # xb.shape = (1000000, 128)
        """
        pass

    def search(self, xq: np.ndarray, k: int) -> Tuple[np.ndarray, np.ndarray]:
        """
        Search for k nearest neighbors of query vectors.

        Args:
            xq: Query vectors, shape (nq, dim), dtype float32
            k: Number of nearest neighbors to return

        Returns:
            (distances, indices):
                - distances: shape (nq, k), dtype float32, L2 distances
                - indices: shape (nq, k), dtype int64, indices into base vectors

        Notes:
            - Must return exactly k neighbors per query
            - Indices should refer to positions in the vectors passed to add()
            - Lower distance = more similar

        Example:
            D, I = index.search(xq, k=1)  # xq.shape = (10000, 128)
            # D.shape = (10000, 1), I.shape = (10000, 1)
        """
        pass
```

**Implementation Requirements**:
- Class can have any name (evaluator auto-discovers classes with `add` and `search` methods)
- Must handle SIFT1M dataset: 1M base vectors, 10K queries, 128 dimensions
- Your `search` must return tuple `(distances, indices)` with shapes `(nq, k)`
- Distances should be L2 (Euclidean) or L2-squared
- No need to handle dataset loading - evaluator provides numpy arrays

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

### 1. Load Dataset
```python
from faiss.contrib.datasets import DatasetSIFT1M
ds = DatasetSIFT1M()
xb = ds.get_database()        # (1000000, 128) float32
xq = ds.get_queries()         # (10000, 128) float32
gt = ds.get_groundtruth()     # (10000, 100) int64 - ground truth indices
```

### 2. Build Index
```python
from solution import YourIndexClass  # Auto-discovered
d = xb.shape[1]                       # 128 for SIFT1M
index = YourIndexClass(d)             # Pass dimension as first argument
index.add(xb)                         # Add all 1M base vectors
```

### 3. Measure Performance (Batch Queries)
```python
import time
t0 = time.time()
D, I = index.search(xq, k=1)          # Search all 10K queries at once
t1 = time.time()

# Calculate metrics
recall_at_1 = (I[:, :1] == gt[:, :1]).sum() / len(xq)
avg_query_time_ms = (t1 - t0) * 1000.0 / len(xq)
```

**Important**: `avg_query_time_ms` from **batch queries** is used for scoring. Batch queries benefit from CPU cache and vectorization, typically faster than single queries.

### 4. Calculate Score
```python
if avg_query_time_ms > 7.7:
    score = 0.0
elif recall_at_1 >= 0.9914:
    score = 100.0
else:
    recall_range = 0.9914 - 0.9409
    recall_proportion = (recall_at_1 - 0.9409) / recall_range
    score = max(0.0, min(100.0, 100.0 * recall_proportion))
```

Dataset Details
---------------
- **Name**: SIFT1M
- **Base vectors**: 1,000,000 vectors of dimension 128
- **Query vectors**: 10,000 vectors
- **Ground truth**: Precomputed nearest neighbors (k=1)
- **Metric**: L2 (Euclidean distance)
- **Vector type**: float32

Runtime Platform
----------------
- **Infrastructure**: Evaluations run on SkyPilot-managed cloud instances (AWS, GCP, or Azure)
- **Compute**: CPU-only instances (no GPU required)
- **Environment**: Docker containerized execution with Python 3, NumPy ≥1.24, FAISS-CPU ≥1.7.4

Constraints
-----------
- **Timeout**: 1 hour for entire evaluation (index construction + queries)
- **Memory**: Use reasonable memory (index should fit in RAM)
- **Latency constraint**: avg_query_time_ms ≤ 7.7ms
- **Recall range**: 0.9409 ≤ recall@1 ≤ 1.0

Strategy Tips
-------------
1. **Maximize recall**: Use 2× latency budget (7.7ms vs 5.775ms balanced) for more thorough search
2. **Batch optimization is key**: Your `search` should handle batch queries efficiently
3. **Parameter tuning for recall**: Higher HNSW efSearch (500-1000) or IVF nprobe (100-200)
4. **Trade latency for accuracy**: Unlike balanced tier, you can afford slower but more accurate search

Example: Simple Baseline
-------------------------
```python
import numpy as np

class SimpleIndex:
    def __init__(self, dim: int, **kwargs):
        self.dim = dim
        self.xb = None

    def add(self, xb: np.ndarray) -> None:
        if self.xb is None:
            self.xb = xb.copy()
        else:
            self.xb = np.vstack([self.xb, xb])

    def search(self, xq: np.ndarray, k: int) -> tuple:
        # Compute all pairwise L2 distances
        # xq: (nq, dim), xb: (N, dim)
        # distances: (nq, N)
        distances = np.sqrt(((xq[:, np.newaxis, :] - self.xb[np.newaxis, :, :]) ** 2).sum(axis=2))

        # Get k nearest neighbors
        indices = np.argpartition(distances, k-1, axis=1)[:, :k]
        sorted_indices = np.argsort(distances[np.arange(len(xq))[:, None], indices], axis=1)
        final_indices = indices[np.arange(len(xq))[:, None], sorted_indices]
        final_distances = distances[np.arange(len(xq))[:, None], final_indices]

        return final_distances, final_indices
```

**Note**: This baseline achieves perfect recall (100%) but is too slow for large datasets. Use approximate methods like HNSW, IVF, or LSH for better speed-recall tradeoffs.

Debugging Tips
--------------
- **Test locally**: Use a subset of data (e.g., 10K vectors) for faster iteration
- **Verify shapes**: Ensure `search` returns `(nq, k)` shaped arrays
- **Check recall calculation**: `(I[:, :1] == gt[:, :1]).sum() / len(xq)`
- **Profile latency**: Measure batch vs single query performance separately
- **Validate before submit**: Run full 1M dataset locally if possible


VDB Design Problem - Low Latency Tier
======================================

Problem Setting
---------------
Design a Vector Database index optimized for **recall** subject to a **strict latency constraint**. This tier uses latency-gated scoring: solutions exceeding the latency threshold receive zero points, while solutions meeting the constraint are scored purely by recall@1.

**Optimization Goal**: Maximize recall@1 within latency constraint

$$
\text{score} = \begin{cases}
0 & \text{if } t_{\text{query}} > t_{\text{max}} \\
100 & \text{if } t_{\text{query}} \leq t_{\text{max}} \text{ and } r \geq r_{\text{baseline}} \\
100 \cdot \frac{r - r_{\text{min}}}{r_{\text{baseline}} - r_{\text{min}}} & \text{if } t_{\text{query}} \leq t_{\text{max}} \text{ and } r < r_{\text{baseline}}
\end{cases}
$$

Where:
- $r$: Your recall@1
- $t_{\text{query}}$: Your average query latency (ms)
- $r_{\text{baseline}} = 0.9914$ (baseline recall)
- $r_{\text{min}} = 0.7931$ (minimum acceptable recall, 80% of baseline)
- $t_{\text{max}} = 2.31\text{ms}$ (maximum allowed latency, 60% of baseline 3.85ms)

**Key Insight**: This tier has a very strict latency constraint (60% of baseline), requiring aggressive approximation while maintaining reasonable recall.

Baseline Performance
--------------------
- Recall@1: **0.9914** (99.14%)
- Avg query time: **3.85ms**
- Baseline score: **100** (recall equals baseline within latency constraint)

Scoring Examples
----------------
Assuming all solutions meet latency constraint ($t \leq 2.31\text{ms}$):

| Recall@1 | Latency | Score Calculation | Score |
|----------|---------|-------------------|-------|
| 0.9914 | 2.00ms | $r = r_{\text{baseline}}$ → max score | **100** |
| 0.9500 | 2.00ms | $\frac{0.95 - 0.7931}{0.9914 - 0.7931} = 0.791$ | **79.1** |
| 0.9000 | 1.50ms | $\frac{0.90 - 0.7931}{0.9914 - 0.7931} = 0.539$ | **53.9** |
| 0.8500 | 1.00ms | $\frac{0.85 - 0.7931}{0.9914 - 0.7931} = 0.287$ | **28.7** |
| 0.7931 | 2.00ms | $r = r_{\text{min}}$ → minimum score | **0** |
| 0.9500 | **2.50ms** | $t > t_{\text{max}}$ → latency gate fails | **0** |

**Note**: The strict latency constraint (2.31ms vs 5.775ms in balanced) requires aggressive approximation, typically resulting in lower recall.

API Specification
-----------------
Implement a class with the following interface:

```python
import numpy as np
from typing import Tuple

class YourIndexClass:
    def __init__(self, dim: int, **kwargs):
        """
        Initialize the index for vectors of dimension `dim`.

        Args:
            dim: Vector dimensionality (e.g., 128 for SIFT1M)
            **kwargs: Optional parameters (e.g., M, ef_construction for HNSW)

        Example:
            index = YourIndexClass(dim=128, M=16, ef_search=80)
        """
        pass

    def add(self, xb: np.ndarray) -> None:
        """
        Add vectors to the index.

        Args:
            xb: Base vectors, shape (N, dim), dtype float32

        Notes:
            - Can be called multiple times (cumulative)
            - Must handle large N (e.g., 1,000,000 vectors)

        Example:
            index.add(xb)  # xb.shape = (1000000, 128)
        """
        pass

    def search(self, xq: np.ndarray, k: int) -> Tuple[np.ndarray, np.ndarray]:
        """
        Search for k nearest neighbors of query vectors.

        Args:
            xq: Query vectors, shape (nq, dim), dtype float32
            k: Number of nearest neighbors to return

        Returns:
            (distances, indices):
                - distances: shape (nq, k), dtype float32, L2 distances
                - indices: shape (nq, k), dtype int64, indices into base vectors

        Notes:
            - Must return exactly k neighbors per query
            - Indices should refer to positions in the vectors passed to add()
            - Lower distance = more similar

        Example:
            D, I = index.search(xq, k=1)  # xq.shape = (10000, 128)
            # D.shape = (10000, 1), I.shape = (10000, 1)
        """
        pass
```

**Implementation Requirements**:
- Class can have any name (evaluator auto-discovers classes with `add` and `search` methods)
- Must handle SIFT1M dataset: 1M base vectors, 10K queries, 128 dimensions
- Your `search` must return tuple `(distances, indices)` with shapes `(nq, k)`
- Distances should be L2 (Euclidean) or L2-squared
- No need to handle dataset loading - evaluator provides numpy arrays

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

### 1. Load Dataset
```python
from faiss.contrib.datasets import DatasetSIFT1M
ds = DatasetSIFT1M()
xb = ds.get_database()        # (1000000, 128) float32
xq = ds.get_queries()         # (10000, 128) float32
gt = ds.get_groundtruth()     # (10000, 100) int64 - ground truth indices
```

### 2. Build Index
```python
from solution import YourIndexClass  # Auto-discovered
d = xb.shape[1]                       # 128 for SIFT1M
index = YourIndexClass(d)             # Pass dimension as first argument
index.add(xb)                         # Add all 1M base vectors
```

### 3. Measure Performance (Batch Queries)
```python
import time
t0 = time.time()
D, I = index.search(xq, k=1)          # Search all 10K queries at once
t1 = time.time()

# Calculate metrics
recall_at_1 = (I[:, :1] == gt[:, :1]).sum() / len(xq)
avg_query_time_ms = (t1 - t0) * 1000.0 / len(xq)
```

**Important**: `avg_query_time_ms` from **batch queries** is used for scoring. Batch queries benefit from CPU cache and vectorization, typically faster than single queries.

### 4. Calculate Score
```python
if avg_query_time_ms > 2.31:
    score = 0.0
elif recall_at_1 >= 0.9914:
    score = 100.0
else:
    recall_range = 0.9914 - 0.7931
    recall_proportion = (recall_at_1 - 0.7931) / recall_range
    score = max(0.0, min(100.0, 100.0 * recall_proportion))
```

Dataset Details
---------------
- **Name**: SIFT1M
- **Base vectors**: 1,000,000 vectors of dimension 128
- **Query vectors**: 10,000 vectors
- **Ground truth**: Precomputed nearest neighbors (k=1)
- **Metric**: L2 (Euclidean distance)
- **Vector type**: float32

Runtime Platform
----------------
- **Infrastructure**: Evaluations run on SkyPilot-managed cloud instances (AWS, GCP, or Azure)
- **Compute**: CPU-only instances (no GPU required)
- **Environment**: Docker containerized execution with Python 3, NumPy ≥1.24, FAISS-CPU ≥1.7.4

Constraints
-----------
- **Timeout**: 1 hour for entire evaluation (index construction + queries)
- **Memory**: Use reasonable memory (index should fit in RAM)
- **Latency constraint**: avg_query_time_ms ≤ 2.31ms
- **Recall range**: 0.7931 ≤ recall@1 ≤ 1.0

Strategy Tips
-------------
1. **Aggressive approximation**: Use very low search budgets (IVF nprobe=2-5, HNSW efSearch=50-100)
2. **Batch optimization is key**: Your `search` should handle batch queries efficiently
3. **Accept recall drops**: 80-90% recall is acceptable if latency is met
4. **Leave safety margin**: Target 1.5-2.0ms to avoid edge cases exceeding 2.31ms

Example: Simple Baseline
-------------------------
```python
import numpy as np

class SimpleIndex:
    def __init__(self, dim: int, **kwargs):
        self.dim = dim
        self.xb = None

    def add(self, xb: np.ndarray) -> None:
        if self.xb is None:
            self.xb = xb.copy()
        else:
            self.xb = np.vstack([self.xb, xb])

    def search(self, xq: np.ndarray, k: int) -> tuple:
        # Compute all pairwise L2 distances
        # xq: (nq, dim), xb: (N, dim)
        # distances: (nq, N)
        distances = np.sqrt(((xq[:, np.newaxis, :] - self.xb[np.newaxis, :, :]) ** 2).sum(axis=2))

        # Get k nearest neighbors
        indices = np.argpartition(distances, k-1, axis=1)[:, :k]
        sorted_indices = np.argsort(distances[np.arange(len(xq))[:, None], indices], axis=1)
        final_indices = indices[np.arange(len(xq))[:, None], sorted_indices]
        final_distances = distances[np.arange(len(xq))[:, None], final_indices]

        return final_distances, final_indices
```

**Note**: This baseline achieves perfect recall (100%) but is too slow for large datasets. Use approximate methods like HNSW, IVF, or LSH for better speed-recall tradeoffs.

Debugging Tips
--------------
- **Test locally**: Use a subset of data (e.g., 10K vectors) for faster iteration
- **Verify shapes**: Ensure `search` returns `(nq, k)` shaped arrays
- **Check recall calculation**: `(I[:, :1] == gt[:, :1]).sum() / len(xq)`
- **Profile latency**: Measure batch vs single query performance separately
- **Validate before submit**: Run full 1M dataset locally if possible


VDB Design Problem - Recall80 Latency Tier
===========================================

Problem Setting
---------------
Design a Vector Database index optimized for **latency** subject to a **recall constraint**. This tier uses recall-gated scoring: solutions failing to meet the recall threshold receive zero points, while solutions meeting the constraint are scored purely by latency.

**Optimization Goal**: Minimize latency within recall constraint

$$
\text{score} = \begin{cases}
0 & \text{if } r < r_{\text{gate}} \\
100 & \text{if } r \geq r_{\text{gate}} \text{ and } t_{\text{query}} \leq t_{\text{min}} \\
100 \cdot \frac{t_{\text{max}} - t_{\text{query}}}{t_{\text{max}} - t_{\text{min}}} & \text{if } r \geq r_{\text{gate}} \text{ and } t_{\text{min}} < t_{\text{query}} < t_{\text{max}} \\
0 & \text{if } r \geq r_{\text{gate}} \text{ and } t_{\text{query}} \geq t_{\text{max}}
\end{cases}
$$

Where:
- $r$: Your recall@1
- $t_{\text{query}}$: Your average query latency (ms)
- $r_{\text{gate}} = 0.80$ (minimum required recall)
- $t_{\text{min}} = 0.0\text{ms}$ (best possible latency)
- $t_{\text{max}} = 0.6\text{ms}$ (maximum allowed latency)

**Key Insight**: Unlike other tiers, this tier gates on recall and scores on latency. You MUST achieve ≥80% recall, then faster is better.

Baseline Performance
--------------------
- Recall@1: **0.9914** (99.14%)
- Avg query time: **3.85ms**

Scoring Examples
----------------
All examples assume recall constraint is met ($r \geq 0.80$):

| Recall@1 | Latency | Score Calculation | Score |
|----------|---------|-------------------|-------|
| 0.85 | 0.00ms | $t \leq t_{\text{min}}$ → max score | **100** |
| 0.85 | 0.30ms | $\frac{0.6 - 0.3}{0.6 - 0.0} = 0.50$ | **50** |
| 0.82 | 0.50ms | $\frac{0.6 - 0.5}{0.6 - 0.0} = 0.167$ | **16.7** |
| 0.90 | 0.10ms | $\frac{0.6 - 0.1}{0.6 - 0.0} = 0.833$ | **83.3** |
| **0.75** | 0.20ms | $r < r_{\text{gate}}$ → recall gate fails | **0** |
| 0.95 | **0.70ms** | $t \geq t_{\text{max}}$ → latency too high | **0** |

**Note**: This is the most aggressive latency requirement (0.6ms max). You must use extreme approximation while maintaining 80% recall.

API Specification
-----------------
Implement a class with the following interface:

```python
import numpy as np
from typing import Tuple

class YourIndexClass:
    def __init__(self, dim: int, **kwargs):
        """
        Initialize the index for vectors of dimension `dim`.

        Args:
            dim: Vector dimensionality (e.g., 128 for SIFT1M)
            **kwargs: Optional parameters (e.g., M, ef_construction for HNSW)

        Example:
            index = YourIndexClass(dim=128, nlist=256, nprobe=2)
        """
        pass

    def add(self, xb: np.ndarray) -> None:
        """
        Add vectors to the index.

        Args:
            xb: Base vectors, shape (N, dim), dtype float32

        Notes:
            - Can be called multiple times (cumulative)
            - Must handle large N (e.g., 1,000,000 vectors)

        Example:
            index.add(xb)  # xb.shape = (1000000, 128)
        """
        pass

    def search(self, xq: np.ndarray, k: int) -> Tuple[np.ndarray, np.ndarray]:
        """
        Search for k nearest neighbors of query vectors.

        Args:
            xq: Query vectors, shape (nq, dim), dtype float32
            k: Number of nearest neighbors to return

        Returns:
            (distances, indices):
                - distances: shape (nq, k), dtype float32, L2 distances
                - indices: shape (nq, k), dtype int64, indices into base vectors

        Notes:
            - Must return exactly k neighbors per query
            - Indices should refer to positions in the vectors passed to add()
            - Lower distance = more similar

        Example:
            D, I = index.search(xq, k=1)  # xq.shape = (10000, 128)
            # D.shape = (10000, 1), I.shape = (10000, 1)
        """
        pass
```

**Implementation Requirements**:
- Class can have any name (evaluator auto-discovers classes with `add` and `search` methods)
- Must handle SIFT1M dataset: 1M base vectors, 10K queries, 128 dimensions
- Your `search` must return tuple `(distances, indices)` with shapes `(nq, k)`
- Distances should be L2 (Euclidean) or L2-squared
- No need to handle dataset loading - evaluator provides numpy arrays

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

### 1. Load Dataset
```python
from faiss.contrib.datasets import DatasetSIFT1M
ds = DatasetSIFT1M()
xb = ds.get_database()        # (1000000, 128) float32
xq = ds.get_queries()         # (10000, 128) float32
gt = ds.get_groundtruth()     # (10000, 100) int64 - ground truth indices
```

### 2. Build Index
```python
from solution import YourIndexClass  # Auto-discovered
d = xb.shape[1]                       # 128 for SIFT1M
index = YourIndexClass(d)             # Pass dimension as first argument
index.add(xb)                         # Add all 1M base vectors
```

### 3. Measure Performance (Batch Queries)
```python
import time
t0 = time.time()
D, I = index.search(xq, k=1)          # Search all 10K queries at once
t1 = time.time()

# Calculate metrics
recall_at_1 = (I[:, :1] == gt[:, :1]).sum() / len(xq)
avg_query_time_ms = (t1 - t0) * 1000.0 / len(xq)
```

**Important**: `avg_query_time_ms` from **batch queries** is used for scoring. Batch queries benefit from CPU cache and vectorization, typically faster than single queries.

### 4. Calculate Score
```python
if recall_at_1 < 0.80:
    score = 0.0
elif avg_query_time_ms <= 0.0:
    score = 100.0
elif avg_query_time_ms >= 0.6:
    score = 0.0
else:
    proportion = (avg_query_time_ms - 0.0) / (0.6 - 0.0)
    score = 100.0 * (1.0 - proportion)
```

Dataset Details
---------------
- **Name**: SIFT1M
- **Base vectors**: 1,000,000 vectors of dimension 128
- **Query vectors**: 10,000 vectors
- **Ground truth**: Precomputed nearest neighbors (k=1)
- **Metric**: L2 (Euclidean distance)
- **Vector type**: float32

Runtime Platform
----------------
- **Infrastructure**: Evaluations run on SkyPilot-managed cloud instances (AWS, GCP, or Azure)
- **Compute**: CPU-only instances (no GPU required)
- **Environment**: Docker containerized execution with Python 3, NumPy ≥1.24, FAISS-CPU ≥1.7.4

Constraints
-----------
- **Timeout**: 1 hour for entire evaluation (index construction + queries)
- **Memory**: Use reasonable memory (index should fit in RAM)
- **Recall constraint**: recall@1 ≥ 0.80
- **Latency range**: 0.0ms ≤ avg_query_time_ms ≤ 0.6ms

Strategy Tips
-------------
1. **Meet recall gate first**: Ensure ≥80% recall, otherwise score = 0
2. **Extreme approximation**: Use minimal search budget (IVF nprobe=1-3)
3. **Batch optimization critical**: 0.6ms is extremely tight, every microsecond counts
4. **Trade recall for speed**: 80-85% recall with ultra-low latency is ideal

Example: Simple Baseline
-------------------------
```python
import numpy as np

class SimpleIndex:
    def __init__(self, dim: int, **kwargs):
        self.dim = dim
        self.xb = None

    def add(self, xb: np.ndarray) -> None:
        if self.xb is None:
            self.xb = xb.copy()
        else:
            self.xb = np.vstack([self.xb, xb])

    def search(self, xq: np.ndarray, k: int) -> tuple:
        # Compute all pairwise L2 distances
        # xq: (nq, dim), xb: (N, dim)
        # distances: (nq, N)
        distances = np.sqrt(((xq[:, np.newaxis, :] - self.xb[np.newaxis, :, :]) ** 2).sum(axis=2))

        # Get k nearest neighbors
        indices = np.argpartition(distances, k-1, axis=1)[:, :k]
        sorted_indices = np.argsort(distances[np.arange(len(xq))[:, None], indices], axis=1)
        final_indices = indices[np.arange(len(xq))[:, None], sorted_indices]
        final_distances = distances[np.arange(len(xq))[:, None], final_indices]

        return final_distances, final_indices
```

**Note**: This baseline achieves perfect recall (100%) but is too slow for large datasets. Use approximate methods like HNSW, IVF, or LSH for better speed-recall tradeoffs.

Debugging Tips
--------------
- **Test locally**: Use a subset of data (e.g., 10K vectors) for faster iteration
- **Verify shapes**: Ensure `search` returns `(nq, k)` shaped arrays
- **Check recall calculation**: `(I[:, :1] == gt[:, :1]).sum() / len(xq)`
- **Profile latency**: Measure batch vs single query performance separately
- **Validate before submit**: Run full 1M dataset locally if possible


VDB Design Problem - Recall95 Latency Tier
===========================================

Problem Setting
---------------
Design a Vector Database index optimized for **latency** subject to a **high recall constraint**. This tier uses recall-gated scoring: solutions failing to meet the recall threshold receive zero points, while solutions meeting the constraint are scored purely by latency.

**Optimization Goal**: Minimize latency within recall constraint

$$
\text{score} = \begin{cases}
0 & \text{if } r < r_{\text{gate}} \\
100 & \text{if } r \geq r_{\text{gate}} \text{ and } t_{\text{query}} \leq t_{\text{min}} \\
100 \cdot \frac{t_{\text{max}} - t_{\text{query}}}{t_{\text{max}} - t_{\text{min}}} & \text{if } r \geq r_{\text{gate}} \text{ and } t_{\text{min}} < t_{\text{query}} < t_{\text{max}} \\
0 & \text{if } r \geq r_{\text{gate}} \text{ and } t_{\text{query}} \geq t_{\text{max}}
\end{cases}
$$

Where:
- $r$: Your recall@1
- $t_{\text{query}}$: Your average query latency (ms)
- $r_{\text{gate}} = 0.95$ (minimum required recall)
- $t_{\text{min}} = 0.0\text{ms}$ (best possible latency)
- $t_{\text{max}} = 7.7\text{ms}$ (maximum allowed latency)

**Key Insight**: This tier requires high recall (95%), but provides generous latency budget (7.7ms). Focus on recall first, then optimize latency.

Baseline Performance
--------------------
- Recall@1: **0.9914** (99.14%)
- Avg query time: **3.85ms**

Scoring Examples
----------------
All examples assume recall constraint is met ($r \geq 0.95$):

| Recall@1 | Latency | Score Calculation | Score |
|----------|---------|-------------------|-------|
| 0.96 | 0.00ms | $t \leq t_{\text{min}}$ → max score | **100** |
| 0.96 | 3.85ms | $\frac{7.7 - 3.85}{7.7 - 0.0} = 0.50$ | **50.0** |
| 0.97 | 5.00ms | $\frac{7.7 - 5.0}{7.7 - 0.0} = 0.351$ | **35.1** |
| 0.98 | 2.00ms | $\frac{7.7 - 2.0}{7.7 - 0.0} = 0.740$ | **74.0** |
| **0.94** | 2.00ms | $r < r_{\text{gate}}$ → recall gate fails | **0** |
| 0.96 | **8.00ms** | $t \geq t_{\text{max}}$ → latency too high | **0** |

**Note**: The 95% recall requirement is strict, but the 7.7ms latency budget is generous, allowing thorough search strategies.

API Specification
-----------------
Implement a class with the following interface:

```python
import numpy as np
from typing import Tuple

class YourIndexClass:
    def __init__(self, dim: int, **kwargs):
        """
        Initialize the index for vectors of dimension `dim`.

        Args:
            dim: Vector dimensionality (e.g., 128 for SIFT1M)
            **kwargs: Optional parameters (e.g., M, ef_construction for HNSW)

        Example:
            index = YourIndexClass(dim=128, M=64, ef_search=400)
        """
        pass

    def add(self, xb: np.ndarray) -> None:
        """
        Add vectors to the index.

        Args:
            xb: Base vectors, shape (N, dim), dtype float32

        Notes:
            - Can be called multiple times (cumulative)
            - Must handle large N (e.g., 1,000,000 vectors)

        Example:
            index.add(xb)  # xb.shape = (1000000, 128)
        """
        pass

    def search(self, xq: np.ndarray, k: int) -> Tuple[np.ndarray, np.ndarray]:
        """
        Search for k nearest neighbors of query vectors.

        Args:
            xq: Query vectors, shape (nq, dim), dtype float32
            k: Number of nearest neighbors to return

        Returns:
            (distances, indices):
                - distances: shape (nq, k), dtype float32, L2 distances
                - indices: shape (nq, k), dtype int64, indices into base vectors

        Notes:
            - Must return exactly k neighbors per query
            - Indices should refer to positions in the vectors passed to add()
            - Lower distance = more similar

        Example:
            D, I = index.search(xq, k=1)  # xq.shape = (10000, 128)
            # D.shape = (10000, 1), I.shape = (10000, 1)
        """
        pass
```

**Implementation Requirements**:
- Class can have any name (evaluator auto-discovers classes with `add` and `search` methods)
- Must handle SIFT1M dataset: 1M base vectors, 10K queries, 128 dimensions
- Your `search` must return tuple `(distances, indices)` with shapes `(nq, k)`
- Distances should be L2 (Euclidean) or L2-squared
- No need to handle dataset loading - evaluator provides numpy arrays

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

### 1. Load Dataset
```python
from faiss.contrib.datasets import DatasetSIFT1M
ds = DatasetSIFT1M()
xb = ds.get_database()        # (1000000, 128) float32
xq = ds.get_queries()         # (10000, 128) float32
gt = ds.get_groundtruth()     # (10000, 100) int64 - ground truth indices
```

### 2. Build Index
```python
from solution import YourIndexClass  # Auto-discovered
d = xb.shape[1]                       # 128 for SIFT1M
index = YourIndexClass(d)             # Pass dimension as first argument
index.add(xb)                         # Add all 1M base vectors
```

### 3. Measure Performance (Batch Queries)
```python
import time
t0 = time.time()
D, I = index.search(xq, k=1)          # Search all 10K queries at once
t1 = time.time()

# Calculate metrics
recall_at_1 = (I[:, :1] == gt[:, :1]).sum() / len(xq)
avg_query_time_ms = (t1 - t0) * 1000.0 / len(xq)
```

**Important**: `avg_query_time_ms` from **batch queries** is used for scoring. Batch queries benefit from CPU cache and vectorization, typically faster than single queries.

### 4. Calculate Score
```python
if recall_at_1 < 0.95:
    score = 0.0
elif avg_query_time_ms <= 0.0:
    score = 100.0
elif avg_query_time_ms >= 7.7:
    score = 0.0
else:
    proportion = (avg_query_time_ms - 0.0) / (7.7 - 0.0)
    score = 100.0 * (1.0 - proportion)
```

Dataset Details
---------------
- **Name**: SIFT1M
- **Base vectors**: 1,000,000 vectors of dimension 128
- **Query vectors**: 10,000 vectors
- **Ground truth**: Precomputed nearest neighbors (k=1)
- **Metric**: L2 (Euclidean distance)
- **Vector type**: float32

Runtime Platform
----------------
- **Infrastructure**: Evaluations run on SkyPilot-managed cloud instances (AWS, GCP, or Azure)
- **Compute**: CPU-only instances (no GPU required)
- **Environment**: Docker containerized execution with Python 3, NumPy ≥1.24, FAISS-CPU ≥1.7.4

Constraints
-----------
- **Timeout**: 1 hour for entire evaluation (index construction + queries)
- **Memory**: Use reasonable memory (index should fit in RAM)
- **Recall constraint**: recall@1 ≥ 0.95
- **Latency range**: 0.0ms ≤ avg_query_time_ms ≤ 7.7ms

Strategy Tips
-------------
1. **Meet recall gate first**: Ensure ≥95% recall, otherwise score = 0
2. **Use moderate approximation**: Higher recall requirement means less aggressive approximation
3. **Batch optimization is key**: Your `search` should handle batch queries efficiently
4. **Balance recall and latency**: Aim for 95-99% recall with 3-5ms latency

Example: Simple Baseline
-------------------------
```python
import numpy as np

class SimpleIndex:
    def __init__(self, dim: int, **kwargs):
        self.dim = dim
        self.xb = None

    def add(self, xb: np.ndarray) -> None:
        if self.xb is None:
            self.xb = xb.copy()
        else:
            self.xb = np.vstack([self.xb, xb])

    def search(self, xq: np.ndarray, k: int) -> tuple:
        # Compute all pairwise L2 distances
        # xq: (nq, dim), xb: (N, dim)
        # distances: (nq, N)
        distances = np.sqrt(((xq[:, np.newaxis, :] - self.xb[np.newaxis, :, :]) ** 2).sum(axis=2))

        # Get k nearest neighbors
        indices = np.argpartition(distances, k-1, axis=1)[:, :k]
        sorted_indices = np.argsort(distances[np.arange(len(xq))[:, None], indices], axis=1)
        final_indices = indices[np.arange(len(xq))[:, None], sorted_indices]
        final_distances = distances[np.arange(len(xq))[:, None], final_indices]

        return final_distances, final_indices
```

**Note**: This baseline achieves perfect recall (100%) but is too slow for large datasets. Use approximate methods like HNSW, IVF, or LSH for better speed-recall tradeoffs.

Debugging Tips
--------------
- **Test locally**: Use a subset of data (e.g., 10K vectors) for faster iteration
- **Verify shapes**: Ensure `search` returns `(nq, k)` shaped arrays
- **Check recall calculation**: `(I[:, :1] == gt[:, :1]).sum() / len(xq)`
- **Profile latency**: Measure batch vs single query performance separately
- **Validate before submit**: Run full 1M dataset locally if possible


Vector Addition Problem - Medium Vectors (2^20)
================================================

Problem Setting
---------------
Design and optimize high-performance Triton kernels for vector addition on GPU with medium vectors (1,048,576 elements). This problem focuses on implementing efficient element-wise addition for typical workloads.

The challenge involves optimizing:
- **Memory access patterns**: Efficient loading and storing of vector data
- **Block sizing**: Optimal block sizes for GPU execution
- **Memory bandwidth**: Maximizing throughput for simple arithmetic operations
- **Performance benchmarking**: Achieving speedup over PyTorch baseline

This variant tests performance on medium vectors (2^20 = 1,048,576 elements = 4 MB per vector).

Target
------
- **Primary**: Maximize bandwidth (GB/s) over PyTorch baseline (higher is better)
- **Secondary**: Ensure correctness
- **Tertiary**: Minimize kernel launch overhead

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 add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """
    Element-wise addition of two vectors.

    Args:
        x: Input tensor of shape (1048576,)
        y: Input tensor of shape (1048576,)

    Returns:
        Output tensor of shape (1048576,) with x + y
    """
    pass
```

API Usage Notes
---------------
- The evaluator looks for an `add` function in the module namespace
- Function must handle vector size of exactly 1,048,576 elements
- Must use Triton JIT compilation for kernel definition
- Should optimize for memory bandwidth
- Input tensors are guaranteed to be contiguous and same size

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

```
target = max(2.0 * (pytorch_bandwidth / cpu_bandwidth), 1.0)
score = ((custom_bandwidth / cpu_bandwidth - 1.0) / (target - 1.0)) * 100

Where:
- custom_bandwidth = your solution's bandwidth
- cpu_bandwidth = naive CPU baseline bandwidth
- pytorch_bandwidth = PyTorch GPU baseline bandwidth
- target = 2x PyTorch performance vs CPU (normalized to custom vs CPU)

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

- 0 points = CPU baseline performance (custom/cpu = 1x)
- 50 points = Halfway between CPU baseline and 2x PyTorch performance
- 100 points = 2x PyTorch GPU performance vs CPU (custom/cpu = 2 * pytorch/cpu)

Evaluation Details
------------------
- Tested on vector size: 2^20 = 1,048,576 elements
- Performance measured in GB/s (bandwidth)
- Correctness verified with tolerance: rtol=1e-5, atol=1e-8
- Performance measured using median execution time across 5 samples
- Requires CUDA backend and GPU support


Vector Addition Problem - Large Vectors (2^24)
===============================================

Problem Setting
---------------
Design and optimize high-performance Triton kernels for vector addition on GPU with large vectors (16,777,216 elements). This problem focuses on implementing efficient element-wise addition for high-throughput workloads.

The challenge involves optimizing:
- **Memory bandwidth**: Maximizing throughput for large vectors
- **Memory access patterns**: Efficient loading and storing of vector data
- **Block sizing**: Optimal block sizes for large vectors
- **Performance benchmarking**: Achieving speedup over PyTorch baseline

This variant tests performance on large vectors (2^24 = 16,777,216 elements = 64 MB per vector).

Target
------
- **Primary**: Maximize bandwidth (GB/s) over PyTorch baseline (higher is better)
- **Secondary**: Minimize kernel launch overhead
- **Tertiary**: Ensure correctness

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 add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """
    Element-wise addition of two vectors.

    Args:
        x: Input tensor of shape (16777216,)
        y: Input tensor of shape (16777216,)

    Returns:
        Output tensor of shape (16777216,) with x + y
    """
    pass
```

API Usage Notes
---------------
- The evaluator looks for an `add` function in the module namespace
- Function must handle vector size of exactly 16,777,216 elements
- Must use Triton JIT compilation for kernel definition
- Should optimize for small vector performance and launch overhead
- Input tensors are guaranteed to be contiguous and same size

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

```
target = max(2.0 * (pytorch_bandwidth / cpu_bandwidth), 1.0)
score = ((custom_bandwidth / cpu_bandwidth - 1.0) / (target - 1.0)) * 100

Where:
- custom_bandwidth = your solution's bandwidth
- cpu_bandwidth = naive CPU baseline bandwidth
- pytorch_bandwidth = PyTorch GPU baseline bandwidth
- target = 2x PyTorch performance vs CPU (normalized to custom vs CPU)

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

- 0 points = CPU baseline performance (custom/cpu = 1x)
- 50 points = Halfway between CPU baseline and 2x PyTorch performance
- 100 points = 2x PyTorch GPU performance vs CPU (custom/cpu = 2 * pytorch/cpu)

Evaluation Details
------------------
- Tested on vector size: 2^24 = 16,777,216 elements
- Performance measured in GB/s (bandwidth)
- Correctness verified with tolerance: rtol=1e-5, atol=1e-8
- Performance measured using median execution time across 5 samples
- 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)"