Iterate Topk

---
description: Run the long-horizon top-k kernel iterate loop autonomously — design 4 variants, JIT-compile, benchmark, generate reports, analyse, repeat. Stops after --max-iterations batches (default 3).
argument-hint: [--max-iterations <N>] [--k <K>] [--num-gpus <G>]
---

You are now running the **top-k kernel iterate loop** autonomously.
Do not ask for confirmation between steps; execute each step in
sequence and loop until the iteration budget is reached.

Unlike `/iterate`, this loop tunes the *top-k selection kernel* used
inside the sparse-attention indexer — a CUDA kernel that, given
per-segment scores, returns the top-K indices. Iteration is fast
(single GPU, seconds-to-minutes per variant) and produces a
`(geomean speedup, recall@189, recall@253)` report per variant.

## Step 0 — parse arguments and activate conda env

Parse `$ARGUMENTS` for `--max-iterations <N>` (default: 3),
`--k <K>` (default: 256), and `--num-gpus <G>` (default: 4 — the
batch size, so behavior matches the historical single-wave run).

- `K` is the top-k size every kernel in this session is tuned for;
  it picks the working folder `vortex_torch/kernels/topk/k_${K}/`
  and is passed to `benchmark.py --k` so reports use the same
  eval points.
- `NUM_GPUS` is the **maximum** number of free GPUs the loop will
  consume per batch. The batch is always 4 variants; with
  `NUM_GPUS < 4` the loop runs `⌈4 / NUM_GPUS⌉` sequential waves
  on the cap-sized GPU pool. With `NUM_GPUS >= 4` (and ≥4 free
  GPUs available) the whole batch fires in a single wave.
```bash
MAX_ITER=3
K=256
NUM_GPUS=4
NEXT=
for arg in $ARGUMENTS; do
  case "$NEXT" in
    max-iterations) MAX_ITER=$arg; NEXT= ;;
    k)              K=$arg;        NEXT= ;;
    num-gpus)       NUM_GPUS=$arg; NEXT= ;;
    *) case "$arg" in
         --max-iterations) NEXT=max-iterations ;;
         --k)              NEXT=k ;;
         --num-gpus)       NEXT=num-gpus ;;
       esac ;;
  esac
done
echo "max_iterations=$MAX_ITER  k=$K  num_gpus=$NUM_GPUS"

Activate the conda environment:

CONDA_BASE=$(conda info --base 2>/dev/null || echo /root/anaconda3)
source "$CONDA_BASE/etc/profile.d/conda.sh"
conda activate vortex_v1
python -c "import sys; print(sys.executable)"

Step 1 — pick your tag and read context (once per session)

Your <tag> is a sanitized lowercase form of your model name (e.g. claude_sonnet_4_6, claude_opus_4_7). Create the per-tag artifact dirs if they don't exist, and scaffold an empty memory_topk.md from the template below on first use of this K:

TAG=<your_tag>
KROOT=vortex_torch/kernels/topk/k_${K}
mkdir -p "$KROOT/configs/$TAG" "$KROOT/sources/$TAG" "$KROOT/reports/$TAG"

if [ ! -f "$KROOT/memory_topk.md" ]; then
  cat > "$KROOT/memory_topk.md" <<EOF
# memory_topk.md — k=${K}

Persistent state for \`/iterate_topk --k ${K}\`. Read at the start of
every session; mutate on every batch launch and every batch
completion. The conversation evaporates; this file does not.

The two reference kernels (against which all proposals are measured)
live at:
- baseline — \`vortex_torch/kernels/topk/configs/sort_default.json\` (CUB BlockRadixSort)
- proposal seed — \`vortex_torch/kernels/topk/configs/radix_default.json\` (8-bit radix)

Benchmark protocol (per proposal): 100 samples × 4 batch sizes ×
4 seq_lens, scores drawn from a 14-distribution rotation, top-k
size **K=${K}**. Recall is reported at \`floor(K*3/4)\` and \`K\`.
Reports land at \`vortex_torch/kernels/topk/k_${K}/reports/<tag>/<batch_X_idY>.md\`.

Quality bar: **min recall@K ≥ 0.98** on every distribution.
Anything below that is structurally broken — log it in §4 and do
not promote.

---

## §1 RUNNING

In-flight batches. One row per launch; remove the row when the
batch lands (move its summary to §2). If a row sits here at session
start, resume waiting on it before launching anything new.

| tag | batch | launched_at | gpu | variants | status |
| --- | ----- | ----------- | --- | -------- | ------ |
| _none_ | | | | | |

---

## §2 Completed batches

One subsection per completed batch (newest first). Each contains
the per-variant comparison table and a 1–3 sentence takeaway.

---

## §3 Hypotheses (pre-registered before launch)

- batch_X_id_Y — <hypothesis> — <novelty axis>

---

## §4 Anti-patterns / broken variants

- <description> — <reason> — <evidence pointer>

---

## §5 Pareto winners (running best by axis)

Updated after every batch. Each row is the best observed *so far*
on its axis; replace when a new variant strictly dominates.

| axis                                          | config | value | notes |
| --------------------------------------------- | ------ | ----- | ----- |
| best geomean speedup                          |        |       |       |
| best worst-case recall@K with speedup ≥ 1.0   |        |       |       |

---

## §6 Insights

One bullet per file read or experiment-confirmed observation. Cite
\`file:line\` when applicable so the insight is checkable.

- **Seed: read prior K buckets first.** Before designing batch 0,
  scan \`vortex_torch/kernels/topk/k_*/memory_topk.md\` for any
  sibling bucket — especially \`k_256/memory_topk.md\` if present.
  Prior experiments expose anti-patterns (§4), alignment / smem
  occupancy caveats (§6), and Pareto winners (§5) that frequently
  transfer across K. Carry the relevant bullets into this file's
  §6 with a "from k_<other>:" prefix instead of re-deriving them.
- **Seed: the algorithm space is wider than radix.** The current
  proposal seed is 8-bit radix (\`radix_topk.cu\`); the baseline is
  CUB \`BlockRadixSort\` (\`sort_topk.cu\`). Other families are
  unexplored and have very different occupancy / recall tradeoffs:
  \`bitonic\` (warp-cooperative sort, no smem candidate buffer),
  hybrid \`radix → bitonic refinement\` for the threshold bin,
  \`heap\` (per-warp min-heap of size K), \`approxTopK\`-style
  atomic-arrival within the threshold bin, and 16-bit-radix
  attacking the bf16-bottom-zero