# Bring Your Own Kernel to FlashInfer-Bench

This guide gives instructions on how to add Definitions, Solutions, capture Workloads, and record Evaluations by walking through each **component of the Trace**, with an end-to-end “apply at runtime” flow.

A **Trace** is an atomic, immutable record of a single benchmark run. It links a specific `Solution` to a specific `Definition`, fixes the exact `workload` (shapes + data), and stores the complete `evaluation`. A folder of Definitions, Solutions, and Traces is your benchmark database.

## Trace Schema (top level)

| Field        | Type   | Required | Description                                       |
| ------------ | ------ | -------- | ------------------------------------------------- |
| `definition` | string | Yes      | The `name` of the `Definition` used in this run.  |
| `solution`   | string | Yes      | The `name` of the `Solution` tested.              |
| `workload`   | object | Yes      | Concrete shapes and input data used for this run. |
| `evaluation` | object | Yes      | Results, logs, and environment snapshot.          |

## Component 1: `definition`

**What it is.** The operator’s contract: axes (const/var), inputs/outputs, constraints, and a correct (not necessarily fast) `reference`.

**Identity rule.** Two kernels are the **same Definition** iff:

* They have the **same axes**,
* Each axis has the **same role** (`const` vs `var`),
* All `const` axes have the **same values**.

**How to add a new kernel Definition.**

1. Refer to schema, choose a `name` (`<type>_<stage>_<axis tokens>`) and `type`; write a clear `description` and helpful `tags`.
2. Specify `axes` with `type: const|var` (+ `value` for const).
3. Add `constraints` that relate axes to inputs (e.g., CSR shapes).
4. Specify `inputs`/`outputs` (names, shapes by axes, dtypes, optional layouts).
5. Provide a correct Python `reference` returning a tuple of outputs.
6. (Optional) Provide minimal tests that run the `reference` on tiny shapes.

## Component 2: `solution`

**What it is.** A concrete implementation of a Definition’s interface (Triton/CUDA/CUTLASS/PyTorch, etc.) plus metadata including target archs, libraries, author (human or LLM).

**Interface.** Your function must take the Definition’s `inputs` and **return** the tuple of `outputs`.

**How to add a Solution.**

1. Add the implementation of the kernel (matching signature).
2. Provide metadata co-located with the code, according to schema.
3. Add unit tests vs `reference` across a representative shapes.

See agent.md (to be added) for our methods to generate Solutions with LLMs.

## Component 3: `workload`

**What it is.** The concrete axes + input data that instantiate a Definition for one run.

| Field    | Description                                   |
| -------- | --------------------------------------------- |
| `axes`   | Map of **var** axis → concrete int value.     |
| `inputs` | Map of **input name** → **actual input**.     |

**How to capture workloads.**

### **Env-vars (zero-code):**

1. **Choose an output dataset root** (optional):

```bash
export FLASHINFER_BENCH_DATASET_PATH=/root/flashinfer-trace
# defaults to ./flashinfer-trace if unset
```

2. **Enable tracing and run your engine or script:**

```bash
export FLASHINFER_BENCH_ENABLE_TRACING=1
python run_engine.py  # your serving or batch script
```

By default, all kernels with a matching **Definition** are traced.

3. **What gets saved & where (default layout):**
```
$FLASHINFER_BENCH_DATASET_PATH/
└── workloads/
    ├── *.jsonl               # workload records (FlashInfer Trace format)
    └── safetensors/          # tensor payloads (when dumped)
```

Writing tensors to file is **async** (background thread) to reduce runtime overhead.

### **Tracing in code (fine-grained control)**

If you want to target a subset of kernels / customize policies:

```python
import flashinfer_bench as fib

# 1) Pick which kernels to trace and how
from flashinfer_bench import TracingConfig

gqa_tracing = TracingConfig(
    tensor_dump_policy="dump_non_float",   # keep scalar and int tensors; skip large float payloads
    filter_policy="shape_only",             # save first occurrence per input-shape signature
)

configs = {
    "gqa_paged_decode_h32_kv4_d128_ps1": gqa_tracing,
    # more kernel definitions...
}

# 2) Enable, run, then finalize
with fib.enable_tracing(dataset_path="/root/flashinfer-trace", tracing_configs=configs):
    run_engine()  # your inference loop
```

**Policies you can use right away:**

* `tensor_dump_policy`: `"dump_all"`, `"dump_none"`, `"dump_non_float"`, or a list of input names to dump.
* `filter_policy`: `"keep_all"`, `"shape_only"`, `"keep_first_k"` (e.g., first k calls), or a custom callable `Workload -> key`.
  These reduce disk/time while keeping representative samples.

## Component 4: `evaluation`

**What it is.** The result bundle for one `(definition, solution, workload)` run.

**How to benchmark to produce Evaluations.**
Run the benchmarker over your `(definition, solution, workload)` triples in the dataset:

CLI:
  ```bash
  flashinfer-bench run --local ./flashinfer-trace --warmup-runs 10 --iterations 50 --save-results
  ```

Use Python API:
### Prepare a `TraceSet` and Run the benchmark

```python
from flashinfer_bench.data import TraceSet
from flashinfer_bench.bench import Benchmark
from flashinfer_bench.bench import BenchmarkConfig

# 1) Build TraceSet (definitions, solutions, workloads)
trace_set = TraceSet(root="./flashinfer-trace")  # scans for definitions, solutions, workloads

# 2) Run the benchmark
bench = Benchmark(trace_set)
bench.run_all(dump_traces=True)   # executes reference + solutions in parallel
```

* **Device pool.** One `MultiProcessRunner` is created per CUDA device.
* **Concurrency.** For each definition and workload, the benchmark:

  * Picks up to `K = min(#devices, #solutions)` runners (round-robin).
  * **Reference phase:** in parallel, calls `runner.run_ref(defn, wl, config)` to build a baseline on each selected runner.

    * If a runner fails during reference, it is removed from the pool and the workload on that runner is skipped.
  * **Solutions phase:** distributes solutions round-robin across the runners that succeeded in the reference phase, calling `runner.run_solution(sol, baseline_handle, config)` in parallel.
* **Status mapping.**

  * Successful run with numerics in tolerance → `PASSED`.
  * Output shape/dtype mismatch → `INCORRECT_SHAPE` / `INCORRECT_DTYPE`.
  * Numeric check fails → `INCORRECT_NUMERICAL`.
  * Runtime fault → `RUNTIME_ERROR`.
  * Build/compile fails → `COMPILE_ERROR`.

Each solution run returns an `Evaluation`; the benchmark immediately stages a `Trace(def_name, workload, sol_name, evaluation)` in memory.

After benchmarking is done, the results can be used to rank solutions, visualize leaderboards, and drive `apply` at runtime.

### Reproducibility

* **`BenchmarkConfig`** controls iteration counts, warmup, tolerances, and timeouts (use your project’s defaults or tune per kernel).
* **Environment snapshot**: runners capture hardware and library versions into `evaluation.environment`.
* **Dead runner handling**: any runner failing the reference is dropped for subsequent work; if all runners fail, a `RuntimeError` is raised.

## Putting it together: Trace lifecycle

1. **Add the Definition**

   * Finalize axes (`const` vs `var`), constraints, I/O shapes, and `reference`.
   * Identity is locked by the axes set/roles/const values.

2. **Add one or more Solutions**

   * Implement the exact interface; return `{output_name: tensor}`.
   * Provide metadata and unit tests vs `reference`.

3. **Capture Workloads**

   * Run with tracing (env-vars or code) over real requests to collect shapes and, when helpful, actual inputs (esp. ragged index tensors).
   * Curate a small but representative set (use `shape_only` or `keep_first_k`).

4. **Benchmark → Emit Traces**

   * For each `(definition, solution, workload)` triple, run the benchmarker to produce one **Trace** JSON with `evaluation`.
   * Store logs and the environment snapshot alongside.

5. **Apply at runtime (end-to-end)**

   * Use runtime substitution to dispatch to the **best** ranked Solution for the current shapes.


## End-to-end “apply” (ties Trace back to serving)

**Decorator form:**

```python
import torch, torch.nn.functional as F
import flashinfer

@flashinfer.apply(lambda A, B: f"gemm_n_{B.shape[0]}_k_{B.shape[1]}")
def gemm(A, B):
    return F.linear(A, B)  # fallback/reference or a simple baseline
```

**Turn on runtime substitution:**

```bash
export FLASHINFER_BENCH_ENABLE_APPLY=1
python serve_or_benchmark.py
```

At call time, `apply` looks up the **Definition** (by name or via the lambda), matches the current **workload** (axes +, when required, data properties), and dispatches to the **best** `Solution` according to your recorded **Traces** (with correctness constraints and numeric tolerances enforced).

### Advanced Usage: Supporting kernels that don’t align with the Definition

Sometimes your production call site can’t be decorated directly—e.g., wrappers that keep internal state across `plan()`/`run()` like `BatchPrefillWithPagedKVCacheWrapper`. In these cases the function you call at runtime doesn’t match the kernel definition’s flat signature, so the decorator can’t attach cleanly. Use the imperative form instead.

#### Imperative `apply(...)` API

Use the function form of `apply` anywhere you call the kernel. It will (1) in **apply** mode: look up the best Solution for the current workload and call it; (2) in **tracing** mode: record the workload, then run the fallback; (3) otherwise: just call the fallback.

```python
import flashinfer

result = flashinfer.apply(
    name: Union[str, Callable[..., str]],
    fallback_function: Callable[..., Any],
    *args,       # All arguments must follow the **kernel definition’s interface
    **kwargs,
)
```

#### Example: stateful paged-attention wrapper → imperative `apply`

In this example, the FlashInfer attention wrapper carries state from `plan()` into `run()`, while the FlashInfer-Bench definition exposes a single `attention(init_params, plan_params, run_params)` entry point. Bridge them with a small monkey-patch that reconstructs the original flow as the fallback:

```python
# Original wrapper shape (state lives across plan/run)
class AttentionWrapper:
    def __init__(self, init_params): ...
    def plan(self, plan_params):
        self.state = compute_state(plan_params)
    def run(self, run_params):
        return call_flashinfer_kernel(run_params, self.state)

# FlashInfer-Bench-side definition interface we want to target:
def attention(init_params, plan_params, run_params):
    return attention_kernel(init_params, plan_params, run_params)
# (covers Q, K, V, page_size, page_indptr, etc.)
```

```python
# Monkey patch to route run() through flashinfer.apply
old_init, old_plan, old_run = (
    AttentionWrapper.__init__,
    AttentionWrapper.plan,
    AttentionWrapper.run,
)

def new_init(self: AttentionWrapper, init_params):
    def fallback(init_params, plan_params, run_params):
        w = AttentionWrapper.__new__(AttentionWrapper)
        old_init(w, init_params)
        old_plan(w, plan_params)
        return old_run(w, run_params)
    self.init_params = init_params
    self._fallback = fallback

def new_plan(self: AttentionWrapper, plan_params):
    self.plan_params = plan_params

def new_run(self: AttentionWrapper, run_params):
    return flashinfer.apply(
        "attention",              # or a lambda resolver if the def varies by shape
        self._fallback,
        self.init_params,
        self.plan_params,
        run_params,
    )

AttentionWrapper.__init__ = new_init
AttentionWrapper.plan = new_plan
AttentionWrapper.run = new_run
```

This preserves wrapper state while letting **apply** choose the best solution (and still trace workloads when enabled).

#### Alternative: avoid monkey-patching (shim inside the class)

If you can edit the wrapper, define a tiny adapter that flattens `(init, plan, run)` into the definition’s signature and call `flashinfer.apply(...)` directly inside `run()`. Same behavior, fewer moving parts.

#### Scope & tips

* Make sure your adapter/fallback **matches the Definition I/O** exactly.
* Group `init_params`, `plan_params`, and `run_params` so they cover the definition’s required tensors (e.g., `Q, K, V, page_size, page_indptr`).
* When definitions vary by shape, pass a **`name` lambda** (e.g., derive hidden size from weights) to resolve the correct Definition at call time.

## Related customization you can enable

* **Apply/trace only selected kernels** via configs (context managers or code APIs), if you don’t want blanket substitution/tracing:

```python
from flashinfer_bench import enable_apply, enable_tracing, ApplyConfig, TracingConfig

apply_cfgs = {
    "gemm_n_4096_k_14336": ApplyConfig(max_atol=1e-5, max_rtol=1e-5),
    "gqa_paged_decode_h32_kv4_d128_ps1": ApplyConfig(),  # defaults OK
}
trace_cfgs = {
    "gqa_paged_decode_h32_kv4_d128_ps1": TracingConfig(
        tensor_dump_policy="dump_non_float",
        filter_policy="shape_only",
    ),
}

with enable_apply(apply_configs=apply_cfgs):
    with enable_tracing(tracing_configs=trace_cfgs):
        run_engine()
```

  This limits substitution/tracing to kernels you care about and mirrors the env-var flow.
