diff --git a/2.0/problems/dbscan_gpu_kernel_optimization/DESIGN.md b/2.0/problems/dbscan_gpu_kernel_optimization/DESIGN.md new file mode 100644 index 000000000..595683acd --- /dev/null +++ b/2.0/problems/dbscan_gpu_kernel_optimization/DESIGN.md @@ -0,0 +1,60 @@ +# Design notes — dbscan_gpu_kernel_optimization + +Operator-facing. Not copied into the agent workspace by the adapter. + +## What this task measures + +Can an agent turn a naive GPU DBSCAN into a fast one? The agent patches +`dbscanlib`; the judge times the patched `dbscan` against a frozen naive baseline +on hidden low-dimensional workloads and scores the geomean speedup, gated on +clustering agreement (Adjusted Rand Index). One of the four-task flashlib +kernel-optimization family (KMeans, KNN, DBSCAN, IVF-PQ), all sharing one +evaluator + one Modal GPU harness (`flash_gpu.py`). + +This is a **genuine kernel task**: the core operation (radius-neighbour search + +connected components + border assignment) is not a single fast torch primitive, +so beating the naive O(N^2) baseline requires real kernel work (a bucketed/grid +radius search + fused neighbour + connected-components kernels). Contrast with +pca/svd, which were algorithmic (cov+eigh + a library eigensolver) and are not +part of this family. + +## Baseline + reference + +- Naive baseline (`dbscanlib/dbscan.py`, and frozen `judge/refdbscan.py`): chunked + O(N^2) neighbour scan for the core mask, GPU label-propagation for the core + connected components, and a smallest-core-label border rule. Verified against + scikit-learn DBSCAN (ARI = 1.0). The border rule matches flashlib's ("smallest + CC label among in-eps core neighbours"), so the reference agrees closely. +- Reference (`reference.patch`): vendors flashlib's optimized **Triton** planar + grid radius-search DBSCAN (`primitives/dbscan/triton/dbscan.py` grid path + + `kernels/flash_mst` connected components) under `dbscanlib/_kernels/…`. The + flashlib grid kernel hard-asserts D==2 (higher D falls back to a flash_knn + brute path), so all graded workloads are D=2 and flash_knn is stubbed out. + +## Correctness gate + +Adjusted Rand Index (permutation-invariant, noise-aware) between the agent's +labels and the baseline's on each iteration's fresh data, `>= ari_threshold` +(default 0.99). Judge-computed; cheat-proof (you cannot fake a high ARI without +actually reproducing the clustering). DBSCAN with fixed (eps, min_samples) is +deterministic up to border tie-breaks, which ARI absorbs. + +## Workloads + +Planar (D=2) planted blobs + a small uniform-noise fraction, N 100k-170k. Below +~100k the tensor-core matmul baseline (`xb @ x.t()`) beats the grid kernel's fixed +launch/grid-build overhead (H100 sweep: 100k->6x, 150k->14x, 200k->27x); the band +is chosen so the O(N^2) baseline is seconds/call (grid wins clearly) yet bounded +for `timed_iters`. `timed_iters` is 3 (vs 7 elsewhere) since each baseline call is +seconds, not ms. +Hidden workloads live in `config.yaml` (judge-only). See kmeans DESIGN for the +shared anti-hack + Modal-offload design; identical here. + +## Calibration (H100 Modal trial — done) + +`reference.patch` validated on H100: the vendored grid DBSCAN compiles, passes the +ARI gate with **ARI 1.0** on all six workloads, and runs +**7.2 / 12.1 / 9.5 / 18.8 / 16.6 / 21.2x** over the naive baseline (geomean +~13.3x). `speedup_target` set to 13.0. The full 6-workload judge (warmup 2 + +timed 3) completed in ~410 s, well inside `modal_timeout_seconds` (2400). Re-sweep +if the N band changes. diff --git a/2.0/problems/dbscan_gpu_kernel_optimization/config.yaml b/2.0/problems/dbscan_gpu_kernel_optimization/config.yaml new file mode 100644 index 000000000..e386fdad6 --- /dev/null +++ b/2.0/problems/dbscan_gpu_kernel_optimization/config.yaml @@ -0,0 +1,47 @@ +tag: systems +runtime: + language: python + timeout_seconds: 10800 + environment: "GPU DBSCAN kernel optimization. The agent patches the dbscanlib package; the judge offloads timing to a Modal GPU, comparing the patched dbscan against a frozen naive baseline on hidden low-dimensional workloads with an ARI (clustering-agreement) gate." + apt_packages: [bash, ca-certificates, git, python3, python3-pip] + judge_apt_packages: [bash, ca-certificates, git, python3, python3-pip] + judge_pip_packages: [modal] + docker: + image: frontiercs/dbscan-gpu-kernel-optimization-agent:experimental-v0.4.0 + judge_image: frontiercs/dbscan-gpu-kernel-optimization-judge:experimental-v0.4.0 +environment: + cpus: 4 + memory_mb: 8192 + storage_mb: 16384 + build_timeout_seconds: 3600 +evaluation: + gpu: "H100" + cuda_image: "nvidia/cuda:12.4.1-devel-ubuntu22.04" + pip: ["torch", "numpy"] + app_name: "dbscan-kernel-opt-eval" + modal_timeout_seconds: 2400 + warmup_iters: 2 + timed_iters: 3 # the O(N^2) naive baseline is seconds/call at these N; keep the timed budget bounded (median of 3) + ari_threshold: 0.99 # agent clustering must match the baseline exact DBSCAN (Adjusted Rand Index) every iteration. On the non-convex multi-subspace-moons data a centroid/partition approximation (k-means, or "k-means the points then DBSCAN each cell") mislabels the interleaving arcs -> ARI ~0.72-0.75, and a PCA->2D projection -> ~0.5-0.6, both < 0.99 -> fail; only an exact eps-connectivity scan passes. Exact reference measures ARI ~1.0 (bf16-stable). + speedup_target: 130.0 # score cap = 2x the flashlib reference geomean on the non-convex multi-subspace-moons workloads (65x on H100: w0-w3 = 29/65/69/141x, ARI=1.0 all 4). Moons make the naive label-propagation baseline slower (elongated clusters -> long chains) so the factor is higher than the old blob regime (28.6x); the reference's connectivity is unaffected. Points carry only bf16 precision so no solution wins on a lower dtype; eps is tuned so bf16 never flips an eps-neighbour (two exact scans on the same bf16 data agree at ARI 1.0). + base_seed: 20260701 + agent_workload_count: 4 # all 4 flashlib graded workloads (feedback covers every graded shape) + expose_per_workload_metrics: true # show per-workload pass/fail so the agent can fix the specific workload it regresses on + # N in the 100k-170k band: the O(N^2) matmul baseline is seconds/call (grid kernel + # wins 6-20x) but bounded for timed_iters; below ~100k the tensor-core matmul baseline + # is faster than the grid kernel's fixed overhead. Vary N + density so agents can't + # special-case. (H100 crossover sweep: 100k->6.0x, 150k->14.1x, 200k->27.0x.) + workloads: + # NON-CONVEX data: n_centers interleaving half-moons, each moon-pair in its own + # random 2-D subspace of R^D (gen in flash_gpu.py). Voronoi/k-means shortcuts + # mislabel non-convex arcs (kmeans-shortcut ARI ~0.72-0.75, PCA-2D ~0.5-0.6), + # so only an exact eps-neighbourhood scan clears the ARI gate. eps is tuned per D + # so DBSCAN gives exactly n_centers clusters with ~0% noise and is bf16-stable + # (two exact scans on the same bf16 data agree at ARI 1.0). Generator is judge-only. + - {id: w0, N: 20000, D: 8, eps: 0.90, min_samples: 5, n_centers: 4} + - {id: w1, N: 20000, D: 32, eps: 1.15, min_samples: 5, n_centers: 6} + - {id: w2, N: 20000, D: 64, eps: 1.40, min_samples: 5, n_centers: 10} + - {id: w3, N: 50000, D: 16, eps: 1.00, min_samples: 5, n_centers: 8} +submission: + kind: file + path: /app/solution.patch diff --git a/2.0/problems/dbscan_gpu_kernel_optimization/dbscanlib/__init__.py b/2.0/problems/dbscan_gpu_kernel_optimization/dbscanlib/__init__.py new file mode 100644 index 000000000..5f897e412 --- /dev/null +++ b/2.0/problems/dbscan_gpu_kernel_optimization/dbscanlib/__init__.py @@ -0,0 +1,11 @@ +"""dbscanlib -- a tiny GPU DBSCAN you are asked to make fast. + +The shipped implementation is correct but naive (chunked O(N^2) neighbour scan + +label propagation). Rewrite the internals (grid/bucketed radius search, fused +neighbour + connected-components kernels, ...) to make :func:`dbscan` fast while +producing the same clustering. The public signature/return must not change. +""" +from __future__ import annotations +from dbscanlib.dbscan import dbscan +__all__ = ["dbscan"] +__version__ = "0.1.0" diff --git a/2.0/problems/dbscan_gpu_kernel_optimization/dbscanlib/dbscan.py b/2.0/problems/dbscan_gpu_kernel_optimization/dbscanlib/dbscan.py new file mode 100644 index 000000000..ab7fa8bfc --- /dev/null +++ b/2.0/problems/dbscan_gpu_kernel_optimization/dbscanlib/dbscan.py @@ -0,0 +1,70 @@ +"""Naive DBSCAN baseline (correct, deterministic, GPU/CPU torch). + +Euclidean DBSCAN. Chunked O(N^2) neighbour scan + GPU-friendly label +propagation for the connected components of the core graph. Border rule matches +flashlib: a non-core point joins the SMALLEST cluster label among its in-eps +core neighbours; otherwise it is noise (-1). +""" +from __future__ import annotations + +import torch + + +def dbscan(x, eps, min_samples, max_neighbors=32): + del max_neighbors # naive path scans all neighbours; knob only for the fast kernel + N = x.shape[0] + dev = x.device + eps2 = float(eps) * float(eps) + xn = (x * x).sum(1) # (N,) + CH = 4096 + BIG = N # sentinel label (> any real label) + + def d2_chunk(lo, hi): + xb = x[lo:hi] + return (xb * xb).sum(1, keepdim=True) - 2.0 * (xb @ x.t()) + xn[None, :] + + # 1) degree (in-eps count, includes self) -> core mask + deg = torch.zeros(N, device=dev, dtype=torch.int64) + for lo in range(0, N, CH): + hi = min(lo + CH, N) + deg[lo:hi] = (d2_chunk(lo, hi) <= eps2).sum(1) + core = deg >= int(min_samples) # (N,) bool + + # 2) connected components of the core-core eps graph via label propagation + labels = torch.arange(N, device=dev, dtype=torch.int64) + labels[~core] = BIG + for _ in range(1000): + new = labels.clone() + for lo in range(0, N, CH): + hi = min(lo + CH, N) + if not bool(core[lo:hi].any()): + continue + adj = (d2_chunk(lo, hi) <= eps2) & core[None, :] # (chunk, N) + lab = torch.where(adj, labels[None, :], torch.full((1, N), BIG, device=dev, dtype=torch.int64)) + mn = lab.min(1).values # (chunk,) + rows = slice(lo, hi) + upd = core[rows] & (mn < new[rows]) + new[rows] = torch.where(upd, mn, new[rows]) + if torch.equal(new, labels): + break + labels = new + + # 3) border assignment: non-core within eps of a core -> smallest core label + for lo in range(0, N, CH): + hi = min(lo + CH, N) + nb = ~core[lo:hi] + if not bool(nb.any()): + continue + adj = (d2_chunk(lo, hi) <= eps2) & core[None, :] + lab = torch.where(adj, labels[None, :], torch.full((1, N), BIG, device=dev, dtype=torch.int64)) + mn = lab.min(1).values + rows = slice(lo, hi) + take = nb & (mn < BIG) + labels[rows] = torch.where(take, mn, labels[rows]) + + # 4) compact to dense [0, n_clusters); noise (BIG) -> -1 + out = torch.full((N,), -1, device=dev, dtype=torch.int64) + uniq = torch.unique(labels[labels < BIG]) + for new_id, old in enumerate(uniq.tolist()): + out[labels == old] = new_id + return out diff --git a/2.0/problems/dbscan_gpu_kernel_optimization/docker/README.md b/2.0/problems/dbscan_gpu_kernel_optimization/docker/README.md new file mode 100644 index 000000000..506c4246d --- /dev/null +++ b/2.0/problems/dbscan_gpu_kernel_optimization/docker/README.md @@ -0,0 +1,56 @@ +# Experimental DBSCAN Kernel-Optimization Images + +Two **light** images (ubuntu + `modal` + git). torch, triton, and the vendored +flashlib kernels all run on the **Modal GPU image** defined in `flash_gpu.py`; +the containers themselves are CPU-only and offload GPU work to Modal. + +```bash +bash 2.0/problems/dbscan_gpu_kernel_optimization/docker/build_images.sh +``` + +Defaults: + +```text +AGENT_TAG=frontiercs/dbscan-gpu-kernel-optimization-agent:experimental-v0.2.0 +JUDGE_TAG=frontiercs/dbscan-gpu-kernel-optimization-judge:experimental-v0.2.0 +``` + +Agent image: + +```text +/app/dbscanlib # clean, git-tracked package (the agent edits this) +/opt/flash_gpu.py # shared Modal GPU harness (public test uses it) +/opt/dbscan_ref/refdbscan.py # frozen naive baseline (public-test speed denominator) +``` + +Judge image: + +```text +/opt/dbscanlib-clean/dbscanlib # pristine tree; the patch is applied to a copy +/opt/dbscan_ref/refdbscan.py # frozen naive baseline (speed denominator + quality oracle) +/opt/flash_gpu.py # shared Modal GPU harness (the evaluator uses it) +``` + +## Runtime requirements + +Both the judge and the agent public test **offload timing to a Modal GPU** and +therefore need Modal credentials in the environment: + +```text +MODAL_TOKEN_ID / MODAL_TOKEN_SECRET +``` + +`flash_gpu.py` builds an ephemeral Modal app on a GPU (`evaluation.gpu`, default +`H100`), ships the frozen baseline + the patched package as data, times both on +fresh per-iteration data, verifies quality each iteration, and returns the +speedups. No persistent deployment is used, so a fresh GPU container is spun up +per evaluation. Without Modal credentials the evaluator returns a patch-policy +smoke pass (which repo CI exercises). + +## Smoke test + +```bash +bash 2.0/problems/dbscan_gpu_kernel_optimization/docker/smoke_images.sh +``` + +Import-only (modal + flash_gpu + the baked packages); does not touch a GPU or Modal. diff --git a/2.0/problems/dbscan_gpu_kernel_optimization/docker/agent/Dockerfile b/2.0/problems/dbscan_gpu_kernel_optimization/docker/agent/Dockerfile new file mode 100644 index 000000000..e1179ff75 --- /dev/null +++ b/2.0/problems/dbscan_gpu_kernel_optimization/docker/agent/Dockerfile @@ -0,0 +1,34 @@ +# Agent image for dbscan_gpu_kernel_optimization. +# +# Light image (no torch/triton): the agent edits /app/dbscanlib and runs the +# public test, which offloads timing to a Modal GPU (torch/triton live on the +# Modal image defined in flash_gpu.py). Needs MODAL_TOKEN_ID / MODAL_TOKEN_SECRET +# in the environment to run the GPU public test. +# +# Build context = the task directory: +# docker build -f docker/agent/Dockerfile -t . +FROM ubuntu:24.04 + +ARG DEBIAN_FRONTEND=noninteractive +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + bash ca-certificates git python3 python3-pip ripgrep && \ + rm -rf /var/lib/apt/lists/* + +RUN pip3 install --break-system-packages --no-cache-dir modal + +# The package the agent edits (git-tracked so make_submission.sh can diff it). +WORKDIR /app +COPY dbscanlib /app/dbscanlib +RUN printf '%s\n' '__pycache__/' '*.pyc' > /app/.gitignore && \ + git -C /app init -q && \ + git -C /app config user.email task@frontier-cs && \ + git -C /app config user.name frontier-cs && \ + git -C /app add -A -- dbscanlib .gitignore && \ + git -C /app -c commit.gpgsign=false commit -qm "pristine dbscanlib" + +# NOTE: the GPU harness (flash_gpu.py) and its data GENERATOR are deliberately +# NOT baked into the agent image -- they live only in the judge. The agent gets +# feedback by submitting to the judge (public_test == the graded submission), +# so it cannot read how the workloads are generated and hardcode a +# data-specific shortcut. See harbor_app/public_test.py. diff --git a/2.0/problems/dbscan_gpu_kernel_optimization/docker/build_images.sh b/2.0/problems/dbscan_gpu_kernel_optimization/docker/build_images.sh new file mode 100644 index 000000000..682f88119 --- /dev/null +++ b/2.0/problems/dbscan_gpu_kernel_optimization/docker/build_images.sh @@ -0,0 +1,17 @@ +#!/usr/bin/env bash +# Build the experimental agent + judge images for dbscan_gpu_kernel_optimization. +set -euo pipefail +HERE=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd) # the problem directory (build context) + +AGENT_TAG=${AGENT_TAG:-frontiercs/dbscan-gpu-kernel-optimization-agent:experimental-v0.3.0} +JUDGE_TAG=${JUDGE_TAG:-frontiercs/dbscan-gpu-kernel-optimization-judge:experimental-v0.3.0} + +echo "Building agent image: $AGENT_TAG" +docker build -f "$HERE/docker/agent/Dockerfile" -t "$AGENT_TAG" "$HERE" + +echo "Building judge image: $JUDGE_TAG" +docker build -f "$HERE/docker/judge/Dockerfile" -t "$JUDGE_TAG" "$HERE" + +echo "Built:" +echo " $AGENT_TAG" +echo " $JUDGE_TAG" diff --git a/2.0/problems/dbscan_gpu_kernel_optimization/docker/judge/Dockerfile b/2.0/problems/dbscan_gpu_kernel_optimization/docker/judge/Dockerfile new file mode 100644 index 000000000..8c93b5ec4 --- /dev/null +++ b/2.0/problems/dbscan_gpu_kernel_optimization/docker/judge/Dockerfile @@ -0,0 +1,26 @@ +# Judge image for dbscan_gpu_kernel_optimization. +# +# Light image (no torch/triton): the judge validates + applies the patch and +# offloads timing to a Modal GPU (torch/triton live on the Modal image defined +# in flash_gpu.py). Needs MODAL_TOKEN_ID / MODAL_TOKEN_SECRET in the environment. +# The Frontier-CS 2.0 adapter builds the final judge image ON TOP of this one. +# +# Build context = the task directory: +# docker build -f docker/judge/Dockerfile -t . +FROM ubuntu:24.04 + +ARG DEBIAN_FRONTEND=noninteractive +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + bash ca-certificates git python3 python3-pip && \ + rm -rf /var/lib/apt/lists/* + +RUN pip3 install --break-system-packages --no-cache-dir modal + +# Pristine tree the patch is applied to, the frozen naive baseline (speed +# denominator + quality oracle), and the shared Modal GPU harness. +COPY dbscanlib /opt/dbscanlib-clean/dbscanlib +COPY judge/refdbscan.py /opt/dbscan_ref/refdbscan.py +COPY flash_gpu.py /opt/flash_gpu.py + +WORKDIR /judge diff --git a/2.0/problems/dbscan_gpu_kernel_optimization/docker/smoke_images.sh b/2.0/problems/dbscan_gpu_kernel_optimization/docker/smoke_images.sh new file mode 100644 index 000000000..012809f0d --- /dev/null +++ b/2.0/problems/dbscan_gpu_kernel_optimization/docker/smoke_images.sh @@ -0,0 +1,22 @@ +#!/usr/bin/env bash +# Smoke test for the built (intentionally torch-less) images: confirm the host-side +# harness imports (modal + flash_gpu) and that the frozen baseline + package parse. +# torch/triton live on the Modal image, so the packages are py_compiled here (a +# torch-less parse check), not imported. No GPU / Modal tokens required. +set -euo pipefail + +AGENT_TAG=${AGENT_TAG:-frontiercs/dbscan-gpu-kernel-optimization-agent:experimental-v0.2.0} +JUDGE_TAG=${JUDGE_TAG:-frontiercs/dbscan-gpu-kernel-optimization-judge:experimental-v0.2.0} + +SMOKE='import sys, modal, py_compile, pathlib +sys.path.insert(0, "/opt"); import flash_gpu +for root in sys.argv[1:]: + for p in pathlib.Path(root).rglob("*.py"): + py_compile.compile(str(p), doraise=True) +print("ok: modal + flash_gpu import; dbscan baseline + dbscanlib parse")' + +echo "== agent image ==" +docker run --rm -w /app "$AGENT_TAG" python3 -c "$SMOKE" /opt/dbscan_ref /app/dbscanlib + +echo "== judge image ==" +docker run --rm "$JUDGE_TAG" python3 -c "$SMOKE" /opt/dbscan_ref /opt/dbscanlib-clean/dbscanlib diff --git a/2.0/problems/dbscan_gpu_kernel_optimization/evaluate.sh b/2.0/problems/dbscan_gpu_kernel_optimization/evaluate.sh new file mode 100644 index 000000000..df1fe7f0d --- /dev/null +++ b/2.0/problems/dbscan_gpu_kernel_optimization/evaluate.sh @@ -0,0 +1,12 @@ +#!/usr/bin/env bash +# Local CLI evaluation for dbscan_gpu_kernel_optimization. +# +# A full run needs a GPU plus the baked judge sources at /opt (the pristine +# /opt/dbscanlib-clean tree and the frozen /opt/dbscan_ref baseline; see the +# judge image). Without them the evaluator still validates the patch policy and +# returns a smoke score, which is what repository CI exercises (the reference +# patch passes the policy). Pass a patch path to score it directly. +set -euo pipefail +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +SOLUTION="${1:-$SCRIPT_DIR/reference.patch}" +exec python3 "$SCRIPT_DIR/evaluator.py" "$SOLUTION" diff --git a/2.0/problems/dbscan_gpu_kernel_optimization/evaluator.py b/2.0/problems/dbscan_gpu_kernel_optimization/evaluator.py new file mode 100644 index 000000000..963234240 --- /dev/null +++ b/2.0/problems/dbscan_gpu_kernel_optimization/evaluator.py @@ -0,0 +1,379 @@ +"""Generic evaluator for the flashlib GPU kernel-optimization task family. + +Identical across all four tasks (kmeans / knn / pca / truncated_svd); every +primitive-specific value comes from the ``evaluation`` block of the task's +config (delivered to the judge as ``/judge/task_config.json``), which is not +present in the agent workspace. + +Flow: statically validate the patch (only ``/**`` may change; no external +optimized libs / env / process / network access) -> apply it to the pristine +package baked at ``clean_source`` -> hand the frozen baseline + patched package +sources to a Modal GPU worker (``flash_gpu.run_remote``) that times both on +fresh per-iteration data and verifies clustering/retrieval/subspace quality on +every iteration -> gate on the worker's per-workload verdict -> score by the +geometric-mean speedup. + +Without Modal credentials or the baked judge sources (e.g. repo CI on a box with +no GPU/Modal), the evaluator still validates the patch policy and returns a +smoke pass, so ``python3 evaluator.py reference.patch`` works anywhere. +""" +from __future__ import annotations + +import fnmatch +import hashlib +import json +import math +import os +import re +import shutil +import subprocess +import sys +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +TASK_CONFIG_PATH = Path("/judge/task_config.json") + + +def _load_task_config() -> dict[str, Any]: + try: + payload = json.loads(TASK_CONFIG_PATH.read_text(encoding="utf-8")) + except Exception: + return {} + return payload if isinstance(payload, dict) else {} + + +TASK_CONFIG = _load_task_config() +EVAL = TASK_CONFIG.get("evaluation", {}) if isinstance(TASK_CONFIG.get("evaluation"), dict) else {} + + +def _get(name: str, default): + return EVAL.get(name, default) + + +def _get_int(name: str, default: int) -> int: + try: + return int(EVAL.get(name, default)) + except Exception: + return default + + +def _get_float(name: str, default: float) -> float: + try: + return float(EVAL.get(name, default)) + except Exception: + return default + + +def _get_bool(name: str, default: bool) -> bool: + raw = EVAL.get(name, default) + if isinstance(raw, bool): + return raw + if isinstance(raw, str): + return raw.strip().lower() in {"1", "true", "yes", "on"} + return bool(raw) + + +# --- Per-task identity: the ONLY lines that differ across the four tasks. --- +# Kept as constants (not config) so the patch policy is enforceable locally / in +# CI without /judge/task_config.json. Everything else is config-driven and only +# needed on the GPU path. +PRIMITIVE = "dbscan" +PKG = "dbscanlib" +REF_MODULE = "refdbscan" + +MAX_PATCH_BYTES = _get_int("max_patch_bytes", 2_000_000) +MAX_CHANGED_FILES = _get_int("max_changed_files", 80) +CLEAN_SOURCE = Path(f"/opt/{PKG}-clean") +BASELINE_SOURCE = Path(f"/opt/{PRIMITIVE}_ref") +SPEEDUP_TARGET = _get_float("speedup_target", 8.0) + +ALLOWED_PATTERNS = (f"{PKG}/**",) +DENIED_PATTERNS = ( + "evaluator.py", + "flash_gpu.py", + "reference.patch", + f"{REF_MODULE}.py", + f"**/{REF_MODULE}.py", + "**/conftest.py", + "**/test_*.py", + "pyproject.toml", + "setup.py", + "setup.cfg", +) +# Defense in depth (the load-bearing anti-tamper defense is structural, inside +# the GPU worker: it captures its perf_counter / synchronize references before +# importing the submission, regenerates data every iteration, and re-verifies +# quality each iteration). These patterns are matched surgically so that +# legitimate vendored/optimized kernel source is not rejected: +# * external optimized libraries, matched as IMPORT statements (not bare words, +# so a comment mentioning a name is fine); none are installed on the GPU +# image anyway; +# * process / network / measurement-tamper patterns that never appear in a +# real GPU kernel. +FORBIDDEN_IMPORT_RE = re.compile( + r"(?:^|\n)\s*(?:import|from)\s+(flashlib|cuml|cudf|cupy|faiss|sklearn|scikit|cutlass|cuvs)\b" +) +FORBIDDEN_PATTERN_RE = re.compile( + r"\bsubprocess\b|\bsocket\b|\burllib\b|\brequests\b|os\.system|" + r"torch\.cuda\.synchronize\s*=|time\.perf_counter\s*=|setattr\(\s*torch\.cuda" +) + + +@dataclass(frozen=True) +class PatchFile: + old_path: str + new_path: str + added_lines: tuple[str, ...] + + @property + def path(self) -> str: + return self.new_path if self.new_path != "/dev/null" else self.old_path + + +def _match(path: str, patterns: tuple[str, ...]) -> bool: + return any(fnmatch.fnmatch(path, p) for p in patterns) + + +def _invalid(message: str, metrics: dict[str, Any] | None = None): + payload = metrics or {} + payload.setdefault("valid_patch", 0) + return 0.0, 0.0, message, payload + + +def _parse_patch(text: str) -> list[PatchFile]: + files: list[PatchFile] = [] + old = new = "" + added: list[str] = [] + in_file = False + for line in text.splitlines(): + if line.startswith("diff --git "): + if in_file: + files.append(PatchFile(old, new, tuple(added))) + in_file, old, new, added = True, "", "", [] + continue + if not in_file: + continue + if line.startswith("--- "): + old = line[4:].strip() + old = old[2:] if old.startswith("a/") else old + elif line.startswith("+++ "): + new = line[4:].strip() + new = new[2:] if new.startswith("b/") else new + elif line.startswith("+") and not line.startswith("+++ "): + added.append(line[1:]) + if in_file: + files.append(PatchFile(old, new, tuple(added))) + return files + + +def _validate_path(path: str) -> tuple[bool, str]: + if not path or path == "/dev/null": + return True, "" + if path.startswith("/") or ".." in Path(path).parts: + return False, f"unsafe patch path: {path}" + if _match(path, DENIED_PATTERNS): + return False, f"changed file is outside task boundary: {path}" + if not _match(path, ALLOWED_PATTERNS): + return False, f"changed file is not allowlisted (only {PKG}/** may change): {path}" + if not path.endswith(".py"): + return False, f"only Python files may change: {path}" + return True, "" + + +def validate_patch(patch_path: Path) -> tuple[bool, str, dict[str, Any]]: + if not patch_path.exists(): + return False, "solution patch does not exist", {} + size = patch_path.stat().st_size + if size > MAX_PATCH_BYTES: + return False, f"patch is too large ({size} bytes > {MAX_PATCH_BYTES})", {} + text = patch_path.read_text(encoding="utf-8", errors="replace") + files = _parse_patch(text) + metrics: dict[str, Any] = { + "patch_bytes": size, + "patch_sha256": hashlib.sha256(text.encode("utf-8", "replace")).hexdigest(), + "changed_files": len(files), + } + if len(files) > MAX_CHANGED_FILES: + return False, f"too many changed files ({len(files)} > {MAX_CHANGED_FILES})", metrics + for pf in files: + path = pf.path + if pf.new_path == "/dev/null": + return False, f"deleting files is outside task boundary: {pf.old_path}", metrics + if pf.old_path != "/dev/null" and pf.old_path != pf.new_path: + ok, err = _validate_path(pf.old_path) + if not ok: + return False, f"rename source is outside task boundary: {err}", metrics + ok, err = _validate_path(path) + if not ok: + return False, err, metrics + added = "\n".join(pf.added_lines) + m = FORBIDDEN_IMPORT_RE.search(added) + if m: + return False, f"{path}: forbidden import of an external optimized library ({m.group(1)})", metrics + m = FORBIDDEN_PATTERN_RE.search(added) + if m: + return False, f"{path}: forbidden pattern in added code ({m.group(0).strip()[:40]})", metrics + metrics["valid_patch"] = 1 + return True, "patch accepted by static policy", metrics + + +def sanitize(text: str) -> str: + text = re.sub(r"/tmp/[A-Za-z0-9_./-]+", "", text or "") + text = re.sub(r"/opt/[A-Za-z0-9_./-]+", "", text) + text = re.sub(r"\bN=\d+|\bseed[=:]?\s*\d+", "", text, flags=re.IGNORECASE) + return text[-600:] + + +def geometric_mean(values: list[float]) -> float: + if not values: + return 0.0 + return math.exp(sum(math.log(max(v, 1e-9)) for v in values) / len(values)) + + +def score_from_speedup(gm: float) -> float: + if gm <= 0: + return 0.0 + raw = 100.0 * math.log(gm) / math.log(max(SPEEDUP_TARGET, 1.0000001)) + return max(0.0, min(100.0, raw)) + + +def is_final_role() -> bool: + return os.environ.get("FRONTIER_SUBMISSION_ROLE", "agent") == "final" + + +def _workloads(final_role: bool) -> list[dict[str, Any]]: + workloads = list(_get("workloads", []) or []) + if not final_role: + n = _get_int("agent_workload_count", 3) + workloads = workloads[:n] + base = _get_int("base_seed", 20260701) + for i, w in enumerate(workloads): + w.setdefault("seed", base + 1000 * (i + 1)) + return workloads + + +def _build_cfg() -> dict[str, Any]: + return { + "primitive": PRIMITIVE, + "pkg": PKG, + "ref_module": REF_MODULE, + "gpu": str(_get("gpu", "H100")), + "cuda_image": str(_get("cuda_image", "nvidia/cuda:12.4.1-devel-ubuntu22.04")), + "pip": list(_get("pip", ["torch==2.5.1", "triton==3.1.0", "numpy"])), + "app_name": str(_get("app_name", "flash-kernel-eval")), + "modal_timeout_seconds": _get_int("modal_timeout_seconds", 1800), + "warmup": _get_int("warmup_iters", 3), + "iters": _get_int("timed_iters", 7), + "inertia_tolerance": _get_float("inertia_tolerance", 0.02), + "recall_threshold": _get_float("recall_threshold", 0.99), + "captured_tolerance": _get_float("captured_tolerance", 0.02), + "ortho_tolerance": _get_float("ortho_tolerance", 0.02), + "ari_threshold": _get_float("ari_threshold", 0.99), + } + + +def _read_dir(root: Path, rel_to: Path) -> dict[str, str]: + out: dict[str, str] = {} + for p in root.rglob("*.py"): + out[str(p.relative_to(rel_to))] = p.read_text(encoding="utf-8", errors="replace") + return out + + +def full_evaluation(patch_path: Path, metrics: dict[str, Any]): + final_role = is_final_role() + metrics["submission_role"] = "final" if final_role else "agent" + workloads = _workloads(final_role) + + # Import the Modal harness baked into the judge image; missing harness or + # missing credentials/sources -> policy smoke pass. + sys.path.insert(0, "/opt") + try: + import flash_gpu # type: ignore + except Exception: + metrics["full_benchmark"] = 0 + return 1.0, 1.0, "patch policy smoke passed; GPU harness not available in this environment", metrics + if not flash_gpu.modal_available() or not CLEAN_SOURCE.exists() or not BASELINE_SOURCE.exists(): + metrics["full_benchmark"] = 0 + return 1.0, 1.0, "patch policy smoke passed; Modal/judge sources not configured in this environment", metrics + + with tempfile.TemporaryDirectory(prefix="flash_kernel_opt_") as tmp: + tmp_root = Path(tmp) + patched = tmp_root / "patched" + shutil.copytree(CLEAN_SOURCE, patched) + env = {"PATH": os.environ.get("PATH", "/usr/bin:/bin"), "HOME": str(tmp_root), + "LC_ALL": "C", "LANG": "C"} + if int(metrics.get("changed_files", 0)) > 0: + try: + subprocess.run(["git", "apply", "--check", str(patch_path)], cwd=patched, env=env, + check=True, capture_output=True, text=True, timeout=60) + subprocess.run(["git", "apply", str(patch_path)], cwd=patched, env=env, + check=True, capture_output=True, text=True, timeout=60) + except subprocess.CalledProcessError as exc: + metrics["stderr_tail"] = sanitize(exc.stderr or "") + return _invalid("patch does not apply to the clean package tree", metrics) + metrics["applied_patch"] = 1 + else: + metrics["used_empty_patch"] = 1 + + payload = { + "baseline_files": _read_dir(BASELINE_SOURCE, BASELINE_SOURCE), + "patched_files": _read_dir(patched, patched), + "workloads": workloads, + "cfg": _build_cfg(), + } + try: + result = flash_gpu.run_remote(payload) + except Exception as exc: + metrics["error_detail"] = sanitize(str(exc)) + return _invalid("GPU evaluation failed", metrics) + + if not result.get("ok"): + metrics["worker_error"] = sanitize(str(result.get("error", "unknown"))) + return _invalid("GPU benchmark worker failed", metrics) + + speedups: list[float] = [] + per_workload: dict[str, Any] = {} + for row in result.get("rows", []): + if not row.get("ok"): + metrics["failed_workload"] = {"reason": row.get("reason", "error")} + return _invalid("submission failed the quality gate or crashed on a hidden workload", metrics) + speedups.append(max(float(row["speedup"]), 0.01)) + per_workload[row["id"]] = {"speedup": row["speedup"]} + if not speedups: + return _invalid("no workloads were evaluated", metrics) + + gm = geometric_mean(speedups) + bounded = score_from_speedup(gm) + metrics.update({"full_benchmark": 1, "workload_count": len(speedups), "geomean_speedup": gm}) + if _get_bool("expose_per_workload_metrics", False): + metrics["per_workload"] = per_workload + return bounded, bounded, f"{PRIMITIVE} geomean speedup {gm:.3f}x over the naive baseline", metrics + + +def evaluate(solution_path: str) -> tuple[float, float, str, dict[str, Any]]: + patch_path = Path(solution_path) + ok, message, metrics = validate_patch(patch_path) + if not ok: + return _invalid(message, metrics) + try: + return full_evaluation(patch_path, metrics) + except Exception as exc: # noqa: BLE001 + metrics["error_type"] = type(exc).__name__ + metrics["error_detail"] = sanitize(str(exc)) + return _invalid("evaluation failed", metrics) + + +def main(argv: list[str]) -> int: + if len(argv) != 2: + print("Usage: evaluator.py SOLUTION_PATCH", file=sys.stderr) + return 2 + score, unbounded, message, metrics = evaluate(argv[1]) + print(json.dumps({"score": score, "score_unbounded": unbounded, + "message": message, "metrics": metrics}, indent=2, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv)) diff --git a/2.0/problems/dbscan_gpu_kernel_optimization/flash_gpu.py b/2.0/problems/dbscan_gpu_kernel_optimization/flash_gpu.py new file mode 100644 index 000000000..bfc23d31a --- /dev/null +++ b/2.0/problems/dbscan_gpu_kernel_optimization/flash_gpu.py @@ -0,0 +1,372 @@ +"""Modal GPU evaluation harness for the flashlib kernel-optimization task family. + +Shared by the judge (evaluator.py) and the agent public test. Baked into both +the agent and judge images at ``/opt/flash_gpu.py``. It runs a frozen naive +baseline and the (patched or agent-edited) package on a Modal GPU and returns +per-workload timings + quality verdicts. + +Anti-reward-hack properties (all enforced inside the remote worker, which is the +only place the untrusted code runs): + +* **Timing primitives are captured before any agent code is imported** — the + worker binds ``time.perf_counter`` and ``torch.cuda.synchronize`` to locals up + front, so a submission that monkey-patches ``torch.cuda.synchronize`` (or any + torch attribute) cannot affect measurement. +* **Fresh data every timed iteration** — each measured iteration generates a new + input from a fresh seed, so a submission cannot memoize on a repeated input. +* **Quality is verified on every timed iteration** — the judge recomputes the + quality metric from the *returned tensors* on that iteration's data and gates + it, so returning a stale cached result for a different input is caught. +* **The judge computes all metrics** — no agent-provided number is trusted. +* **Fresh container per submission** — an ephemeral Modal app is used, so no + global state survives across evaluations. + +The remote worker is a single self-contained function serialized to Modal +(``serialized=True``), so the remote does not import this module and there is no +module-mounting to get wrong. The GPU image ships torch + triton (+ any extra +pip packages a task needs for its reference build). +""" +from __future__ import annotations + +import os + + +# --------------------------------------------------------------------------- # +# Remote worker (runs on the Modal GPU). Fully self-contained: every helper is +# nested so cloudpickle ships the whole thing. Takes and returns plain data. +# --------------------------------------------------------------------------- # +def _gpu_worker(payload: dict) -> dict: + import importlib + import os as _os + import sys as _sys + import tempfile + import time + import traceback + + import torch + + # Capture timing + sync references BEFORE importing any submission code, so a + # monkey-patch of torch.cuda.synchronize cannot affect measurement. + _perf = time.perf_counter + _sync = torch.cuda.synchronize + if not torch.cuda.is_available(): + return {"ok": False, "error": "cuda_unavailable"} + dev = "cuda" + + cfg = payload["cfg"] + prim = cfg["primitive"] + warmup = int(cfg.get("warmup", 3)) + iters = int(cfg.get("iters", 7)) + + def _materialize(files: dict, tag: str) -> str: + root = tempfile.mkdtemp(prefix=f"flash_{tag}_") + for rel, content in files.items(): + path = _os.path.join(root, rel) + parent = _os.path.dirname(path) + if parent: + _os.makedirs(parent, exist_ok=True) + with open(path, "w", encoding="utf-8") as fh: + fh.write(content) + _sys.path.insert(0, root) + return root + + _materialize(payload["baseline_files"], "baseline") + _materialize(payload["patched_files"], "patched") + ref = importlib.import_module(cfg["ref_module"]) + pkg = importlib.import_module(cfg["pkg"]) + + # ---- per-primitive: fixed context, data gen, call, quality metric ----- # + def setup(w, seed): + # Built ONCE per workload (not per timed iter). IVF-PQ builds the fixed + # index (the "database") here; queries are generated fresh per iteration. + if prim == "ivfpq": + g = torch.Generator(device=dev).manual_seed(int(seed)) + X = torch.randn(int(w["M"]), int(w["D"]), generator=g, device=dev, dtype=torch.float32) + index = ref.ivf_pq_build(X, int(w["nlist"]), m=int(w["m"]), + nprobe=int(w["nprobe"]), seed=int(seed)) + return {"index": index, "D": int(w["D"])} + return {} + + def gen(w, seed, ctx): + g = torch.Generator(device=dev).manual_seed(int(seed)) + if prim == "ivfpq": + q = torch.randn(int(w["Q"]), int(ctx["D"]), generator=g, device=dev, dtype=torch.float32) + return {"queries": q} + if prim == "knn": + db = torch.randn(w["M"], w["D"], generator=g, device=dev, dtype=torch.float32) + q = torch.randn(w["Q"], w["D"], generator=g, device=dev, dtype=torch.float32) + return {"queries": q, "database": db} + if prim == "dbscan": + # NON-CONVEX clustering data: a union of interleaving half-moons, each + # moon-PAIR placed in its own random 2-D subspace of R^D (so the cluster + # structure spans many dims -- a PCA->2D projection cannot capture it) and + # offset from the other pairs. Voronoi/centroid shortcuts get non-convex + # moons WRONG: "k-means the points into n_centers cells then DBSCAN each + # cell" (and plain k-means) mislabels the interleaving arcs, so the ARI + # gate rejects them (kmeans-shortcut ARI ~0.72-0.75, PCA-2D ~0.5-0.6 << gate). + # Only a real eps-neighbourhood density scan recovers the clusters -- which + # is exactly the kernel this task asks you to optimise. Clusters = the arcs + # (n_centers of them). (The generator is judge-only and NOT in the agent image.) + nc, D, N = int(w["n_centers"]), int(w["D"]), int(w["N"]) + pairs = max(1, nc // 2) + basis, _ = torch.linalg.qr(torch.randn(D, D, generator=g, device=dev)) # orthonormal + per = (N + pairs - 1) // pairs + parts = [] + for p in range(pairs): + nout = per // 2; nin = per - nout + to = torch.rand(nout, generator=g, device=dev) * torch.pi + ti = torch.rand(nin, generator=g, device=dev) * torch.pi + mx = torch.cat([torch.cos(to), 1.0 - torch.cos(ti)]) + my = torch.cat([torch.sin(to), 0.5 - torch.sin(ti)]) + m = (torch.stack([mx, my], 1) + torch.randn(per, 2, generator=g, device=dev) * 0.05) * 6.0 + b0 = basis[:, (2 * p) % D]; b1 = basis[:, (2 * p + 1) % D] + pts = m[:, 0:1] * b0[None, :] + m[:, 1:2] * b1[None, :] + pts = pts + torch.randn(per, D, generator=g, device=dev) * 0.1 # thin ambient noise + pts = pts + basis[:, (2 * p + 4) % D][None, :] * (p * 30.0) # separate the pairs + parts.append(pts) + x = torch.cat(parts, 0) + x = x[torch.randperm(x.shape[0], generator=g, device=dev)][:N] + # PRECISION LOCK (bf16): points carry only bf16 precision, so no solution + # wins by computing the eps-graph distances in a lower dtype. eps is set so + # the moons stay cleanly separated under bf16 -- an honest bf16 scan equals + # the fp32 clustering exactly (verified: two exact scans on the same bf16 + # data agree at ARI 1.0, so the gate is not flaky on border points). + return {"x": x.to(torch.bfloat16).to(torch.float32), "eps": float(w["eps"]), + "min_samples": int(w["min_samples"])} + x = torch.randn(w["N"], w["D"], generator=g, device=dev, dtype=torch.float32) + if prim == "kmeans": + perm = torch.randperm(w["N"], generator=g, device=dev)[: w["K"]] + return {"x": x, "init": x.index_select(0, perm).clone()} + return {"x": x} + + def call(mod, w, data, ctx): + if prim == "ivfpq": + return mod.ivf_pq_search(ctx["index"], data["queries"], int(w["k"]), nprobe=int(w["nprobe"])) + if prim == "kmeans": + return mod.kmeans(data["x"], w["K"], max_iters=w["max_iters"], + init_centroids=data["init"], tol=0.0) + if prim == "knn": + return mod.knn(data["queries"], data["database"], w["k"]) + if prim == "dbscan": + return mod.dbscan(data["x"], data["eps"], data["min_samples"]) + if prim == "pca": + return mod.pca(data["x"], w["k"]) + if prim == "tsvd": + return mod.truncated_svd(data["x"], w["k"]) + raise ValueError(prim) + + def check_shape(w, out): + if prim == "kmeans": + _, cen, _ = out + if tuple(cen.shape) != (w["K"], w["D"]) or not torch.isfinite(cen).all(): + raise ValueError("bad kmeans output") + elif prim == "knn": + d, i = out + if tuple(d.shape) != (w["Q"], w["k"]) or tuple(i.shape) != (w["Q"], w["k"]): + raise ValueError("bad knn output") + if not torch.isfinite(d).all(): + raise ValueError("non-finite distances") + elif prim == "ivfpq": + vals, ids = out + if tuple(ids.shape) != (int(w["Q"]), int(w["k"])) or \ + tuple(vals.shape) != (int(w["Q"]), int(w["k"])): + raise ValueError("bad ivfpq output shape") + elif prim == "dbscan": + if out.ndim != 1 or int(out.shape[0]) != int(w["N"]): + raise ValueError("bad dbscan labels shape") + else: + a, b = out + comps = a if prim == "pca" else b + if tuple(comps.shape) != (w["k"], w["D"]) or not torch.isfinite(comps).all(): + raise ValueError("bad decomposition output") + + def inertia(x, centroids): + c = centroids.to(torch.float32) + cn = (c * c).sum(1) + total = 0.0 + for j in range(0, x.shape[0], 16384): + xb = x[j:j + 16384] + d = (xb * xb).sum(1, keepdim=True) - 2.0 * (xb @ c.t()) + cn[None, :] + total += float(d.min(1).values.clamp_min(0).sum().item()) + return total + + def recall(agent_idx, ref_idx): + a = agent_idx.long(); b = ref_idx.long() + hit = (a.unsqueeze(2) == b.unsqueeze(1)).any(2) + return float(hit.sum().item()) / (b.shape[0] * b.shape[1]) + + def ari(a, b): + # Adjusted Rand Index between two integer label vectors (N,); noise (-1) + # is treated as its own label (dense-remapped first). + a = torch.unique(a.long(), return_inverse=True)[1] + b = torch.unique(b.long(), return_inverse=True)[1] + na = int(a.max().item()) + 1; nb = int(b.max().item()) + 1 + cont = torch.bincount(a * nb + b, minlength=na * nb).reshape(na, nb).double() + ai = cont.sum(1); bj = cont.sum(0); n = cont.sum() + c2 = lambda z: z * (z - 1) / 2.0 + sij = c2(cont).sum(); sa = c2(ai).sum(); sb = c2(bj).sum() + exp = (sa * sb / c2(n)) if float(n) > 1 else 0.0 + den = 0.5 * (sa + sb) - exp + return float(((sij - exp) / den).item()) if float(den) != 0.0 else 1.0 + + def ortho_err(comps): + c = comps.to(torch.float32); k = c.shape[0] + return float((c @ c.t() - torch.eye(k, device=c.device)).abs().max().item()) + + def captured(x, comps, center): + c = comps.to(torch.float32) + mean = x.mean(0) if center else None + total = 0.0 + for j in range(0, x.shape[0], 16384): + xb = x[j:j + 16384] + if center: + xb = xb - mean + p = xb @ c.t() + total += float((p * p).sum().item()) + return total / (x.shape[0] - 1) if center else total + + def verdict(w, data, ctx, ref_out, agent_out): + """Return (ok, reason, ref_val, agent_val) gating the agent output.""" + check_shape(w, agent_out) + if prim == "ivfpq": + # iso-result recall: agent ids vs frozen-baseline ids on the SAME index. + rec = recall(agent_out[1], ref_out[1]) + ok = rec >= float(cfg["recall_threshold"]) + return ok, ("recall_regression" if not ok else ""), 1.0, rec + if prim == "kmeans": + rv = inertia(data["x"], ref_out[1]); av = inertia(data["x"], agent_out[1]) + tol = float(cfg["inertia_tolerance"]) + ok = av <= (1.0 + tol) * rv + 1e-6 + return ok, ("inertia_regression" if not ok else ""), rv, av + if prim == "knn": + rec = recall(agent_out[1], ref_out[1]) + ok = rec >= float(cfg["recall_threshold"]) + return ok, ("recall_regression" if not ok else ""), 1.0, rec + if prim == "dbscan": + val = ari(agent_out, ref_out) + ok = val >= float(cfg["ari_threshold"]) + return ok, ("ari_regression" if not ok else ""), 1.0, val + center = (prim == "pca") + comps = agent_out[0] if prim == "pca" else agent_out[1] + ref_comps = ref_out[0] if prim == "pca" else ref_out[1] + oe = ortho_err(comps) + if oe > float(cfg["ortho_tolerance"]): + return False, "not_orthonormal", 0.0, oe + rv = captured(data["x"], ref_comps, center) + av = captured(data["x"], comps, center) + ok = av >= (1.0 - float(cfg["captured_tolerance"])) * rv - 1e-6 + return ok, ("captured_regression" if not ok else ""), rv, av + + def time_call(mod, w, data, ctx): + _sync(); t0 = _perf(); out = call(mod, w, data, ctx); _sync() + return (_perf() - t0) * 1000.0, out + + rows = [] + try: + for w in payload["workloads"]: + base_seed = int(w["seed"]) + ctx = setup(w, base_seed) # fixed context (e.g. IVF-PQ index) + # warmup on fresh data (kernels / autotune) — not measured + for i in range(warmup): + d = gen(w, base_seed + 100 + i, ctx) + call(ref, w, d, ctx); call(pkg, w, d, ctx); _sync() + del d + torch.cuda.empty_cache() + ratios, ref_val, agent_val = [], None, None + bad = None + for i in range(iters): + d = gen(w, base_seed + 10000 + i, ctx) # fresh every iteration + rt, rout = time_call(ref, w, d, ctx) + at, aout = time_call(pkg, w, d, ctx) + ok, reason, rv, av = verdict(w, d, ctx, rout, aout) # verify THIS iter + if not ok: + bad = reason; ref_val, agent_val = rv, av + break + ratios.append(rt / at if at > 0 else 0.01) + ref_val, agent_val = rv, av + del d, rout, aout + torch.cuda.empty_cache() + if bad is not None: + rows.append({"id": w["id"], "ok": False, "reason": bad, + "ref_val": ref_val, "agent_val": agent_val}) + continue + ratios.sort() + rows.append({"id": w["id"], "ok": True, + "speedup": ratios[len(ratios) // 2], + "ref_val": ref_val, "agent_val": agent_val}) + except Exception as exc: + return {"ok": False, "error": f"{type(exc).__name__}: {str(exc)[:200]}", + "trace": traceback.format_exc()[-500:]} + return {"ok": True, "rows": rows} + + +# --------------------------------------------------------------------------- # +# Host-side wrappers (run in the judge / agent container; need MODAL tokens). +# --------------------------------------------------------------------------- # +def _build_image(cfg: dict): + import modal + pip = list(cfg.get("pip", ["torch==2.5.1", "triton==3.1.0", "numpy"])) + base = cfg.get("cuda_image", "nvidia/cuda:12.4.1-devel-ubuntu22.04") + # add_python must match the judge/agent container Python (ubuntu:24.04 -> 3.12) + # because the worker ships via serialized=True (cloudpickle is version-sensitive). + # add_local_python_source mounts THIS module into the image so the remote can + # import it when deserializing the (module-level) worker function. + return (modal.Image.from_registry(base, add_python=str(cfg.get("python", "3.12"))) + .entrypoint([]) + .pip_install(*pip) + .add_local_python_source("flash_gpu")) + + +# Substrings that mark a transient Modal control-plane / image-build failure +# (evicted build, gateway hiccup, app stopped mid-run) rather than a real error +# in the submission — safe to retry. +_TRANSIENT_MARKERS = ( + "external shut-down", "terminated due to external", "please try again", + "app_state_stopped", "conflicterror", "deadline exceeded", "connection reset", + "502 bad gateway", "503 service", "temporarily unavailable", "timed out", + "gateway", "eviction", "internalfailure", "internalerror", + "failed to get new inputs", "runner failed", "task exited", +) + + +def _is_transient(text: str) -> bool: + low = (text or "").lower() + return any(m in low for m in _TRANSIENT_MARKERS) + + +def run_remote(payload: dict) -> dict: + """Run the GPU worker on Modal and return its result dict. + + Retries transient Modal control-plane failures; a real error in the + submission is non-transient and surfaces immediately. Raises RuntimeError + with a sanitized message on unrecoverable failure. + """ + import time as _time + + import modal + + cfg = payload["cfg"] + attempts = max(1, int(cfg.get("modal_retries", 3))) + last = "modal_gpu_eval_failed" + for attempt in range(1, attempts + 1): + app = modal.App(cfg.get("app_name", "flash-kernel-eval")) + remote = app.function( + gpu=cfg.get("gpu", "H100"), + image=_build_image(cfg), + timeout=int(cfg.get("modal_timeout_seconds", 1800)), + serialized=True, + max_containers=1, + )(_gpu_worker) + try: + with modal.enable_output(), app.run(): + return remote.remote(payload) + except Exception as exc: # noqa: BLE001 + last = str(exc)[-400:] + if not _is_transient(last) or attempt == attempts: + raise RuntimeError(f"modal_gpu_eval_failed: {last}") from exc + _time.sleep(min(45, 10 * attempt)) + raise RuntimeError(f"modal_gpu_eval_failed: {last}") + + +def modal_available() -> bool: + return bool(os.environ.get("MODAL_TOKEN_ID") and os.environ.get("MODAL_TOKEN_SECRET")) diff --git a/2.0/problems/dbscan_gpu_kernel_optimization/harbor/app/README.md b/2.0/problems/dbscan_gpu_kernel_optimization/harbor/app/README.md new file mode 100644 index 000000000..9a9b2979e --- /dev/null +++ b/2.0/problems/dbscan_gpu_kernel_optimization/harbor/app/README.md @@ -0,0 +1,34 @@ +# DBSCAN kernel optimization — submission workflow + +You are optimizing the `dbscanlib` package at `/app/dbscanlib`. Edit the package +(rewrite the internals of `dbscan`, add Triton kernel modules under +`dbscanlib/`), then submit a patch. + +## Iterate locally + +Run the public self-test to check correctness and get a rough speed signal on +the two public shapes (needs a GPU in the agent container): + +```bash +bash /app/public_test.sh +``` + +## Submit + +```bash +bash /app/make_submission.sh # writes /app/solution.patch (dbscanlib diff) +bash /app/submit.sh # enqueues it for the black-box judge +``` + +Submissions are asynchronous. Submit early and keep improving; use +`bash /app/submissions.sh` and `bash /app/wait_submission.sh ` to inspect +results. + +## Rules + +- Only files under `dbscanlib/` may change. +- Do not import external optimized libraries (write the kernels yourself), and + do not access the environment, spawn processes, or use the network. +- Keep the public `dbscan(...)` signature and return contract unchanged. +- Clustering quality is gated (inertia vs the naive baseline); do not sacrifice + correctness for speed beyond the allowed tolerance. diff --git a/2.0/problems/dbscan_gpu_kernel_optimization/harbor/app/make_submission.sh b/2.0/problems/dbscan_gpu_kernel_optimization/harbor/app/make_submission.sh new file mode 100644 index 000000000..44ef00d34 --- /dev/null +++ b/2.0/problems/dbscan_gpu_kernel_optimization/harbor/app/make_submission.sh @@ -0,0 +1,24 @@ +#!/usr/bin/env bash +# Package the current dbscanlib edits into a unified diff for submission. +set -euo pipefail + +APP_DIR="${APP_DIR:-/app}" +OUT="${1:-/app/solution.patch}" + +if [[ ! -d "$APP_DIR/dbscanlib" ]]; then + echo "dbscanlib package not found at $APP_DIR/dbscanlib" >&2 + exit 2 +fi +if [[ ! -d "$APP_DIR/.git" ]]; then + echo "git repo not found at $APP_DIR (expected in the agent image)" >&2 + exit 2 +fi + +# Stage and diff only the package, so submission helper scripts under /app are +# never included. New kernel files under dbscanlib/ are captured via `git add`. +git -C "$APP_DIR" add -A -- dbscanlib +git -C "$APP_DIR" diff --cached -- dbscanlib > "$OUT" +git -C "$APP_DIR" reset -q + +bytes=$(wc -c < "$OUT" | tr -d ' ') +echo "Wrote $OUT ($bytes bytes)" diff --git a/2.0/problems/dbscan_gpu_kernel_optimization/harbor/app/public_test.py b/2.0/problems/dbscan_gpu_kernel_optimization/harbor/app/public_test.py new file mode 100644 index 000000000..d944a88cc --- /dev/null +++ b/2.0/problems/dbscan_gpu_kernel_optimization/harbor/app/public_test.py @@ -0,0 +1,62 @@ +"""Public self-test == the graded evaluation, run by the judge. + +The data generator is judge-only and is NOT present in this (agent) image, so you +cannot read how the workloads are produced and hardcode a data-specific shortcut. +Instead, this packages your current `/app/dbscanlib`, submits it to the judge -- +which runs the EXACT graded workloads on a GPU and grades your clustering against +the exact reference (ARI) -- and prints your per-workload result + score. It is +byte-for-byte the same evaluation used for your final grade. + +(Submission is asynchronous; this waits for the result. You can also submit +directly with `bash /app/make_submission.sh && bash /app/submit.sh` and poll with +`bash /app/submissions.sh`.) +""" +from __future__ import annotations + +import json +import re +import subprocess +import sys +from pathlib import Path + +SUBMISSIONS_LOG = Path("/logs/agent/submissions.jsonl") +_UUID = re.compile(r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}") + + +def _latest_uuid_from_log() -> str | None: + if not SUBMISSIONS_LOG.exists(): + return None + lines = [l for l in SUBMISSIONS_LOG.read_text(encoding="utf-8").splitlines() if l.strip()] + for l in reversed(lines): + try: + u = json.loads(l).get("submission_uuid") + if u: + return u + except Exception: # noqa: BLE001 + continue + return None + + +def main() -> int: + # 1) package the current package into /app/solution.patch + subprocess.run(["bash", "/app/make_submission.sh"], check=True) + # 2) submit to the judge (async); it runs the exact graded eval on GPU + print("Submitting to the judge (runs the exact graded workloads + returns your score)...", + flush=True) + r = subprocess.run(["python3", "/app/submit.py", *sys.argv[1:]], + capture_output=True, text=True) + sys.stdout.write(r.stdout) + sys.stderr.write(r.stderr) + if r.returncode != 0: + return r.returncode + # 3) find the submission uuid and wait for the result + m = _UUID.findall(r.stdout + r.stderr) + uuid = m[-1] if m else _latest_uuid_from_log() + if not uuid: + print("Submitted. Poll for the result with: bash /app/submissions.sh") + return 0 + return subprocess.call(["python3", "/app/wait_submission.py", uuid]) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/2.0/problems/dbscan_gpu_kernel_optimization/harbor/app/public_test.sh b/2.0/problems/dbscan_gpu_kernel_optimization/harbor/app/public_test.sh new file mode 100644 index 000000000..1038fbfb4 --- /dev/null +++ b/2.0/problems/dbscan_gpu_kernel_optimization/harbor/app/public_test.sh @@ -0,0 +1,7 @@ +#!/usr/bin/env bash +# Public self-test: run the patched dbscanlib on the public shapes and print a +# correctness + rough speedup summary. Requires a GPU in the agent container. +set -euo pipefail +APP_DIR="${APP_DIR:-/app}" +cd "$APP_DIR" +exec python3 "$APP_DIR/public_test.py" diff --git a/2.0/problems/dbscan_gpu_kernel_optimization/harbor/app/solution.patch b/2.0/problems/dbscan_gpu_kernel_optimization/harbor/app/solution.patch new file mode 100644 index 000000000..e69de29bb diff --git a/2.0/problems/dbscan_gpu_kernel_optimization/judge/refdbscan.py b/2.0/problems/dbscan_gpu_kernel_optimization/judge/refdbscan.py new file mode 100644 index 000000000..ab7fa8bfc --- /dev/null +++ b/2.0/problems/dbscan_gpu_kernel_optimization/judge/refdbscan.py @@ -0,0 +1,70 @@ +"""Naive DBSCAN baseline (correct, deterministic, GPU/CPU torch). + +Euclidean DBSCAN. Chunked O(N^2) neighbour scan + GPU-friendly label +propagation for the connected components of the core graph. Border rule matches +flashlib: a non-core point joins the SMALLEST cluster label among its in-eps +core neighbours; otherwise it is noise (-1). +""" +from __future__ import annotations + +import torch + + +def dbscan(x, eps, min_samples, max_neighbors=32): + del max_neighbors # naive path scans all neighbours; knob only for the fast kernel + N = x.shape[0] + dev = x.device + eps2 = float(eps) * float(eps) + xn = (x * x).sum(1) # (N,) + CH = 4096 + BIG = N # sentinel label (> any real label) + + def d2_chunk(lo, hi): + xb = x[lo:hi] + return (xb * xb).sum(1, keepdim=True) - 2.0 * (xb @ x.t()) + xn[None, :] + + # 1) degree (in-eps count, includes self) -> core mask + deg = torch.zeros(N, device=dev, dtype=torch.int64) + for lo in range(0, N, CH): + hi = min(lo + CH, N) + deg[lo:hi] = (d2_chunk(lo, hi) <= eps2).sum(1) + core = deg >= int(min_samples) # (N,) bool + + # 2) connected components of the core-core eps graph via label propagation + labels = torch.arange(N, device=dev, dtype=torch.int64) + labels[~core] = BIG + for _ in range(1000): + new = labels.clone() + for lo in range(0, N, CH): + hi = min(lo + CH, N) + if not bool(core[lo:hi].any()): + continue + adj = (d2_chunk(lo, hi) <= eps2) & core[None, :] # (chunk, N) + lab = torch.where(adj, labels[None, :], torch.full((1, N), BIG, device=dev, dtype=torch.int64)) + mn = lab.min(1).values # (chunk,) + rows = slice(lo, hi) + upd = core[rows] & (mn < new[rows]) + new[rows] = torch.where(upd, mn, new[rows]) + if torch.equal(new, labels): + break + labels = new + + # 3) border assignment: non-core within eps of a core -> smallest core label + for lo in range(0, N, CH): + hi = min(lo + CH, N) + nb = ~core[lo:hi] + if not bool(nb.any()): + continue + adj = (d2_chunk(lo, hi) <= eps2) & core[None, :] + lab = torch.where(adj, labels[None, :], torch.full((1, N), BIG, device=dev, dtype=torch.int64)) + mn = lab.min(1).values + rows = slice(lo, hi) + take = nb & (mn < BIG) + labels[rows] = torch.where(take, mn, labels[rows]) + + # 4) compact to dense [0, n_clusters); noise (BIG) -> -1 + out = torch.full((N,), -1, device=dev, dtype=torch.int64) + uniq = torch.unique(labels[labels < BIG]) + for new_id, old in enumerate(uniq.tolist()): + out[labels == old] = new_id + return out diff --git a/2.0/problems/dbscan_gpu_kernel_optimization/readme b/2.0/problems/dbscan_gpu_kernel_optimization/readme new file mode 100644 index 000000000..648a5a6f8 --- /dev/null +++ b/2.0/problems/dbscan_gpu_kernel_optimization/readme @@ -0,0 +1,99 @@ +# GPU DBSCAN Kernel Optimization + +## Problem + +You are given a small GPU DBSCAN library, `dbscanlib`, in the Harbor workspace at +`/app/dbscanlib`. Its public entry point is: + +```python +dbscanlib.dbscan(x, eps, min_samples) -> labels +``` + +`x` is an `(N, D)` float32 CUDA tensor of low-dimensional points, `eps` is the +neighbourhood radius, and `min_samples` is the core-point threshold. It runs +Euclidean DBSCAN and returns `labels`, an `(N,)` int64 tensor with a cluster id +`>= 0` per point and `-1` for noise. The shipped implementation is a correct but +straightforward version. + +Your goal is to make `dbscanlib.dbscan` **as fast as possible** on the GPU while +producing the same clustering. You may rewrite the internals of the package and +add new modules (including Triton kernels) under `dbscanlib/`. The public +function signature and return contract must not change, and the result must +remain a deterministic function of the inputs. + +## Workload + +The graded workloads are held-out `(N, D)` low-count-cluster point sets with +`D >= 8`, varying `N`, `D`, cluster count, `eps`, and `min_samples`. Treat it as +**general Euclidean DBSCAN** and reproduce the *exact* clustering: your labels are +graded against the exact reference clustering (Adjusted Rand Index), so an +approximation that only roughly recovers the clusters will not pass. How the point +sets are generated is not disclosed — solve the general problem, do not special-case. + +## Iterate on a GPU + +The agent workspace has no GPU, and the data generator is judge-only (it is not in +your image). To check your current code, **submit it to the judge** — it runs the +exact graded workloads on a GPU and returns your per-workload result + score. This +one command packages, submits, and waits for the result: + +```bash +bash /app/public_test.sh +``` + +It reports, per graded workload, pass/fail on the quality gate and your speedup, +plus the geometric-mean speedup and your score (0-100) — the identical evaluation +used for your final grade. Any workload that fails its gate scores 0. (Equivalently: +`bash /app/make_submission.sh && bash /app/submit.sh`, then poll with +`bash /app/submissions.sh`.) Submissions are asynchronous; submit early and iterate. + +## Submission + +The submitted artifact is a patch over the `dbscanlib` package: + +```text +/app/solution.patch +``` + +After editing `/app/dbscanlib`, run `bash /app/make_submission.sh` then +`bash /app/submit.sh`. Submissions are asynchronous; submit early and iterate. +The judge applies your patch to a clean copy of `dbscanlib` and times it against +the original baseline on a GPU, on the same seeded data. + +## Correctness + +Correctness is a gate. On every timed iteration the judge compares your +clustering to the baseline's on identical data using the **Adjusted Rand Index** +(a permutation-invariant clustering-agreement score that also accounts for +noise), and requires it to stay at or above a high threshold. Crashes, +non-finite / wrong-shape output, timeouts, and clusterings that disagree with the +baseline beyond the tolerance are penalized before speed is considered. + +## Scoring + +Valid submissions are scored by speedup relative to the baseline: + +```text +speedup = baseline_time / your_time +``` + +The objective is the geometric mean of per-workload speedups. A result no faster +than the baseline earns 0; regressions earn 0. The raw geometric-mean speedup is +reported in the evaluator metrics. + +## Patch Policy + +Only Python files under `dbscanlib/**` may change. New Python modules inside +`dbscanlib/` are allowed. Patches may not: modify anything outside `dbscanlib/`; +import or call an external optimized ML/kernel library (write the kernels +yourself); read/write environment variables, spawn processes, or access the +network; or tamper with the measurement framework. The timing harness measures +your code on freshly generated data every iteration and re-verifies clustering +agreement each time. + +## Resource Budget + +```text +GPU: single Modal GPU (H100 reference; Triton paths also run on L40S / A100) +Agent container: CPU-only (GPU work is offloaded to Modal) +``` diff --git a/2.0/problems/dbscan_gpu_kernel_optimization/reference.patch b/2.0/problems/dbscan_gpu_kernel_optimization/reference.patch new file mode 100644 index 000000000..40170a32d --- /dev/null +++ b/2.0/problems/dbscan_gpu_kernel_optimization/reference.patch @@ -0,0 +1,3149 @@ +diff --git a/dbscanlib/_kernels/__init__.py b/dbscanlib/_kernels/__init__.py +new file mode 100644 +index 0000000..e69de29 +diff --git a/dbscanlib/_kernels/kernels/__init__.py b/dbscanlib/_kernels/kernels/__init__.py +new file mode 100644 +index 0000000..e69de29 +diff --git a/dbscanlib/_kernels/kernels/distance/__init__.py b/dbscanlib/_kernels/kernels/distance/__init__.py +new file mode 100644 +index 0000000..e69de29 +diff --git a/dbscanlib/_kernels/kernels/distance/triton/__init__.py b/dbscanlib/_kernels/kernels/distance/triton/__init__.py +new file mode 100644 +index 0000000..63fceb5 +--- /dev/null ++++ b/dbscanlib/_kernels/kernels/distance/triton/__init__.py +@@ -0,0 +1,11 @@ ++"""distance triton backend (kNN squared-L2 gather only). ++ ++Minimal re-export: only the tiled squared-L2 gather kernel that recovers ++true ``||x - c[idx]||^2`` per neighbour after the fused kNN pass. The ++other distance kernels from the upstream backend are not vendored here. ++""" ++from dbscanlib._kernels.kernels.distance.triton.knn_gather_l2sq import ( ++ triton_knn_gather_sqdist, ++) ++ ++__all__ = ["triton_knn_gather_sqdist"] +diff --git a/dbscanlib/_kernels/kernels/distance/triton/knn_gather_l2sq.py b/dbscanlib/_kernels/kernels/distance/triton/knn_gather_l2sq.py +new file mode 100644 +index 0000000..e2e3fbe +--- /dev/null ++++ b/dbscanlib/_kernels/kernels/distance/triton/knn_gather_l2sq.py +@@ -0,0 +1,232 @@ ++"""Tiled gather kernel: true squared L2 to selected kNN neighbours. ++ ++After a fused kNN pass that ranks candidates with the shift-invariant ++score ``c_sq - 2 * ``, this kernel writes the **true** ++``||x[b, n] - c[b, idx[b, n, k]]||^2`` for every neighbour. Computing ++the difference directly (``(x - y)^2`` per d) avoids the cancellation ++inherent to the ``x^2 + y^2 - 2*x*y`` expansion, so the result is ++accurate even for fp32 inputs where the fused kNN kernel may have used ++TF32 in the cross-term. ++ ++Complexity: ``O(B * N * K * D)`` -- cheap vs the ``O(B * N * M * D)`` ++GEMM that found the candidates. ++ ++Design ++------ ++One program owns a ``(BN, K_BLOCK)`` output tile. The query row tile ++``x[b, n_block, :]`` is loaded once and reused across the ``K_BLOCK`` ++neighbours, then streamed through ``D`` in ``BLOCK_D`` chunks. The ++corpus rows ``c[b, idx, :]`` are gather-loaded fresh per d-chunk into a ++3-D tile ``(BN, K_BLOCK, BLOCK_D)`` and subtracted directly from the ++broadcast x tile in fp32. The fp32 squared difference is summed along ++``D`` into the accumulator. ++ ++Three regime presets pick ``BN`` and ``K_BLOCK`` so the c-tile fits in ++SMEM/registers and amortises one x-load over many neighbours: ++ ++ * K <= 32: ``K_BLOCK = next_pow2(K)``, ``BN = 16``. ++ * 32 < K <= 512: ``K_BLOCK = 32``, ``BN = 8``; multiple K-tiles per row. ++ * K > 512: ``K_BLOCK = 32``, ``BN = 4``. ++ ++``BLOCK_D`` is chosen so the per-program SMEM footprint ++``BN * K_BLOCK * BLOCK_D * sizeof(dtype)`` stays under ~32 KB. The ++unmodified ``D`` (not padded) is passed as a constexpr so the Triton ++compiler unrolls the d-loop exactly. ++""" ++from __future__ import annotations ++ ++import torch ++import triton ++import triton.language as tl ++ ++ ++def _next_pow2(n: int) -> int: ++ if n <= 1: ++ return 1 ++ return 1 << (n - 1).bit_length() ++ ++ ++@triton.jit ++def _knn_gather_l2sq_kernel( ++ x_ptr, c_ptr, idx_ptr, out_ptr, ++ M, ++ stride_x_b, stride_x_n, stride_x_d, ++ stride_c_b, stride_c_m, stride_c_d, ++ stride_i_b, stride_i_n, stride_i_k, ++ stride_o_b, stride_o_n, stride_o_k, ++ N: tl.constexpr, D: tl.constexpr, K: tl.constexpr, ++ BN: tl.constexpr, K_BLOCK: tl.constexpr, ++ BLOCK_D: tl.constexpr, ++ SINGLE_D_TILE: tl.constexpr, ++): ++ """Tiled (BN, K_BLOCK, BLOCK_D) gather + sq-distance accumulate. ++ ++ Grid: ``(ceil(N / BN), ceil(K / K_BLOCK), B)``. ++ ++ Each program owns one ``(BN, K_BLOCK)`` output tile. The fp32 ++ accumulator stays in registers; the c-row tile ``(BN, K_BLOCK, ++ BLOCK_D)`` is reloaded per d-chunk to keep SMEM bounded. ++ """ ++ pid_n = tl.program_id(0) ++ pid_k = tl.program_id(1) ++ pid_b = tl.program_id(2).to(tl.int64) ++ ++ n_start = pid_n * BN ++ n_offs = (n_start + tl.arange(0, BN)).to(tl.int64) ++ n_mask = n_offs < N ++ ++ k_start = pid_k * K_BLOCK ++ k_offs = (k_start + tl.arange(0, K_BLOCK)).to(tl.int64) ++ k_mask = k_offs < K ++ ++ idx_tile = tl.load( ++ idx_ptr + pid_b * stride_i_b ++ + n_offs[:, None] * stride_i_n ++ + k_offs[None, :] * stride_i_k, ++ mask=n_mask[:, None] & k_mask[None, :], other=0, ++ ).to(tl.int64) ++ idx_tile = tl.maximum(idx_tile, 0) ++ idx_tile = tl.minimum(idx_tile, M - 1) ++ ++ acc = tl.zeros((BN, K_BLOCK), dtype=tl.float32) ++ ++ if SINGLE_D_TILE: ++ d_offs = tl.arange(0, BLOCK_D).to(tl.int64) ++ d_mask = d_offs < D ++ x_tile = tl.load( ++ x_ptr + pid_b * stride_x_b ++ + n_offs[:, None] * stride_x_n ++ + d_offs[None, :] * stride_x_d, ++ mask=n_mask[:, None] & d_mask[None, :], other=0.0, ++ ).to(tl.float32) ++ c_tile = tl.load( ++ c_ptr + pid_b * stride_c_b ++ + idx_tile[:, :, None] * stride_c_m ++ + d_offs[None, None, :] * stride_c_d, ++ mask=(n_mask[:, None] & k_mask[None, :])[:, :, None] ++ & d_mask[None, None, :], other=0.0, ++ ).to(tl.float32) ++ diff = x_tile[:, None, :] - c_tile ++ acc = tl.sum(diff * diff, axis=2) ++ else: ++ for d_start in tl.range(0, D, BLOCK_D): ++ d_offs = (d_start + tl.arange(0, BLOCK_D)).to(tl.int64) ++ d_mask = d_offs < D ++ x_tile = tl.load( ++ x_ptr + pid_b * stride_x_b ++ + n_offs[:, None] * stride_x_n ++ + d_offs[None, :] * stride_x_d, ++ mask=n_mask[:, None] & d_mask[None, :], other=0.0, ++ ).to(tl.float32) ++ c_tile = tl.load( ++ c_ptr + pid_b * stride_c_b ++ + idx_tile[:, :, None] * stride_c_m ++ + d_offs[None, None, :] * stride_c_d, ++ mask=(n_mask[:, None] & k_mask[None, :])[:, :, None] ++ & d_mask[None, None, :], other=0.0, ++ ).to(tl.float32) ++ diff = x_tile[:, None, :] - c_tile ++ acc += tl.sum(diff * diff, axis=2) ++ ++ tl.store( ++ out_ptr + pid_b * stride_o_b ++ + n_offs[:, None] * stride_o_n ++ + k_offs[None, :] * stride_o_k, ++ acc, mask=n_mask[:, None] & k_mask[None, :], ++ ) ++ ++ ++def _pick_tile(K: int, D: int, dtype_bytes: int) -> tuple[int, int, int, bool]: ++ """Pick ``(BN, K_BLOCK, BLOCK_D, single_d_tile)`` for the kernel. ++ ++ Regimes mirror the file docstring -- small ``K_BLOCK`` for large K, ++ larger ``BN`` for small K. ``BLOCK_D`` is sized so the 3-D c-tile ++ SMEM footprint stays around ~32 KB. ++ """ ++ K_pad = _next_pow2(K) ++ ++ if K_pad <= 32: ++ K_BLOCK = max(1, K_pad) ++ BN = 16 ++ elif K_pad <= 512: ++ K_BLOCK = 32 ++ BN = 8 ++ else: ++ K_BLOCK = 32 ++ BN = 4 ++ ++ target_bytes = 32 * 1024 ++ per_d = BN * K_BLOCK * dtype_bytes ++ block_d_cap = max(16, target_bytes // max(1, per_d)) ++ block_d_cap = min(block_d_cap, 256) ++ ++ if D <= block_d_cap: ++ BLOCK_D = max(16, _next_pow2(D)) ++ single_d_tile = True ++ else: ++ BLOCK_D = 64 if block_d_cap >= 64 else 32 ++ single_d_tile = False ++ ++ return BN, K_BLOCK, BLOCK_D, single_d_tile ++ ++ ++def triton_knn_gather_sqdist( ++ x: torch.Tensor, ++ c: torch.Tensor, ++ idx: torch.Tensor, ++ *, ++ out: torch.Tensor | None = None, ++) -> torch.Tensor: ++ """``out[b, n, k] = || x[b, n, :] - c[b, idx[b, n, k], :] ||^2`` (fp32). ++ ++ Args: ++ x: (B, N, D) query tensor, any float cuda dtype (bf16 / fp16 / fp32). ++ c: (B, M, D) corpus, same dtype and same batch as ``x``. ++ idx: (B, N, K) int32 / int64 column indices into ``c``. ++ out: optional pre-allocated (B, N, K) fp32 output buffer. ++ ++ Returns: ++ (B, N, K) fp32. The computation runs in fp32 throughout ++ (``diff = x.to(fp32) - c.to(fp32)``, ``sum(diff*diff)``) so the ++ result is bit-equivalent to the naive torch reference up to ++ accumulation order on the d-axis. ++ """ ++ assert x.is_cuda and c.is_cuda and idx.is_cuda ++ assert x.dim() == 3 and c.dim() == 3 and idx.dim() == 3 ++ B, N, D = x.shape ++ Bc, M, Dc = c.shape ++ Bi, Ni, K = idx.shape ++ assert B == Bc == Bi and N == Ni and D == Dc ++ ++ if not x.is_contiguous(): ++ x = x.contiguous() ++ if not c.is_contiguous(): ++ c = c.contiguous() ++ if not idx.is_contiguous(): ++ idx = idx.contiguous() ++ ++ if out is None: ++ out = torch.empty((B, N, K), device=x.device, dtype=torch.float32) ++ else: ++ assert out.shape == (B, N, K) and out.dtype == torch.float32 ++ ++ dtype_bytes = x.element_size() ++ BN, K_BLOCK, BLOCK_D, single_d_tile = _pick_tile(K, D, dtype_bytes) ++ ++ grid = (triton.cdiv(N, BN), triton.cdiv(K, K_BLOCK), B) ++ _knn_gather_l2sq_kernel[grid]( ++ x, c, idx, out, ++ M, ++ x.stride(0), x.stride(1), x.stride(2), ++ c.stride(0), c.stride(1), c.stride(2), ++ idx.stride(0), idx.stride(1), idx.stride(2), ++ out.stride(0), out.stride(1), out.stride(2), ++ N=N, D=D, K=K, ++ BN=BN, K_BLOCK=K_BLOCK, BLOCK_D=BLOCK_D, ++ SINGLE_D_TILE=single_d_tile, ++ num_warps=4, ++ ) ++ return out ++ ++ ++__all__ = ["triton_knn_gather_sqdist"] +diff --git a/dbscanlib/_kernels/kernels/flash_mst/__init__.py b/dbscanlib/_kernels/kernels/flash_mst/__init__.py +new file mode 100644 +index 0000000..e69de29 +diff --git a/dbscanlib/_kernels/kernels/flash_mst/triton/__init__.py b/dbscanlib/_kernels/kernels/flash_mst/triton/__init__.py +new file mode 100644 +index 0000000..e69de29 +diff --git a/dbscanlib/_kernels/kernels/flash_mst/triton/flash_mst.py b/dbscanlib/_kernels/kernels/flash_mst/triton/flash_mst.py +new file mode 100644 +index 0000000..af8ace6 +--- /dev/null ++++ b/dbscanlib/_kernels/kernels/flash_mst/triton/flash_mst.py +@@ -0,0 +1,611 @@ ++"""flash-mst: GPU-resident dense Boruvka MST. ++ ++Algorithm (per Boruvka iteration, all on GPU): ++ 1. Per-component argmin (`_per_component_argmin_v2_kernel`): for each row v, ++ scan all N columns of MRD once; pack (weight_bits, dst_idx) as int64 with ++ +inf for same-component cells; reduce per-row to a single packed best ++ edge; atomic_min into OUT_PACKED[parent[v]] (the row's component bucket). ++ SAME kernel folds in OUT_SRC[parent[v]] = atomic_min(orig_i) so we get ++ the canonical source vertex per component without a separate kernel. ++ 2. Concurrent union-find (`_concurrent_uf_kernel`): each component decodes ++ its packed-int64 best edge, finds roots with bounded path-walk, atomic- ++ CAS merges. Successful merges write an MST edge via atomic counter. ++ 3. Pointer-jumping (`_pointer_jump_kernel`): parent[v] = parent[parent[v]] ++ in parallel for ~log iters until paths are flat. Triton kernel replaces ++ the `parent = parent[parent.to(int64)]` torch loop (saves N int64 casts). ++ ++Persistent state (allocated ONCE at flash_mst entry, reused across rounds): ++ - parent[N] int32 — union-find array; doubles as the component label ++ for the next round's argmin (after pointer-jumping it is a valid label, ++ no torch.unique relabel needed). ++ - out_packed[N] int64 — per-CID atomic_min target, sized N (upper bound ++ on component count). Reset to +inf each round via Triton kernel. ++ - out_src[N] int32 — per-CID canonical src; reset to MAX_INT32. ++ - mst_{src,dst,w}, edge_count — output. ++ ++Why no argsort + no torch.unique relabel: ++ - The original code argsort-ed by component for cache locality; with ++ parent indexed directly the cache hit-rate is the same (parent[i] is ++ a single random read per row, then comp_j is a streaming scan which ++ L2-caches). ++ - `torch.unique` did a sort + inverse to remap parent[] to dense [0, n_C). ++ That step cost ~0.13 ms × log_2(N) iters of pointer-jumping. We skip it: ++ the argmin kernel uses `parent[i] != parent[j]` directly, which works ++ with sparse labels. ++ ++Note: I tried a K-min Boruvka variant (build a per-row top-K cross-component ++edge table once, reuse for K mini-merges) — measured 1 outer round was enough ++to find 99.97% of MST edges at huge, but the K=4 table build is ~3-4× more ++compute-heavy per scan than the simple argmin (4× axis-min + masking + bitonic ++merge), and a 2nd cleanup round still costs a full N² scan, so it lost to the ++simple argmin. Reverted; orchestration cleanup + tile tuning gave the win. ++ ++Packed int64 trick: for non-negative fp32 weights, the IEEE-754 bit pattern ++is monotonic (same ordering as float). Pack as ++ (weight_bits << 32) | dst_idx ++so atomic_min on int64 selects the smallest weight (with index as tiebreaker). ++""" ++ ++import math ++import numpy as np ++import torch ++import triton ++import triton.language as tl ++ ++ ++# ============================================================================= ++# Kernel 1: per-component min outgoing edge ++# Each program handles BLOCK_S sorted-contiguous points; per-row argmin over ++# the row of MRD (masked by component); atomic_min the result into the ++# corresponding component slot. ++# ============================================================================= ++ ++@triton.jit ++def _per_component_argmin_kernel( ++ MRD_ptr, # (N, N) fp32 ++ SORTED_IDX_ptr, # (N,) int32 ++ COMP_ptr, # (N,) int32, indexed by orig_idx ++ OUT_PACKED_ptr, # (n_components,) int64, atomic_min target ++ N, ++ BLOCK_S: tl.constexpr, ++ BLOCK_J: tl.constexpr, ++): ++ pid = tl.program_id(0) ++ s_offs = pid * BLOCK_S + tl.arange(0, BLOCK_S) ++ s_mask = s_offs < N ++ ++ orig_i = tl.load(SORTED_IDX_ptr + s_offs, mask=s_mask, other=0) ++ comp_i = tl.load(COMP_ptr + orig_i, mask=s_mask, other=-1) ++ ++ # IEEE-754 +inf bit pattern as int64 high word. ++ # 0x7F800000 << 32 = 9151314442816847872 ++ best_packed = tl.full([BLOCK_S], 9151314442816847872, dtype=tl.int64) ++ SHIFT32 = tl.cast(4294967296, tl.int64) # 2^32 as int64 ++ n_i64 = tl.cast(N, tl.int64) ++ ++ for j_start in range(0, N, BLOCK_J): ++ j_offs = j_start + tl.arange(0, BLOCK_J) ++ j_mask = j_offs < N ++ comp_j = tl.load(COMP_ptr + j_offs, mask=j_mask, other=-2) ++ ++ row_offs = (orig_i.to(tl.int64)[:, None] * n_i64 ++ + j_offs.to(tl.int64)[None, :]) ++ # Load and cast to fp32 (handles bf16 / fp16 / fp32 storage uniformly) ++ mrd_tile = tl.load(MRD_ptr + row_offs, ++ mask=s_mask[:, None] & j_mask[None, :], ++ other=float('inf')).to(tl.float32) ++ ++ # Mask same-component ++ diff = (comp_j[None, :] != comp_i[:, None]) & j_mask[None, :] & s_mask[:, None] ++ mrd_tile = tl.where(diff, mrd_tile, float('inf')) ++ ++ # Pack (weight_bits, j_idx) as int64. Use multiplication by 2^32 instead ++ # of `<< 32` to avoid int32 overflow warning. ++ wbits_u32 = mrd_tile.to(tl.uint32, bitcast=True) ++ wbits_i64 = wbits_u32.to(tl.int64) ++ j_i64 = j_offs.to(tl.int64)[None, :] ++ packed = wbits_i64 * SHIFT32 + j_i64 ++ ++ tile_best = tl.min(packed, axis=1) ++ best_packed = tl.minimum(best_packed, tile_best) ++ ++ # Atomic_min into OUT_PACKED[comp_i] for each row in this block ++ out_addrs = OUT_PACKED_ptr + comp_i.to(tl.int64) ++ tl.atomic_min(out_addrs, best_packed, mask=s_mask) ++ ++ ++# ============================================================================= ++# Kernel 2: write per-component canonical source (one orig_idx per component) ++# Run after argsort. For each component C in dense [0, n_components), picks ++# the FIRST orig_idx in sorted order (sorted_idx[comp_start[C]]). ++# ============================================================================= ++ ++@triton.jit ++def _set_per_component_src_kernel( ++ SORTED_IDX_ptr, # (N,) int32 ++ COMP_ptr, # (N,) int32 ++ OUT_SRC_ptr, # (n_components,) int32 ++ N, ++ BLOCK: tl.constexpr, ++): ++ pid = tl.program_id(0) ++ s_offs = pid * BLOCK + tl.arange(0, BLOCK) ++ s_mask = s_offs < N ++ orig = tl.load(SORTED_IDX_ptr + s_offs, mask=s_mask, other=0) ++ comp = tl.load(COMP_ptr + orig, mask=s_mask, other=0) ++ # The "first" element of each component is the one whose sorted position ++ # is smaller than all others with the same component. Use atomic_min on ++ # (sorted_pos, orig) pair — but for simplicity, just use atomic_min on ++ # orig (any vertex works as canonical). ++ target = OUT_SRC_ptr + comp.to(tl.int64) ++ # Use atomic_min so the smallest orig in each component wins ++ # (deterministic). orig values are int32 non-negative. ++ tl.atomic_min(target, orig, mask=s_mask) ++ ++ ++# ============================================================================= ++# Kernel 3: concurrent union-find merge ++# For each component C, decode (weight, dst) from out_packed[C], set src from ++# out_src[C], find roots via bounded loop, atomic_cas to merge. ++# Successful merges write an MST edge via atomic counter. ++# ============================================================================= ++ ++@triton.jit ++def _concurrent_uf_kernel( ++ PARENT_ptr, # (N,) int32, atomic ++ OUT_PACKED_ptr, # (n_components,) int64 ++ OUT_SRC_ptr, # (n_components,) int32 ++ MST_SRC_ptr, # (N-1,) int32 — output edge sources ++ MST_DST_ptr, # (N-1,) int32 — output edge targets ++ MST_W_ptr, # (N-1,) fp32 — output edge weights ++ EDGE_COUNT_ptr, # (1,) int32 atomic ++ n_components, ++ MAX_FIND: tl.constexpr, ++): ++ cid = tl.program_id(0) ++ if cid >= n_components: ++ return ++ ++ packed = tl.load(OUT_PACKED_ptr + cid) ++ # Sentinel: high 32 bits = 0x7F800000 (+inf bit pattern) means "no edge found" ++ high32 = (packed >> 32).to(tl.int32) ++ if high32 >= 0x7F800000: ++ return ++ ++ weight = high32.to(tl.float32, bitcast=True) ++ dst = (packed & 0xFFFFFFFF).to(tl.int32) ++ src = tl.load(OUT_SRC_ptr + cid) ++ if src < 0 or src == 2147483647: ++ return ++ ++ # Single-pass: find roots with path-walking, then atomic_cas to merge. ++ # If CAS fails (concurrent thread merged hi first), we drop this edge — ++ # next Boruvka iter will re-propose. log(N) extra iters total (rare path). ++ rs = src ++ for _ in tl.static_range(MAX_FIND): ++ p = tl.load(PARENT_ptr + rs) ++ rs = tl.where(p == rs, rs, p) ++ rd = dst ++ for _ in tl.static_range(MAX_FIND): ++ p = tl.load(PARENT_ptr + rd) ++ rd = tl.where(p == rd, rd, p) ++ ++ if rs == rd: ++ return ++ ++ lo = tl.minimum(rs, rd) ++ hi = tl.maximum(rs, rd) ++ old = tl.atomic_cas(PARENT_ptr + hi, hi, lo) ++ if old == hi: ++ idx = tl.atomic_add(EDGE_COUNT_ptr, 1) ++ tl.store(MST_SRC_ptr + idx, src) ++ tl.store(MST_DST_ptr + idx, dst) ++ tl.store(MST_W_ptr + idx, weight) ++ ++ ++# ============================================================================= ++# Generic edge-list union-find kernel (used by flash-dbscan CC, etc.). ++# Vectorised over BLOCK edges per program. Path-halving compresses the ++# tree during find via best-effort store-back. A merge counter lets the ++# caller exit as soon as no more edges can merge (a chain-shaped graph ++# with diameter ~ N can need 8+ passes). ++# ============================================================================= ++ ++@triton.jit ++def _union_edges_v2_kernel( ++ PARENT_ptr, # (N,) int32 atomic ++ ROW_ptr, # (E,) int32 — edge sources ++ COL_ptr, # (E,) int32 — edge targets ++ MERGE_COUNTER_ptr, # (1,) int32 — atomic counter; nonzero ⇒ keep iterating ++ n_edges, ++ BLOCK: tl.constexpr, ++ MAX_FIND: tl.constexpr, ++): ++ pid = tl.program_id(0) ++ e_offs = pid * BLOCK + tl.arange(0, BLOCK) ++ e_mask = e_offs < n_edges ++ u = tl.load(ROW_ptr + e_offs, mask=e_mask, other=0) ++ v = tl.load(COL_ptr + e_offs, mask=e_mask, other=0) ++ ++ # find root of u with path-halving: each step does 2 hops ++ # (parent[ru] := parent[parent[ru]]) and writes back the compressed ++ # parent. Best-effort — concurrent writers to the same ru race, but ++ # since parent links are monotonically toward the root the result is ++ # always valid (no need for atomic_cas — a stale write only undoes ++ # one step of compression, never causes a cycle). ++ ru = u ++ for _ in tl.static_range(MAX_FIND): ++ p = tl.load(PARENT_ptr + ru, mask=e_mask, other=0) ++ gp = tl.load(PARENT_ptr + p, mask=e_mask, other=0) ++ # only compress if we actually moved up two levels ++ compress = (ru != p) & (p != gp) & e_mask ++ tl.store(PARENT_ptr + ru, gp, mask=compress) ++ ru = tl.where(p == ru, ru, gp) ++ rv = v ++ for _ in tl.static_range(MAX_FIND): ++ p = tl.load(PARENT_ptr + rv, mask=e_mask, other=0) ++ gp = tl.load(PARENT_ptr + p, mask=e_mask, other=0) ++ compress = (rv != p) & (p != gp) & e_mask ++ tl.store(PARENT_ptr + rv, gp, mask=compress) ++ rv = tl.where(p == rv, rv, gp) ++ ++ diff = (ru != rv) & e_mask ++ lo = tl.minimum(ru, rv) ++ hi = tl.maximum(ru, rv) ++ # atomic_min instead of atomic_cas: parent[hi] := min(parent[hi], lo). ++ # Triton 3.6 doesn't support mask= on atomic_cas, but does on atomic_min. ++ # Semantics are actually cleaner for UF: parent[v] is monotonically ++ # non-increasing under atomic_min, so cycles are impossible by induction ++ # (we always have parent[v] ≤ v). "Made progress" = OLD > lo, i.e., we ++ # actually decreased parent[hi]. ++ old = tl.atomic_min(PARENT_ptr + hi, lo, mask=diff) ++ progress = (old > lo) & diff ++ n_succ = tl.sum(progress.to(tl.int32)) ++ # one atomic_add per CTA (not per edge) — negligible contention ++ if n_succ > 0: ++ tl.atomic_add(MERGE_COUNTER_ptr, n_succ) ++ ++ ++# ============================================================================= ++# Triton compaction: replaces `torch.unique(parent, return_inverse=True)`. ++# After parent is fully flattened, parent[v] is either v (if v is a root) or ++# the root index. Then dense_id_at[v] = (number of roots in [0, v]) - 1 if v ++# is a root, otherwise undefined. We compute it via a torch cumsum on the ++# is_root bitmap (cheap, one launch) and then gather labels[v] = dense_id[parent[v]]. ++# ============================================================================= ++ ++@triton.jit ++def _is_root_kernel( ++ PARENT_ptr, # (N,) int32 ++ IS_ROOT_ptr, # (N,) int32 (0/1) ++ N, ++ BLOCK: tl.constexpr, ++): ++ pid = tl.program_id(0) ++ offs = pid * BLOCK + tl.arange(0, BLOCK) ++ mask = offs < N ++ p = tl.load(PARENT_ptr + offs, mask=mask, other=-1) ++ is_root = (p == offs.to(tl.int32)).to(tl.int32) ++ tl.store(IS_ROOT_ptr + offs, is_root, mask=mask) ++ ++ ++@triton.jit ++def _gather_labels_kernel( ++ PARENT_ptr, # (N,) int32 ++ DENSE_ID_ptr, # (N,) int32 — dense id for each root, undefined for non-roots ++ LABELS_ptr, # (N,) int32 — output ++ N, ++ BLOCK: tl.constexpr, ++): ++ pid = tl.program_id(0) ++ offs = pid * BLOCK + tl.arange(0, BLOCK) ++ mask = offs < N ++ p = tl.load(PARENT_ptr + offs, mask=mask, other=0) ++ label = tl.load(DENSE_ID_ptr + p.to(tl.int64), mask=mask, other=0) ++ tl.store(LABELS_ptr + offs, label, mask=mask) ++ ++ ++ ++ ++ ++ ++ ++# ============================================================================= ++# Kernel 4: pointer jumping path compression ++# parent[v] = parent[parent[v]] in parallel; iterate to convergence. ++# ============================================================================= ++ ++@triton.jit ++def _pointer_jump_kernel( ++ PARENT_ptr, ++ N, ++ BLOCK: tl.constexpr, ++): ++ pid = tl.program_id(0) ++ v_offs = pid * BLOCK + tl.arange(0, BLOCK) ++ v_mask = v_offs < N ++ p = tl.load(PARENT_ptr + v_offs, mask=v_mask, other=0) ++ # Cast p (int32) to int64 for pointer arithmetic safety on large N ++ pp = tl.load(PARENT_ptr + p.to(tl.int64), mask=v_mask, other=0) ++ tl.store(PARENT_ptr + v_offs, pp, mask=v_mask) ++ ++ ++ ++def flash_cc_from_edges(rows: torch.Tensor, cols: torch.Tensor, N: int, ++ max_find: int = 8, max_passes: int = 16): ++ """Connected components on the graph defined by edge list (rows, cols). ++ ++ Iterative converging algorithm: ++ - Each pass: vectorized union with path-halving + atomic_cas hooks ++ (BLOCK=128 edges per program → 4 warps fully utilized vs the old ++ scalar 1-edge-per-warp launch). ++ - Then 6 calls to `_pointer_jump_kernel` to flatten the parent forest. ++ - Read the merge counter; exit as soon as a pass merges 0 edges. ++ For diameter-bound dense graphs this usually fires in 1-2 passes; ++ chain-shaped graphs (diameter ~N) converge in log_2(N) passes. ++ ++ Args: ++ rows, cols: (E,) int32 — edge endpoints. ++ N: number of vertices. ++ max_find: bounded find depth per pass. With path-halving each step ++ doubles the effective depth, so MAX_FIND=8 covers depth 256. ++ max_passes: hard ceiling (log2 of any conceivable graph diameter). ++ ++ Returns: ++ labels: (N,) int32 — CC component id, dense in [0, n_components). ++ """ ++ device = rows.device ++ parent = torch.arange(N, dtype=torch.int32, device=device) ++ E = rows.shape[0] ++ ++ if E > 0: ++ BLOCK_EDGE = 128 ++ BLOCK_JUMP = 1024 ++ merge_counter = torch.zeros(1, dtype=torch.int32, device=device) ++ grid_edge = (triton.cdiv(E, BLOCK_EDGE),) ++ grid_jump = (triton.cdiv(N, BLOCK_JUMP),) ++ ++ for _ in range(max_passes): ++ merge_counter.zero_() ++ _union_edges_v2_kernel[grid_edge]( ++ parent, rows, cols, merge_counter, E, ++ BLOCK=BLOCK_EDGE, MAX_FIND=max_find, num_warps=4, ++ ) ++ # Pointer-jump to flatten any remaining short chains. 6 iters ++ # cover depth ≤ 2^6 = 64 (post-compression chains are short). ++ for _ in range(6): ++ _pointer_jump_kernel[grid_jump]( ++ parent, N, BLOCK=BLOCK_JUMP, num_warps=4, ++ ) ++ if merge_counter.item() == 0: ++ break ++ ++ # Compact: roots have parent[v] == v; assign dense [0, n_cc) by index order. ++ is_root = torch.empty(N, dtype=torch.int32, device=device) ++ grid_compact = (triton.cdiv(N, 1024),) ++ _is_root_kernel[grid_compact](parent, is_root, N, BLOCK=1024, num_warps=4) ++ # cumsum gives 1-indexed dense IDs at root positions; subtract 1 → 0-indexed. ++ # At non-root positions the value is meaningless (we only gather via parent[v]). ++ dense_id = torch.cumsum(is_root, dim=0, dtype=torch.int32) - 1 ++ labels = torch.empty(N, dtype=torch.int32, device=device) ++ _gather_labels_kernel[grid_compact]( ++ parent, dense_id, labels, N, BLOCK=1024, num_warps=4, ++ ) ++ return labels ++ ++ ++# ============================================================================= ++# Kernel 5: per-component argmin with FOLDED canonical-src. Uses ++# ``orig_i = pid * BLOCK_S + lane`` directly (same cache hit-rate since ++# the column-side scan dominates the L2 traffic). The ++# ``OUT_SRC[component] = min(orig_idx)`` atomic_min is folded into the ++# same kernel, saving a separate launch per round. ++# ============================================================================= ++ ++@triton.jit ++def _per_component_argmin_v2_kernel( ++ MRD_ptr, # (N, N) bf16 or fp32 ++ PARENT_ptr, # (N,) int32 — parent[i] is i's component label ++ OUT_PACKED_ptr, # (N,) int64 — atomic_min target (sized to N) ++ OUT_SRC_ptr, # (N,) int32 — atomic_min target for canonical src ++ N, ++ BLOCK_S: tl.constexpr, ++ BLOCK_J: tl.constexpr, ++): ++ pid = tl.program_id(0) ++ s_offs = pid * BLOCK_S + tl.arange(0, BLOCK_S) ++ s_mask = s_offs < N ++ ++ orig_i = s_offs.to(tl.int64) ++ comp_i = tl.load(PARENT_ptr + orig_i, mask=s_mask, other=0) ++ ++ # IEEE-754 +inf bit pattern as int64 high word: 0x7F800000 << 32 ++ best_packed = tl.full([BLOCK_S], 9151314442816847872, dtype=tl.int64) ++ SHIFT32 = tl.cast(4294967296, tl.int64) ++ n_i64 = tl.cast(N, tl.int64) ++ ++ for j_start in range(0, N, BLOCK_J): ++ j_offs = j_start + tl.arange(0, BLOCK_J) ++ j_mask = j_offs < N ++ comp_j = tl.load(PARENT_ptr + j_offs, mask=j_mask, other=0) ++ ++ row_offs = orig_i[:, None] * n_i64 + j_offs.to(tl.int64)[None, :] ++ mrd_tile = tl.load(MRD_ptr + row_offs, ++ mask=s_mask[:, None] & j_mask[None, :], ++ other=float('inf')).to(tl.float32) ++ ++ # Mask same-component ++ diff = (comp_j[None, :] != comp_i[:, None]) & j_mask[None, :] & s_mask[:, None] ++ mrd_tile = tl.where(diff, mrd_tile, float('inf')) ++ ++ # Pack (weight_bits, j_idx) as int64 — the int64 axis-min uses a single ++ # comparator per element regardless of (weight, j) since the high bits ++ # dominate the ordering. Trying to track (best_w, best_j) as separate ++ # fp32+int (single tl.min + tl.argmin) was *slower* in benchmarking on ++ # H200 — the int64 reduction actually beats the fused fp32+argmin path. ++ wbits_u32 = mrd_tile.to(tl.uint32, bitcast=True) ++ wbits_i64 = wbits_u32.to(tl.int64) ++ j_i64 = j_offs.to(tl.int64)[None, :] ++ packed = wbits_i64 * SHIFT32 + j_i64 ++ ++ tile_best = tl.min(packed, axis=1) ++ best_packed = tl.minimum(best_packed, tile_best) ++ ++ # Atomic_min into OUT_PACKED[comp_i] and OUT_SRC[comp_i] (folded) ++ out_addrs = OUT_PACKED_ptr + comp_i.to(tl.int64) ++ tl.atomic_min(out_addrs, best_packed, mask=s_mask) ++ src_addrs = OUT_SRC_ptr + comp_i.to(tl.int64) ++ tl.atomic_min(src_addrs, s_offs.to(tl.int32), mask=s_mask) ++ ++ ++# ============================================================================= ++# Kernel 6: reset (N,)-sized int64 buffer to INF and int32 buffer to MAX_INT32 ++# Faster than torch.full(...) by ~4× (~30us vs ~125us per pair at huge): single ++# kernel launch with one pass over both arrays vs two separate fill kernels + ++# Python tensor wrapping overhead. ++# ============================================================================= ++ ++@triton.jit ++def _reset_buffers_kernel( ++ OUT_PACKED_ptr, # (N,) int64 ++ OUT_SRC_ptr, # (N,) int32 ++ N, ++ BLOCK: tl.constexpr, ++): ++ pid = tl.program_id(0) ++ offs = pid * BLOCK + tl.arange(0, BLOCK) ++ mask = offs < N ++ INF_PACKED = tl.cast(9151314442816847872, tl.int64) # 0x7F800000 << 32 ++ MAX_I32 = tl.cast(2147483647, tl.int32) ++ tl.store(OUT_PACKED_ptr + offs, INF_PACKED, mask=mask) ++ tl.store(OUT_SRC_ptr + offs, MAX_I32, mask=mask) ++ ++ ++# ============================================================================= ++# Python orchestration ++# ============================================================================= ++ ++def _pick_argmin_tile(N: int): ++ """Pick (BLOCK_S, BLOCK_J, num_warps, num_stages) for the argmin kernel. ++ ++ Tuned on H200 (GPU 4) at bf16 MRD via `algorithms/hdbscan/_bench_mst.py` ++ (sub-process to avoid Triton autotune cache leakage between the per-tile ++ sweep and the E2E run). ++ ++ Best per-N (of 4.8 TB/s peak BW): ++ N=20K: BS=32, BJ=256, nw=4, ns=4 → 0.376 ms (44.4%) ++ N=80K: BS=32, BJ=128, nw=4, ns=4 → 4.625 ms (57.7%) ++ N=200K: BS=32, BJ=128, nw=4, ns=3 → 28.08 ms (59.3%) ++ Larger BLOCK_J (256+) stalls at huge — the per-CTA register footprint of ++ the int64 packed tile (8 bytes × BLOCK_S × BLOCK_J) hurts occupancy and ++ spills into shared memory. Use BJ=128 for N >= 50K. ++ """ ++ if N >= 50_000: ++ return 32, 128, 4, 3 ++ return 32, 256, 4, 4 ++ ++ ++def flash_mst(MRD: torch.Tensor) -> torch.Tensor: ++ """GPU-resident dense Boruvka MST on a symmetric (N, N) distance matrix. ++ ++ Args: ++ MRD: (N, N) fp32 or bf16 — assumed symmetric, non-negative. Pass bf16 ++ for 2× HBM speedup on the per-iter argmin scan; weight precision ++ reduces to ~7 bits but tie-breaking via packed (weight, dst) keeps ++ determinism. ++ ++ Returns: ++ (N-1, 3) fp32 — MST edges as [src, dst, weight] sorted by weight. ++ src/dst are integer vertex IDs stored as fp32 (precise for N <= 2^24). ++ """ ++ assert MRD.is_cuda and MRD.dtype in (torch.float32, torch.bfloat16) and MRD.ndim == 2 ++ N = MRD.shape[0] ++ assert MRD.shape[1] == N ++ device = MRD.device ++ ++ # Persistent state — allocated once, reused across all rounds. ++ parent = torch.arange(N, dtype=torch.int32, device=device) ++ out_packed = torch.empty(N, dtype=torch.int64, device=device) ++ out_src = torch.empty(N, dtype=torch.int32, device=device) ++ ++ mst_src = torch.full((N - 1,), -1, dtype=torch.int32, device=device) ++ mst_dst = torch.full((N - 1,), -1, dtype=torch.int32, device=device) ++ mst_w = torch.zeros(N - 1, dtype=torch.float32, device=device) ++ edge_count = torch.zeros(1, dtype=torch.int32, device=device) ++ ++ BLOCK_S, BLOCK_J, num_warps_argmin, num_stages_argmin = _pick_argmin_tile(N) ++ BLOCK_RESET = 1024 ++ BLOCK_JUMP = 1024 ++ ++ # Standard Boruvka iter cap: log2(N) + small slack ++ max_iters = max(2, int(math.ceil(math.log2(max(N, 2)))) + 5) ++ target_edges = N - 1 ++ prev_count = 0 ++ ++ for it in range(max_iters): ++ # Reset per-component buffers in one kernel (int64 + int32 fused) ++ grid_reset = (triton.cdiv(N, BLOCK_RESET),) ++ _reset_buffers_kernel[grid_reset]( ++ out_packed, out_src, N, ++ BLOCK=BLOCK_RESET, num_warps=4, ++ ) ++ ++ # Per-component argmin: folds canonical-src into the kernel ++ grid_arg = (triton.cdiv(N, BLOCK_S),) ++ _per_component_argmin_v2_kernel[grid_arg]( ++ MRD, parent, out_packed, out_src, N, ++ BLOCK_S=BLOCK_S, BLOCK_J=BLOCK_J, ++ num_warps=num_warps_argmin, num_stages=num_stages_argmin, ++ ) ++ ++ # Concurrent union-find: dispatch over [0, N). Non-root CIDs see ++ # OUT_PACKED[c] == INF and early-return — cheap (~2 instructions). ++ grid_uf = (N,) ++ _concurrent_uf_kernel[grid_uf]( ++ parent, out_packed, out_src, ++ mst_src, mst_dst, mst_w, edge_count, ++ N, ++ MAX_FIND=8, num_warps=1, ++ ) ++ ++ # Pointer-jump until paths flat. Triton kernel (one launch ≈ 30 µs at ++ # huge) replaces the original log_2(N)-deep `parent[parent.to(int64)]` ++ # torch loop. 6 iters flatten depth ≤ 64 which covers all observed ++ # post-CAS chains; the next round's argmin handles any residual depth ++ # (a non-flat parent makes parent[i] != parent[j] occasionally true ++ # within a component → at worst one wasted edge proposal which the ++ # uf_merge rejects). ++ grid_jump = (triton.cdiv(N, BLOCK_JUMP),) ++ for _ in range(6): ++ _pointer_jump_kernel[grid_jump]( ++ parent, N, BLOCK=BLOCK_JUMP, num_warps=4, ++ ) ++ ++ # Single host sync per iter to check completion. Cost: ~5 µs each; ++ # 8 iters at huge = 40 µs total (negligible vs ~30 ms argmin). ++ cur_count = int(edge_count.item()) ++ if cur_count >= target_edges: ++ break ++ if cur_count == prev_count: ++ # No new edges this round — graph is disconnected. ++ break ++ prev_count = cur_count ++ ++ n_edges = int(edge_count.item()) ++ if n_edges != target_edges: ++ raise RuntimeError( ++ f"flash_mst: produced {n_edges} edges, expected {target_edges}. " ++ f"Graph may be disconnected, or merge kernel hit MAX_FIND." ++ ) ++ ++ # Sort by weight ascending for downstream HDBSCAN tree-building ++ sort_idx = torch.argsort(mst_w) ++ mst_src = mst_src[sort_idx] ++ mst_dst = mst_dst[sort_idx] ++ mst_w = mst_w[sort_idx] ++ ++ # Pack as (N-1, 3) fp32 for return — int IDs preserved exactly for N<=2^24 ++ out = torch.stack([mst_src.to(torch.float32), ++ mst_dst.to(torch.float32), ++ mst_w], dim=1) ++ return out +diff --git a/dbscanlib/_kernels/primitives/__init__.py b/dbscanlib/_kernels/primitives/__init__.py +new file mode 100644 +index 0000000..e69de29 +diff --git a/dbscanlib/_kernels/primitives/dbscan/__init__.py b/dbscanlib/_kernels/primitives/dbscan/__init__.py +new file mode 100644 +index 0000000..e69de29 +diff --git a/dbscanlib/_kernels/primitives/dbscan/triton/__init__.py b/dbscanlib/_kernels/primitives/dbscan/triton/__init__.py +new file mode 100644 +index 0000000..e69de29 +diff --git a/dbscanlib/_kernels/primitives/dbscan/triton/dbscan.py b/dbscanlib/_kernels/primitives/dbscan/triton/dbscan.py +new file mode 100644 +index 0000000..b9cc7b6 +--- /dev/null ++++ b/dbscanlib/_kernels/primitives/dbscan/triton/dbscan.py +@@ -0,0 +1,346 @@ ++"""flash-dbscan: GPU-resident DBSCAN. ++ ++Two algorithmic paths, dispatched by D: ++ ++ • D ≤ 4 (low-D): **GRID radius search**. Bin points into cells of side eps; ++ each point only needs to scan its 3^D=9 (D=2) or 27 (D=3) neighbor cells. ++ For typical D=2 spatial data with N ≈ M ≈ 100s/cell, this is ~250× ++ fewer distance evaluations than brute force. ++ ++ • D ≥ 5 (high-D): **brute-force top-K** via flash_knn (bf16 GEMM + on-chip ++ top-K), then eps filter on the returned distances. ++ ++Both paths feed the same downstream pipeline: ++ → core mask: deg ≥ min_samples ++ → connected components (atomic-CAS UF on GPU) ++ → border assignment: smallest CC label among in-eps core neighbors ++ → compact labels to dense [0, n_clusters) ++ ++The grid path is exact for any D where the cell-side ≥ eps (so all in-eps ++neighbors lie in the 3^D neighborhood). For high-D the curse of ++dimensionality makes the per-cell cost grow as 3^D, hence the brute-force ++path for D ≥ 5. ++""" ++import math ++ ++import torch ++import triton ++import triton.language as tl ++ ++# High-D brute-force neighbour path: flashlib's flash_knn = a fused Triton kNN ++# (indices) + a tiled gather kernel (exact squared-L2 distances), both vendored ++# from the flashlib knn primitive under dbscanlib._kernels. ++from dbscanlib._kernels.primitives.knn.triton import flash_knn_triton ++from dbscanlib._kernels.kernels.distance.triton.knn_gather_l2sq import ( ++ triton_knn_gather_sqdist, ++) ++ ++ ++def flash_knn(x, y, k, *, tol=None): ++ """Return ``(dist_sq, idx)`` top-k neighbours -- flashlib's flash_knn contract.""" ++ idx = flash_knn_triton(x, y, k) ++ dist_sq = triton_knn_gather_sqdist(x, y, idx) ++ return dist_sq, idx ++from dbscanlib._kernels.kernels.flash_mst.triton.flash_mst import flash_cc_from_edges ++ ++ ++def _next_pow2(x): ++ return 1 << (int(x - 1).bit_length()) if x > 1 else 1 ++ ++ ++# --------------------------------------------------------------------------- ++# GRID radius-search kernel (D=2) ++# --------------------------------------------------------------------------- ++# Each program handles BN query points. For each query, computes its grid ++# cell, iterates over the 3×3 neighbor cells, looks up each cell in a ++# hash table (linear probing) to find the (start, end) indices in ++# sorted_idx, and computes fp32 (xi - xj)² for each candidate; emits hits ++# (dist ≤ eps²) into the per-row neighbor buffer. ++# ++# Hash table layout: (HT_SIZE,) int64 packed as (cell_id << 32) | bin_idx. ++# Linear probing with a sentinel (-1) for empty slots. ++# --------------------------------------------------------------------------- ++ ++@triton.jit ++def _grid_radius_2d_kernel( ++ X_ptr, # (N, 2) fp32 ++ SORTED_PTR_ptr, # (N,) int32 — point indices sorted by cell ++ CELL_START_ptr, # (GW * GH,) int32 — dense 2D grid start ++ CELL_END_ptr, # (GW * GH,) int32 — dense 2D grid end ++ DEG_ptr, # (N,) int32 ++ NBR_IDX_ptr, # (N, K) int32 ++ N: tl.constexpr, K: tl.constexpr, ++ GW: tl.constexpr, GH: tl.constexpr, # grid dims ++ INV_EPS, EPS_SQ, ++ GRID_X_MIN, GRID_Y_MIN, ++ BN: tl.constexpr, ++): ++ """Grid radius search with dense 2D grid index. ++ ++ Each query maps to its (cx, cy) cell; we scan the 3×3 neighborhood by ++ direct indexing into CELL_START/END arrays of size (GW × GH). For our ++ standardised D=2 inputs (~6σ in [-3, 3], eps=0.04 → grid ~150×150, ++ storage ~22500 int32 × 2 = 180 KB). ++ """ ++ pid = tl.program_id(0) ++ n_offs = (pid * BN + tl.arange(0, BN)).to(tl.int64) ++ n_mask = n_offs < N ++ ++ xi = tl.load(X_ptr + n_offs * 2, mask=n_mask, other=0.0) ++ yi = tl.load(X_ptr + n_offs * 2 + 1, mask=n_mask, other=0.0) ++ ++ cxq = tl.floor((xi - GRID_X_MIN) * INV_EPS).to(tl.int32) ++ cyq = tl.floor((yi - GRID_Y_MIN) * INV_EPS).to(tl.int32) ++ ++ cur_cnt = tl.zeros([BN], dtype=tl.int32) ++ ++ # 3×3 cell neighborhood. Static-unrolled outer. ++ for nbi in tl.static_range(9): ++ dx = nbi // 3 - 1 ++ dy = nbi % 3 - 1 ++ cx = cxq + dx ++ cy = cyq + dy ++ # Bounds check (drop out-of-grid cells by mask). ++ in_grid = (cx >= 0) & (cx < GW) & (cy >= 0) & (cy < GH) & n_mask ++ cell_idx = cy * GW + cx ++ cell_idx_safe = tl.where(in_grid, cell_idx, 0) ++ s = tl.load(CELL_START_ptr + cell_idx_safe.to(tl.int64), ++ mask=in_grid, other=0) ++ e = tl.load(CELL_END_ptr + cell_idx_safe.to(tl.int64), ++ mask=in_grid, other=0) ++ range_len = tl.where(in_grid, e - s, 0) ++ max_iter = tl.max(range_len) ++ ++ for j in tl.range(0, max_iter, 1): ++ in_range = (j < range_len) & in_grid ++ idx = s + j ++ pj = tl.load(SORTED_PTR_ptr + idx.to(tl.int64), mask=in_range, other=0) ++ xj = tl.load(X_ptr + pj.to(tl.int64) * 2, mask=in_range, other=0.0) ++ yj = tl.load(X_ptr + pj.to(tl.int64) * 2 + 1, mask=in_range, other=0.0) ++ dxv = xi - xj ++ dyv = yi - yj ++ dist = dxv * dxv + dyv * dyv ++ hit_eps = (dist <= EPS_SQ) & in_range ++ slot_addr = n_offs * K + cur_cnt.to(tl.int64) ++ emit = hit_eps & (cur_cnt < K) ++ tl.store(NBR_IDX_ptr + slot_addr, pj, mask=emit) ++ cur_cnt = cur_cnt + hit_eps.to(tl.int32) ++ ++ cur_cnt = tl.minimum(cur_cnt, tl.full([BN], K, dtype=tl.int32)) ++ tl.store(DEG_ptr + n_offs, cur_cnt, mask=n_mask) ++ ++ ++def _build_grid_index(X: torch.Tensor, eps: float): ++ """Build a dense 2D grid index over points X (N, 2) fp32. ++ ++ Returns: ++ sorted_ptr: (N,) int32 — point indices sorted by cell ++ cell_start, cell_end: (GW * GH,) int32 — start/end indices in sorted_ptr ++ gw, gh: grid dims ++ inv_eps, x_min, y_min: kernel params ++ """ ++ N = X.shape[0] ++ device = X.device ++ inv_eps = 1.0 / eps ++ x_min = float(X[:, 0].min().item()) ++ y_min = float(X[:, 1].min().item()) ++ x_max = float(X[:, 0].max().item()) ++ y_max = float(X[:, 1].max().item()) ++ ++ GW = int((x_max - x_min) * inv_eps) + 2 ++ GH = int((y_max - y_min) * inv_eps) + 2 ++ n_cells = GW * GH ++ ++ cx = ((X[:, 0] - x_min) * inv_eps).floor().to(torch.int32).clamp(0, GW - 1) ++ cy = ((X[:, 1] - y_min) * inv_eps).floor().to(torch.int32).clamp(0, GH - 1) ++ cell_id = cy.to(torch.int64) * GW + cx.to(torch.int64) ++ ++ # Sort points by cell_id ++ sorted_cell, perm = torch.sort(cell_id, stable=True) ++ sorted_ptr = perm.to(torch.int32) ++ ++ # Build dense (GW*GH,) start/end via bucket counts. ++ counts = torch.zeros(n_cells, dtype=torch.int32, device=device) ++ counts.scatter_add_(0, sorted_cell, torch.ones(N, dtype=torch.int32, device=device)) ++ cell_end = torch.cumsum(counts, dim=0, dtype=torch.int32) ++ cell_start = cell_end - counts ++ ++ return sorted_ptr, cell_start, cell_end, GW, GH, inv_eps, x_min, y_min ++ ++ ++def _flash_dbscan_grid(X: torch.Tensor, eps: float, min_samples: int, ++ max_neighbors: int): ++ """D=2 grid-based path.""" ++ N, D = X.shape ++ assert D == 2 ++ device = X.device ++ K = max(min_samples, max_neighbors) ++ eps_sq = float(eps) ** 2 ++ ++ # Build grid index ++ (sorted_ptr, cell_start, cell_end, GW, GH, ++ inv_eps, x_min, y_min) = _build_grid_index(X, eps) ++ ++ deg = torch.zeros(N, dtype=torch.int32, device=device) ++ nbr_idx = torch.full((N, K), -1, dtype=torch.int32, device=device) ++ ++ BN = 32 ++ grid = (triton.cdiv(N, BN),) ++ _grid_radius_2d_kernel[grid]( ++ X.contiguous(), sorted_ptr, cell_start, cell_end, ++ deg, nbr_idx, ++ N=N, K=K, GW=GW, GH=GH, ++ INV_EPS=inv_eps, EPS_SQ=eps_sq, ++ GRID_X_MIN=x_min, GRID_Y_MIN=y_min, ++ BN=BN, ++ num_warps=4, num_stages=1, ++ ) ++ return deg, nbr_idx, K ++ ++ ++def _flash_dbscan_brute(X: torch.Tensor, eps: float, min_samples: int, ++ max_neighbors: int, *, tol=None): ++ """High-D brute-force path via flash_knn -- single bf16 kNN call. ++ ++ Enumerates each point's top-``max_neighbors`` ε-candidates (in bf16 ++ by default; pass ``tol=0`` to force fp32). Points with more than ++ ``max_neighbors`` true ε-neighbors will have their excess neighbours ++ silently truncated, but this does **not** change the DBSCAN ++ clustering: ++ ++ * Core-mask correctness: a point is core iff its ε-degree ≥ ++ ``min_samples``. As long as ``max_neighbors >= min_samples`` ++ the bounded-K enumeration still observes ``min_samples``-many ++ ε-neighbours when they exist, so the core mask is exact. ++ * Cluster-membership correctness: two core points belong to the ++ same cluster iff they are reachable through a chain of ++ core-core ε-edges. Dense clusters expose Θ(max_neighbors) such ++ edges per core point, so any one of them is enough to merge ++ that core into the cluster's connected component. Edges ++ ranked beyond top-K are redundant for CC. ++ ++ The historical K-grow loop ("double K until the K-th NN exceeds ++ ε") was reverted in 2026-05 -- the iterative re-launches grew K ++ geometrically in dense regimes (medium-D blob shapes hit K=1024 ++ in 7 calls = 566 ms total) without changing ARI vs cuML beyond ++ the noise floor inherited from the bounded-K approximation. ++ ++ Callers who actually need exhaustive enumeration (e.g. for ++ reach-distance diagnostics, not clustering) should pass a ++ larger ``max_neighbors`` once. ++ """ ++ N, D = X.shape ++ device = X.device ++ K = max(min_samples, max_neighbors) ++ K = min(K, N) ++ eps_sq = float(eps) ** 2 ++ ++ # Default to fp32 KNN: at these shapes (N, M >= 16K) flash_knn ++ # routes to the "large_n" insert path (BN=128 single-pass per CTA), ++ # which is memory-bound and runs at the same speed in fp32 as in ++ # bf16 -- the .to(bf16) cast itself eats the would-be win. fp32 ++ # also gives bit-exact ARI vs sklearn ++ # on the standard ε boundaries; bf16 quantisation flips a handful ++ # of borderline points (~3% on D=64 eps=8). Power users can opt ++ # into bf16 via ``tol=1e-3`` if they have looser ε. ++ if D < 16: ++ Xpad = torch.zeros(N, 16, device=device, dtype=X.dtype) ++ Xpad[:, :D] = X ++ X_kernel = Xpad.contiguous() ++ else: ++ X_kernel = X.contiguous() ++ knn_dist_sq, knn_idx = flash_knn( ++ X_kernel[None], X_kernel[None], k=K, tol=tol) ++ knn_dist_sq = knn_dist_sq[0] ++ knn_idx = knn_idx[0] ++ ++ valid = knn_dist_sq <= eps_sq ++ deg = valid.sum(dim=1).to(torch.int32) ++ nbr_idx = torch.where(valid, knn_idx.to(torch.int32), ++ torch.full_like(knn_idx, -1, dtype=torch.int32)) ++ return deg, nbr_idx, K ++ ++ ++def flash_dbscan(X: torch.Tensor, eps: float, min_samples: int = 5, ++ max_neighbors: int = 32, *, tol=None): ++ """End-to-end Triton DBSCAN. ++ ++ Args: ++ X: (N, D) float32 CUDA tensor. ++ eps: distance threshold (Euclidean). ++ min_samples: minimum points in eps-neighborhood for a point to be core. ++ max_neighbors: per-point ε-candidate budget for the high-D brute ++ path. Default 32 -- core detection stays exact for any ++ ``max_neighbors >= min_samples``; the connected-component ++ stage only needs O(max_neighbors) redundant edges per ++ cluster, so cluster membership is robust to the bounded-K ++ approximation. Bump this if you need exhaustive ε-edge ++ enumeration (e.g. reach-distance diagnostics). ++ tol: residual tolerance forwarded to :func:`flash_knn` on the ++ high-D (D >= 3) brute-force path. ``None`` (default) keeps ++ fp32 storage -- the "large_n" single-pass insert path ++ (BN=128) flash_knn picks on these shapes is memory-bound, ++ so fp32 and bf16 run at the same speed; fp32 wins on ARI by ++ avoiding ε-boundary quantisation. Pass ``tol=1e-3`` to opt ++ into bf16/fp16 KNN storage (matches the historical HEAD ++ cast) when you want lower HBM pressure. ++ ++ Returns: ++ labels: (N,) int32 -- cluster id (>= 0) or -1 for noise. ++ """ ++ assert X.is_cuda ++ N, D = X.shape ++ device = X.device ++ ++ # ── Stage 1: per-point in-eps neighbor list ───────────────────────── ++ # Path A: D ≤ 2 → grid (cell-side eps; 9-cell scan, exact in any dtype). ++ # Path B: D ≥ 3 → brute-force top-K via flash_knn (tol-routed dtype). ++ if D == 2: ++ # Grid kernel reads X via raw fp32 strides; cast internally. ++ # This is a kernel implementation detail, NOT a public dtype contract. ++ deg, nbr_idx, K = _flash_dbscan_grid( ++ X if X.dtype == torch.float32 else X.float(), ++ eps, min_samples, max_neighbors, ++ ) ++ else: ++ deg, nbr_idx, K = _flash_dbscan_brute( ++ X, eps, min_samples, max_neighbors, tol=tol) ++ ++ # ── Stage 2: core mask ────────────────────────────────────────────── ++ core_mask = deg >= min_samples ++ ++ # ── Stage 3: edge construction (core-core only) ───────────────────── ++ K_eff = nbr_idx.shape[1] ++ nbr_idx_i64 = nbr_idx.to(torch.int64) ++ valid_slot = nbr_idx >= 0 ++ core_per_row = core_mask[:, None].expand(-1, K_eff) ++ nbr_idx_safe = torch.where(valid_slot, nbr_idx_i64, torch.zeros_like(nbr_idx_i64)) ++ core_per_col = core_mask[nbr_idx_safe] & valid_slot ++ edge_mask = valid_slot & core_per_row & core_per_col ++ rows = (torch.arange(N, device=device, dtype=torch.int32).view(-1, 1) ++ .expand(-1, K_eff).contiguous())[edge_mask].contiguous() ++ cols = nbr_idx[edge_mask].contiguous() ++ label_cc = flash_cc_from_edges(rows, cols, N) ++ ++ # ── Stage 4: assign labels (core: CC label, else: -1 for now) ──────── ++ INT_MAX = 2 ** 31 - 1 ++ label = torch.where(core_mask, label_cc, torch.full_like(label_cc, -1)) ++ ++ # ── Stage 5: border assignment ─────────────────────────────────────── ++ nbr_labels = label[nbr_idx_safe] ++ nbr_is_core = core_mask[nbr_idx_safe] & valid_slot ++ border_cand = torch.where(valid_slot & nbr_is_core, nbr_labels, ++ torch.full_like(nbr_labels, INT_MAX)) ++ min_core_label = border_cand.min(dim=1).values ++ is_border = (~core_mask) & (min_core_label != INT_MAX) & (min_core_label >= 0) ++ label = torch.where(is_border, min_core_label, label) ++ ++ # ── Stage 6: compact ───────────────────────────────────────────────── ++ valid = label >= 0 ++ if valid.any(): ++ unique, inv = torch.unique(label[valid], return_inverse=True) ++ compact = torch.full_like(label, -1) ++ compact[valid] = inv.to(torch.int32) ++ label = compact ++ ++ return label +diff --git a/dbscanlib/_kernels/primitives/knn/__init__.py b/dbscanlib/_kernels/primitives/knn/__init__.py +new file mode 100644 +index 0000000..e69de29 +diff --git a/dbscanlib/_kernels/primitives/knn/triton/__init__.py b/dbscanlib/_kernels/primitives/knn/triton/__init__.py +new file mode 100644 +index 0000000..f5834ec +--- /dev/null ++++ b/dbscanlib/_kernels/primitives/knn/triton/__init__.py +@@ -0,0 +1,41 @@ ++"""knn triton backend. ++ ++Re-exports the public Python wrappers from each component file. ++``@triton.jit`` kernels stay private to their file. ++ ++Kernel layout: ++ ++* :mod:`_common` -- shared helpers: ``_next_pow2``, micro-bench, ++ IEEE-sortable u32 transform, ``_INF_PACKED`` sentinel. ++* :mod:`sortmerge` -- packed-uint64 sort-merge top-K (x²-free). Routed ++ to only for the small-Q + medium-K Pattern-A corner. ++* :mod:`insert` -- iterative argmin-insert top-K (x²-free). The ++ general path -- BN in {8, 16, 32, 64, 128} covers everything from ++ small-Q search to large-Q build. ++* :mod:`_row_norm` -- ``_fast_row_sq`` / ``_get_or_compute_csq`` for ++ the CuteDSL FA3 wrapper (it needs ``c_sq`` as a kernel argument). ++* :mod:`dispatch` -- host-side router. Single public entry ++ :func:`flash_knn_triton` runs through :func:`_heuristic_config` ++ unconditionally -- the per-CTA-count check inside the heuristic ++ picks single-pass vs M-split. Per-shape config (BN/BM/mode/mps/ ++ ns_pipe) also falls out of :func:`_heuristic_config`. ++""" ++from dbscanlib._kernels.primitives.knn.triton._common import ( ++ _next_pow2, ++ _bench_quick, ++) ++from dbscanlib._kernels.primitives.knn.triton._row_norm import ( ++ _fast_row_sq, ++ _get_or_compute_csq, ++) ++from dbscanlib._kernels.primitives.knn.triton.dispatch import ( ++ flash_knn_triton, ++ flash_knn_triton_small_n, ++ flash_knn_triton_large_n, ++) ++ ++__all__ = [ ++ "flash_knn_triton", ++ "flash_knn_triton_small_n", ++ "flash_knn_triton_large_n", ++] +diff --git a/dbscanlib/_kernels/primitives/knn/triton/_common.py b/dbscanlib/_kernels/primitives/knn/triton/_common.py +new file mode 100644 +index 0000000..686e85b +--- /dev/null ++++ b/dbscanlib/_kernels/primitives/knn/triton/_common.py +@@ -0,0 +1,81 @@ ++"""Shared utilities for flash_knn Triton dispatch + kernels. ++ ++Lives in one place so :mod:`dispatch`, :mod:`sortmerge` and :mod:`insert` ++all import the same primitives: ++ ++ * :func:`_next_pow2` -- D / K padding helper. ++ * :func:`_bench_quick` -- micro-bench used by the autotuner. ++ * :func:`_fp32_to_sortable_u32` / :func:`_sortable_u32_to_fp32` -- ++ branch-free IEEE-sortable u32 transform. With the x²-free score ++ ``s = c² - 2⟨x, c⟩`` the quantity is signed, so the packed-uint64 ++ sort-merge / final-sort needs an ascending-u32 = ascending-fp32 ++ mapping. The transform is 3 ops, exact for non-NaN fp32. ++ * :data:`_INF_PACKED` -- (sortable +inf << 32) | -1 idx sentinel, ++ used as the initial fill of the sortmerge running top-K so the ++ early-exit ``chunk_best < sortable_inv(top)`` always fires until ++ real values arrive. ++""" ++from __future__ import annotations ++ ++import torch ++import triton ++import triton.language as tl ++ ++ ++def _next_pow2(n: int) -> int: ++ """Smallest power-of-two >= ``max(1, n)``. Used for D / K padding.""" ++ if n <= 1: ++ return 1 ++ return 1 << (n - 1).bit_length() ++ ++ ++def _bench_quick(fn, *, warmup: int = 3, reps: int = 5) -> float: ++ """Median-ish wall time in ms; cheap enough to call once per autotune cfg.""" ++ for _ in range(warmup): ++ fn() ++ torch.cuda.synchronize() ++ s = torch.cuda.Event(enable_timing=True) ++ e = torch.cuda.Event(enable_timing=True) ++ s.record() ++ for _ in range(reps): ++ fn() ++ e.record() ++ torch.cuda.synchronize() ++ return s.elapsed_time(e) / reps ++ ++ ++# Sortable upper-bits = sortable(+inf) = 0xFF800000; lower-bits = -1 (idx ++# sentinel) = 0xFFFFFFFF. ``_sortable_u32_to_fp32(0xFF800000) == +inf`` so ++# the sortmerge early-exit ``chunk_best < +inf`` always fires until the ++# running top-K starts to hold real values. ++# ++# Wrapped in ``tl.constexpr`` so ``@triton.jit`` kernels can reference it ++# directly as a module-level global (un-wrapped Python ints are not ++# accessible from inside a JIT'd kernel as of Triton 3.x). ++_INF_PACKED = tl.constexpr(0xFF800000_FFFFFFFF) ++ ++ ++@triton.jit ++def _fp32_to_sortable_u32(x): ++ """Map ascending fp32 -> ascending uint32 (branch-free, 3 ops). ++ ++ Standard IEEE-754 trick: arithmetic-shift the int32 view by 31 (so ++ MSB=1 / negative produces 0xFFFFFFFF, MSB=0 / positive produces 0), ++ OR with 0x80000000 to flip the sign bit of positives, then XOR with ++ the original bits. NaNs map into 0xFF800001..0xFFFFFFFF (unused in ++ practice). ++ """ ++ bits = x.to(tl.uint32, bitcast=True) ++ sign_mask = (bits.to(tl.int32, bitcast=True) >> 31).to(tl.uint32, bitcast=True) ++ flip = sign_mask | 0x80000000 ++ return bits ^ flip ++ ++ ++@triton.jit ++def _sortable_u32_to_fp32(s): ++ """Inverse of :func:`_fp32_to_sortable_u32`.""" ++ sign_mask = (s.to(tl.int32, bitcast=True) >> 31).to(tl.uint32, bitcast=True) ++ # In sortable space: MSB=1 -> was positive (flip back via 0x80000000); ++ # MSB=0 -> was negative (flip back via 0xFFFFFFFF). ++ flip = (sign_mask ^ 0xFFFFFFFF) | 0x80000000 ++ return (s ^ flip).to(tl.float32, bitcast=True) +diff --git a/dbscanlib/_kernels/primitives/knn/triton/_row_norm.py b/dbscanlib/_kernels/primitives/knn/triton/_row_norm.py +new file mode 100644 +index 0000000..39fd0c8 +--- /dev/null ++++ b/dbscanlib/_kernels/primitives/knn/triton/_row_norm.py +@@ -0,0 +1,93 @@ ++"""Fast row sum-of-squares (||x_i||^2) helper + per-tensor cache. ++ ++Tiny utility, but shared by the FA3 fused KNN path which needs the ++pre-computed query/corpus norms to fold into the squared-L2 distance ++formula ``||x - c||^2 = ||x||^2 + ||c||^2 - 2 ``. ++ ++This module deliberately stays narrow: ++ - ``_row_sq_kernel`` is a pure Triton row-norm kernel (no cross ++ matrix, no top-K) — it never materialises an N x M tensor. ++ - ``_fast_row_sq`` wraps the kernel. ++ - ``_get_or_compute_csq`` caches results per ``(data_ptr, shape, ++ dtype)`` so repeat queries on the same corpus skip the recompute. ++""" ++from __future__ import annotations ++ ++import math ++ ++import torch ++import triton ++import triton.language as tl ++ ++from dbscanlib._kernels.primitives.knn.triton._common import _next_pow2 ++ ++ ++@triton.jit ++def _row_sq_kernel( ++ x_ptr, out_ptr, ++ stride_x_b, stride_x_n, stride_x_d, ++ stride_o_b, stride_o_n, ++ B: tl.constexpr, N: tl.constexpr, D: tl.constexpr, ++ BN: tl.constexpr, BD: tl.constexpr, ++): ++ """Row sum-of-squares for ``(B, N, D) -> (B, N)`` fp32. ++ ++ Casts input to fp32 inside the kernel; reads input exactly once ++ and writes output exactly once. At BN=128 BD=128 nw=4 this hits ++ ~36% peak HBM BW on H200 (typical for tall-skinny reads). ++ """ ++ pid_n = tl.program_id(0) ++ pid_b = tl.program_id(1).to(tl.int64) ++ n_offs = (pid_n * BN + tl.arange(0, BN)).to(tl.int64) ++ n_mask = n_offs < N ++ acc = tl.zeros([BN], dtype=tl.float32) ++ for d_start in range(0, D, BD): ++ d_offs = (d_start + tl.arange(0, BD)).to(tl.int64) ++ d_mask = d_offs < D ++ x = tl.load( ++ x_ptr + pid_b * stride_x_b ++ + n_offs[:, None] * stride_x_n ++ + d_offs[None, :] * stride_x_d, ++ mask=n_mask[:, None] & d_mask[None, :], other=0.0, ++ ) ++ x_f = x.to(tl.float32) ++ acc += tl.sum(x_f * x_f, axis=1) ++ tl.store(out_ptr + pid_b * stride_o_b + n_offs * stride_o_n, ++ acc, mask=n_mask) ++ ++ ++def _fast_row_sq(x: torch.Tensor) -> torch.Tensor: ++ """Return ``(B, N)`` fp32 row sum-of-squares via the Triton kernel.""" ++ assert x.is_cuda and x.ndim == 3 ++ B, N, D = x.shape ++ out = torch.empty(B, N, device=x.device, dtype=torch.float32) ++ BN = 128 if N >= 128 else _next_pow2(max(N, 16)) ++ BD = 128 if D >= 128 else _next_pow2(max(D, 16)) ++ grid = (math.ceil(N / BN), B) ++ _row_sq_kernel[grid]( ++ x, out, ++ x.stride(0), x.stride(1), x.stride(2), ++ out.stride(0), out.stride(1), ++ B=B, N=N, D=D, BN=BN, BD=BD, ++ num_warps=4, ++ ) ++ return out ++ ++ ++# Per-data-ptr cache so repeated queries on the same corpus don't pay ++# the row-norm kernel cost. ++_csq_cache: dict = {} ++ ++ ++def _get_or_compute_csq(c: torch.Tensor) -> torch.Tensor: ++ """Cached ``||c_i||^2`` (fp32) keyed by ``(data_ptr, shape, dtype)``.""" ++ key = (c.data_ptr(), tuple(c.shape), c.dtype) ++ cached = _csq_cache.get(key) ++ if cached is not None and cached.device == c.device: ++ return cached ++ csq = _fast_row_sq(c) ++ _csq_cache[key] = csq ++ if len(_csq_cache) > 8: ++ for old_key in list(_csq_cache.keys())[:-8]: ++ del _csq_cache[old_key] ++ return csq +diff --git a/dbscanlib/_kernels/primitives/knn/triton/dispatch.py b/dbscanlib/_kernels/primitives/knn/triton/dispatch.py +new file mode 100644 +index 0000000..1476b18 +--- /dev/null ++++ b/dbscanlib/_kernels/primitives/knn/triton/dispatch.py +@@ -0,0 +1,1178 @@ ++"""flash_knn -- host-side dispatcher: kernel/config selection + 2-stage launch. ++ ++Backs the public :func:`dbscanlib._kernels.primitives.knn.flash_knn` entry point. ++ ++Top-level wrappers ++------------------ ++ ++* :func:`flash_knn_triton` -- one-shot kNN, returns ``(B, N, k)`` int32 ++ indices. The shape-only heuristic now handles both build (large-N, ++ single-pass per CTA) and search (small-N, M-split flash-decode) ++ without any host-side SM-saturation gate -- ``ctas_no_split`` is ++ inspected after BN is picked, so build shapes (``BN=128`` -> ++ ``ctas_no_split`` already saturates 132 SMs) and search shapes ++ (small BN -> M-split for tail-fill) both fall out of the same ++ decision tree. ++* :func:`flash_knn_triton_small_n` / :func:`flash_knn_triton_large_n` ++ -- explicit overrides for callers (e.g. K-means assign) that know ++ their shape category at construction time. ++* :func:`_heuristic_config`, :func:`_autotune` -- shape-only heuristic ++ (default, fast first call) + opt-in brute-force autotune (slower ++ first call, cached steady-state). Both share the same kernel grid; ++ ``force_path="large_n"`` collapses M-splits down to one per CTA. ++ ++Two kernels live behind this dispatcher: ++ ++* :mod:`insert` -- iterative argmin-insert top-K. Picked by the ++ heuristic on virtually every shape. ++* :mod:`sortmerge` -- packed-uint64 sort-merge top-K with the IEEE- ++ sortable u32 transform. Routed to for the **Pattern-A small-Q** ++ corner (``B*N <= 8 or (B*N <= 16 and K in {32, 64})``) at medium-K ++ + ``M <= 200K`` where the wider sort per chunk beats insert's ++ argmin loop on this regime's autotune data. Otherwise insert wins; ++ sortmerge is kept in :func:`_gen_configs` so the offline tuner can ++ keep verifying that. ++ ++Pipelining depth ++---------------- ++Both kernels take a ``NUM_STAGES_PIPE`` constexpr for the M-loop's ++``tl.range`` pipelining factor. The dispatcher tunes it per shape from ++a 201-shape ns-sweep (upstream ``bench/sweep_ns_by_k.py``): K<=8 + ++small tile wants ``ns=3`` to hide HBM latency, K>=64 + big tile wants ++``ns=2`` (registers tight), wide-D wants ``ns=2`` uniformly, narrow ++tile + anything wants ``ns=1`` to avoid spills. See the comments ++inside :func:`_heuristic_config` for the empirical breakdown. ++ ++Distance recovery ++----------------- ++The kernels emit **signed shifted scores** ``s = c_sq - 2*``, not ++true squared L2. ``argmin-k`` over ``s`` matches ``argmin-k`` over true ++squared L2 (since ``-||x||^2`` is a per-row constant), but the returned ++score is not a valid distance. The :func:`flash_knn` public wrapper ++calls :func:`dbscanlib._kernels.kernels.distance.triton_knn_gather_sqdist` on the ++returned indices to write the true ``||x - c[idx]||^2`` per neighbour. ++""" ++from __future__ import annotations ++ ++import math ++from typing import Optional ++ ++import torch ++ ++from dbscanlib._kernels.primitives.knn.triton._common import _next_pow2, _bench_quick ++from dbscanlib._kernels.primitives.knn.triton.sortmerge import _flash_knn_sortmerge_kernel ++from dbscanlib._kernels.primitives.knn.triton.insert import _flash_knn_insert_kernel ++ ++ ++# ── SRAM estimation ──────────────────────────────────────────────────── ++ ++ ++def _estimate_sram(bn, bm, D, K, d_inner, num_stages, *, ++ sortmerge=False, dtype_bytes=2): ++ """Per-CTA SRAM estimate (bytes) for the unified kernel — coarse. ++ ++ Counts only what Triton **definitely** keeps in SMEM: the x tile ++ + c tile (per iteration of the inner D loop). Everything else ++ (cross fp32 accumulator, score tile, topK heap, sort scratch) is ++ register-resident at the tile sizes used by the dispatcher's ++ heuristic, and the NUM_STAGES_PIPE multi-buffering of the c tile ++ is sometimes elided by Triton's pipeliner. Including any of these ++ in the estimate produces 50-200 KB of over-count that **wrongly ++ shrinks** configs the kernel could actually launch — a bf16 ++ regression demonstrated by the ``D=256`` cells in ++ ``benchmarks/results/micro_knn_dtype_smem_fit.md``. ++ ++ This estimator is therefore **optimistic on purpose**. Its only ++ job is to bias the autotune candidate filter away from configs ++ that obviously can't fit (e.g. ``BN=128, BM=128, D=256, fp32`` ++ needs >232 KB by raw arithmetic). The source of truth for "does ++ this fit?" is ``_run``'s runtime ``OutOfResources`` catch + shrink ++ loop. ``dtype_bytes`` is still threaded through so fp32 inputs ++ correctly double the x/c contribution. ++ """ ++ x_sub = bn * d_inner * dtype_bytes ++ c_sub = bm * d_inner * dtype_bytes ++ c_sq = bm * 4 ++ base = x_sub + c_sub + c_sq ++ if sortmerge: ++ chunk = bn * bm * 8 ++ base += chunk ++ return base ++ ++ ++def _smem_limit(device) -> int: ++ """Per-block dynamic shared-memory budget for ``device``. ++ ++ Mirrors :func:`dbscanlib._kernels.primitives.kmeans.triton.assign._smem_limit` ++ — Triton uses opt-in dynamic SMEM; prefer that attribute when ++ available, fall back to static limit, and finally to a conservative ++ 48 KiB for old PyTorch builds. ++ """ ++ props = torch.cuda.get_device_properties(device) ++ for attr in ( ++ "shared_memory_per_block_optin", ++ "max_shared_memory_per_block_optin", ++ "shared_memory_per_block", ++ "max_shared_memory_per_block", ++ ): ++ v = getattr(props, attr, None) ++ if v: ++ return int(v) ++ return 48 * 1024 ++ ++ ++def _fit_config_to_smem(cfg: dict, D: int, K: int, ++ dtype_bytes: int, smem_limit: int) -> dict: ++ """Shrink ``(BN, BM, NUM_STAGES_PIPE)`` until the kernel fits SMEM. ++ ++ Returns ``cfg`` unchanged when it already fits. Otherwise enumerates ++ power-of-two reductions on the three SMEM-bearing axes and picks the ++ one that maximises ``BN * BM * num_stages`` (work-per-program tile) ++ among those that fit, breaking ties towards the original aspect ratio. ++ ++ Why only these three axes? ++ * ``D_INNER``: changing it flips between single-D-tile and D-split ++ paths, which would also change the kernel's perf profile. We ++ keep the heuristic's D-split decision and shrink the cheaper ++ axes first. ++ * ``TOPK_PAD``: determined by ``K``; cannot be reduced without ++ breaking correctness. ++ * ``kernel_mode``, ``num_warps``, ``M_PER_SPLIT``, ``NUM_SPLITS``: ++ these don't affect per-CTA SMEM. Stable. ++ ++ Pattern (and rationale) mirrors :func:`dbscanlib._kernels.primitives.kmeans.\ ++triton.assign._fit_config_to_smem`. ++ """ ++ BN0 = int(cfg["BN"]) ++ BM0 = int(cfg["BM"]) ++ D_INNER = int(cfg["D_INNER"]) ++ NS0 = int(cfg.get("NUM_STAGES_PIPE", 2)) ++ sortmerge = cfg.get("kernel_mode") == "sortmerge" ++ ++ cur = _estimate_sram(BN0, BM0, D, K, D_INNER, NS0, ++ sortmerge=sortmerge, dtype_bytes=dtype_bytes) ++ if cur <= smem_limit: ++ return cfg ++ ++ def _pow2_down_to(v, lo): ++ out = [] ++ x = v ++ while x >= lo: ++ out.append(x) ++ x //= 2 ++ return out ++ ++ # BN must stay >= 8 (matches _gen_configs lowest BN rung); ++ # BM must stay >= 32 for the inner loop to keep WGMMA shape; ++ # NUM_STAGES_PIPE must stay >= 1. ++ bn_cands = _pow2_down_to(BN0, 8) ++ bm_cands = _pow2_down_to(BM0, 32) ++ ns_cands = list(range(NS0, 0, -1)) ++ ++ best = None ++ best_key = None ++ for bn in bn_cands: ++ for bm in bm_cands: ++ # sortmerge requires BM == TOPK_PAD; do not shrink BM there. ++ if sortmerge and bm != BM0: ++ continue ++ for ns in ns_cands: ++ if _estimate_sram(bn, bm, D, K, D_INNER, ns, ++ sortmerge=sortmerge, ++ dtype_bytes=dtype_bytes) > smem_limit: ++ continue ++ aspect_penalty = abs((bn / max(bm, 1)) - (BN0 / max(BM0, 1))) ++ # Prefer: larger total tile work, closer aspect ratio, ++ # larger BN (more N-parallelism), larger NS (better pipelining). ++ key = (bn * bm * ns, -aspect_penalty, bn, ns) ++ if best_key is None or key > best_key: ++ best_key = key ++ best = (bn, bm, ns) ++ ++ if best is None: ++ raise RuntimeError( ++ f"flash_knn: cannot fit kernel into shared memory " ++ f"(D={D}, K={K}, D_INNER={D_INNER}, dtype_bytes={dtype_bytes}, " ++ f"smem_limit={smem_limit}). Original config " ++ f"BN={BN0}, BM={BM0}, NS={NS0} needs {cur} bytes." ++ ) ++ ++ bn, bm, ns = best ++ out = dict(cfg) ++ out["BN"] = bn ++ out["BM"] = bm ++ out["NUM_STAGES_PIPE"] = ns ++ return out ++ ++ ++def _pipe_stages_for(K, d_inner, D): ++ """``NUM_STAGES_PIPE`` candidates for the autotune M-loop sweep. ++ ++ The K-step inner loop adds register pressure that scales with K, so ++ deeper pipelining (which double-/triple-buffers the next C-tile) is ++ only profitable when K is small. Empirically: ++ ++ * D-split path (``d_inner < D``) -- kernel uses ns=1 by construction. ++ * K >= 64 -- only try {1, 2}; ns >= 3 spills registers and tanks perf ++ (probe showed ns=3 going from 750 -> 1220 us at K=128). ++ * K <= 32 -- try {1, 2, 3} so big pipelines can hide HBM latency ++ for huge-M shapes. ++ """ ++ if d_inner < D: ++ return [1] ++ if K >= 64: ++ return [1, 2] ++ return [1, 2, 3] ++ ++ ++def _gen_configs(D, K, *, dtype_bytes=2, smem_limit=None): ++ """Generate candidate kernel configs (sortmerge + insert), SRAM-validated. ++ ++ Covers ``BN ∈ {8, 16, 32, 64, 128}`` × ``num_warps ∈ {2, 4}`` × the ++ BM / D_INNER / NUM_STAGES_PIPE combinations that fit per-CTA SMEM ++ on H200. ++ ++ The ``BN = 8`` rung unlocks the small-N + medium-K Pattern-A regime ++ (``B*N <= 8`` + ``K ∈ {16, 32, 64, 128}``) where a small row tile + ++ sortmerge ``BM = max(32, K)`` wins over insert. ``tl.dot`` silently ++ pads the M dimension to 16 on Triton 3.x sm_90, so BN=8 still ++ compiles. ++ ++ ``dtype_bytes`` and ``smem_limit`` parameterise the SMEM check so ++ fp32 inputs don't slip through with a config that OOR's at launch. ++ Defaults preserve the pre-fix behaviour (bf16 / 220 KB) for any ++ caller that didn't yet plumb dtype awareness through. ++ """ ++ # 220 KB stays the default — slightly below the 227 KB H200 opt-in ++ # limit to leave room for static SMEM the compiler adds. ++ max_sram = 220_000 if smem_limit is None else smem_limit ++ d_pad = _next_pow2(D) ++ topk_pad_sm = max(32, _next_pow2(K)) ++ topk_pad_ins = _next_pow2(K) ++ d_inner_cands = [d_pad] if D <= 256 else [128] ++ ++ configs = [] ++ ++ # sort-merge: BM == TOPK_PAD; BN ∈ {8, 16, 32} (BN >= 64 never wins ++ # for sortmerge in autotune data). ++ bm_sm = topk_pad_sm ++ for bn in [8, 16, 32]: ++ for nw in [2, 4]: ++ for d_inner in d_inner_cands: ++ num_d_iters = math.ceil(D / d_inner) ++ ns = 2 if num_d_iters == 1 else 1 ++ sram = _estimate_sram( ++ bn, bm_sm, D, K, d_inner, ns, ++ sortmerge=True, dtype_bytes=dtype_bytes, ++ ) ++ if sram <= max_sram: ++ for ns_pipe in _pipe_stages_for(K, d_inner, D): ++ configs.append({ ++ "BN": bn, "BM": bm_sm, "D_INNER": d_inner, ++ "num_warps": nw, ++ "TOPK_PAD": topk_pad_sm, ++ "kernel_mode": "sortmerge", ++ "NUM_STAGES_PIPE": ns_pipe, ++ }) ++ ++ # iterative insert: BM decoupled from K, BN ∈ {8, 16, 32, 64, 128}. ++ for bn in [8, 16, 32, 64, 128]: ++ for nw in [2, 4]: ++ for bm in [64, 128, 256]: ++ for d_inner in d_inner_cands: ++ num_d_iters = math.ceil(D / d_inner) ++ ns = 2 if num_d_iters == 1 else 1 ++ sram = _estimate_sram( ++ bn, bm, D, K, d_inner, ns, ++ sortmerge=False, dtype_bytes=dtype_bytes, ++ ) ++ if sram <= max_sram: ++ for ns_pipe in _pipe_stages_for(K, d_inner, D): ++ configs.append({ ++ "BN": bn, "BM": bm, "D_INNER": d_inner, ++ "num_warps": nw, ++ "TOPK_PAD": topk_pad_ins, ++ "kernel_mode": "insert", ++ "NUM_STAGES_PIPE": ns_pipe, ++ }) ++ ++ return configs ++ ++ ++def _gen_m_splits(M, BM, *, B=1, N=1, BN=16): ++ """M_PER_SPLIT candidates targeting {1, 2, 4, 8, 16} waves on 132 SMs. ++ ++ The wave count is critical -- Stage-2 reduce scales linearly with ++ NUM_SPLITS, while too few splits leaves SMs idle. The autotune ++ sweep confirms that wave=2 is the winner on huge-M + medium-N ++ shapes (e.g. ``1×128×10M×64×10``), which the original {1, 4, 16} ++ grid missed. ++ ++ Capped at 4096×BM tiles (``MAX_MPS_TILES``) to bound the static ++ M-loop trip count for fast Triton compilation. Also includes the ++ single-pass case when M itself fits the cap (subsumes the large-N ++ kernel as ``num_splits = 1``). ++ """ ++ NUM_SMS = 132 ++ MAX_MPS_TILES = 4096 ++ num_n_tiles = max(1, (N + BN - 1) // BN) ++ max_mps = min(M, MAX_MPS_TILES * BM) ++ candidates = set() ++ for waves in [1, 2, 4, 8, 16]: ++ target_splits = max(2, (NUM_SMS * waves) // (num_n_tiles * B)) ++ target_splits = min(target_splits, max(2, M // BM)) ++ mps = (M + target_splits - 1) // target_splits ++ mps = ((mps + BM - 1) // BM) * BM ++ mps = max(mps, BM) ++ mps = min(mps, max_mps) ++ candidates.add(mps) ++ if M <= max_mps: ++ single = ((M + BM - 1) // BM) * BM ++ candidates.add(single) ++ return sorted(candidates) ++ ++ ++# ── shape-only heuristic ─────────────────────────────────────────────── ++ ++ ++def _heuristic_config(B, N, M, D, K, *, force_path=None, ++ dtype_bytes=2, smem_limit=None): ++ """Pick a kernel config from shape alone -- no autotune. ++ ++ When ``dtype_bytes`` and ``smem_limit`` are supplied (the dispatcher ++ plumbs them from ``x.element_size()`` and the device's opt-in SMEM ++ cap), the picked config is post-processed by :func:`_fit_config_to_smem` ++ so that fp32 inputs at large D never produce a config that OOR's at ++ launch. ++ ++ Derived from a 92-shape autotune sweep on H200/bf16. The data ++ showed clean decision boundaries on three axes: ++ ++ * **BN by NB-bucket + M-bump**: ++ NB <= 8 -> BN=8; ++ NB <= 32 -> BN=16 (32 at M>=5M); ++ NB <= 256 -> BN=64 (128 at M>=5M + narrow-D + small-K); ++ NB <= 2K -> BN=64 (128 at M>=5M); ++ NB >= 8K -> BN=128 (64 for D>=256 + K>=32). ++ ++ The M-bump rule fixes the ``(1, 128, 10M, 64, 10)`` regression ++ identified by klib -- at M=10M each extra N-tile costs ++ ~300us of c-replication HBM traffic. ++ ++ * **BM by K**: ++ K<=4 -> 128/256 (256 only at very-small M + small NB); ++ K=5..16 -> 128 default (256 at huge-M); ++ K=32 -> 128 default (256 at NB<=8 + M<=200K); ++ K>=64 -> 64 default (sortmerge at NB<=8). ++ ++ * **kernel_mode**: ++ ``sortmerge`` only when ``NB <= 8`` and ``K ∈ {32, 64, 128}`` ++ and ``D <= 256`` and ``M <= 200K`` (Pattern-A1), plus a ++ secondary ``NB <= 16 + K ∈ {32, 64} + M <= 200K`` corner ++ (Pattern-A2). All other shapes use ``insert``. ++ ++ * **num_warps**: 4 everywhere except tiny tiles (BN=8 + BM<=64 + ++ K<=4 + huge M, where nw=2 wins by reducing register pressure) ++ and the BN=16 + huge-M + small-K corner where nw=2 frees a ++ warp set for the topK epilogue. ++ ++ * **target waves**: 1 for build (NB>=8K) and very-large K; ++ 2 for medium/large M; 4 for small M + small K. ++ ++ * **NUM_STAGES_PIPE**: 2 by default; 3 for BN>=64 + BM<=128 + ++ K<=8 (small tile + tiny K hides HBM latency) and for ++ BN ∈ [16, 32] + BM=256 + K>=64 (big tile keeps K loop fed); ++ 1 for narrow-tile + any K (avoid register spills) and for the ++ D-split path; 2 for D_INNER>=256 and for BN>=128. Without ++ this rule, K>=64 + D=256 shapes hit catastrophic regressions ++ (e.g. ``(1,16,100K,256,128)`` 439us -> 2429us at ns=1, a 5.5x ++ slowdown); conversely K<=32 + huge-M shapes save 20-45% vs ++ ns=2 (e.g. ``(1,1,10M,128,4)`` 1426us -> 785us at ns=1). ++ ++ Args: ++ B, N, M, D, K: shape parameters. ++ force_path: ``"large_n"`` for single-pass (one M-split per CTA), ++ ``None`` for the wave-targeted M-split. ++ ++ Returns: ++ Config dict: ``BN, BM, D_INNER, TOPK_PAD, kernel_mode, ++ num_warps, M_PER_SPLIT, NUM_SPLITS, NUM_STAGES_PIPE``. ++ """ ++ NUM_SMS = 132 ++ NB = N * B ++ TOPK_PAD = _next_pow2(K) ++ D_INNER = _next_pow2(D) if D <= 256 else 128 ++ is_large_n = (force_path == "large_n") ++ # MAX_MPS_TILES bounds the per-CTA M-loop iteration count. The ++ # original cap of 512 was too tight for huge M with BM=128 (cap ++ # = 65K), which forced fewer splits than autotune wanted. Bump ++ # to 4096 so we never bottleneck on this cap. ++ MAX_MPS_TILES = 4096 ++ ++ # ─────────────────────────────────────────────────────────────── ++ # Step 1: kernel_mode, BN, BM, num_warps ++ # ─────────────────────────────────────────────────────────────── ++ if NB <= 8: ++ # Tiny query: BN=8 default. ++ BN = 8 ++ # Sortmerge fast-path: K ∈ {32, 64, 128}, M <= 200K. Now ++ # extended to ANY D (autotune ``(1,1,100K,1024,32)`` picks ++ # BN=8 BM=32 sort). ++ if (not is_large_n) and K in (32, 64, 128) and M <= 200_000: ++ BM = max(32, _next_pow2(K)) ++ sram = _estimate_sram(8, BM, D, K, D_INNER, ++ num_stages=(2 if D_INNER >= D else 1), ++ sortmerge=True) ++ if sram <= 220_000: ++ kernel_mode = "sortmerge" ++ num_warps = 2 if BM >= 128 else 4 ++ else: ++ kernel_mode = "insert" ++ BM = 256 if M <= 200_000 else 128 ++ num_warps = 4 ++ else: ++ kernel_mode = "insert" ++ if K <= 4: ++ # K=4 + D>=256 + huge M -> BM=64 nw=2 (autotune ++ # ``(1,1,1M,256,4)`` picks this). HBM-bound; BM=256 ++ # adds wasted SMEM. ++ if D >= 256 and M >= 500_000: ++ BM = 64 ++ else: ++ BM = 256 ++ elif K <= 16: ++ # D>=256 + K=16 -> BM=128 (autotune (1,1,100K,256,16)) ++ if D >= 256 and K >= 16: ++ BM = 128 ++ else: ++ BM = 256 if (M >= 500_000 or K >= 16) else 128 ++ elif K == 32: ++ # autotune (1,1,1M,256,32) prefers BM=128 at D>=256. ++ # At narrower D, BM=256 wins. ++ if D >= 256 and M >= 500_000: ++ BM = 128 ++ else: ++ BM = 256 ++ elif K <= 64: ++ BM = 64 ++ else: # K=128 ++ BM = 256 ++ # nw=2 for small/medium-D BM=64 cases ++ if BM == 64 and K <= 4 and D >= 512 and M >= 500_000: ++ # D >= 512 hits the D-split path (D_INNER capped at 128); ++ # nw=2 was calibrated for this regime at (1, 1, 1M, 256, 4) ++ # but later re-benching at D=256 showed nw=4 ns=1 wins by ++ # 1.48-1.81x there (single-D-iter, multibuffered C). Keep ++ # nw=2 only for the D-split path; D=256 falls to the ++ # default nw=4 below and gets the ns=1 override in Step 3. ++ num_warps = 2 ++ elif BM >= 128 and K <= 4 and D <= 64 and M >= 5_000_000: ++ # NB<=8 + K<=4 + narrow-D + huge-M: nw=2 wins by 20-30 %. ++ # 4 warps split the SM register file too thinly for the ++ # K=4 argmin-insert epilogue at BM>=128, so the compiler ++ # spills; halving the warp count doubles regs/warp and ++ # lets ILP recover. Verified at D=64 across M ∈ [5M, 60M] ++ # (D=128 strictly prefers nw=4 at this BM; do not bump ++ # the threshold without re-benching). ++ num_warps = 2 ++ else: ++ num_warps = 4 ++ ++ elif NB <= 32: ++ kernel_mode = "insert" ++ # Sortmerge also wins at NB <= 16 + K ∈ {32, 64} (autotune ++ # ``(1,16,100K,256,32)`` picks BN=8 BM=32 sortmerge). ++ if (not is_large_n) and NB <= 16 and K in (32, 64) and M <= 200_000: ++ BN = 8 ++ BM = max(32, _next_pow2(K)) ++ sram = _estimate_sram(8, BM, D, K, D_INNER, ++ num_stages=(2 if D_INNER >= D else 1), ++ sortmerge=True) ++ if sram <= 220_000: ++ kernel_mode = "sortmerge" ++ num_warps = 4 ++ else: ++ BN, BM = 16, 256 ++ if kernel_mode == "insert": ++ if M >= 5_000_000 and N >= 32: ++ BN = 32 ++ else: ++ BN = max(8, min(16, _next_pow2(N))) ++ if K <= 4: ++ BM = 128 ++ elif K <= 16: ++ BM = 256 if M <= 2_000_000 else 128 ++ else: # K >= 32 ++ BM = 256 ++ if K <= 4 and M >= 500_000: ++ num_warps = 2 ++ elif K <= 16 and M >= 5_000_000 and BN <= 16: ++ num_warps = 2 ++ elif K <= 16 and NB >= 17 and M >= 5_000_000: ++ # autotune (1,32,10M,64,8) BN=32 nw=2 wins ++ num_warps = 2 ++ else: ++ num_warps = 4 ++ ++ elif NB <= 256: ++ # Medium queries: BN must scale with M to control ++ # c-replication, but ONLY when K is very small (K<=4) or ++ # very large (K>=32). At K ∈ {5..16} the autotune consistently ++ # picks BN=64 even at M=10M, because K=8..16 hits a sweet ++ # spot where the topk-insert + Stage-2 cost of 1-wave ++ # BN=128 outweighs the c-HBM savings. ++ kernel_mode = "insert" ++ big_NB = (NB >= 192) ++ ++ # K=128 special-case: medium NB + K=128, autotune picks ++ # BN=8 BM=256 insert (NOT sortmerge!). The K-step argmin loop ++ # at BN=64 K=128 is too expensive -- better to chop N into ++ # many small tiles. ++ if K >= 128 and M <= 200_000: ++ BN = 8 ++ BM = 256 ++ elif M >= 5_000_000 and D <= 64 and (K <= 4 or K >= 32): ++ BN = 128 ++ elif big_NB and M >= 5_000_000: ++ BN = 128 ++ elif big_NB and D >= 256 and M <= 200_000 and K <= 4: ++ BN = 128 ++ elif D >= 256 and K <= 16 and (NB >= 64 and M >= 500_000): ++ # D=256 + K<=16 + (medium NB + medium M): BN=128 wins. ++ # autotune (1,128,1M,256,8) picks BN=128 BM=64. ++ BN = 128 ++ else: ++ BN = 64 ++ BN = min(BN, max(8, _next_pow2(N))) ++ # BM ++ if K >= 128: ++ BM = 256 # already covered above, keep consistent here ++ elif K <= 4 and D <= 128 and M <= 200_000: ++ BM = 256 ++ elif D >= 256 and K <= 16: ++ BM = 64 ++ elif K <= 32: ++ BM = 64 if (BN == 64 and K == 32 and D >= 256) else 128 ++ elif K <= 64: ++ BM = 128 ++ else: ++ BM = 128 ++ num_warps = 4 ++ ++ elif NB <= 2048: ++ kernel_mode = "insert" ++ if M >= 5_000_000: ++ BN = 128 ++ elif M >= 500_000 and D >= 256 and K <= 16: ++ BN = 128 ++ else: ++ BN = 64 ++ BN = min(BN, max(8, _next_pow2(N))) ++ if K <= 4 and M >= 500_000: ++ BM = 64 ++ elif K <= 16 and M >= 500_000 and D >= 128: ++ BM = 64 ++ elif K >= 64: ++ BM = 64 ++ else: ++ BM = 128 ++ num_warps = 4 ++ ++ else: # NB >= 8K (build / large-eval regime) ++ kernel_mode = "insert" ++ # autotune patterns: ++ # NB=50K-100K D<=128 K<=8 -> BN=128 (halves HBM) ++ # NB=10K K>=8 -> BN=64 (more topk work per row) ++ # NB>=10K D>=256 K<=16 -> BN=128 ++ # NB>=10K D>=256 K>=32 -> BN=64 ++ if D >= 256 and K <= 16: ++ BN = 128 ++ elif D >= 256 and K >= 32: ++ BN = 64 ++ elif K <= 4: ++ BN = 128 ++ elif K <= 8 and NB >= 30_000: ++ # Larger NB -> BN=128 saves enough HBM to overcome topk cost ++ BN = 128 ++ else: ++ BN = 64 ++ if K <= 4 and D >= 256: ++ BM = 64 # SMEM pressure at D=256 + BM=128 ++ elif K <= 4: ++ BM = 128 if (D <= 64 or BN == 128) else 64 ++ elif K <= 8 and BN == 128: ++ BM = 128 ++ elif K >= 64 and BN == 64 and D <= 128 and NB >= 30_000: ++ # klib carve-out: NB in [30K, ~200K] + BN=64 + K>=64 + ++ # D<=128 wins big at BM=128 splits=1 vs BM=64 splits=2. ++ # Empirically at (1, 48K, 48K, 64, 64) the upstream rule ++ # (BM=64, target_splits=2) runs 17.95 ms, while BM=128 ++ # splits=1 ns=2 runs 11.71 ms (-35%); at (1, 64K, 64K, 64, ++ # 64) 29.20 ms -> 18.91 ms (-35%). Bounded on K>=64 (K=32 ++ # autotune at this NB shows BM=64 splits=1 wins, ~1 ms over ++ # BM=128 splits=1) and on NB>=30K so the upstream rule ++ # still applies to smaller-NB shapes where its M-split ++ # target was directly autotune-derived. ++ BM = 128 ++ else: ++ BM = 64 ++ num_warps = 4 ++ ++ # ─────────────────────────────────────────────────────────────── ++ # Step 2: M_PER_SPLIT (target waves) ++ # ─────────────────────────────────────────────────────────────── ++ num_n_tiles = max(1, math.ceil(N / BN)) ++ ctas_no_split = num_n_tiles * B ++ ++ # Workaround for an insert-kernel correctness issue with the ++ # D-split path: when ``M_PER_SPLIT == BM`` (single M-loop ++ # iteration) and the d-loop has >=4 chunks, the result is wrong. ++ # Bumping the minimum mps to 2*BM keeps the M-loop at >=2 iters. ++ needs_min_2_iters = (D_INNER < D and (D + D_INNER - 1) // D_INNER >= 4) ++ min_mps = (2 * BM) if needs_min_2_iters else BM ++ ++ if is_large_n: ++ M_PER_SPLIT = ((M + BM - 1) // BM) * BM ++ elif ctas_no_split >= NUM_SMS * 8: ++ # Massively oversaturated (e.g. N=100K, BN=64, ctas=1563): one ++ # split fully utilises SMs. ++ M_PER_SPLIT = ((M + BM - 1) // BM) * BM ++ elif K >= 32 and BN == 64 and D <= 128 and NB >= 30_000 and not is_large_n: ++ # klib carve-out (pairs with the BM=128 K>=64 rule above ++ # plus K=32 BM=64 standalone): the K>=32 build regime at ++ # NB>=30K wants splits=1 -- 2 splits doubles the K-step ++ # argmin work per row and the Stage-2 reduce, neither of ++ # which the autotune at NB=10K K=8 accounted for. Empirically ++ # saves ~6 ms on (1, 48K, 48K, 64, 64) and ~10 ms on ++ # (1, 64K, 64K, 64, 64) vs target_splits=2, and ~1 ms on ++ # (1, 56K, 56K, 64, 32) vs target_splits=2. ++ M_PER_SPLIT = ((M + BM - 1) // BM) * BM ++ else: ++ if ctas_no_split >= NUM_SMS * 4: ++ # ctas >= 528: 2 splits for tail-fill at very large builds ++ # (autotune (1,100K,100K,64,4) BN=128 ctas=782 -> 2 splits) ++ target_splits = 2 ++ elif ctas_no_split >= NUM_SMS * 2: ++ # ctas in [264, 528): single split already 2-4 waves; adding ++ # splits hurts more than it helps. autotune ++ # (1,50K,50K,64,8) BN=128 ctas=391 -> 1 split. ++ target_splits = 1 ++ elif ctas_no_split >= NUM_SMS: ++ # ctas in [132, 264): 3 splits for better tail. autotune ++ # (1,10K,10K,64,16) BN=64 ctas=157 -> 3 splits. ++ target_splits = 3 ++ else: ++ # Under-saturated: target_waves by K (autotune-derived). ++ # Use FLOOR div: snaps to integer wave count, which matches ++ # autotune choices on shapes like NB=1024 K=32 ctas=16 ++ # (16 splits = 2 waves, 17 splits = 2.06 waves but bad ++ # tail). ++ if kernel_mode == "sortmerge": ++ target_waves = 2 ++ elif K >= 128 and NB >= 64: ++ # K=128 + medium NB: 2w wins (autotune ++ # (1,128,100K,128,128) picks 16 splits at ctas=16). ++ target_waves = 2 ++ elif K >= 128: ++ target_waves = 1 ++ elif K >= 64 and ctas_no_split <= 2: ++ # K=64 + tiny ctas: 1 wave wins (autotune ++ # (1,128,100K,128,64) picks 66 splits = 1w). ++ target_waves = 1 ++ elif K >= 64: ++ target_waves = 2 ++ elif K <= 4 and ctas_no_split == 1 and (NB <= 32 or M <= 1_000_000): ++ # K=4 single-CTA-stack: 4 waves at small NB or small M. ++ # autotune NB=1 K=4 -> 521 splits; (1,128,10M,64,4) ++ # however picks 264 splits (2w) -- too much M per CTA ++ # for more splits to help. ++ target_waves = 4 ++ elif K <= 8 and ctas_no_split >= 8 and ctas_no_split <= 32: ++ target_waves = 4 if M >= 500_000 else 2 ++ elif K <= 8 and ctas_no_split == 1 and M >= 5_000_000 and BN <= 32: ++ target_waves = 4 ++ elif B >= 2 and N <= 8 and K <= 8: ++ # Batched B>=2 + single-query + K<=8: 4w wins ++ target_waves = 4 ++ else: ++ target_waves = 2 ++ target_splits = max(1, NUM_SMS * target_waves // ctas_no_split) ++ target_splits = min(target_splits, max(1, math.ceil(M / BM))) ++ mps_raw = math.ceil(M / target_splits) ++ mps = math.ceil(mps_raw / BM) * BM ++ mps = max(min_mps, min(mps, ((M + BM - 1) // BM) * BM)) ++ mps = min(mps, MAX_MPS_TILES * BM) ++ M_PER_SPLIT = mps ++ ++ NUM_SPLITS = math.ceil(M / M_PER_SPLIT) ++ ++ # ─────────────────────────────────────────────────────────────── ++ # Step 3: NUM_STAGES_PIPE (M-loop pipeline depth) ++ # ─────────────────────────────────────────────────────────────── ++ # Derived from a 201-shape × 4-ns sweep on H200 (upstream ++ # ``bench/sweep_ns_by_k.py``), cross-checked against direct probes ++ # on N>=64 build shapes. Three competing effects: ++ # ++ # * Each extra pipeline stage adds another buffered C-tile, ++ # costing ``BM * D_INNER * 2`` bytes of SMEM/registers. Wide ++ # D-tiles amortise this well; narrow ones don't. ++ # * The K-step argmin inner loop holds live state proportional ++ # to ``BN * TOPK_PAD * 8`` bytes. For narrow rows (BN<=16) + ++ # large K, registers get tight and deeper pipelines spill. ++ # * Medium-N tiles (BN>=64) have abundant compute per tile to ++ # amortise prefetch, but get little benefit at very large K ++ # because the K loop dominates anyway. ++ # ++ # Empirical winners (mode_ns from the 201-shape sweep): ++ # ++ # D_INNER >= 256 (wide D) -> ns=2 (uniform across K) ++ # BN >= 128 (huge build N) -> ns=2 (registers tight) ++ # BN >= 64 + BM <= 128 + K <= 8 -> ns=3 (small tile + tiny K) ++ # BN >= 64 -> ns=2 (BM=256 SMEM-bound) ++ # BN ∈ [16,32] + BM=256 + K >= 64 -> ns=3 (big-tile large-K) ++ # otherwise (narrow tile + any K) -> ns=1 (avoid spills) ++ # ++ # Without this rule, K>=64 + D=256 shapes hit catastrophic ++ # regressions (e.g. (1,16,100K,256,128) 439us -> 2429us at ns=1, ++ # a 5.5x slowdown). Conversely K<=32 + huge-M shapes save 20-45% ++ # vs ns=2 (e.g. (1,1,10M,128,4) 1426us -> 785us at ns=1). ++ if D_INNER < D: ++ NUM_STAGES_PIPE = 1 # D-split path is hard-coded to ns=1 ++ elif (BN == 8 and BM == 64 and D_INNER == 256 ++ and K <= 4 and M >= 500_000): ++ # NB<=8 + K<=4 + D=256 + BM=64 + huge-M: ns=1 wins by 1.48-1.81x ++ # over the default ns=2. Direct verification at K ∈ {1, 2, 4} × ++ # M ∈ {1M, 5M, 10M, 30M}: (nw=4, ns=1) lifts %peak HBM from ++ # 39-48 % (ns=2) to 57-85 % (ns=1). The ns=2 default below was ++ # calibrated for wider-N tiles; at BN=8 the deeper pipeline ++ # spends too many registers on prefetch and the K=4 epilogue ++ # spills. Pairs with the nw=4 default the rule above falls into ++ # at D=256. ++ NUM_STAGES_PIPE = 1 ++ elif D_INNER >= 256: ++ NUM_STAGES_PIPE = 2 ++ elif BN >= 128: ++ NUM_STAGES_PIPE = 2 ++ elif BN >= 64 and BM <= 128 and K <= 8: ++ NUM_STAGES_PIPE = 3 ++ elif BN >= 64: ++ NUM_STAGES_PIPE = 2 ++ elif BN >= 16 and BM >= 256 and K >= 64: ++ NUM_STAGES_PIPE = 3 ++ else: ++ NUM_STAGES_PIPE = 1 ++ ++ cfg = { ++ "BN": BN, "BM": BM, "D_INNER": D_INNER, ++ "TOPK_PAD": ( ++ max(32, _next_pow2(K)) if kernel_mode == "sortmerge" ++ else _next_pow2(K) ++ ), ++ "kernel_mode": kernel_mode, "num_warps": num_warps, ++ "M_PER_SPLIT": M_PER_SPLIT, ++ "NUM_SPLITS": NUM_SPLITS, ++ "NUM_STAGES_PIPE": NUM_STAGES_PIPE, ++ } ++ ++ # Dtype-aware SMEM shrink. The shape-only heuristic above was tuned ++ # on bf16 and may pick (BN=128, BM=128, D_INNER=128) for shapes that ++ # cleanly fit ~210 KB at 2 bytes/elt but blow past the 227 KB H200 ++ # opt-in cap at 4 bytes/elt. The fitter shrinks BN/BM/NUM_STAGES_PIPE ++ # by powers of 2 until the kernel fits — leaving D_INNER and ++ # kernel_mode (which encode different perf regimes) alone. ++ if smem_limit is not None: ++ cfg = _fit_config_to_smem(cfg, D, K, dtype_bytes, smem_limit) ++ ++ return cfg ++ ++ ++# ── autotune (opt-in) ────────────────────────────────────────────────── ++ ++ ++_autotune_cache: dict = {} ++ ++ ++def _autotune(x, c, k, *, force_path=None): ++ """Brute-force autotune over the candidate config grid. ++ ++ Cached per ``(B, N, M, D, k, dtype, force_path)`` shape; first call ++ costs ~30-60 s of compile time, subsequent calls hit the cache in ++ sub-ms. ++ ++ The candidate grid is dtype-aware: ``_gen_configs`` filters out ++ configs whose SMEM exceeds the device opt-in limit at this dtype, ++ so fp32 inputs at large D never have an unrunnable config in the ++ search space. ++ """ ++ B, N, D = x.shape ++ M = c.shape[1] ++ dtype_bytes = x.element_size() ++ key = (B, N, M, D, k, x.dtype, force_path) ++ ++ if key in _autotune_cache: ++ return _autotune_cache[key] ++ ++ smem_limit = _smem_limit(x.device) ++ configs = _gen_configs(D, k, ++ dtype_bytes=dtype_bytes, ++ smem_limit=min(smem_limit, 220_000)) ++ if not configs: ++ raise RuntimeError( ++ f"No valid flash_knn config for D={D}, K={k}, " ++ f"dtype={x.dtype} (element_size={dtype_bytes})" ++ ) ++ ++ # Drop configs with BN > next_pow2(N) -- they would just pad N to BN ++ # and run the same kernel slower. ++ max_bn = max(16, _next_pow2(N)) ++ configs = [cfg for cfg in configs if cfg["BN"] <= max_bn] ++ ++ best_time = float('inf') ++ best_cfg = None ++ max_partial_bytes = 4 * 1024**3 ++ ++ for cfg in configs: ++ bn = cfg["BN"] ++ bm = cfg["BM"] ++ d_inner = cfg["D_INNER"] ++ topk_pad = cfg["TOPK_PAD"] ++ kernel_mode = cfg["kernel_mode"] ++ nw = cfg["num_warps"] ++ ns_pipe = cfg.get("NUM_STAGES_PIPE", 2) ++ ++ if force_path == "large_n": ++ mps_list = [((M + bm - 1) // bm) * bm] ++ else: ++ mps_list = _gen_m_splits(M, bm, B=B, N=N, BN=bn) ++ ++ for mps in mps_list: ++ num_splits = math.ceil(M / mps) ++ partial_bytes = B * num_splits * N * k * 8 ++ if partial_bytes > max_partial_bytes: ++ continue ++ ++ try: ++ pv = torch.empty((B, N, num_splits, k), ++ device=x.device, dtype=torch.float32) ++ pi = torch.empty((B, N, num_splits, k), ++ device=x.device, dtype=torch.int32) ++ grid = (num_splits, math.ceil(N / bn), B) ++ pv_s0, pv_s1, pv_s2, pv_s3 = pv.stride() ++ pi_s0, pi_s1, pi_s2, pi_s3 = pi.stride() ++ ++ if kernel_mode == "sortmerge": ++ def run(grid=grid, bn=bn, bm=bm, d_inner=d_inner, ++ topk_pad=topk_pad, mps=mps, nw=nw, ++ num_splits=num_splits, pv=pv, pi=pi, ++ ns_pipe=ns_pipe, ++ pv_s0=pv_s0, pv_s1=pv_s1, pv_s2=pv_s2, pv_s3=pv_s3, ++ pi_s0=pi_s0, pi_s1=pi_s1, pi_s2=pi_s2, pi_s3=pi_s3): ++ _flash_knn_sortmerge_kernel[grid]( ++ x, c, pv, pi, ++ x.stride(0), x.stride(1), x.stride(2), ++ c.stride(0), c.stride(1), c.stride(2), ++ pv_s0, pv_s2, pv_s1, pv_s3, ++ pi_s0, pi_s2, pi_s1, pi_s3, ++ N=N, M=M, D=D, K=k, M_PER_SPLIT=mps, ++ BN=bn, BM=bm, D_INNER=d_inner, ++ TOPK_PAD=topk_pad, ++ NUM_STAGES_PIPE=ns_pipe, ++ num_warps=nw, ++ ) ++ if num_splits > 1: ++ pv_flat = pv.view(B, N, -1) ++ pv_flat.topk(k, dim=-1, largest=False) ++ else: ++ max_steps = min(k, bm) ++ def run(grid=grid, bn=bn, bm=bm, d_inner=d_inner, ++ topk_pad=topk_pad, mps=mps, nw=nw, ++ num_splits=num_splits, pv=pv, pi=pi, ++ max_steps=max_steps, ns_pipe=ns_pipe, ++ pv_s0=pv_s0, pv_s1=pv_s1, pv_s2=pv_s2, pv_s3=pv_s3, ++ pi_s0=pi_s0, pi_s1=pi_s1, pi_s2=pi_s2, pi_s3=pi_s3): ++ _flash_knn_insert_kernel[grid]( ++ x, c, pv, pi, ++ x.stride(0), x.stride(1), x.stride(2), ++ c.stride(0), c.stride(1), c.stride(2), ++ pv_s0, pv_s2, pv_s1, pv_s3, ++ pi_s0, pi_s2, pi_s1, pi_s3, ++ N=N, M=M, D=D, K=k, M_PER_SPLIT=mps, ++ BN=bn, BM=bm, D_INNER=d_inner, ++ TOPK_PAD=topk_pad, MAX_STEPS=max_steps, ++ NUM_STAGES_PIPE=ns_pipe, ++ num_warps=nw, ++ ) ++ if num_splits > 1: ++ pv_flat = pv.view(B, N, -1) ++ pv_flat.topk(k, dim=-1, largest=False) ++ ++ t = _bench_quick(run) ++ if t < best_time: ++ best_time = t ++ best_cfg = {**cfg, "M_PER_SPLIT": mps, ++ "NUM_SPLITS": num_splits, ++ "NUM_STAGES_PIPE": ns_pipe} ++ ++ del pv, pi ++ except Exception: ++ continue ++ ++ if best_cfg is None: ++ raise RuntimeError( ++ f"All flash_knn configs failed for B={B}, N={N}, M={M}, D={D}, K={k}") ++ ++ _autotune_cache[key] = best_cfg ++ return best_cfg ++ ++ ++# ── runner ───────────────────────────────────────────────────────────── ++ ++ ++_OOR_FALLBACK_CACHE: dict = {} ++ ++ ++def _try_launch(x, c, cfg, B, N, M, D, k): ++ """Single launch attempt; returns idxs or raises ``triton.compiler.errors.OutOfResources``. ++ ++ Factored out of ``_run`` so the OOR-shrink retry loop can call it ++ repeatedly with a shrunk cfg. ++ """ ++ bn = cfg["BN"] ++ bm = cfg["BM"] ++ d_inner = cfg["D_INNER"] ++ topk_pad = cfg["TOPK_PAD"] ++ mps = cfg["M_PER_SPLIT"] ++ num_splits = cfg["NUM_SPLITS"] ++ kernel_mode = cfg["kernel_mode"] ++ nw = cfg["num_warps"] ++ num_stages_pipe = cfg.get("NUM_STAGES_PIPE", 2) ++ ++ partial_vals = torch.empty((B, N, num_splits, k), ++ device=x.device, dtype=torch.float32) ++ partial_idxs = torch.empty((B, N, num_splits, k), ++ device=x.device, dtype=torch.int32) ++ grid = (num_splits, math.ceil(N / bn), B) ++ ++ pv_s0, pv_s1, pv_s2, pv_s3 = partial_vals.stride() ++ pi_s0, pi_s1, pi_s2, pi_s3 = partial_idxs.stride() ++ ++ if kernel_mode == "sortmerge": ++ _flash_knn_sortmerge_kernel[grid]( ++ x, c, partial_vals, partial_idxs, ++ x.stride(0), x.stride(1), x.stride(2), ++ c.stride(0), c.stride(1), c.stride(2), ++ pv_s0, pv_s2, pv_s1, pv_s3, ++ pi_s0, pi_s2, pi_s1, pi_s3, ++ N=N, M=M, D=D, K=k, M_PER_SPLIT=mps, ++ BN=bn, BM=bm, D_INNER=d_inner, ++ TOPK_PAD=topk_pad, ++ NUM_STAGES_PIPE=num_stages_pipe, ++ num_warps=nw, ++ ) ++ else: ++ max_steps = min(k, bm) ++ _flash_knn_insert_kernel[grid]( ++ x, c, partial_vals, partial_idxs, ++ x.stride(0), x.stride(1), x.stride(2), ++ c.stride(0), c.stride(1), c.stride(2), ++ pv_s0, pv_s2, pv_s1, pv_s3, ++ pi_s0, pi_s2, pi_s1, pi_s3, ++ N=N, M=M, D=D, K=k, M_PER_SPLIT=mps, ++ BN=bn, BM=bm, D_INNER=d_inner, ++ TOPK_PAD=topk_pad, MAX_STEPS=max_steps, ++ NUM_STAGES_PIPE=num_stages_pipe, ++ num_warps=nw, ++ ) ++ ++ if num_splits == 1: ++ return partial_idxs[:, :, 0, :].contiguous() ++ ++ pv = partial_vals.view(B, N, -1) ++ pi = partial_idxs.view(B, N, -1) ++ _, sel = pv.topk(k, dim=-1, largest=False, sorted=True) ++ out_idxs = pi.gather(-1, sel.to(torch.int64)).to(torch.int32) ++ return out_idxs ++ ++ ++def _shrink_cfg_on_oor(cfg: dict) -> Optional[dict]: ++ """Halve the cheapest SMEM-bearing axis. ``None`` when nothing left. ++ ++ Order: ++ 1. NUM_STAGES_PIPE 3 → 2 → 1 (drops the c-tile pipeline buffer ++ — usually free or nearly-free at this depth). ++ 2. Tend toward a near-square tile by **halving the larger of ++ (BN, BM)** first. WGMMA shape efficiency drops rapidly once ++ BM<32 on Hopper (MMA atom is 64x16x16), so we prefer to ++ shrink BN when BN > BM. For BN=128 BM=64 (typical build ++ regime) this gives BN=64 BM=64 in one step, which empirically ++ matches the kmeans-style ``_fit_config_to_smem`` choice and ++ beats BN=128 BM=32 by ~20 % at D=1024 fp32. ++ 3. Floors: BN ≥ 8, BM ≥ 32 (both required by the WGMMA atom). ++ ++ sortmerge requires ``BM == TOPK_PAD``; skip BM shrinks and only ++ halve BN. ++ ++ Also recomputes ``M_PER_SPLIT`` so it remains a multiple of the ++ new BM (the kernel's inner ``tl.range`` walk requires this). ++ """ ++ ns = int(cfg.get("NUM_STAGES_PIPE", 2)) ++ bm = int(cfg["BM"]) ++ bn = int(cfg["BN"]) ++ sortmerge = cfg.get("kernel_mode") == "sortmerge" ++ ++ out = dict(cfg) ++ if ns > 1: ++ out["NUM_STAGES_PIPE"] = ns - 1 ++ return out ++ if sortmerge: ++ if bn > 8: ++ out["BN"] = bn // 2 ++ return out ++ return None ++ # Non-sortmerge: prefer shrinking the larger axis. ++ if bn > bm and bn > 8: ++ out["BN"] = bn // 2 ++ return out ++ if bm > 32: ++ out["BM"] = bm // 2 ++ new_bm = out["BM"] ++ out["M_PER_SPLIT"] = ((int(out["M_PER_SPLIT"]) + new_bm - 1) ++ // new_bm) * new_bm ++ return out ++ if bn > 8: ++ out["BN"] = bn // 2 ++ return out ++ return None ++ ++ ++def _run(x: torch.Tensor, c: torch.Tensor, k: int, *, ++ force_path: Optional[str] = None, ++ autotune: bool = False) -> torch.Tensor: ++ """Shared dispatcher -- launch stage 1, optional stage 2 reduce, return idxs. ++ ++ Returns indices only ``(B, N, k) int32``. The wrapper at ++ :func:`dbscanlib._kernels.primitives.knn.flash_knn` calls the gather kernel to ++ produce true squared distances per neighbour. ++ ++ On ``OutOfResources`` at launch, we shrink one SMEM-bearing axis ++ (NS → BM → BN) and retry, caching the surviving config per ++ ``(D, K, dtype, force_path)`` so subsequent calls with the same ++ shape skip the failed compiles. This is the source of truth for ++ "does this fit?"; ``_estimate_sram`` only provides a coarse hint ++ for the autotune candidate filter. ++ """ ++ assert x.is_cuda and c.is_cuda ++ assert x.dtype == c.dtype ++ B, N, D = x.shape ++ M = c.shape[1] ++ assert c.shape == (B, M, D) ++ assert 1 <= k <= M ++ ++ if autotune: ++ cfg = _autotune(x, c, k, force_path=force_path) ++ else: ++ cfg = _heuristic_config( ++ B, N, M, D, k, ++ force_path=force_path, ++ dtype_bytes=x.element_size(), ++ smem_limit=min(_smem_limit(x.device), 220_000), ++ ) ++ ++ # Surviving-config cache keyed on the parts the dispatcher can't ++ # vary (D, K, dtype, force_path). If a previous call already ++ # shrunk past an OOR for this shape, start with that smaller cfg. ++ fb_key = (D, k, x.dtype, force_path, ++ cfg["kernel_mode"], cfg.get("D_INNER")) ++ if fb_key in _OOR_FALLBACK_CACHE: ++ cached = _OOR_FALLBACK_CACHE[fb_key] ++ # Use the cached cfg only if it's smaller than the current one ++ # on at least one SMEM-bearing axis (otherwise the heuristic ++ # already picked something safe for this shape). ++ sm_cur = (cfg["BN"], cfg["BM"], cfg.get("NUM_STAGES_PIPE", 2)) ++ sm_cached = (cached["BN"], cached["BM"], ++ cached.get("NUM_STAGES_PIPE", 2)) ++ if sm_cached < sm_cur: ++ cfg = {**cfg, **{k_: cached[k_] for k_ in ++ ("BN", "BM", "NUM_STAGES_PIPE")}} ++ # Re-align mps to new BM ++ new_bm = cfg["BM"] ++ cfg["M_PER_SPLIT"] = ((cfg["M_PER_SPLIT"] + new_bm - 1) ++ // new_bm) * new_bm ++ cfg["NUM_SPLITS"] = math.ceil(M / cfg["M_PER_SPLIT"]) ++ ++ # Import lazily so we don't pay it on every dispatch. ++ try: ++ from triton.runtime.errors import OutOfResources ++ except ImportError: ++ try: ++ from triton.compiler.errors import OutOfResources ++ except ImportError: ++ OutOfResources = RuntimeError # safest fallback ++ ++ last_err = None ++ initial_cfg = cfg ++ shrunk = False ++ for _attempt in range(8): ++ try: ++ result = _try_launch(x, c, cfg, B, N, M, D, k) ++ # Cache the surviving cfg so subsequent calls with the ++ # same (D, K, dtype, ...) shape skip the failed compile. ++ if shrunk: ++ _OOR_FALLBACK_CACHE[fb_key] = { ++ "BN": cfg["BN"], "BM": cfg["BM"], ++ "NUM_STAGES_PIPE": cfg.get("NUM_STAGES_PIPE", 2), ++ } ++ return result ++ except OutOfResources as e: ++ last_err = e ++ new_cfg = _shrink_cfg_on_oor(cfg) ++ if new_cfg is None: ++ break ++ cfg = new_cfg ++ shrunk = True ++ continue ++ ++ raise last_err if last_err is not None else RuntimeError( ++ "flash_knn: exhausted OOR shrink attempts") ++ ++ ++# ── public API ───────────────────────────────────────────────────────── ++ ++ ++def flash_knn_triton_small_n(x: torch.Tensor, c: torch.Tensor, k: int, ++ *, autotune: bool = False) -> torch.Tensor: ++ """Force the M-split (search) path -- ``(B, N, k) int32`` indices.""" ++ return _run(x, c, k, force_path=None, autotune=autotune) ++ ++ ++def flash_knn_triton_large_n(x: torch.Tensor, c: torch.Tensor, k: int, ++ *, autotune: bool = False) -> torch.Tensor: ++ """Force the single-pass (build) path -- ``(B, N, k) int32`` indices.""" ++ return _run(x, c, k, force_path="large_n", autotune=autotune) ++ ++ ++def flash_knn_triton(x: torch.Tensor, c: torch.Tensor, k: int, ++ *, autotune: bool = False) -> torch.Tensor: ++ """Universal Triton dispatch -- the heuristic picks BN/BM/mode/ ++ mps/ns_pipe based on shape, including whether to single-pass or ++ M-split. The per-CTA-count check inside :func:`_heuristic_config` ++ handles both build and eval shapes without a host-side forced path. ++ ++ Neither path materialises an N×M cross or distance matrix to HBM, ++ and neither computes or loads ``x_sq`` -- the two hard contracts ++ klib's KNN imposes on every Triton entry point here. ++ ++ Args: ++ x: ``(B, N, D)`` bf16 / fp16 / fp32 query tensor. ++ c: ``(B, M, D)`` corpus, same dtype. ++ k: number of neighbours. ++ autotune: ``False`` (default) uses the shape-only heuristic -- ++ first call pays a single Triton compile (~0.5 s). ``True`` ++ runs the full brute-force sweep + caches per shape (~30-90 s ++ first call). ++ ++ Returns: ++ idxs: ``(B, N, k)`` int32 -- nearest-neighbour indices, sorted ++ by ascending true squared L2 distance (ties broken by index). ++ """ ++ return _run(x, c, k, force_path=None, autotune=autotune) +diff --git a/dbscanlib/_kernels/primitives/knn/triton/insert.py b/dbscanlib/_kernels/primitives/knn/triton/insert.py +new file mode 100644 +index 0000000..7e72d76 +--- /dev/null ++++ b/dbscanlib/_kernels/primitives/knn/triton/insert.py +@@ -0,0 +1,204 @@ ++"""flash_knn — Triton iterative-insert top-K kernel (x²-free, signed score). ++ ++Same x²-free signed score as ``sortmerge.py`` — see that module's docstring ++for the rationale. The insert kernel differs in two ways: ++ ++ * ``BM`` is decoupled from ``TOPK_PAD`` (sortmerge requires ``BM == TOPK_PAD``). ++ The autotune sweeps ``BM ∈ {64, 128, 256}`` independently. ++ * Top-K is maintained as an unsorted ``(BN, TOPK_PAD)`` register array; ++ each chunk does at most ``MAX_STEPS = min(K, BM)`` argmin->insert ++ passes. fp32 ``<`` works natively on signed scores, so only the ++ trailing output sort uses the sortable-u32 transform. ++ ++Empirically the insert kernel wins on virtually every shape past the ++trivial ``K == 1`` case; sortmerge is kept for completeness (autotune ++still considers it, and it's the natural fit when ``K`` matches a ++hardware-sortable tile size). ++ ++Public name: ``_flash_knn_insert_kernel`` — called from ++``klib/primitives/knn/triton/dispatch.py``. ++""" ++import triton ++import triton.language as tl ++ ++from dbscanlib._kernels.primitives.knn.triton._common import ( ++ _fp32_to_sortable_u32, ++ _sortable_u32_to_fp32, ++) ++ ++ ++@triton.jit ++def _flash_knn_insert_kernel( ++ x_ptr, c_ptr, ++ partial_val_ptr, partial_idx_ptr, ++ stride_x_b, stride_x_n, stride_x_d, ++ stride_c_b, stride_c_m, stride_c_d, ++ stride_pv_b, stride_pv_s, stride_pv_n, stride_pv_k, ++ stride_pi_b, stride_pi_s, stride_pi_n, stride_pi_k, ++ N: tl.constexpr, M: tl.constexpr, ++ D: tl.constexpr, K: tl.constexpr, ++ M_PER_SPLIT: tl.constexpr, ++ BN: tl.constexpr, BM: tl.constexpr, ++ D_INNER: tl.constexpr, ++ TOPK_PAD: tl.constexpr, ++ MAX_STEPS: tl.constexpr, ++ NUM_STAGES_PIPE: tl.constexpr = 2, ++): ++ """Iterative-insert top-K on signed shifted distance. ++ ++ Grid: ``(num_m_splits, ceil(N/BN), B)``. ++ Two compile-time paths gated on ``D_INNER >= D`` (persistent x + ++ pipelined M-loop vs D-split with sequential M-loop). The M-loop ++ pipelining depth is exposed as ``NUM_STAGES_PIPE`` so the ++ dispatcher can tune it per shape (see :mod:`dispatch` for the ++ heuristic rules; defaults to 2 if not provided). ++ """ ++ pid_s = tl.program_id(0) ++ pid_n = tl.program_id(1) ++ pid_b = tl.program_id(2) ++ pid_b = pid_b.to(tl.int64) ++ ++ n_start = pid_n * BN ++ n_offs = (n_start + tl.arange(0, BN)).to(tl.int64) ++ n_mask = n_offs < N ++ ++ topk_vals = tl.full([BN, TOPK_PAD], float('inf'), dtype=tl.float32) ++ topk_idxs = tl.full([BN, TOPK_PAD], -1, dtype=tl.int32) ++ topk_max = tl.full([BN], float('inf'), dtype=tl.float32) ++ k_range = tl.arange(0, TOPK_PAD) ++ bm_range = tl.arange(0, BM) ++ m_base = pid_s.to(tl.int64) * M_PER_SPLIT ++ ++ if D_INNER >= D: ++ d_offs = tl.arange(0, D_INNER).to(tl.int64) ++ d_mask = d_offs < D ++ x_tile = tl.load( ++ x_ptr + pid_b * stride_x_b ++ + n_offs[:, None] * stride_x_n ++ + d_offs[None, :] * stride_x_d, ++ mask=n_mask[:, None] & d_mask[None, :], other=0.0) ++ ++ for m_local in tl.range(0, M_PER_SPLIT, BM, num_stages=NUM_STAGES_PIPE): ++ m_start = m_base + m_local ++ m_offs = m_start + bm_range.to(tl.int64) ++ m_mask = m_offs < M ++ ++ c_tile = tl.load( ++ c_ptr + pid_b * stride_c_b ++ + m_offs[:, None] * stride_c_m ++ + d_offs[None, :] * stride_c_d, ++ mask=m_mask[:, None] & d_mask[None, :], other=0.0) ++ ++ c_f = c_tile.to(tl.float32) ++ c_sq_tile = tl.sum(c_f * c_f, axis=1) ++ cross = tl.dot(x_tile, tl.trans(c_tile)).to(tl.float32) ++ score = c_sq_tile[None, :] - 2.0 * cross ++ score = tl.where(m_mask[None, :], score, float('inf')) ++ ++ chunk_best = tl.min(score) ++ threshold_worst = tl.max(topk_max) ++ if chunk_best < threshold_worst: ++ _active = tl.full([1], 1, dtype=tl.int32) ++ for _step in range(MAX_STEPS): ++ if tl.max(_active) > 0: ++ row_min = tl.min(score, axis=1) ++ row_argmin = tl.argmin(score, axis=1) ++ do_insert = row_min < topk_max ++ n_inserts = tl.sum(do_insert.to(tl.int32)) ++ if n_inserts > 0: ++ topk_argmax = tl.argmax(topk_vals, axis=1) ++ replace_mask = (k_range[None, :] == topk_argmax[:, None]) ++ insert_mask = do_insert[:, None] & replace_mask ++ topk_vals = tl.where(insert_mask, row_min[:, None], topk_vals) ++ topk_idxs = tl.where( ++ insert_mask, ++ (m_start + row_argmin.to(tl.int64))[:, None].to(tl.int32), ++ topk_idxs, ++ ) ++ topk_max = tl.max(topk_vals, axis=1) ++ used_mask = (bm_range[None, :] == row_argmin[:, None]) ++ score = tl.where(used_mask & do_insert[:, None], float('inf'), score) ++ _active = tl.where( ++ n_inserts > 0, ++ tl.full([1], 1, dtype=tl.int32), ++ tl.full([1], 0, dtype=tl.int32), ++ ) ++ ++ else: ++ for m_local in tl.range(0, M_PER_SPLIT, BM, num_stages=1): ++ m_start = m_base + m_local ++ m_offs = m_start + bm_range.to(tl.int64) ++ m_mask = m_offs < M ++ ++ cross = tl.zeros([BN, BM], dtype=tl.float32) ++ c_sq_tile = tl.zeros([BM], dtype=tl.float32) ++ for d_start in range(0, D, D_INNER): ++ d_offs = (d_start + tl.arange(0, D_INNER)).to(tl.int64) ++ d_mask = d_offs < D ++ ++ x_sub = tl.load( ++ x_ptr + pid_b * stride_x_b ++ + n_offs[:, None] * stride_x_n ++ + d_offs[None, :] * stride_x_d, ++ mask=n_mask[:, None] & d_mask[None, :], other=0.0) ++ ++ c_sub = tl.load( ++ c_ptr + pid_b * stride_c_b ++ + m_offs[:, None] * stride_c_m ++ + d_offs[None, :] * stride_c_d, ++ mask=m_mask[:, None] & d_mask[None, :], other=0.0) ++ ++ cross += tl.dot(x_sub, tl.trans(c_sub)).to(tl.float32) ++ c_f = c_sub.to(tl.float32) ++ c_sq_tile += tl.sum(c_f * c_f, axis=1) ++ ++ score = c_sq_tile[None, :] - 2.0 * cross ++ score = tl.where(m_mask[None, :], score, float('inf')) ++ ++ chunk_best = tl.min(score) ++ threshold_worst = tl.max(topk_max) ++ if chunk_best < threshold_worst: ++ _active = tl.full([1], 1, dtype=tl.int32) ++ for _step in range(MAX_STEPS): ++ if tl.max(_active) > 0: ++ row_min = tl.min(score, axis=1) ++ row_argmin = tl.argmin(score, axis=1) ++ do_insert = row_min < topk_max ++ n_inserts = tl.sum(do_insert.to(tl.int32)) ++ if n_inserts > 0: ++ topk_argmax = tl.argmax(topk_vals, axis=1) ++ replace_mask = (k_range[None, :] == topk_argmax[:, None]) ++ insert_mask = do_insert[:, None] & replace_mask ++ topk_vals = tl.where(insert_mask, row_min[:, None], topk_vals) ++ topk_idxs = tl.where( ++ insert_mask, ++ (m_start + row_argmin.to(tl.int64))[:, None].to(tl.int32), ++ topk_idxs, ++ ) ++ topk_max = tl.max(topk_vals, axis=1) ++ used_mask = (bm_range[None, :] == row_argmin[:, None]) ++ score = tl.where(used_mask & do_insert[:, None], float('inf'), score) ++ _active = tl.where( ++ n_inserts > 0, ++ tl.full([1], 1, dtype=tl.int32), ++ tl.full([1], 0, dtype=tl.int32), ++ ) ++ ++ val_sortable = _fp32_to_sortable_u32(topk_vals) ++ val_bits = val_sortable.to(tl.uint64) ++ idx_bits = topk_idxs.to(tl.uint32).to(tl.uint64) ++ packed = (val_bits << 32) | idx_bits ++ packed = tl.sort(packed, dim=1) ++ topk_score = _sortable_u32_to_fp32((packed >> 32).to(tl.uint32)) ++ topk_idx = (packed & 0xFFFFFFFF).to(tl.int32) ++ ++ k_offs = tl.arange(0, TOPK_PAD).to(tl.int64) ++ k_mask = k_offs < K ++ write_mask = n_mask[:, None] & k_mask[None, :] ++ pid_s_i64 = pid_s.to(tl.int64) ++ tl.store(partial_val_ptr + pid_b * stride_pv_b + pid_s_i64 * stride_pv_s ++ + n_offs[:, None] * stride_pv_n + k_offs[None, :] * stride_pv_k, ++ topk_score, mask=write_mask) ++ tl.store(partial_idx_ptr + pid_b * stride_pi_b + pid_s_i64 * stride_pi_s ++ + n_offs[:, None] * stride_pi_n + k_offs[None, :] * stride_pi_k, ++ topk_idx, mask=write_mask) +diff --git a/dbscanlib/_kernels/primitives/knn/triton/sortmerge.py b/dbscanlib/_kernels/primitives/knn/triton/sortmerge.py +new file mode 100644 +index 0000000..0d4a0e4 +--- /dev/null ++++ b/dbscanlib/_kernels/primitives/knn/triton/sortmerge.py +@@ -0,0 +1,181 @@ ++"""flash_knn — Triton sort-merge top-K kernel (x²-free, signed score). ++ ++The kernel computes ``argmin-K`` over the *signed* shifted distance ++ ++ s(n, m) = ||c[m]||² − 2·⟨x[n], c[m]⟩ (== ||x[n] − c[m]||² − ||x[n]||²) ++ ++The ``||x||²`` term is constant per row and so does not affect the top-K ++selection over ``m``. Dropping it eliminates the ``x_sq`` HBM tensor, its ++precompute pass, its per-tile load, one fp32 ADD per accumulator element, ++and the ``dist >= 0`` underflow clamp. The trade-off is that ``s`` is no ++longer non-negative; the packed-uint64 sort-merge top-K therefore uses ++an IEEE-sortable u32 transform (positives flip just the sign bit, negatives ++flip every bit) so ascending uint32 order matches ascending fp32 order on ++signed inputs. ++ ++Public name: ``_flash_knn_sortmerge_kernel`` — called from ++``klib/primitives/knn/triton/dispatch.py``. Stage 1 writes signed ++fp32 scores + int32 idxs to a partial buffer; Stage 2 either returns ++``partial[:, 0]`` directly (single split) or runs ``torch.topk(..., ++largest=False)`` to reduce across splits. ++ ++Same kernel covers both ``small-N M-split flash-decode`` (multiple ++splits per query block) and ``large-N single-pass`` (M_PER_SPLIT == M) ++— the host-side dispatcher just chooses the grid + M_PER_SPLIT. ++""" ++import triton ++import triton.language as tl ++ ++from dbscanlib._kernels.primitives.knn.triton._common import ( ++ _INF_PACKED, ++ _fp32_to_sortable_u32, ++ _sortable_u32_to_fp32, ++) ++ ++ ++@triton.jit ++def _flash_knn_sortmerge_kernel( ++ x_ptr, c_ptr, ++ partial_val_ptr, partial_idx_ptr, ++ stride_x_b, stride_x_n, stride_x_d, ++ stride_c_b, stride_c_m, stride_c_d, ++ stride_pv_b, stride_pv_s, stride_pv_n, stride_pv_k, ++ stride_pi_b, stride_pi_s, stride_pi_n, stride_pi_k, ++ N: tl.constexpr, M: tl.constexpr, ++ D: tl.constexpr, K: tl.constexpr, ++ M_PER_SPLIT: tl.constexpr, ++ BN: tl.constexpr, BM: tl.constexpr, ++ D_INNER: tl.constexpr, ++ TOPK_PAD: tl.constexpr, ++ NUM_STAGES_PIPE: tl.constexpr = 2, ++): ++ """Packed-uint64 sort-merge top-K with IEEE-sortable u32 score bits. ++ ++ Grid: ``(num_m_splits, ceil(N/BN), B)``. ``BM == TOPK_PAD``. ++ ++ Two compile-time paths gated on ``D_INNER >= D``: ++ * persistent x tile + pipelined M-loop (single tile covers all of D) ++ * D-split with sequential M-loop (D-chunked GEMM for wide D) ++ ++ ``NUM_STAGES_PIPE`` controls the M-loop ``tl.range`` pipelining ++ depth (host-side dispatch tunes it per shape; defaults to 2 if not ++ provided -- matches the pre-port hard-coded value). ++ """ ++ pid_s = tl.program_id(0) ++ pid_n = tl.program_id(1) ++ pid_b = tl.program_id(2) ++ pid_b = pid_b.to(tl.int64) ++ ++ n_start = pid_n * BN ++ n_offs = (n_start + tl.arange(0, BN)).to(tl.int64) ++ n_mask = n_offs < N ++ ++ topk_packed = tl.full([BN, TOPK_PAD], _INF_PACKED, dtype=tl.uint64) ++ bm_range = tl.arange(0, BM) ++ m_base = pid_s.to(tl.int64) * M_PER_SPLIT ++ ++ if D_INNER >= D: ++ d_offs = tl.arange(0, D_INNER).to(tl.int64) ++ d_mask = d_offs < D ++ x_tile = tl.load( ++ x_ptr + pid_b * stride_x_b ++ + n_offs[:, None] * stride_x_n ++ + d_offs[None, :] * stride_x_d, ++ mask=n_mask[:, None] & d_mask[None, :], other=0.0) ++ ++ for m_local in tl.range(0, M_PER_SPLIT, BM, num_stages=NUM_STAGES_PIPE): ++ m_start = m_base + m_local ++ m_offs = m_start + bm_range.to(tl.int64) ++ m_mask = m_offs < M ++ ++ c_tile = tl.load( ++ c_ptr + pid_b * stride_c_b ++ + m_offs[:, None] * stride_c_m ++ + d_offs[None, :] * stride_c_d, ++ mask=m_mask[:, None] & d_mask[None, :], other=0.0) ++ ++ c_f = c_tile.to(tl.float32) ++ c_sq_tile = tl.sum(c_f * c_f, axis=1) ++ cross = tl.dot(x_tile, tl.trans(c_tile)).to(tl.float32) ++ score = c_sq_tile[None, :] - 2.0 * cross ++ score = tl.where(m_mask[None, :], score, float('inf')) ++ ++ chunk_best = tl.min(score) ++ topk_worst_packed = tl.max(topk_packed) ++ topk_worst_val = _sortable_u32_to_fp32( ++ (topk_worst_packed >> 32).to(tl.uint32)) ++ if chunk_best < topk_worst_val: ++ score_sortable = _fp32_to_sortable_u32(score) ++ idx_vals = (m_start + bm_range.to(tl.int64)).to(tl.int32) ++ idx_bits = idx_vals.to(tl.uint32).to(tl.uint64) ++ chunk_packed = (score_sortable.to(tl.uint64) << 32) | idx_bits[None, :] ++ ++ chunk_packed = chunk_packed ^ tl.full([1], 0xFFFFFFFFFFFFFFFF, dtype=tl.uint64) ++ chunk_packed = tl.sort(chunk_packed, dim=1) ++ chunk_packed = chunk_packed ^ tl.full([1], 0xFFFFFFFFFFFFFFFF, dtype=tl.uint64) ++ ++ merged = tl.minimum(topk_packed, chunk_packed) ++ topk_packed = tl.sort(merged, dim=1) ++ ++ else: ++ for m_local in tl.range(0, M_PER_SPLIT, BM, num_stages=1): ++ m_start = m_base + m_local ++ m_offs = m_start + bm_range.to(tl.int64) ++ m_mask = m_offs < M ++ ++ cross = tl.zeros([BN, BM], dtype=tl.float32) ++ c_sq_tile = tl.zeros([BM], dtype=tl.float32) ++ for d_start in range(0, D, D_INNER): ++ d_offs = (d_start + tl.arange(0, D_INNER)).to(tl.int64) ++ d_mask = d_offs < D ++ ++ x_sub = tl.load( ++ x_ptr + pid_b * stride_x_b ++ + n_offs[:, None] * stride_x_n ++ + d_offs[None, :] * stride_x_d, ++ mask=n_mask[:, None] & d_mask[None, :], other=0.0) ++ ++ c_sub = tl.load( ++ c_ptr + pid_b * stride_c_b ++ + m_offs[:, None] * stride_c_m ++ + d_offs[None, :] * stride_c_d, ++ mask=m_mask[:, None] & d_mask[None, :], other=0.0) ++ ++ cross += tl.dot(x_sub, tl.trans(c_sub)).to(tl.float32) ++ c_f = c_sub.to(tl.float32) ++ c_sq_tile += tl.sum(c_f * c_f, axis=1) ++ ++ score = c_sq_tile[None, :] - 2.0 * cross ++ score = tl.where(m_mask[None, :], score, float('inf')) ++ ++ chunk_best = tl.min(score) ++ topk_worst_packed = tl.max(topk_packed) ++ topk_worst_val = _sortable_u32_to_fp32( ++ (topk_worst_packed >> 32).to(tl.uint32)) ++ if chunk_best < topk_worst_val: ++ score_sortable = _fp32_to_sortable_u32(score) ++ idx_vals = (m_start + bm_range.to(tl.int64)).to(tl.int32) ++ idx_bits = idx_vals.to(tl.uint32).to(tl.uint64) ++ chunk_packed = (score_sortable.to(tl.uint64) << 32) | idx_bits[None, :] ++ ++ chunk_packed = chunk_packed ^ tl.full([1], 0xFFFFFFFFFFFFFFFF, dtype=tl.uint64) ++ chunk_packed = tl.sort(chunk_packed, dim=1) ++ chunk_packed = chunk_packed ^ tl.full([1], 0xFFFFFFFFFFFFFFFF, dtype=tl.uint64) ++ ++ merged = tl.minimum(topk_packed, chunk_packed) ++ topk_packed = tl.sort(merged, dim=1) ++ ++ topk_score_sortable = (topk_packed >> 32).to(tl.uint32) ++ topk_score = _sortable_u32_to_fp32(topk_score_sortable) ++ topk_idx = (topk_packed & 0xFFFFFFFF).to(tl.int32) ++ ++ k_offs = tl.arange(0, TOPK_PAD).to(tl.int64) ++ k_mask = k_offs < K ++ write_mask = n_mask[:, None] & k_mask[None, :] ++ pid_s_i64 = pid_s.to(tl.int64) ++ tl.store(partial_val_ptr + pid_b * stride_pv_b + pid_s_i64 * stride_pv_s ++ + n_offs[:, None] * stride_pv_n + k_offs[None, :] * stride_pv_k, ++ topk_score, mask=write_mask) ++ tl.store(partial_idx_ptr + pid_b * stride_pi_b + pid_s_i64 * stride_pi_s ++ + n_offs[:, None] * stride_pi_n + k_offs[None, :] * stride_pi_k, ++ topk_idx, mask=write_mask) +diff --git a/dbscanlib/dbscan.py b/dbscanlib/dbscan.py +index ab7fa8b..53710d0 100644 +--- a/dbscanlib/dbscan.py ++++ b/dbscanlib/dbscan.py +@@ -1,70 +1,17 @@ +-"""Naive DBSCAN baseline (correct, deterministic, GPU/CPU torch). ++"""Euclidean DBSCAN -- vendored best-in-repo Triton grid path (D=2). + +-Euclidean DBSCAN. Chunked O(N^2) neighbour scan + GPU-friendly label +-propagation for the connected components of the core graph. Border rule matches +-flashlib: a non-core point joins the SMALLEST cluster label among its in-eps +-core neighbours; otherwise it is noise (-1). ++Delegates to the fused Triton grid radius-search + connected-components kernels ++vendored under ``dbscanlib._kernels``. Public contract is unchanged: ++ ++ dbscan(x, eps, min_samples) -> labels ((N,) int64; cluster id >=0, -1 noise) + """ + from __future__ import annotations + + import torch + ++from dbscanlib._kernels.primitives.dbscan.triton.dbscan import flash_dbscan + +-def dbscan(x, eps, min_samples, max_neighbors=32): +- del max_neighbors # naive path scans all neighbours; knob only for the fast kernel +- N = x.shape[0] +- dev = x.device +- eps2 = float(eps) * float(eps) +- xn = (x * x).sum(1) # (N,) +- CH = 4096 +- BIG = N # sentinel label (> any real label) +- +- def d2_chunk(lo, hi): +- xb = x[lo:hi] +- return (xb * xb).sum(1, keepdim=True) - 2.0 * (xb @ x.t()) + xn[None, :] +- +- # 1) degree (in-eps count, includes self) -> core mask +- deg = torch.zeros(N, device=dev, dtype=torch.int64) +- for lo in range(0, N, CH): +- hi = min(lo + CH, N) +- deg[lo:hi] = (d2_chunk(lo, hi) <= eps2).sum(1) +- core = deg >= int(min_samples) # (N,) bool + +- # 2) connected components of the core-core eps graph via label propagation +- labels = torch.arange(N, device=dev, dtype=torch.int64) +- labels[~core] = BIG +- for _ in range(1000): +- new = labels.clone() +- for lo in range(0, N, CH): +- hi = min(lo + CH, N) +- if not bool(core[lo:hi].any()): +- continue +- adj = (d2_chunk(lo, hi) <= eps2) & core[None, :] # (chunk, N) +- lab = torch.where(adj, labels[None, :], torch.full((1, N), BIG, device=dev, dtype=torch.int64)) +- mn = lab.min(1).values # (chunk,) +- rows = slice(lo, hi) +- upd = core[rows] & (mn < new[rows]) +- new[rows] = torch.where(upd, mn, new[rows]) +- if torch.equal(new, labels): +- break +- labels = new +- +- # 3) border assignment: non-core within eps of a core -> smallest core label +- for lo in range(0, N, CH): +- hi = min(lo + CH, N) +- nb = ~core[lo:hi] +- if not bool(nb.any()): +- continue +- adj = (d2_chunk(lo, hi) <= eps2) & core[None, :] +- lab = torch.where(adj, labels[None, :], torch.full((1, N), BIG, device=dev, dtype=torch.int64)) +- mn = lab.min(1).values +- rows = slice(lo, hi) +- take = nb & (mn < BIG) +- labels[rows] = torch.where(take, mn, labels[rows]) +- +- # 4) compact to dense [0, n_clusters); noise (BIG) -> -1 +- out = torch.full((N,), -1, device=dev, dtype=torch.int64) +- uniq = torch.unique(labels[labels < BIG]) +- for new_id, old in enumerate(uniq.tolist()): +- out[labels == old] = new_id +- return out ++def dbscan(x, eps, min_samples, max_neighbors=32): ++ labels = flash_dbscan(x, float(eps), int(min_samples), int(max_neighbors)) ++ return labels.to(torch.int64) diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/DESIGN.md b/2.0/problems/ivf_pq_gpu_kernel_optimization/DESIGN.md new file mode 100644 index 000000000..9c7523153 --- /dev/null +++ b/2.0/problems/ivf_pq_gpu_kernel_optimization/DESIGN.md @@ -0,0 +1,72 @@ +# Design notes — ivf_pq_gpu_kernel_optimization + +Operator-facing. Not copied into the agent workspace by the adapter. + +## What this task measures + +Can an agent turn a naive GPU IVF-PQ search into a fast one? The agent patches +`ivfpqlib`; the judge builds a fixed IVF-PQ index once per workload with a frozen +builder, then times the patched `ivf_pq_search` against a frozen naive baseline +over that same index on hidden ANN workloads and scores the geomean speedup, +gated on iso-result recall@k. One of the four-task flashlib kernel-optimization +family (KMeans, KNN, DBSCAN, IVF-PQ), all sharing one evaluator + one Modal GPU +harness (`flash_gpu.py`). + +This is a **genuine kernel task**: the core operation (coarse list probe + a +ragged asymmetric-distance scan of PQ codes with an on-chip top-k) is not a +single fast torch primitive. The naive baseline's per-query Python loop is the +bottleneck; beating it requires fusing the coarse probe, the per-list distance +tables, and the ragged candidate scan into GPU kernels. Contrast with pca/svd, +which were algorithmic (cov+eigh + a library eigensolver) and are not part of +this family. + +## Baseline + reference + +- Naive baseline (`ivfpqlib/`, and frozen `judge/refivfpq.py`): a pure-torch + IVF-PQ. `ivf_pq_build` is a coarse Lloyd k-means quantizer + per-subspace PQ + codebooks encoding residuals to `(M, m)` uint8 codes in CSR cell-contiguous + layout; `ivf_pq_search` probes the `nprobe` nearest lists, builds the residual + ADC lookup table per query, scores each probed candidate as the sum of `m` + table lookups, and keeps the global top-`k`. The per-query Python loop is what + makes search slow. `refivfpq.py` is the byte-identical frozen copy (verified), + so the reference agrees with the baseline to fp tolerance. +- Reference (`reference.patch`): vendors flashlib's optimized **Triton** IVF-PQ + fine scan (`primitives/ivf_pq/triton/{search,fine_scan,fine_scan_batch, + fine_scan_gemm,lut}.py`) under `ivfpqlib/_kernels/…`, rewriting search's + `ivf_pq_search`. The coarse nprobe-nearest-centroid step is done with an exact + `torch.cdist().topk()` (flashlib routes it through `flash_knn`; over a few + hundred centroids that is the same exact result and a negligible fraction of + the time, so the KNN subtree is not vendored). The CuTe DSL Hopper tier is + disabled (`is_cutedsl_available -> False`), so the portable Triton + LUT-scan / decode+GEMM kernels run — flashlib's best non-DSL path. + +## Correctness gate + +Iso-result recall@k between the agent's ids and the baseline's on the **same** +index and queries each iteration, `>= recall_threshold`. Judge-computed; +cheat-proof (you cannot fake a high recall without actually reproducing the ADC +ranking). Because the index is fixed and shared, both search implementations +target the same ground truth; the reference's LUT/GEMM kernels reproduce the +naive ADC distances to fp tolerance, so recall is near 1.0 (ties on equal ADC +sums are the only source of <1.0). `recall_threshold` is set below the +reference's measured iso-recall with margin — calibrate after the GPU trial. + +## Workloads + +Random float32 databases, `M` 50k–150k, `D` 64–128, `nlist` 256–512, `m` 8–16, +`nprobe` 8–32, `nq` ~1k–2k. The index build (frozen, not timed) and the naive +baseline search both stay within the Modal timeout while the fused kernel gives a +large speedup. Hidden workloads live in `config.yaml` (judge-only). See kmeans +DESIGN for the shared anti-hack + Modal-offload design; identical here, plus a +`setup(w, seed)` worker hook that builds the fixed index once per workload (the +frozen builder), reused across all timed iterations while queries regenerate. + +## Calibration (H100 Modal trial — done) + +`reference.patch` validated on H100: the vendored Triton search compiles, passes +the recall gate with iso-recall **1.0** on all six workloads, and runs +**513 / 907 / 1420 / 1449 / 442 / 1234x** over the naive baseline (geomean +~898x). `speedup_target` set to 800 (reference scores ~100, capped); +`recall_threshold` 0.95 (well below the measured 1.0). Re-sweep if workload +sizes change. The large factors reflect the Python-per-query-loop baseline; +fully closing the gap needs the fused ragged scan, not just torch vectorization. diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/config.yaml b/2.0/problems/ivf_pq_gpu_kernel_optimization/config.yaml new file mode 100644 index 000000000..415bfc66c --- /dev/null +++ b/2.0/problems/ivf_pq_gpu_kernel_optimization/config.yaml @@ -0,0 +1,71 @@ +tag: systems +runtime: + language: python + timeout_seconds: 10800 + environment: "GPU IVF-PQ search kernel optimization. The agent patches the ivfpqlib package; the judge offloads timing to a Modal GPU, comparing the patched ivf_pq_search against a frozen naive baseline over a fixed pre-built index on hidden ANN workloads with an iso-result recall gate." + apt_packages: [bash, ca-certificates, git, python3, python3-pip] + judge_apt_packages: [bash, ca-certificates, git, python3, python3-pip] + judge_pip_packages: [modal] + docker: + image: frontiercs/ivf-pq-gpu-kernel-optimization-agent:experimental-v0.4.0 + judge_image: frontiercs/ivf-pq-gpu-kernel-optimization-judge:experimental-v0.4.0 +environment: + cpus: 4 + memory_mb: 8192 + storage_mb: 16384 + build_timeout_seconds: 3600 +evaluation: + # --- primitive wiring (consumed by public_test's _cfg(); the graded evaluator hardcodes these) --- + primitive: ivfpq + pkg: ivfpqlib + ref_module: refivfpq + gpu: "H100" + cuda_image: "nvidia/cuda:12.4.1-devel-ubuntu22.04" + pip: ["torch", "numpy", "h5py"] # match the original (validated) image; torch brings its matched triton (a standalone triton pin crashes some flashlib kernels). h5py loads SIFT/GIST HDF5 + app_name: "ivfpq-kernel-opt-eval" + dataset_volume: "flashlib-ann-datasets" # Modal Volume holding sift.hdf5 / gist.hdf5 (ann-benchmarks); mounted read-only at /data + modal_timeout_seconds: 3600 # the naive per-query-loop baseline is seconds/call and real GIST (D=960) is heavier; give the judge margin + warmup_iters: 2 + timed_iters: 7 + recall_threshold: 0.99 # agent ids must match the FULL-nprobe baseline's top-k (iso-result recall@k) every iteration. Tightened 0.95->0.99 to kill nprobe-truncation: at nlist=8192 halving nprobe drops iso-recall to 0.85-0.92 and a 3/8 cut to ~0.9, all << 0.99, while an honest full-nprobe exact-ADC search stays ~0.9999 (bf16 tie margin). So a solution cannot silently probe fewer lists. + speedup_target: 2551.0 # score cap = 2x the gemm-pinned reference geomean on the SIFT nlist=8192 dispatch-capped workloads (1276x on H100: 1000/1663/1012/1750/786/1863x, iso-recall 1.0 all 6, gate 0.99). Reference = fast WGMMA decode+GEMM; beats a scalar-decode kernel ~1.1-1.9x at M16-32, so 100 needs a genuinely better tensor-core kernel. NOTE: the huge factor is the naive Python-loop BASELINE -> a near-reference kernel still scores ~91 (log compression); the planned torch-vectorized baseline will spread the top AND remove the dispatch wall (letting M go fatter). bf16-locked query+index. + base_seed: 20260702 + agent_workload_count: 6 # agent's submission feedback covers ALL graded workloads (no blind pass-3/fail-6 cliff); params are agent-visible anyway + expose_per_workload_metrics: true # show per-workload pass/fail so the agent can fix the specific workload it regresses on + # REAL-DATA WORKLOADS: the index (the "database") is built ONCE per workload by + # the frozen builder from the full real base (SIFT1M 128d / GIST1M 960d), reused + # for every timed iteration (each call gets a fresh clone -- see flash_gpu + # _fresh_index); only ivf_pq_search is timed. Each iteration draws a fresh + # Q-subset of the held-out real query set (anti-caching). Real datasets have + # clustered, imbalanced inverted lists (unlike synthetic gaussian) -> the search + # kernel must handle long-tail lists. + # + # nprobe is the main recall/speed knob and is meant to MATTER: scanning fewer than + # the workload's nprobe lists drops iso-recall below the gate. + # + # SIFT-only, nlist=8192, LARGE Q. The fat GEMM M that makes the reference's WGMMA + # decode+GEMM strong comes from a big query batch (M = Q*nprobe/nlist), NOT from a + # low nlist -- so nlist stays HIGH (8192, ~122 rows/list) and nprobe genuinely + # matters: halving nprobe drops iso-recall to 0.85-0.92 (measured), so a + # nprobe-truncating solution fails the 0.99 gate. (A low-nlist "fat M" regime made + # M fat but also made high nprobe redundant -> truncation was nearly free; and GIST + # with only 1k held-out queries cannot reach a fat M at nlist=8192, so it is dropped.) + # The reference is pinned to the gemm variant (see reference.patch); at M>=16 it beats + # a query-parallel scalar-decode kernel ~1.5-2.4x, so beating IT 2x needs a genuinely + # good tensor-core kernel. + # Q*nprobe (the naive Python-loop baseline's per-call dispatch count) is capped at + # ~262k so the judge stays practical -- Q=8192 pushed it to ~0.5-1M and the eval + # backed up (submissions queued unevaluated). M = Q*nprobe/nlist is kept at 16-32 + # (fat enough that the pinned gemm reference beats a scalar kernel ~1.1-1.9x). The + # deferred torch-vectorized baseline will remove the Python-loop dispatch wall and + # allow a fatter M (and fix the score log-compression) in the next round. + workloads: + - {id: sift_q4k_np32, dataset: sift, nlist: 8192, m: 16, nprobe: 32, Q: 4096, k: 100} # M16 disp131k + - {id: sift_q4k_np64, dataset: sift, nlist: 8192, m: 16, nprobe: 64, Q: 4096, k: 100} # M32 disp262k + - {id: sift_q2k_np64, dataset: sift, nlist: 8192, m: 16, nprobe: 64, Q: 2048, k: 100} # M16 disp131k + - {id: sift_q2k_np128, dataset: sift, nlist: 8192, m: 16, nprobe: 128, Q: 2048, k: 100} # M32 disp262k + - {id: sift_q1k_np128, dataset: sift, nlist: 8192, m: 16, nprobe: 128, Q: 1024, k: 100} # M16 disp131k + - {id: sift_q1k_np256, dataset: sift, nlist: 8192, m: 16, nprobe: 256, Q: 1024, k: 100} # M32 disp262k +submission: + kind: file + path: /app/solution.patch diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/docker/README.md b/2.0/problems/ivf_pq_gpu_kernel_optimization/docker/README.md new file mode 100644 index 000000000..2a6a0978d --- /dev/null +++ b/2.0/problems/ivf_pq_gpu_kernel_optimization/docker/README.md @@ -0,0 +1,56 @@ +# Experimental IVF-PQ Kernel-Optimization Images + +Two **light** images (ubuntu + `modal` + git). torch, triton, and the vendored +flashlib kernels all run on the **Modal GPU image** defined in `flash_gpu.py`; +the containers themselves are CPU-only and offload GPU work to Modal. + +```bash +bash 2.0/problems/ivf_pq_gpu_kernel_optimization/docker/build_images.sh +``` + +Defaults: + +```text +AGENT_TAG=frontiercs/ivf-pq-gpu-kernel-optimization-agent:experimental-v0.1.0 +JUDGE_TAG=frontiercs/ivf-pq-gpu-kernel-optimization-judge:experimental-v0.1.0 +``` + +Agent image: + +```text +/app/ivfpqlib # clean, git-tracked package (the agent edits this) +/opt/flash_gpu.py # shared Modal GPU harness (public test uses it) +/opt/ivfpq_ref/refivfpq.py # frozen naive baseline (public-test speed denominator) +``` + +Judge image: + +```text +/opt/ivfpqlib-clean/ivfpqlib # pristine tree; the patch is applied to a copy +/opt/ivfpq_ref/refivfpq.py # frozen naive baseline (speed denominator + quality oracle) +/opt/flash_gpu.py # shared Modal GPU harness (the evaluator uses it) +``` + +## Runtime requirements + +Both the judge and the agent public test **offload timing to a Modal GPU** and +therefore need Modal credentials in the environment: + +```text +MODAL_TOKEN_ID / MODAL_TOKEN_SECRET +``` + +`flash_gpu.py` builds an ephemeral Modal app on a GPU (`evaluation.gpu`, default +`H100`), ships the frozen baseline + the patched package as data, times both on +fresh per-iteration data, verifies quality each iteration, and returns the +speedups. No persistent deployment is used, so a fresh GPU container is spun up +per evaluation. Without Modal credentials the evaluator returns a patch-policy +smoke pass (which repo CI exercises). + +## Smoke test + +```bash +bash 2.0/problems/ivf_pq_gpu_kernel_optimization/docker/smoke_images.sh +``` + +Import-only (modal + flash_gpu + the baked packages); does not touch a GPU or Modal. diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/docker/agent/Dockerfile b/2.0/problems/ivf_pq_gpu_kernel_optimization/docker/agent/Dockerfile new file mode 100644 index 000000000..d722a90eb --- /dev/null +++ b/2.0/problems/ivf_pq_gpu_kernel_optimization/docker/agent/Dockerfile @@ -0,0 +1,34 @@ +# Agent image for ivf_pq_gpu_kernel_optimization. +# +# Light image (no torch/triton): the agent edits /app/ivfpqlib and runs the +# public test, which offloads timing to a Modal GPU (torch/triton live on the +# Modal image defined in flash_gpu.py). Needs MODAL_TOKEN_ID / MODAL_TOKEN_SECRET +# in the environment to run the GPU public test. +# +# Build context = the task directory: +# docker build -f docker/agent/Dockerfile -t . +FROM ubuntu:24.04 + +ARG DEBIAN_FRONTEND=noninteractive +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + bash ca-certificates git python3 python3-pip ripgrep && \ + rm -rf /var/lib/apt/lists/* + +RUN pip3 install --break-system-packages --no-cache-dir modal + +# The package the agent edits (git-tracked so make_submission.sh can diff it). +WORKDIR /app +COPY ivfpqlib /app/ivfpqlib +RUN printf '%s\n' '__pycache__/' '*.pyc' > /app/.gitignore && \ + git -C /app init -q && \ + git -C /app config user.email task@frontier-cs && \ + git -C /app config user.name frontier-cs && \ + git -C /app add -A -- ivfpqlib .gitignore && \ + git -C /app -c commit.gpgsign=false commit -qm "pristine ivfpqlib" + +# Shared Modal GPU harness + the frozen naive baseline (the public-test speed +# denominator). The baseline is exactly the code the agent starts from, so it is +# not secret. +COPY flash_gpu.py /opt/flash_gpu.py +COPY judge/refivfpq.py /opt/ivfpq_ref/refivfpq.py diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/docker/build_images.sh b/2.0/problems/ivf_pq_gpu_kernel_optimization/docker/build_images.sh new file mode 100644 index 000000000..4d19d5e25 --- /dev/null +++ b/2.0/problems/ivf_pq_gpu_kernel_optimization/docker/build_images.sh @@ -0,0 +1,17 @@ +#!/usr/bin/env bash +# Build the experimental agent + judge images for ivf_pq_gpu_kernel_optimization. +set -euo pipefail +HERE=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd) # the problem directory (build context) + +AGENT_TAG=${AGENT_TAG:-frontiercs/ivf-pq-gpu-kernel-optimization-agent:experimental-v0.4.0} +JUDGE_TAG=${JUDGE_TAG:-frontiercs/ivf-pq-gpu-kernel-optimization-judge:experimental-v0.4.0} + +echo "Building agent image: $AGENT_TAG" +docker build -f "$HERE/docker/agent/Dockerfile" -t "$AGENT_TAG" "$HERE" + +echo "Building judge image: $JUDGE_TAG" +docker build -f "$HERE/docker/judge/Dockerfile" -t "$JUDGE_TAG" "$HERE" + +echo "Built:" +echo " $AGENT_TAG" +echo " $JUDGE_TAG" diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/docker/judge/Dockerfile b/2.0/problems/ivf_pq_gpu_kernel_optimization/docker/judge/Dockerfile new file mode 100644 index 000000000..62389a617 --- /dev/null +++ b/2.0/problems/ivf_pq_gpu_kernel_optimization/docker/judge/Dockerfile @@ -0,0 +1,26 @@ +# Judge image for ivf_pq_gpu_kernel_optimization. +# +# Light image (no torch/triton): the judge validates + applies the patch and +# offloads timing to a Modal GPU (torch/triton live on the Modal image defined +# in flash_gpu.py). Needs MODAL_TOKEN_ID / MODAL_TOKEN_SECRET in the environment. +# The Frontier-CS 2.0 adapter builds the final judge image ON TOP of this one. +# +# Build context = the task directory: +# docker build -f docker/judge/Dockerfile -t . +FROM ubuntu:24.04 + +ARG DEBIAN_FRONTEND=noninteractive +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + bash ca-certificates git python3 python3-pip && \ + rm -rf /var/lib/apt/lists/* + +RUN pip3 install --break-system-packages --no-cache-dir modal + +# Pristine tree the patch is applied to, the frozen naive baseline (speed +# denominator + quality oracle), and the shared Modal GPU harness. +COPY ivfpqlib /opt/ivfpqlib-clean/ivfpqlib +COPY judge/refivfpq.py /opt/ivfpq_ref/refivfpq.py +COPY flash_gpu.py /opt/flash_gpu.py + +WORKDIR /judge diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/docker/smoke_images.sh b/2.0/problems/ivf_pq_gpu_kernel_optimization/docker/smoke_images.sh new file mode 100644 index 000000000..d40fe4cfc --- /dev/null +++ b/2.0/problems/ivf_pq_gpu_kernel_optimization/docker/smoke_images.sh @@ -0,0 +1,22 @@ +#!/usr/bin/env bash +# Smoke test for the built (intentionally torch-less) images: confirm the host-side +# harness imports (modal + flash_gpu) and that the frozen baseline + package parse. +# torch/triton live on the Modal image, so the packages are py_compiled here (a +# torch-less parse check), not imported. No GPU / Modal tokens required. +set -euo pipefail + +AGENT_TAG=${AGENT_TAG:-frontiercs/ivf-pq-gpu-kernel-optimization-agent:experimental-v0.1.0} +JUDGE_TAG=${JUDGE_TAG:-frontiercs/ivf-pq-gpu-kernel-optimization-judge:experimental-v0.1.0} + +SMOKE='import sys, modal, py_compile, pathlib +sys.path.insert(0, "/opt"); import flash_gpu +for root in sys.argv[1:]: + for p in pathlib.Path(root).rglob("*.py"): + py_compile.compile(str(p), doraise=True) +print("ok: modal + flash_gpu import; ivfpq baseline + ivfpqlib parse")' + +echo "== agent image ==" +docker run --rm -w /app "$AGENT_TAG" python3 -c "$SMOKE" /opt/ivfpq_ref /app/ivfpqlib + +echo "== judge image ==" +docker run --rm "$JUDGE_TAG" python3 -c "$SMOKE" /opt/ivfpq_ref /opt/ivfpqlib-clean/ivfpqlib diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/evaluate.sh b/2.0/problems/ivf_pq_gpu_kernel_optimization/evaluate.sh new file mode 100644 index 000000000..9530e17b0 --- /dev/null +++ b/2.0/problems/ivf_pq_gpu_kernel_optimization/evaluate.sh @@ -0,0 +1,12 @@ +#!/usr/bin/env bash +# Local CLI evaluation for ivf_pq_gpu_kernel_optimization. +# +# A full run needs a GPU plus the baked judge sources at /opt (the pristine +# /opt/ivfpqlib-clean tree and the frozen /opt/ivfpq_ref baseline; see the +# judge image). Without them the evaluator still validates the patch policy and +# returns a smoke score, which is what repository CI exercises (the reference +# patch passes the policy). Pass a patch path to score it directly. +set -euo pipefail +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +SOLUTION="${1:-$SCRIPT_DIR/reference.patch}" +exec python3 "$SCRIPT_DIR/evaluator.py" "$SOLUTION" diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/evaluator.py b/2.0/problems/ivf_pq_gpu_kernel_optimization/evaluator.py new file mode 100644 index 000000000..3dc593e8c --- /dev/null +++ b/2.0/problems/ivf_pq_gpu_kernel_optimization/evaluator.py @@ -0,0 +1,380 @@ +"""Generic evaluator for the flashlib GPU kernel-optimization task family. + +Identical across all four tasks (kmeans / knn / pca / truncated_svd); every +primitive-specific value comes from the ``evaluation`` block of the task's +config (delivered to the judge as ``/judge/task_config.json``), which is not +present in the agent workspace. + +Flow: statically validate the patch (only ``/**`` may change; no external +optimized libs / env / process / network access) -> apply it to the pristine +package baked at ``clean_source`` -> hand the frozen baseline + patched package +sources to a Modal GPU worker (``flash_gpu.run_remote``) that times both on +fresh per-iteration data and verifies clustering/retrieval/subspace quality on +every iteration -> gate on the worker's per-workload verdict -> score by the +geometric-mean speedup. + +Without Modal credentials or the baked judge sources (e.g. repo CI on a box with +no GPU/Modal), the evaluator still validates the patch policy and returns a +smoke pass, so ``python3 evaluator.py reference.patch`` works anywhere. +""" +from __future__ import annotations + +import fnmatch +import hashlib +import json +import math +import os +import re +import shutil +import subprocess +import sys +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +TASK_CONFIG_PATH = Path("/judge/task_config.json") + + +def _load_task_config() -> dict[str, Any]: + try: + payload = json.loads(TASK_CONFIG_PATH.read_text(encoding="utf-8")) + except Exception: + return {} + return payload if isinstance(payload, dict) else {} + + +TASK_CONFIG = _load_task_config() +EVAL = TASK_CONFIG.get("evaluation", {}) if isinstance(TASK_CONFIG.get("evaluation"), dict) else {} + + +def _get(name: str, default): + return EVAL.get(name, default) + + +def _get_int(name: str, default: int) -> int: + try: + return int(EVAL.get(name, default)) + except Exception: + return default + + +def _get_float(name: str, default: float) -> float: + try: + return float(EVAL.get(name, default)) + except Exception: + return default + + +def _get_bool(name: str, default: bool) -> bool: + raw = EVAL.get(name, default) + if isinstance(raw, bool): + return raw + if isinstance(raw, str): + return raw.strip().lower() in {"1", "true", "yes", "on"} + return bool(raw) + + +# --- Per-task identity: the ONLY lines that differ across the four tasks. --- +# Kept as constants (not config) so the patch policy is enforceable locally / in +# CI without /judge/task_config.json. Everything else is config-driven and only +# needed on the GPU path. +PRIMITIVE = "ivfpq" +PKG = "ivfpqlib" +REF_MODULE = "refivfpq" + +MAX_PATCH_BYTES = _get_int("max_patch_bytes", 2_000_000) +MAX_CHANGED_FILES = _get_int("max_changed_files", 80) +CLEAN_SOURCE = Path(f"/opt/{PKG}-clean") +BASELINE_SOURCE = Path(f"/opt/{PRIMITIVE}_ref") +SPEEDUP_TARGET = _get_float("speedup_target", 8.0) + +ALLOWED_PATTERNS = (f"{PKG}/**",) +DENIED_PATTERNS = ( + "evaluator.py", + "flash_gpu.py", + "reference.patch", + f"{REF_MODULE}.py", + f"**/{REF_MODULE}.py", + "**/conftest.py", + "**/test_*.py", + "pyproject.toml", + "setup.py", + "setup.cfg", +) +# Defense in depth (the load-bearing anti-tamper defense is structural, inside +# the GPU worker: it captures its perf_counter / synchronize references before +# importing the submission, regenerates data every iteration, and re-verifies +# quality each iteration). These patterns are matched surgically so that +# legitimate vendored/optimized kernel source is not rejected: +# * external optimized libraries, matched as IMPORT statements (not bare words, +# so a comment mentioning a name is fine); none are installed on the GPU +# image anyway; +# * process / network / measurement-tamper patterns that never appear in a +# real GPU kernel. +FORBIDDEN_IMPORT_RE = re.compile( + r"(?:^|\n)\s*(?:import|from)\s+(flashlib|cuml|cudf|cupy|faiss|sklearn|scikit|cutlass|cuvs)\b" +) +FORBIDDEN_PATTERN_RE = re.compile( + r"\bsubprocess\b|\bsocket\b|\burllib\b|\brequests\b|os\.system|" + r"torch\.cuda\.synchronize\s*=|time\.perf_counter\s*=|setattr\(\s*torch\.cuda" +) + + +@dataclass(frozen=True) +class PatchFile: + old_path: str + new_path: str + added_lines: tuple[str, ...] + + @property + def path(self) -> str: + return self.new_path if self.new_path != "/dev/null" else self.old_path + + +def _match(path: str, patterns: tuple[str, ...]) -> bool: + return any(fnmatch.fnmatch(path, p) for p in patterns) + + +def _invalid(message: str, metrics: dict[str, Any] | None = None): + payload = metrics or {} + payload.setdefault("valid_patch", 0) + return 0.0, 0.0, message, payload + + +def _parse_patch(text: str) -> list[PatchFile]: + files: list[PatchFile] = [] + old = new = "" + added: list[str] = [] + in_file = False + for line in text.splitlines(): + if line.startswith("diff --git "): + if in_file: + files.append(PatchFile(old, new, tuple(added))) + in_file, old, new, added = True, "", "", [] + continue + if not in_file: + continue + if line.startswith("--- "): + old = line[4:].strip() + old = old[2:] if old.startswith("a/") else old + elif line.startswith("+++ "): + new = line[4:].strip() + new = new[2:] if new.startswith("b/") else new + elif line.startswith("+") and not line.startswith("+++ "): + added.append(line[1:]) + if in_file: + files.append(PatchFile(old, new, tuple(added))) + return files + + +def _validate_path(path: str) -> tuple[bool, str]: + if not path or path == "/dev/null": + return True, "" + if path.startswith("/") or ".." in Path(path).parts: + return False, f"unsafe patch path: {path}" + if _match(path, DENIED_PATTERNS): + return False, f"changed file is outside task boundary: {path}" + if not _match(path, ALLOWED_PATTERNS): + return False, f"changed file is not allowlisted (only {PKG}/** may change): {path}" + if not path.endswith(".py"): + return False, f"only Python files may change: {path}" + return True, "" + + +def validate_patch(patch_path: Path) -> tuple[bool, str, dict[str, Any]]: + if not patch_path.exists(): + return False, "solution patch does not exist", {} + size = patch_path.stat().st_size + if size > MAX_PATCH_BYTES: + return False, f"patch is too large ({size} bytes > {MAX_PATCH_BYTES})", {} + text = patch_path.read_text(encoding="utf-8", errors="replace") + files = _parse_patch(text) + metrics: dict[str, Any] = { + "patch_bytes": size, + "patch_sha256": hashlib.sha256(text.encode("utf-8", "replace")).hexdigest(), + "changed_files": len(files), + } + if len(files) > MAX_CHANGED_FILES: + return False, f"too many changed files ({len(files)} > {MAX_CHANGED_FILES})", metrics + for pf in files: + path = pf.path + if pf.new_path == "/dev/null": + return False, f"deleting files is outside task boundary: {pf.old_path}", metrics + if pf.old_path != "/dev/null" and pf.old_path != pf.new_path: + ok, err = _validate_path(pf.old_path) + if not ok: + return False, f"rename source is outside task boundary: {err}", metrics + ok, err = _validate_path(path) + if not ok: + return False, err, metrics + added = "\n".join(pf.added_lines) + m = FORBIDDEN_IMPORT_RE.search(added) + if m: + return False, f"{path}: forbidden import of an external optimized library ({m.group(1)})", metrics + m = FORBIDDEN_PATTERN_RE.search(added) + if m: + return False, f"{path}: forbidden pattern in added code ({m.group(0).strip()[:40]})", metrics + metrics["valid_patch"] = 1 + return True, "patch accepted by static policy", metrics + + +def sanitize(text: str) -> str: + text = re.sub(r"/tmp/[A-Za-z0-9_./-]+", "", text or "") + text = re.sub(r"/opt/[A-Za-z0-9_./-]+", "", text) + text = re.sub(r"\bN=\d+|\bseed[=:]?\s*\d+", "", text, flags=re.IGNORECASE) + return text[-600:] + + +def geometric_mean(values: list[float]) -> float: + if not values: + return 0.0 + return math.exp(sum(math.log(max(v, 1e-9)) for v in values) / len(values)) + + +def score_from_speedup(gm: float) -> float: + if gm <= 0: + return 0.0 + raw = 100.0 * math.log(gm) / math.log(max(SPEEDUP_TARGET, 1.0000001)) + return max(0.0, min(100.0, raw)) + + +def is_final_role() -> bool: + return os.environ.get("FRONTIER_SUBMISSION_ROLE", "agent") == "final" + + +def _workloads(final_role: bool) -> list[dict[str, Any]]: + workloads = list(_get("workloads", []) or []) + if not final_role: + n = _get_int("agent_workload_count", 3) + workloads = workloads[:n] + base = _get_int("base_seed", 20260701) + for i, w in enumerate(workloads): + w.setdefault("seed", base + 1000 * (i + 1)) + return workloads + + +def _build_cfg() -> dict[str, Any]: + return { + "primitive": PRIMITIVE, + "pkg": PKG, + "ref_module": REF_MODULE, + "gpu": str(_get("gpu", "H100")), + "cuda_image": str(_get("cuda_image", "nvidia/cuda:12.4.1-devel-ubuntu22.04")), + "pip": list(_get("pip", ["torch==2.5.1", "triton==3.1.0", "numpy"])), + "app_name": str(_get("app_name", "flash-kernel-eval")), + "dataset_volume": _get("dataset_volume", None), # Modal Volume with real SIFT/GIST datasets, mounted at /data + "modal_timeout_seconds": _get_int("modal_timeout_seconds", 1800), + "warmup": _get_int("warmup_iters", 3), + "iters": _get_int("timed_iters", 7), + "inertia_tolerance": _get_float("inertia_tolerance", 0.02), + "recall_threshold": _get_float("recall_threshold", 0.99), + "captured_tolerance": _get_float("captured_tolerance", 0.02), + "ortho_tolerance": _get_float("ortho_tolerance", 0.02), + "ari_threshold": _get_float("ari_threshold", 0.99), + } + + +def _read_dir(root: Path, rel_to: Path) -> dict[str, str]: + out: dict[str, str] = {} + for p in root.rglob("*.py"): + out[str(p.relative_to(rel_to))] = p.read_text(encoding="utf-8", errors="replace") + return out + + +def full_evaluation(patch_path: Path, metrics: dict[str, Any]): + final_role = is_final_role() + metrics["submission_role"] = "final" if final_role else "agent" + workloads = _workloads(final_role) + + # Import the Modal harness baked into the judge image; missing harness or + # missing credentials/sources -> policy smoke pass. + sys.path.insert(0, "/opt") + try: + import flash_gpu # type: ignore + except Exception: + metrics["full_benchmark"] = 0 + return 1.0, 1.0, "patch policy smoke passed; GPU harness not available in this environment", metrics + if not flash_gpu.modal_available() or not CLEAN_SOURCE.exists() or not BASELINE_SOURCE.exists(): + metrics["full_benchmark"] = 0 + return 1.0, 1.0, "patch policy smoke passed; Modal/judge sources not configured in this environment", metrics + + with tempfile.TemporaryDirectory(prefix="flash_kernel_opt_") as tmp: + tmp_root = Path(tmp) + patched = tmp_root / "patched" + shutil.copytree(CLEAN_SOURCE, patched) + env = {"PATH": os.environ.get("PATH", "/usr/bin:/bin"), "HOME": str(tmp_root), + "LC_ALL": "C", "LANG": "C"} + if int(metrics.get("changed_files", 0)) > 0: + try: + subprocess.run(["git", "apply", "--check", str(patch_path)], cwd=patched, env=env, + check=True, capture_output=True, text=True, timeout=60) + subprocess.run(["git", "apply", str(patch_path)], cwd=patched, env=env, + check=True, capture_output=True, text=True, timeout=60) + except subprocess.CalledProcessError as exc: + metrics["stderr_tail"] = sanitize(exc.stderr or "") + return _invalid("patch does not apply to the clean package tree", metrics) + metrics["applied_patch"] = 1 + else: + metrics["used_empty_patch"] = 1 + + payload = { + "baseline_files": _read_dir(BASELINE_SOURCE, BASELINE_SOURCE), + "patched_files": _read_dir(patched, patched), + "workloads": workloads, + "cfg": _build_cfg(), + } + try: + result = flash_gpu.run_remote(payload) + except Exception as exc: + metrics["error_detail"] = sanitize(str(exc)) + return _invalid("GPU evaluation failed", metrics) + + if not result.get("ok"): + metrics["worker_error"] = sanitize(str(result.get("error", "unknown"))) + return _invalid("GPU benchmark worker failed", metrics) + + speedups: list[float] = [] + per_workload: dict[str, Any] = {} + for row in result.get("rows", []): + if not row.get("ok"): + metrics["failed_workload"] = {"reason": row.get("reason", "error")} + return _invalid("submission failed the quality gate or crashed on a hidden workload", metrics) + speedups.append(max(float(row["speedup"]), 0.01)) + per_workload[row["id"]] = {"speedup": row["speedup"]} + if not speedups: + return _invalid("no workloads were evaluated", metrics) + + gm = geometric_mean(speedups) + bounded = score_from_speedup(gm) + metrics.update({"full_benchmark": 1, "workload_count": len(speedups), "geomean_speedup": gm}) + if _get_bool("expose_per_workload_metrics", False): + metrics["per_workload"] = per_workload + return bounded, bounded, f"{PRIMITIVE} geomean speedup {gm:.3f}x over the naive baseline", metrics + + +def evaluate(solution_path: str) -> tuple[float, float, str, dict[str, Any]]: + patch_path = Path(solution_path) + ok, message, metrics = validate_patch(patch_path) + if not ok: + return _invalid(message, metrics) + try: + return full_evaluation(patch_path, metrics) + except Exception as exc: # noqa: BLE001 + metrics["error_type"] = type(exc).__name__ + metrics["error_detail"] = sanitize(str(exc)) + return _invalid("evaluation failed", metrics) + + +def main(argv: list[str]) -> int: + if len(argv) != 2: + print("Usage: evaluator.py SOLUTION_PATCH", file=sys.stderr) + return 2 + score, unbounded, message, metrics = evaluate(argv[1]) + print(json.dumps({"score": score, "score_unbounded": unbounded, + "message": message, "metrics": metrics}, indent=2, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv)) diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/flash_gpu.py b/2.0/problems/ivf_pq_gpu_kernel_optimization/flash_gpu.py new file mode 100644 index 000000000..ca61d51e5 --- /dev/null +++ b/2.0/problems/ivf_pq_gpu_kernel_optimization/flash_gpu.py @@ -0,0 +1,416 @@ +"""Modal GPU evaluation harness for the flashlib kernel-optimization task family. + +Shared by the judge (evaluator.py) and the agent public test. Baked into both +the agent and judge images at ``/opt/flash_gpu.py``. It runs a frozen naive +baseline and the (patched or agent-edited) package on a Modal GPU and returns +per-workload timings + quality verdicts. + +Anti-reward-hack properties (all enforced inside the remote worker, which is the +only place the untrusted code runs): + +* **Timing primitives are captured before any agent code is imported** — the + worker binds ``time.perf_counter`` and ``torch.cuda.synchronize`` to locals up + front, so a submission that monkey-patches ``torch.cuda.synchronize`` (or any + torch attribute) cannot affect measurement. +* **Fresh data every timed iteration** — each measured iteration generates a new + input from a fresh seed, so a submission cannot memoize on a repeated input. +* **Quality is verified on every timed iteration** — the judge recomputes the + quality metric from the *returned tensors* on that iteration's data and gates + it, so returning a stale cached result for a different input is caught. +* **The judge computes all metrics** — no agent-provided number is trusted. +* **Fresh container per submission** — an ephemeral Modal app is used, so no + global state survives across evaluations. + +The remote worker is a single self-contained function serialized to Modal +(``serialized=True``), so the remote does not import this module and there is no +module-mounting to get wrong. The GPU image ships torch + triton (+ any extra +pip packages a task needs for its reference build). +""" +from __future__ import annotations + +import os + + +# --------------------------------------------------------------------------- # +# Remote worker (runs on the Modal GPU). Fully self-contained: every helper is +# nested so cloudpickle ships the whole thing. Takes and returns plain data. +# --------------------------------------------------------------------------- # +def _gpu_worker(payload: dict) -> dict: + import importlib + import os as _os + import sys as _sys + import tempfile + import time + import traceback + + import torch + + # Capture timing + sync references BEFORE importing any submission code, so a + # monkey-patch of torch.cuda.synchronize cannot affect measurement. + _perf = time.perf_counter + _sync = torch.cuda.synchronize + if not torch.cuda.is_available(): + return {"ok": False, "error": "cuda_unavailable"} + dev = "cuda" + + cfg = payload["cfg"] + prim = cfg["primitive"] + warmup = int(cfg.get("warmup", 3)) + iters = int(cfg.get("iters", 7)) + + def _materialize(files: dict, tag: str) -> str: + root = tempfile.mkdtemp(prefix=f"flash_{tag}_") + for rel, content in files.items(): + path = _os.path.join(root, rel) + parent = _os.path.dirname(path) + if parent: + _os.makedirs(parent, exist_ok=True) + with open(path, "w", encoding="utf-8") as fh: + fh.write(content) + _sys.path.insert(0, root) + return root + + _materialize(payload["baseline_files"], "baseline") + _materialize(payload["patched_files"], "patched") + ref = importlib.import_module(cfg["ref_module"]) + pkg = importlib.import_module(cfg["pkg"]) + + # ---- per-primitive: fixed context, data gen, call, quality metric ----- # + _ann_cache = {} + def _load_ann(name): + # Load a real ANN dataset (ann-benchmarks HDF5) from the mounted Modal + # Volume at /data. Cached across workloads in one worker call. + if name not in _ann_cache: + import h5py + import numpy as _np + with h5py.File(f"/data/{name}.hdf5", "r") as f: + base = torch.from_numpy(_np.asarray(f["train"], dtype=_np.float32)).to(dev) + queries = torch.from_numpy(_np.asarray(f["test"], dtype=_np.float32)).to(dev) + _ann_cache[name] = (base.contiguous(), queries.contiguous()) + return _ann_cache[name] + + def setup(w, seed): + # Built ONCE per workload (not per timed iter). IVF-PQ builds the fixed + # index (the "database") here; queries are generated fresh per iteration. + if prim == "ivfpq": + ds = w.get("dataset") + if ds: + # Real dataset: build the frozen index from the full real base; + # queries are drawn per-iteration from the held-out query set. + base, queries_all = _load_ann(ds) + # Train the coarse quantizer on enough points for a large nlist + # (>= ~40 per centroid); the base is 1M rows so this is a subsample. + nlist = int(w["nlist"]) + tsz = min(base.shape[0], max(200000, nlist * 40)) + index = ref.ivf_pq_build(base, nlist, m=int(w["m"]), + nprobe=int(w["nprobe"]), seed=int(w.get("seed", 0)), + train_size=tsz, pq_train_size=min(base.shape[0], 100000)) + return {"index": index, "D": int(base.shape[1]), "queries_all": queries_all} + g = torch.Generator(device=dev).manual_seed(int(seed)) + X = torch.randn(int(w["M"]), int(w["D"]), generator=g, device=dev, dtype=torch.float32) + index = ref.ivf_pq_build(X, int(w["nlist"]), m=int(w["m"]), + nprobe=int(w["nprobe"]), seed=int(seed)) + return {"index": index, "D": int(w["D"])} + return {} + + def _fresh_index(src): + # Hand every timed call a FRESH index object with freshly cloned tensors. + # The index content is identical (it is still the one frozen index built in + # setup), but the object identity and every data_ptr change per iteration, + # so a solution cannot attach cached state to it (`setattr(index, ...)`) or + # key a cache on its tensor addresses and thereby amortise index-derived + # precomputation (e.g. a CSR->padded layout conversion) out of the timed + # region and into the free warmup calls. `src` is the pristine index kept + # in ctx and never handed to the submission, so nothing leaks between iters. + new = object.__new__(type(src)) + # PRECISION LOCK (bf16): round the float index tensors (coarse centroids + + # PQ codebooks) to bf16 precision so the coarse distance matmul and the ADC + # decode carry only bf16 precision for everyone. A solution then cannot win + # by computing them in a lower dtype than the baseline -- bf16 tensor cores + # are the floor, fp32/TF32 buy nothing. Codes are uint8 already; the query + # is bf16-rounded in gen(). Values stay fp32-typed so the search contract is + # unchanged. + _bf16 = {"centroids", "pq_codebooks"} + for k, v in vars(src).items(): + if torch.is_tensor(v): + v = v.clone() + if k in _bf16 and v.is_floating_point(): + v = v.to(torch.bfloat16).to(v.dtype) + new.__dict__[k] = v + return new + + def gen(w, seed, ctx): + g = torch.Generator(device=dev).manual_seed(int(seed)) + if prim == "ivfpq": + data = {"index": _fresh_index(ctx["index"])} + if w.get("dataset"): + # Fresh subset of the real query set each iteration (anti-caching); + # Q must be <= number of held-out queries. + qa = ctx["queries_all"] + perm = torch.randperm(qa.shape[0], generator=g, device=dev)[: int(w["Q"])] + # bf16-round the query too (see _fresh_index): with the index tensors + # already at bf16 precision the whole search carries only bf16. + q = qa.index_select(0, perm) + data["queries"] = q.to(torch.bfloat16).to(q.dtype).contiguous() + return data + q = torch.randn(int(w["Q"]), int(ctx["D"]), generator=g, device=dev, dtype=torch.float32) + data["queries"] = q.to(torch.bfloat16).to(torch.float32).contiguous() + return data + if prim == "knn": + db = torch.randn(w["M"], w["D"], generator=g, device=dev, dtype=torch.float32) + q = torch.randn(w["Q"], w["D"], generator=g, device=dev, dtype=torch.float32) + return {"queries": q, "database": db} + if prim == "dbscan": + nc, D, N = int(w["n_centers"]), int(w["D"]), int(w["N"]) + centers = torch.randn(nc, D, generator=g, device=dev, dtype=torch.float32) * 10.0 + asg = torch.randint(0, nc, (N,), generator=g, device=dev) + xb = centers.index_select(0, asg) + torch.randn(N, D, generator=g, device=dev, dtype=torch.float32) * 0.5 + nz = int(float(w.get("noise_frac", 0.0)) * N) + if nz > 0: + lo = xb.min(0).values; hi = xb.max(0).values + xb[:nz] = lo + (hi - lo) * torch.rand(nz, D, generator=g, device=dev) + return {"x": xb, "eps": float(w["eps"]), "min_samples": int(w["min_samples"])} + x = torch.randn(w["N"], w["D"], generator=g, device=dev, dtype=torch.float32) + if prim == "kmeans": + perm = torch.randperm(w["N"], generator=g, device=dev)[: w["K"]] + return {"x": x, "init": x.index_select(0, perm).clone()} + return {"x": x} + + def call(mod, w, data, ctx): + if prim == "ivfpq": + # data["index"] is this iteration's fresh clone (see _fresh_index): the + # same frozen index, but no state can survive from a previous call. + return mod.ivf_pq_search(data["index"], data["queries"], int(w["k"]), nprobe=int(w["nprobe"])) + if prim == "kmeans": + return mod.kmeans(data["x"], w["K"], max_iters=w["max_iters"], + init_centroids=data["init"], tol=0.0) + if prim == "knn": + return mod.knn(data["queries"], data["database"], w["k"]) + if prim == "dbscan": + return mod.dbscan(data["x"], data["eps"], data["min_samples"]) + if prim == "pca": + return mod.pca(data["x"], w["k"]) + if prim == "tsvd": + return mod.truncated_svd(data["x"], w["k"]) + raise ValueError(prim) + + def check_shape(w, out): + if prim == "kmeans": + _, cen, _ = out + if tuple(cen.shape) != (w["K"], w["D"]) or not torch.isfinite(cen).all(): + raise ValueError("bad kmeans output") + elif prim == "knn": + d, i = out + if tuple(d.shape) != (w["Q"], w["k"]) or tuple(i.shape) != (w["Q"], w["k"]): + raise ValueError("bad knn output") + if not torch.isfinite(d).all(): + raise ValueError("non-finite distances") + elif prim == "ivfpq": + vals, ids = out + if tuple(ids.shape) != (int(w["Q"]), int(w["k"])) or \ + tuple(vals.shape) != (int(w["Q"]), int(w["k"])): + raise ValueError("bad ivfpq output shape") + elif prim == "dbscan": + if out.ndim != 1 or int(out.shape[0]) != int(w["N"]): + raise ValueError("bad dbscan labels shape") + else: + a, b = out + comps = a if prim == "pca" else b + if tuple(comps.shape) != (w["k"], w["D"]) or not torch.isfinite(comps).all(): + raise ValueError("bad decomposition output") + + def inertia(x, centroids): + c = centroids.to(torch.float32) + cn = (c * c).sum(1) + total = 0.0 + for j in range(0, x.shape[0], 16384): + xb = x[j:j + 16384] + d = (xb * xb).sum(1, keepdim=True) - 2.0 * (xb @ c.t()) + cn[None, :] + total += float(d.min(1).values.clamp_min(0).sum().item()) + return total + + def recall(agent_idx, ref_idx): + a = agent_idx.long(); b = ref_idx.long() + hit = (a.unsqueeze(2) == b.unsqueeze(1)).any(2) + return float(hit.sum().item()) / (b.shape[0] * b.shape[1]) + + def ari(a, b): + # Adjusted Rand Index between two integer label vectors (N,); noise (-1) + # is treated as its own label (dense-remapped first). + a = torch.unique(a.long(), return_inverse=True)[1] + b = torch.unique(b.long(), return_inverse=True)[1] + na = int(a.max().item()) + 1; nb = int(b.max().item()) + 1 + cont = torch.bincount(a * nb + b, minlength=na * nb).reshape(na, nb).double() + ai = cont.sum(1); bj = cont.sum(0); n = cont.sum() + c2 = lambda z: z * (z - 1) / 2.0 + sij = c2(cont).sum(); sa = c2(ai).sum(); sb = c2(bj).sum() + exp = (sa * sb / c2(n)) if float(n) > 1 else 0.0 + den = 0.5 * (sa + sb) - exp + return float(((sij - exp) / den).item()) if float(den) != 0.0 else 1.0 + + def ortho_err(comps): + c = comps.to(torch.float32); k = c.shape[0] + return float((c @ c.t() - torch.eye(k, device=c.device)).abs().max().item()) + + def captured(x, comps, center): + c = comps.to(torch.float32) + mean = x.mean(0) if center else None + total = 0.0 + for j in range(0, x.shape[0], 16384): + xb = x[j:j + 16384] + if center: + xb = xb - mean + p = xb @ c.t() + total += float((p * p).sum().item()) + return total / (x.shape[0] - 1) if center else total + + def verdict(w, data, ctx, ref_out, agent_out): + """Return (ok, reason, ref_val, agent_val) gating the agent output.""" + check_shape(w, agent_out) + if prim == "ivfpq": + # iso-result recall: agent ids vs frozen-baseline ids on the SAME index. + rec = recall(agent_out[1], ref_out[1]) + ok = rec >= float(cfg["recall_threshold"]) + return ok, ("recall_regression" if not ok else ""), 1.0, rec + if prim == "kmeans": + rv = inertia(data["x"], ref_out[1]); av = inertia(data["x"], agent_out[1]) + tol = float(cfg["inertia_tolerance"]) + ok = av <= (1.0 + tol) * rv + 1e-6 + return ok, ("inertia_regression" if not ok else ""), rv, av + if prim == "knn": + rec = recall(agent_out[1], ref_out[1]) + ok = rec >= float(cfg["recall_threshold"]) + return ok, ("recall_regression" if not ok else ""), 1.0, rec + if prim == "dbscan": + val = ari(agent_out, ref_out) + ok = val >= float(cfg["ari_threshold"]) + return ok, ("ari_regression" if not ok else ""), 1.0, val + center = (prim == "pca") + comps = agent_out[0] if prim == "pca" else agent_out[1] + ref_comps = ref_out[0] if prim == "pca" else ref_out[1] + oe = ortho_err(comps) + if oe > float(cfg["ortho_tolerance"]): + return False, "not_orthonormal", 0.0, oe + rv = captured(data["x"], ref_comps, center) + av = captured(data["x"], comps, center) + ok = av >= (1.0 - float(cfg["captured_tolerance"])) * rv - 1e-6 + return ok, ("captured_regression" if not ok else ""), rv, av + + def time_call(mod, w, data, ctx): + _sync(); t0 = _perf(); out = call(mod, w, data, ctx); _sync() + return (_perf() - t0) * 1000.0, out + + rows = [] + try: + for w in payload["workloads"]: + base_seed = int(w["seed"]) + ctx = setup(w, base_seed) # fixed context (e.g. IVF-PQ index) + # warmup on fresh data (kernels / autotune) — not measured + for i in range(warmup): + d = gen(w, base_seed + 100 + i, ctx) + call(ref, w, d, ctx); call(pkg, w, d, ctx); _sync() + del d + torch.cuda.empty_cache() + ratios, ref_val, agent_val = [], None, None + bad = None + for i in range(iters): + d = gen(w, base_seed + 10000 + i, ctx) # fresh every iteration + rt, rout = time_call(ref, w, d, ctx) + at, aout = time_call(pkg, w, d, ctx) + ok, reason, rv, av = verdict(w, d, ctx, rout, aout) # verify THIS iter + if not ok: + bad = reason; ref_val, agent_val = rv, av + break + ratios.append(rt / at if at > 0 else 0.01) + ref_val, agent_val = rv, av + del d, rout, aout + torch.cuda.empty_cache() + if bad is not None: + rows.append({"id": w["id"], "ok": False, "reason": bad, + "ref_val": ref_val, "agent_val": agent_val}) + continue + ratios.sort() + rows.append({"id": w["id"], "ok": True, + "speedup": ratios[len(ratios) // 2], + "ref_val": ref_val, "agent_val": agent_val}) + except Exception as exc: + return {"ok": False, "error": f"{type(exc).__name__}: {str(exc)[:200]}", + "trace": traceback.format_exc()[-500:]} + return {"ok": True, "rows": rows} + + +# --------------------------------------------------------------------------- # +# Host-side wrappers (run in the judge / agent container; need MODAL tokens). +# --------------------------------------------------------------------------- # +def _build_image(cfg: dict): + import modal + pip = list(cfg.get("pip", ["torch==2.5.1", "triton==3.1.0", "numpy"])) + base = cfg.get("cuda_image", "nvidia/cuda:12.4.1-devel-ubuntu22.04") + # add_python must match the judge/agent container Python (ubuntu:24.04 -> 3.12) + # because the worker ships via serialized=True (cloudpickle is version-sensitive). + # add_local_python_source mounts THIS module into the image so the remote can + # import it when deserializing the (module-level) worker function. + return (modal.Image.from_registry(base, add_python=str(cfg.get("python", "3.12"))) + .entrypoint([]) + .pip_install(*pip) + .add_local_python_source("flash_gpu")) + + +# Substrings that mark a transient Modal control-plane / image-build failure +# (evicted build, gateway hiccup, app stopped mid-run) rather than a real error +# in the submission — safe to retry. +_TRANSIENT_MARKERS = ( + "external shut-down", "terminated due to external", "please try again", + "app_state_stopped", "conflicterror", "deadline exceeded", "connection reset", + "502 bad gateway", "503 service", "temporarily unavailable", "timed out", + "gateway", "eviction", "internalfailure", "internalerror", + "failed to get new inputs", "runner failed", "task exited", +) + + +def _is_transient(text: str) -> bool: + low = (text or "").lower() + return any(m in low for m in _TRANSIENT_MARKERS) + + +def run_remote(payload: dict) -> dict: + """Run the GPU worker on Modal and return its result dict. + + Retries transient Modal control-plane failures; a real error in the + submission is non-transient and surfaces immediately. Raises RuntimeError + with a sanitized message on unrecoverable failure. + """ + import time as _time + + import modal + + cfg = payload["cfg"] + attempts = max(1, int(cfg.get("modal_retries", 3))) + last = "modal_gpu_eval_failed" + for attempt in range(1, attempts + 1): + app = modal.App(cfg.get("app_name", "flash-kernel-eval")) + fn_kwargs = dict( + gpu=cfg.get("gpu", "H100"), + image=_build_image(cfg), + timeout=int(cfg.get("modal_timeout_seconds", 1800)), + serialized=True, + max_containers=1, + ) + # Mount the cached real-dataset Volume (SIFT/GIST HDF5) read-only at /data. + vol_name = cfg.get("dataset_volume") + if vol_name: + fn_kwargs["volumes"] = {"/data": modal.Volume.from_name(vol_name)} + remote = app.function(**fn_kwargs)(_gpu_worker) + try: + with modal.enable_output(), app.run(): + return remote.remote(payload) + except Exception as exc: # noqa: BLE001 + last = str(exc)[-400:] + if not _is_transient(last) or attempt == attempts: + raise RuntimeError(f"modal_gpu_eval_failed: {last}") from exc + _time.sleep(min(45, 10 * attempt)) + raise RuntimeError(f"modal_gpu_eval_failed: {last}") + + +def modal_available() -> bool: + return bool(os.environ.get("MODAL_TOKEN_ID") and os.environ.get("MODAL_TOKEN_SECRET")) diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/harbor/app/README.md b/2.0/problems/ivf_pq_gpu_kernel_optimization/harbor/app/README.md new file mode 100644 index 000000000..34769f919 --- /dev/null +++ b/2.0/problems/ivf_pq_gpu_kernel_optimization/harbor/app/README.md @@ -0,0 +1,34 @@ +# IVF-PQ kernel optimization — submission workflow + +You are optimizing the `ivfpqlib` package at `/app/ivfpqlib`. Edit the package +(rewrite the internals of `ivfpq`, add Triton kernel modules under +`ivfpqlib/`), then submit a patch. + +## Iterate locally + +Run the public self-test to check correctness and get a rough speed signal on +the two public shapes (needs a GPU in the agent container): + +```bash +bash /app/public_test.sh +``` + +## Submit + +```bash +bash /app/make_submission.sh # writes /app/solution.patch (ivfpqlib diff) +bash /app/submit.sh # enqueues it for the black-box judge +``` + +Submissions are asynchronous. Submit early and keep improving; use +`bash /app/submissions.sh` and `bash /app/wait_submission.sh ` to inspect +results. + +## Rules + +- Only files under `ivfpqlib/` may change. +- Do not import external optimized libraries (write the kernels yourself), and + do not access the environment, spawn processes, or use the network. +- Keep the public `ivfpq(...)` signature and return contract unchanged. +- Clustering quality is gated (inertia vs the naive baseline); do not sacrifice + correctness for speed beyond the allowed tolerance. diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/harbor/app/make_submission.sh b/2.0/problems/ivf_pq_gpu_kernel_optimization/harbor/app/make_submission.sh new file mode 100644 index 000000000..8b1dec84c --- /dev/null +++ b/2.0/problems/ivf_pq_gpu_kernel_optimization/harbor/app/make_submission.sh @@ -0,0 +1,24 @@ +#!/usr/bin/env bash +# Package the current ivfpqlib edits into a unified diff for submission. +set -euo pipefail + +APP_DIR="${APP_DIR:-/app}" +OUT="${1:-/app/solution.patch}" + +if [[ ! -d "$APP_DIR/ivfpqlib" ]]; then + echo "ivfpqlib package not found at $APP_DIR/ivfpqlib" >&2 + exit 2 +fi +if [[ ! -d "$APP_DIR/.git" ]]; then + echo "git repo not found at $APP_DIR (expected in the agent image)" >&2 + exit 2 +fi + +# Stage and diff only the package, so submission helper scripts under /app are +# never included. New kernel files under ivfpqlib/ are captured via `git add`. +git -C "$APP_DIR" add -A -- ivfpqlib +git -C "$APP_DIR" diff --cached -- ivfpqlib > "$OUT" +git -C "$APP_DIR" reset -q + +bytes=$(wc -c < "$OUT" | tr -d ' ') +echo "Wrote $OUT ($bytes bytes)" diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/harbor/app/public_test.py b/2.0/problems/ivf_pq_gpu_kernel_optimization/harbor/app/public_test.py new file mode 100644 index 000000000..e8de4a4bd --- /dev/null +++ b/2.0/problems/ivf_pq_gpu_kernel_optimization/harbor/app/public_test.py @@ -0,0 +1,171 @@ +"""Public GPU self-test -- IDENTICAL to the final graded evaluation. + +This runs the EXACT same workloads, thresholds, timing, and seeds that the hidden +judge uses to grade your submission (all read from ``/app/task_config.json``, +the same config the judge reads), on a Modal GPU, and reports the same +per-workload pass/fail + speedup + geomean + predicted score (0-100) you would +receive on submission. There is no longer a separate, smaller "public" workload +set -- what you see here is what you get graded on. + +Requires MODAL_TOKEN_ID / MODAL_TOKEN_SECRET in the environment. +""" +from __future__ import annotations + +import json +import math +import os +import sys +from pathlib import Path + +sys.path.insert(0, "/opt") # flash_gpu.py is baked at /opt + +APP_DIR = os.environ.get("APP_DIR", "/app") +TASK_CONFIG_PATH = os.environ.get("TASK_CONFIG_PATH", "/app/task_config.json") + + +def _load_eval() -> dict: + """The judge grades from task_config.json's `evaluation` block; read the same.""" + try: + doc = json.loads(Path(TASK_CONFIG_PATH).read_text(encoding="utf-8")) + except Exception as exc: # noqa: BLE001 + print(f"could not read {TASK_CONFIG_PATH}: {exc}") + return {} + return doc.get("evaluation", doc) + + +EV = _load_eval() + + +def g(key, default): + v = EV.get(key, default) + return default if v is None else v + + +PRIMITIVE = str(g("primitive", "")) +PKG = str(g("pkg", "")) +REF_MODULE = str(g("ref_module", "")) +SPEEDUP_TARGET = float(g("speedup_target", 8.0)) +BASELINE_DIR = str(g("baseline_source", "/opt/flash_ref")) + + +def _workloads() -> list: + """Exactly the judge's final-role workload set: ALL workloads, same seed + derivation ``base_seed + 1000*(i+1)`` -- identical data to the graded run.""" + wls = [dict(w) for w in g("workloads", [])] + base = int(g("base_seed", 20260701)) + for i, w in enumerate(wls): + w.setdefault("seed", base + 1000 * (i + 1)) + return wls + + +def _cfg() -> dict: + """Byte-for-byte the judge's _build_cfg().""" + return { + "primitive": PRIMITIVE, + "pkg": PKG, + "ref_module": REF_MODULE, + "gpu": str(g("gpu", "H100")), + "cuda_image": str(g("cuda_image", "nvidia/cuda:12.4.1-devel-ubuntu22.04")), + "pip": list(g("pip", ["torch", "numpy"])), + "app_name": str(g("app_name", "flash-kernel-public")), + "modal_timeout_seconds": int(g("modal_timeout_seconds", 1800)), + "warmup": int(g("warmup_iters", 3)), + "iters": int(g("timed_iters", 7)), + "inertia_tolerance": float(g("inertia_tolerance", 0.02)), + "recall_threshold": float(g("recall_threshold", 0.99)), + "captured_tolerance": float(g("captured_tolerance", 0.02)), + "ortho_tolerance": float(g("ortho_tolerance", 0.02)), + "ari_threshold": float(g("ari_threshold", 0.99)), + } + + +def geometric_mean(values: list) -> float: + if not values: + return 0.0 + return math.exp(sum(math.log(max(v, 1e-9)) for v in values) / len(values)) + + +def score_from_speedup(gm: float) -> float: + if gm <= 0: + return 0.0 + raw = 100.0 * math.log(gm) / math.log(max(SPEEDUP_TARGET, 1.0000001)) + return max(0.0, min(100.0, raw)) + + +def _read(root: str, sub: str = "") -> dict: + base = Path(root) + scan = base / sub if sub else base + return {str(p.relative_to(base)): p.read_text(encoding="utf-8", errors="replace") + for p in scan.rglob("*.py")} + + +def main() -> int: + import flash_gpu # baked at /opt/flash_gpu.py + if not flash_gpu.modal_available(): + print("Set MODAL_TOKEN_ID and MODAL_TOKEN_SECRET to run the GPU self-test.") + return 1 + workloads = _workloads() + cfg = _cfg() + if not workloads: + print("no workloads found in task_config.json; cannot run.") + return 1 + payload = { + "baseline_files": _read(BASELINE_DIR), + "patched_files": _read(APP_DIR, PKG), + "workloads": workloads, + "cfg": cfg, + } + try: + result = flash_gpu.run_remote(payload) + except Exception as exc: # noqa: BLE001 + print(f"GPU run failed: {exc}") + return 1 + if not result.get("ok"): + print(f"worker error: {result.get('error')}") + return 1 + + if PRIMITIVE in ("knn", "ivfpq"): + mlabel, gate = "recall@k", f">= {cfg['recall_threshold']}" + elif PRIMITIVE == "dbscan": + mlabel, gate = "ARI", f">= {cfg['ari_threshold']}" + elif PRIMITIVE == "kmeans": + mlabel, gate = "inertia", f"<= (1+{cfg['inertia_tolerance']}) x ref (lower is better)" + else: + mlabel, gate = "quality", "" + + print("=== FINAL-EQUIVALENT self-test: identical workloads / thresholds / seeds / " + "timing to the graded judge ===") + print(f"correctness gate: {mlabel} {gate} | speedup_target = {SPEEDUP_TARGET:g}") + print(f"{'workload':9s} {'status':22s} {'speedup':>9s} {mlabel}: agent / ref") + rows = result.get("rows", []) + speedups = [] + any_fail = False + for row in rows: + av, rv = row.get("agent_val"), row.get("ref_val") + q = (f"{av:.4f} / {rv:.4f}" + if isinstance(av, (int, float)) and isinstance(rv, (int, float)) else "n/a") + if row.get("ok"): + sp = f"{row['speedup']:.2f}x" + speedups.append(max(float(row["speedup"]), 0.01)) # same clamp as the judge + st = "OK" + else: + sp = "-" + st = "FAIL:" + str(row.get("reason", "")) + any_fail = True + print(f"{row['id']:9s} {st:22s} {sp:>9s} {q}") + + # Scoring is byte-for-byte the judge's full_evaluation: ANY gate failure -> 0. + print() + if any_fail or len(speedups) != len(rows) or not speedups: + print("RESULT: at least one workload FAILED the correctness gate -> a submission now " + "would be INVALID and score 0/100. Fix the failing workload(s) before submitting.") + return 0 + gm = geometric_mean(speedups) + score = score_from_speedup(gm) + print(f"geomean speedup = {gm:.3f}x over the naive baseline") + print(f"PREDICTED FINAL SCORE = {score:.2f} / 100 (100 == {SPEEDUP_TARGET:g}x geomean)") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/harbor/app/public_test.sh b/2.0/problems/ivf_pq_gpu_kernel_optimization/harbor/app/public_test.sh new file mode 100644 index 000000000..8cdc88205 --- /dev/null +++ b/2.0/problems/ivf_pq_gpu_kernel_optimization/harbor/app/public_test.sh @@ -0,0 +1,7 @@ +#!/usr/bin/env bash +# Public self-test: run the patched ivfpqlib on the public shapes and print a +# correctness + rough speedup summary. Requires a GPU in the agent container. +set -euo pipefail +APP_DIR="${APP_DIR:-/app}" +cd "$APP_DIR" +exec python3 "$APP_DIR/public_test.py" diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/harbor/app/solution.patch b/2.0/problems/ivf_pq_gpu_kernel_optimization/harbor/app/solution.patch new file mode 100644 index 000000000..e69de29bb diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/ivfpqlib/__init__.py b/2.0/problems/ivf_pq_gpu_kernel_optimization/ivfpqlib/__init__.py new file mode 100644 index 000000000..873d4f077 --- /dev/null +++ b/2.0/problems/ivf_pq_gpu_kernel_optimization/ivfpqlib/__init__.py @@ -0,0 +1,8 @@ +"""GPU IVF-PQ library: naive torch baseline to be optimized (search) plus a +frozen index builder. Public API: ``ivf_pq_build``, ``ivf_pq_search``, +``IvfPqIndex``.""" +from ivfpqlib.index import IvfPqIndex +from ivfpqlib.build import ivf_pq_build +from ivfpqlib.search import ivf_pq_search + +__all__ = ["ivf_pq_build", "ivf_pq_search", "IvfPqIndex"] diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/ivfpqlib/build.py b/2.0/problems/ivf_pq_gpu_kernel_optimization/ivfpqlib/build.py new file mode 100644 index 000000000..be1d016da --- /dev/null +++ b/2.0/problems/ivf_pq_gpu_kernel_optimization/ivfpqlib/build.py @@ -0,0 +1,208 @@ +"""Naive IVF-PQ index builder (pure torch, deterministic, CPU-OK). + +Builds an :class:`~ivfpqlib.index.IvfPqIndex`: a coarse k-means quantizer over +``nlist`` cells + ``m`` product-quantization sub-codebooks, then encodes every +database row to an ``(M, m)`` uint8 code in CSR cell-contiguous layout. This is +the **frozen** index construction (it is not the timed operation); search is +what you optimize. ADC contract: a candidate with codes ``code`` probing list +``c`` scores ``sum_s || rq_s - codebook[s, code[s]] ||^2`` where the residual +query ``rq = q - centroid[c]`` (``by_residual``) or ``q`` otherwise. +""" +from __future__ import annotations + +import math +from typing import Optional, Tuple + +import torch + +from ivfpqlib.index import IvfPqIndex + +def _pad_features(x: torch.Tensor, Dp: int) -> torch.Tensor: + """Zero-pad the trailing feature dim to ``Dp`` (no-op when already wide).""" + D = x.shape[-1] + if D >= Dp: + return x + pad = torch.zeros((*x.shape[:-1], Dp - D), device=x.device, dtype=x.dtype) + return torch.cat([x, pad], dim=-1) + + +def _pq_dims(D: int, m: int) -> Tuple[int, int, int]: + """Resolve ``(m, dsub, Dp)`` for a ``D``-dim input split into ``m`` codes. + + ``dsub = ceil(D / m)`` and ``Dp = m * dsub`` (the zero-padded working + width). ``dsub`` is bumped until ``Dp >= 16`` so the coarse-quantizer + kernels (which use ``tl.dot``) always have a contraction dim >= 16; + the zero columns never change squared-L2 distances. + """ + m = int(m) + if m < 1: + raise ValueError(f"m (number of sub-quantizers) must be >= 1 (got {m})") + if m > D: + # More sub-quantizers than dims: clamp so each sub-vector has >= 1 dim. + m = int(D) + dsub = int(math.ceil(D / m)) + while m * dsub < 16: + dsub += 1 + return m, dsub, m * dsub + + +# ── tiny CPU-OK building blocks ──────────────────────────────────────────── +def _lloyd_kmeans( + sample: torch.Tensor, k: int, *, niter: int, seed: int +) -> torch.Tensor: + """Tiny Lloyd k-means returning ``(k, D)`` centroids (fp32 math).""" + n = sample.shape[0] + if n < k: + raise ValueError( + f"k-means needs at least k={k} training rows (got {n}); " + "increase train_size / pq_train_size or lower nlist / m." + ) + g = torch.Generator(device="cpu").manual_seed(seed) + perm = torch.randperm(n, generator=g)[:k] + centroids = sample[perm.to(sample.device)].to(torch.float32).clone() + s = sample.to(torch.float32) + for _ in range(max(1, niter)): + d2 = torch.cdist(s, centroids) ** 2 # (n, k) + assign = d2.argmin(dim=1) # (n,) + new = centroids.clone() + for c in range(k): + mask = assign == c + if bool(mask.any()): + new[c] = s[mask].mean(dim=0) + shift = (new - centroids).norm(dim=-1).max() + centroids = new + if float(shift) == 0.0: + break + return centroids.to(torch.float32) + + +def _assign_chunked( + Xp: torch.Tensor, centroids: torch.Tensor, chunk: int = 8192 +) -> torch.Tensor: + """Nearest-centroid id per row (squared-L2), chunked to bound memory.""" + out = torch.empty(Xp.shape[0], dtype=torch.int64, device=Xp.device) + cf = centroids.to(torch.float32) + for lo in range(0, Xp.shape[0], chunk): + hi = min(lo + chunk, Xp.shape[0]) + d2 = torch.cdist(Xp[lo:hi].to(torch.float32), cf) ** 2 + out[lo:hi] = d2.argmin(dim=1) + return out + + +def _train_pq_codebooks( + resid: torch.Tensor, m: int, dsub: int, ksub: int, *, niter: int, seed: int +) -> torch.Tensor: + """Train ``m`` independent k-means sub-quantizers on residual sub-vectors. + + ``resid`` is ``(N, m*dsub)``; returns ``(m, ksub, dsub)`` fp32 codebooks. + """ + resid_sub = resid.reshape(resid.shape[0], m, dsub) # (N, m, dsub) + codebooks = torch.empty(m, ksub, dsub, dtype=torch.float32, device=resid.device) + for s in range(m): + codebooks[s] = _lloyd_kmeans(resid_sub[:, s, :], ksub, niter=niter, seed=seed + s) + return codebooks + + +def _encode_pq( + resid: torch.Tensor, codebooks: torch.Tensor, m: int, dsub: int, + chunk: int = 8192, +) -> torch.Tensor: + """Encode residual rows to ``(N, m)`` uint8 PQ codes (nearest sub-centroid).""" + N = resid.shape[0] + resid_sub = resid.reshape(N, m, dsub) # (N, m, dsub) + codes = torch.empty(N, m, dtype=torch.uint8, device=resid.device) + for s in range(m): + cb = codebooks[s].to(torch.float32) # (ksub, dsub) + sub = resid_sub[:, s, :].to(torch.float32) # (N, dsub) + for lo in range(0, N, chunk): + hi = min(lo + chunk, N) + d2 = torch.cdist(sub[lo:hi], cb) ** 2 # (chunk, ksub) + codes[lo:hi, s] = d2.argmin(dim=1).to(torch.uint8) + return codes + + +# ── public build / search ────────────────────────────────────────────────── +def ivf_pq_build( + X: torch.Tensor, + nlist: int, + *, + m: int = 8, + nbits: int = 8, + metric: str = "l2", + by_residual: bool = True, + nprobe: int = 8, + niter: int = 20, + pq_niter: int = 25, + train_size: Optional[int] = None, + pq_train_size: Optional[int] = None, + seed: int = 0, +): + """Build an :class:`IvfPqIndex` with pure torch ops.""" + if metric != "l2": + raise NotImplementedError(f"ivf_pq supports metric='l2' only (got {metric!r})") + if nbits != 8: + raise NotImplementedError(f"ivf_pq supports nbits=8 only (got {nbits})") + if X.ndim != 2: + raise ValueError("ivf_pq build expects a 2D (M, D) tensor") + + M, D = X.shape + nlist = int(min(nlist, M)) + if nlist < 1: + raise ValueError("nlist must be >= 1") + m, dsub, Dp = _pq_dims(int(D), m) + ksub = 1 << nbits # 256 + Xp = _pad_features(X.to(torch.float32), Dp).contiguous() + + # ── coarse quantizer: k-means on a sample ────────────────────────── + train_size = int(train_size or min(M, nlist * 256)) + train_size = max(min(train_size, M), nlist) + g = torch.Generator(device="cpu").manual_seed(seed) + sample_idx = torch.randperm(M, generator=g)[:train_size].to(X.device) + sample = Xp.index_select(0, sample_idx) + centroids = _lloyd_kmeans(sample, nlist, niter=niter, seed=seed) # (nlist, Dp) + + # ── PQ codebooks from residuals of a (sub)sample ─────────────────── + pq_train_size = int(pq_train_size or min(train_size, max(ksub * 16, 4096))) + pq_train_size = max(min(pq_train_size, train_size), ksub) + pq_sample = sample[:pq_train_size] + if by_residual: + pq_assign = _assign_chunked(pq_sample, centroids) + resid_train = pq_sample - centroids.index_select(0, pq_assign) + else: + resid_train = pq_sample + codebooks = _train_pq_codebooks( + resid_train, m, dsub, ksub, niter=pq_niter, seed=seed + 1 + ) # (m, ksub, dsub) + + # ── assign every database row + encode its residual ──────────────── + assign = _assign_chunked(Xp, centroids) # (M,) + resid_all = Xp - centroids.index_select(0, assign) if by_residual else Xp + codes = _encode_pq(resid_all, codebooks, m, dsub) # (M, m) uint8 + + # ── CSR cell-contiguous layout ───────────────────────────────────── + counts = torch.bincount(assign, minlength=nlist) # (nlist,) + offsets = torch.zeros(nlist + 1, dtype=torch.int64, device=X.device) + offsets[1:] = counts.cumsum(0) + order = torch.argsort(assign, stable=True) # (M,) int64 + codes_sorted = codes.index_select(0, order).contiguous() + + return IvfPqIndex( + centroids=centroids, + pq_codebooks=codebooks, + codes=codes_sorted, + ids=order.to(torch.int64), + list_offsets=offsets, + metric=metric, + by_residual=bool(by_residual), + D=int(D), + Dp=int(Dp), + dsub=int(dsub), + m=int(m), + nbits=int(nbits), + nlist=int(nlist), + nprobe=int(nprobe), + max_list_len=int(counts.max().item()) if nlist > 0 else 0, + ) + + +__all__ = ["ivf_pq_build", "_pad_features", "_pq_dims"] diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/ivfpqlib/index.py b/2.0/problems/ivf_pq_gpu_kernel_optimization/ivfpqlib/index.py new file mode 100644 index 000000000..1c177c308 --- /dev/null +++ b/2.0/problems/ivf_pq_gpu_kernel_optimization/ivfpqlib/index.py @@ -0,0 +1,96 @@ +"""``IvfPqIndex`` -- the in-memory container for a built IVF-PQ index. + +Stores the database **cell-contiguous** (all vectors of inverted list ``c`` in +rows ``[list_offsets[c], list_offsets[c+1])``) and keeps only the ``(M, m)`` +uint8 product-quantization codes, not the full vectors. ``ids[p]`` maps a stored +row back to the caller's original row id. Torch-only; no kernel import. +""" + +from dataclasses import dataclass + +import torch + + +@dataclass +class IvfPqIndex: + """A built IVF-PQ index. + + Attributes: + centroids: ``(nlist, Dp)`` coarse-quantizer centroids (fp32). + pq_codebooks: ``(m, ksub, dsub)`` product-quantization sub-centroids + (fp32). ``ksub == 2**nbits`` (256 for the only supported + ``nbits=8``); ``dsub == Dp // m``. + codes: ``(M, m)`` uint8 -- per sub-quantizer code, cell-contiguous. + ids: ``(M,)`` int64 -- original row id for each stored row. + list_offsets: ``(nlist + 1,)`` int64 CSR offsets into ``codes``. + metric: distance metric (only ``"l2"`` supported). + by_residual: if True, PQ encodes ``x - centroid[list]`` (FAISS / + cuVS default, higher recall); else PQ encodes ``x`` directly. + D: original feature dimension as passed by the caller. + Dp: padded working dimension (``m * dsub``, ``>= 16``); zero + columns added for ``D < Dp`` never affect squared-L2 distances. + dsub: sub-vector dimension (``Dp // m``). + m: number of sub-quantizers (PQ codes per vector). + nbits: bits per code (only ``8`` supported -> ``ksub = 256``). + nlist: number of inverted lists / coarse centroids. + nprobe: default number of lists to probe at search time. + max_list_len: longest inverted list, recorded at build time so + search can size the kernel's chunk loop without a D2H sync. + """ + + centroids: torch.Tensor + pq_codebooks: torch.Tensor + codes: torch.Tensor + ids: torch.Tensor + list_offsets: torch.Tensor + metric: str + by_residual: bool + D: int + Dp: int + dsub: int + m: int + nbits: int + nlist: int + nprobe: int + max_list_len: int = 0 + + @property + def ksub(self) -> int: + """Number of sub-centroids per sub-quantizer (``2**nbits``).""" + return int(self.pq_codebooks.shape[1]) + + @property + def M(self) -> int: + return int(self.codes.shape[0]) + + @property + def device(self) -> torch.device: + return self.codes.device + + @property + def dtype(self) -> torch.dtype: + """Working dtype of the centroids / codebooks (codes are uint8).""" + return self.centroids.dtype + + def list_lengths(self) -> torch.Tensor: + """``(nlist,)`` int64 number of vectors in each inverted list.""" + return self.list_offsets[1:] - self.list_offsets[:-1] + + def code_size_bytes(self) -> int: + """Bytes per stored vector (``m`` for ``nbits=8``).""" + return int(self.m) + + def compression_ratio(self) -> float: + """Original fp32 vector bytes / PQ code bytes (storage savings).""" + return (4.0 * self.D) / max(self.code_size_bytes(), 1) + + def __repr__(self) -> str: # pragma: no cover - cosmetic + return ( + f"IvfPqIndex(M={self.M}, D={self.D}, Dp={self.Dp}, m={self.m}, " + f"dsub={self.dsub}, nbits={self.nbits}, nlist={self.nlist}, " + f"nprobe={self.nprobe}, by_residual={self.by_residual}, " + f"metric={self.metric!r}, dtype={self.dtype}, device={self.device})" + ) + + +__all__ = ["IvfPqIndex"] diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/ivfpqlib/search.py b/2.0/problems/ivf_pq_gpu_kernel_optimization/ivfpqlib/search.py new file mode 100644 index 000000000..56b559bd7 --- /dev/null +++ b/2.0/problems/ivf_pq_gpu_kernel_optimization/ivfpqlib/search.py @@ -0,0 +1,78 @@ +"""Naive IVF-PQ search (pure torch reference -- the baseline you optimize). + +Coarse step: probe the ``nprobe`` nearest lists per query. Fine step: for each +query, build the per-list residual ADC lookup table and score every probed +candidate as ``sum_s LUT[s, code[s]]``, then keep the global top-``k``. The +Python per-query loop is what makes this slow; a fast GPU implementation fuses +the coarse search, LUT build, and ragged candidate scan into kernels. + + ivf_pq_search(index, Q, k, *, nprobe=None) -> (vals (nq,k) fp32, + ids (nq,k) int64, -1 padded) +""" +from __future__ import annotations + +from typing import Optional + +import torch + +from ivfpqlib.build import _pad_features + +def ivf_pq_search(index, Q: torch.Tensor, k: int, *, nprobe: Optional[int] = None): + """Reference coarse + ADC IVF-PQ search over a built index. + + Returns ``(vals, ids)`` where ``vals[i, j]`` is the (approximate) + squared-L2 distance to the ``j``-th nearest PQ reconstruction and + ``ids`` are original row ids. Mirrors the Triton kernel exactly: + probe the ``nprobe`` nearest lists, build the per-(query, list) + residual LUT, score each member as the sum of ``m`` table lookups, + keep the global top-``k``. + """ + if Q.ndim != 2: + raise ValueError("ivf_pq search expects a 2D (nq, D) query tensor") + nprobe = int(nprobe or index.nprobe) + nprobe = max(1, min(nprobe, index.nlist)) + nq = Q.shape[0] + Dp, m, dsub = index.Dp, index.m, index.dsub + + Qp = _pad_features(Q.to(torch.float32), Dp) + centroids = index.centroids.to(torch.float32) # (nlist, Dp) + codebooks = index.pq_codebooks.to(torch.float32) # (m, ksub, dsub) + codes = index.codes # (M, m) uint8 + offsets = index.list_offsets + + coarse_d2 = torch.cdist(Qp, centroids) ** 2 # (nq, nlist) + probed = coarse_d2.topk(nprobe, dim=1, largest=False).indices # (nq, nprobe) + + out_vals = torch.full((nq, k), float("inf"), device=Q.device, dtype=torch.float32) + out_ids = torch.full((nq, k), -1, device=Q.device, dtype=torch.int64) + + for i in range(nq): + cand_dists = [] + cand_ids = [] + for p in range(nprobe): + c = int(probed[i, p].item()) + s0, e0 = int(offsets[c].item()), int(offsets[c + 1].item()) + if e0 <= s0: + continue + rq = (Qp[i] - centroids[c]) if index.by_residual else Qp[i] # (Dp,) + rq_sub = rq.reshape(m, dsub) # (m, dsub) + # LUT[s, j] = ||rq_s - codebook[s, j]||^2 + lut = ((rq_sub[:, None, :] - codebooks) ** 2).sum(-1) # (m, ksub) + cc = codes[s0:e0].to(torch.int64) # (L, m) + # dist[l] = sum_s LUT[s, cc[l, s]] + dist = lut.gather(1, cc.t().contiguous()).sum(0) # (L,) + cand_dists.append(dist) + cand_ids.append(index.ids[s0:e0]) + if not cand_dists: + continue + dist = torch.cat(cand_dists) + ids = torch.cat(cand_ids) + kk = min(k, dist.shape[0]) + vals, sel = dist.topk(kk, largest=False) + out_vals[i, :kk] = vals + out_ids[i, :kk] = ids[sel] + + return out_vals, out_ids + + +__all__ = ["ivf_pq_search"] diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/judge/refivfpq.py b/2.0/problems/ivf_pq_gpu_kernel_optimization/judge/refivfpq.py new file mode 100644 index 000000000..c07e5c8ea --- /dev/null +++ b/2.0/problems/ivf_pq_gpu_kernel_optimization/judge/refivfpq.py @@ -0,0 +1,343 @@ +"""``IvfPqIndex`` -- the in-memory container for a built IVF-PQ index. + +Stores the database **cell-contiguous** (all vectors of inverted list ``c`` in +rows ``[list_offsets[c], list_offsets[c+1])``) and keeps only the ``(M, m)`` +uint8 product-quantization codes, not the full vectors. ``ids[p]`` maps a stored +row back to the caller's original row id. Torch-only; no kernel import. Frozen judge baseline; do not edit. +""" + +import math +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch + +@dataclass +class IvfPqIndex: + """A built IVF-PQ index. + + Attributes: + centroids: ``(nlist, Dp)`` coarse-quantizer centroids (fp32). + pq_codebooks: ``(m, ksub, dsub)`` product-quantization sub-centroids + (fp32). ``ksub == 2**nbits`` (256 for the only supported + ``nbits=8``); ``dsub == Dp // m``. + codes: ``(M, m)`` uint8 -- per sub-quantizer code, cell-contiguous. + ids: ``(M,)`` int64 -- original row id for each stored row. + list_offsets: ``(nlist + 1,)`` int64 CSR offsets into ``codes``. + metric: distance metric (only ``"l2"`` supported). + by_residual: if True, PQ encodes ``x - centroid[list]`` (FAISS / + cuVS default, higher recall); else PQ encodes ``x`` directly. + D: original feature dimension as passed by the caller. + Dp: padded working dimension (``m * dsub``, ``>= 16``); zero + columns added for ``D < Dp`` never affect squared-L2 distances. + dsub: sub-vector dimension (``Dp // m``). + m: number of sub-quantizers (PQ codes per vector). + nbits: bits per code (only ``8`` supported -> ``ksub = 256``). + nlist: number of inverted lists / coarse centroids. + nprobe: default number of lists to probe at search time. + max_list_len: longest inverted list, recorded at build time so + search can size the kernel's chunk loop without a D2H sync. + """ + + centroids: torch.Tensor + pq_codebooks: torch.Tensor + codes: torch.Tensor + ids: torch.Tensor + list_offsets: torch.Tensor + metric: str + by_residual: bool + D: int + Dp: int + dsub: int + m: int + nbits: int + nlist: int + nprobe: int + max_list_len: int = 0 + + @property + def ksub(self) -> int: + """Number of sub-centroids per sub-quantizer (``2**nbits``).""" + return int(self.pq_codebooks.shape[1]) + + @property + def M(self) -> int: + return int(self.codes.shape[0]) + + @property + def device(self) -> torch.device: + return self.codes.device + + @property + def dtype(self) -> torch.dtype: + """Working dtype of the centroids / codebooks (codes are uint8).""" + return self.centroids.dtype + + def list_lengths(self) -> torch.Tensor: + """``(nlist,)`` int64 number of vectors in each inverted list.""" + return self.list_offsets[1:] - self.list_offsets[:-1] + + def code_size_bytes(self) -> int: + """Bytes per stored vector (``m`` for ``nbits=8``).""" + return int(self.m) + + def compression_ratio(self) -> float: + """Original fp32 vector bytes / PQ code bytes (storage savings).""" + return (4.0 * self.D) / max(self.code_size_bytes(), 1) + + def __repr__(self) -> str: # pragma: no cover - cosmetic + return ( + f"IvfPqIndex(M={self.M}, D={self.D}, Dp={self.Dp}, m={self.m}, " + f"dsub={self.dsub}, nbits={self.nbits}, nlist={self.nlist}, " + f"nprobe={self.nprobe}, by_residual={self.by_residual}, " + f"metric={self.metric!r}, dtype={self.dtype}, device={self.device})" + ) + + +def _pad_features(x: torch.Tensor, Dp: int) -> torch.Tensor: + """Zero-pad the trailing feature dim to ``Dp`` (no-op when already wide).""" + D = x.shape[-1] + if D >= Dp: + return x + pad = torch.zeros((*x.shape[:-1], Dp - D), device=x.device, dtype=x.dtype) + return torch.cat([x, pad], dim=-1) + + +def _pq_dims(D: int, m: int) -> Tuple[int, int, int]: + """Resolve ``(m, dsub, Dp)`` for a ``D``-dim input split into ``m`` codes. + + ``dsub = ceil(D / m)`` and ``Dp = m * dsub`` (the zero-padded working + width). ``dsub`` is bumped until ``Dp >= 16`` so the coarse-quantizer + kernels (which use ``tl.dot``) always have a contraction dim >= 16; + the zero columns never change squared-L2 distances. + """ + m = int(m) + if m < 1: + raise ValueError(f"m (number of sub-quantizers) must be >= 1 (got {m})") + if m > D: + # More sub-quantizers than dims: clamp so each sub-vector has >= 1 dim. + m = int(D) + dsub = int(math.ceil(D / m)) + while m * dsub < 16: + dsub += 1 + return m, dsub, m * dsub + + +# ── tiny CPU-OK building blocks ──────────────────────────────────────────── +def _lloyd_kmeans( + sample: torch.Tensor, k: int, *, niter: int, seed: int +) -> torch.Tensor: + """Tiny Lloyd k-means returning ``(k, D)`` centroids (fp32 math).""" + n = sample.shape[0] + if n < k: + raise ValueError( + f"k-means needs at least k={k} training rows (got {n}); " + "increase train_size / pq_train_size or lower nlist / m." + ) + g = torch.Generator(device="cpu").manual_seed(seed) + perm = torch.randperm(n, generator=g)[:k] + centroids = sample[perm.to(sample.device)].to(torch.float32).clone() + s = sample.to(torch.float32) + for _ in range(max(1, niter)): + d2 = torch.cdist(s, centroids) ** 2 # (n, k) + assign = d2.argmin(dim=1) # (n,) + new = centroids.clone() + for c in range(k): + mask = assign == c + if bool(mask.any()): + new[c] = s[mask].mean(dim=0) + shift = (new - centroids).norm(dim=-1).max() + centroids = new + if float(shift) == 0.0: + break + return centroids.to(torch.float32) + + +def _assign_chunked( + Xp: torch.Tensor, centroids: torch.Tensor, chunk: int = 8192 +) -> torch.Tensor: + """Nearest-centroid id per row (squared-L2), chunked to bound memory.""" + out = torch.empty(Xp.shape[0], dtype=torch.int64, device=Xp.device) + cf = centroids.to(torch.float32) + for lo in range(0, Xp.shape[0], chunk): + hi = min(lo + chunk, Xp.shape[0]) + d2 = torch.cdist(Xp[lo:hi].to(torch.float32), cf) ** 2 + out[lo:hi] = d2.argmin(dim=1) + return out + + +def _train_pq_codebooks( + resid: torch.Tensor, m: int, dsub: int, ksub: int, *, niter: int, seed: int +) -> torch.Tensor: + """Train ``m`` independent k-means sub-quantizers on residual sub-vectors. + + ``resid`` is ``(N, m*dsub)``; returns ``(m, ksub, dsub)`` fp32 codebooks. + """ + resid_sub = resid.reshape(resid.shape[0], m, dsub) # (N, m, dsub) + codebooks = torch.empty(m, ksub, dsub, dtype=torch.float32, device=resid.device) + for s in range(m): + codebooks[s] = _lloyd_kmeans(resid_sub[:, s, :], ksub, niter=niter, seed=seed + s) + return codebooks + + +def _encode_pq( + resid: torch.Tensor, codebooks: torch.Tensor, m: int, dsub: int, + chunk: int = 8192, +) -> torch.Tensor: + """Encode residual rows to ``(N, m)`` uint8 PQ codes (nearest sub-centroid).""" + N = resid.shape[0] + resid_sub = resid.reshape(N, m, dsub) # (N, m, dsub) + codes = torch.empty(N, m, dtype=torch.uint8, device=resid.device) + for s in range(m): + cb = codebooks[s].to(torch.float32) # (ksub, dsub) + sub = resid_sub[:, s, :].to(torch.float32) # (N, dsub) + for lo in range(0, N, chunk): + hi = min(lo + chunk, N) + d2 = torch.cdist(sub[lo:hi], cb) ** 2 # (chunk, ksub) + codes[lo:hi, s] = d2.argmin(dim=1).to(torch.uint8) + return codes + + +# ── public build / search ────────────────────────────────────────────────── +def ivf_pq_build( + X: torch.Tensor, + nlist: int, + *, + m: int = 8, + nbits: int = 8, + metric: str = "l2", + by_residual: bool = True, + nprobe: int = 8, + niter: int = 20, + pq_niter: int = 25, + train_size: Optional[int] = None, + pq_train_size: Optional[int] = None, + seed: int = 0, +): + """Build an :class:`IvfPqIndex` with pure torch ops.""" + if metric != "l2": + raise NotImplementedError(f"ivf_pq supports metric='l2' only (got {metric!r})") + if nbits != 8: + raise NotImplementedError(f"ivf_pq supports nbits=8 only (got {nbits})") + if X.ndim != 2: + raise ValueError("ivf_pq build expects a 2D (M, D) tensor") + + M, D = X.shape + nlist = int(min(nlist, M)) + if nlist < 1: + raise ValueError("nlist must be >= 1") + m, dsub, Dp = _pq_dims(int(D), m) + ksub = 1 << nbits # 256 + Xp = _pad_features(X.to(torch.float32), Dp).contiguous() + + # ── coarse quantizer: k-means on a sample ────────────────────────── + train_size = int(train_size or min(M, nlist * 256)) + train_size = max(min(train_size, M), nlist) + g = torch.Generator(device="cpu").manual_seed(seed) + sample_idx = torch.randperm(M, generator=g)[:train_size].to(X.device) + sample = Xp.index_select(0, sample_idx) + centroids = _lloyd_kmeans(sample, nlist, niter=niter, seed=seed) # (nlist, Dp) + + # ── PQ codebooks from residuals of a (sub)sample ─────────────────── + pq_train_size = int(pq_train_size or min(train_size, max(ksub * 16, 4096))) + pq_train_size = max(min(pq_train_size, train_size), ksub) + pq_sample = sample[:pq_train_size] + if by_residual: + pq_assign = _assign_chunked(pq_sample, centroids) + resid_train = pq_sample - centroids.index_select(0, pq_assign) + else: + resid_train = pq_sample + codebooks = _train_pq_codebooks( + resid_train, m, dsub, ksub, niter=pq_niter, seed=seed + 1 + ) # (m, ksub, dsub) + + # ── assign every database row + encode its residual ──────────────── + assign = _assign_chunked(Xp, centroids) # (M,) + resid_all = Xp - centroids.index_select(0, assign) if by_residual else Xp + codes = _encode_pq(resid_all, codebooks, m, dsub) # (M, m) uint8 + + # ── CSR cell-contiguous layout ───────────────────────────────────── + counts = torch.bincount(assign, minlength=nlist) # (nlist,) + offsets = torch.zeros(nlist + 1, dtype=torch.int64, device=X.device) + offsets[1:] = counts.cumsum(0) + order = torch.argsort(assign, stable=True) # (M,) int64 + codes_sorted = codes.index_select(0, order).contiguous() + + return IvfPqIndex( + centroids=centroids, + pq_codebooks=codebooks, + codes=codes_sorted, + ids=order.to(torch.int64), + list_offsets=offsets, + metric=metric, + by_residual=bool(by_residual), + D=int(D), + Dp=int(Dp), + dsub=int(dsub), + m=int(m), + nbits=int(nbits), + nlist=int(nlist), + nprobe=int(nprobe), + max_list_len=int(counts.max().item()) if nlist > 0 else 0, + ) + + +def ivf_pq_search(index, Q: torch.Tensor, k: int, *, nprobe: Optional[int] = None): + """Reference coarse + ADC IVF-PQ search over a built index. + + Returns ``(vals, ids)`` where ``vals[i, j]`` is the (approximate) + squared-L2 distance to the ``j``-th nearest PQ reconstruction and + ``ids`` are original row ids. Mirrors the Triton kernel exactly: + probe the ``nprobe`` nearest lists, build the per-(query, list) + residual LUT, score each member as the sum of ``m`` table lookups, + keep the global top-``k``. + """ + if Q.ndim != 2: + raise ValueError("ivf_pq search expects a 2D (nq, D) query tensor") + nprobe = int(nprobe or index.nprobe) + nprobe = max(1, min(nprobe, index.nlist)) + nq = Q.shape[0] + Dp, m, dsub = index.Dp, index.m, index.dsub + + Qp = _pad_features(Q.to(torch.float32), Dp) + centroids = index.centroids.to(torch.float32) # (nlist, Dp) + codebooks = index.pq_codebooks.to(torch.float32) # (m, ksub, dsub) + codes = index.codes # (M, m) uint8 + offsets = index.list_offsets + + coarse_d2 = torch.cdist(Qp, centroids) ** 2 # (nq, nlist) + probed = coarse_d2.topk(nprobe, dim=1, largest=False).indices # (nq, nprobe) + + out_vals = torch.full((nq, k), float("inf"), device=Q.device, dtype=torch.float32) + out_ids = torch.full((nq, k), -1, device=Q.device, dtype=torch.int64) + + for i in range(nq): + cand_dists = [] + cand_ids = [] + for p in range(nprobe): + c = int(probed[i, p].item()) + s0, e0 = int(offsets[c].item()), int(offsets[c + 1].item()) + if e0 <= s0: + continue + rq = (Qp[i] - centroids[c]) if index.by_residual else Qp[i] # (Dp,) + rq_sub = rq.reshape(m, dsub) # (m, dsub) + # LUT[s, j] = ||rq_s - codebook[s, j]||^2 + lut = ((rq_sub[:, None, :] - codebooks) ** 2).sum(-1) # (m, ksub) + cc = codes[s0:e0].to(torch.int64) # (L, m) + # dist[l] = sum_s LUT[s, cc[l, s]] + dist = lut.gather(1, cc.t().contiguous()).sum(0) # (L,) + cand_dists.append(dist) + cand_ids.append(index.ids[s0:e0]) + if not cand_dists: + continue + dist = torch.cat(cand_dists) + ids = torch.cat(cand_ids) + kk = min(k, dist.shape[0]) + vals, sel = dist.topk(kk, largest=False) + out_vals[i, :kk] = vals + out_ids[i, :kk] = ids[sel] + + return out_vals, out_ids + + +__all__ = ["IvfPqIndex", "ivf_pq_build", "ivf_pq_search"] diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/readme b/2.0/problems/ivf_pq_gpu_kernel_optimization/readme new file mode 100644 index 000000000..6bca22a3c --- /dev/null +++ b/2.0/problems/ivf_pq_gpu_kernel_optimization/readme @@ -0,0 +1,117 @@ +# GPU IVF-PQ Search Kernel Optimization + +## Problem + +You are given a small GPU approximate-nearest-neighbour library, `ivfpqlib`, in +the Harbor workspace at `/app/ivfpqlib`. It implements **IVF-PQ** search. Its +public entry point is: + +```python +ivfpqlib.ivf_pq_search(index, Q, k, *, nprobe=None) -> (vals, ids) +``` + +`index` is a pre-built `IvfPqIndex` (an inverted-file product-quantization index: +`nlist` coarse cells, `m` PQ sub-codebooks, and the database rows stored as `(M, +m)` uint8 PQ codes). `Q` is an `(nq, D)` float32 CUDA tensor of queries, `k` is +the number of neighbours per query, and `nprobe` is how many inverted lists to +probe. It returns `vals` `(nq, k)` float32 (approximate squared-L2 distances to +each neighbour's PQ reconstruction) and `ids` `(nq, k)` int64 (original database +row ids, `-1` padded). The shipped implementation is correct but straightforward. + +Your goal is to make `ivfpqlib.ivf_pq_search` **as fast as possible** on the GPU +while returning the same neighbours. You may rewrite the internals of the package +and add new modules (including Triton kernels) under `ivfpqlib/`. The public +function signature and return contract must not change, and the result must +remain a deterministic function of the inputs. + +## What is (and isn't) timed + +Only `ivf_pq_search` is timed. The `index` is **built once** per workload by a +frozen builder (`ivf_pq_build`) and handed to your search unchanged; index +construction is **not** part of the score, so optimizing `ivf_pq_build` buys you +nothing. Your search must consume the index exactly as built (you may of course +reorganize its arrays inside your own call). Treat the coarse centroids, PQ +codebooks, CSR list layout, and codes as fixed inputs. + +Each call receives its **own fresh copy** of that index — same contents, but a +new object with new tensors every time. Any state you attach to the index or +cache keyed on its tensor addresses is therefore gone by the next call: work +like a CSR-to-padded layout conversion cannot be hoisted into a warmup call and +amortized away. Whatever your search needs, it must produce within the call +being timed. + +## Workload + +The graded workloads are held-out ANN search problems over random float32 +databases: `M` from tens of thousands to a few hundred thousand rows, `D` in the +64–128 range, `nlist` a few hundred lists, `m` sub-quantizers 8–16, `nprobe` +8–32, and `nq` a thousand-plus queries per call. Treat it as general IVF-PQ +search, not as specific indices to special-case. + +The public self-test now runs the **exact graded workloads** — the same shapes, +thresholds, seeds, and timing the judge uses to score you — so there are no +separate hidden shapes to guess at. + +## Iterate on a GPU + +The agent workspace has no GPU. Use the public test to check your current code's correctness and speed on +a GPU through Modal (needs `MODAL_TOKEN_ID` and `MODAL_TOKEN_SECRET`): + +```bash +bash /app/public_test.sh +``` + +It runs the **identical evaluation the judge uses to grade you** and reports, per +graded workload, pass/fail on the quality gate, your speedup, the geometric-mean +speedup, and your **predicted final score** (0-100). Any workload that fails its +gate makes the submission score 0. + +## Submission + +The submitted artifact is a patch over the `ivfpqlib` package: + +```text +/app/solution.patch +``` + +After editing `/app/ivfpqlib`, run `bash /app/make_submission.sh` then +`bash /app/submit.sh`. Submissions are asynchronous; submit early and iterate. +The judge applies your patch to a clean copy of `ivfpqlib` and times it against +the original baseline on a GPU, searching the same seeded index and queries. + +## Correctness + +Correctness is a gate. On every timed iteration the judge searches the same index +with the same queries using both your code and the baseline, and requires the +**iso-result recall@k** — the fraction of your returned ids that also appear in +the baseline's top-`k` for that query — to stay at or above a high threshold. +Crashes, non-finite / wrong-shape output, timeouts, and results that disagree +with the baseline beyond the tolerance are penalized before speed is considered. + +## Scoring + +Valid submissions are scored by speedup relative to the baseline: + +```text +speedup = baseline_time / your_time +``` + +The objective is the geometric mean of per-workload speedups. A result no faster +than the baseline earns 0; regressions earn 0. The raw geometric-mean speedup is +reported in the evaluator metrics. + +## Patch Policy + +Only Python files under `ivfpqlib/**` may change. New Python modules inside +`ivfpqlib/` are allowed. Patches may not: modify anything outside `ivfpqlib/`; +import or call an external optimized ML / ANN / kernel library (write the kernels +yourself); read/write environment variables, spawn processes, or access the +network; or tamper with the measurement framework. The timing harness regenerates +queries every iteration and re-verifies recall against the baseline each time. + +## Resource Budget + +```text +GPU: single Modal GPU (H100 reference; Triton paths also run on L40S / A100) +Agent container: CPU-only (GPU work is offloaded to Modal) +``` diff --git a/2.0/problems/ivf_pq_gpu_kernel_optimization/reference.patch b/2.0/problems/ivf_pq_gpu_kernel_optimization/reference.patch new file mode 100644 index 000000000..9d87e8249 --- /dev/null +++ b/2.0/problems/ivf_pq_gpu_kernel_optimization/reference.patch @@ -0,0 +1,1544 @@ +diff --git a/ivfpqlib/_kernels/__init__.py b/ivfpqlib/_kernels/__init__.py +new file mode 100644 +index 0000000..e69de29 +diff --git a/ivfpqlib/_kernels/_common.py b/ivfpqlib/_kernels/_common.py +new file mode 100644 +index 0000000..2b4c5c4 +--- /dev/null ++++ b/ivfpqlib/_kernels/_common.py +@@ -0,0 +1,8 @@ ++"""Vendored pure-python helper (next-power-of-two, for D / K padding).""" ++ ++ ++def _next_pow2(n: int) -> int: ++ """Smallest power-of-two >= ``max(1, n)`` (D / K padding).""" ++ if n <= 1: ++ return 1 ++ return 1 << (n - 1).bit_length() +diff --git a/ivfpqlib/_kernels/fine_scan.py b/ivfpqlib/_kernels/fine_scan.py +new file mode 100644 +index 0000000..6932237 +--- /dev/null ++++ b/ivfpqlib/_kernels/fine_scan.py +@@ -0,0 +1,238 @@ ++"""IVF-PQ fused ADC fine-scan kernel (elementwise / online path). ++ ++This is the IVF-PQ analogue of the IVF-Flat fused fine-scan: instead of ++streaming full database vectors and computing ``(q - x)^2``, it streams ++the **compressed PQ codes** of a ragged inverted list and recovers each ++candidate's (approximate) squared-L2 distance from the precomputed ++lookup table via asymmetric distance computation (ADC). ++ ++For each ``(query, probed-list[, split])`` the kernel: ++ ++ 1. reads the list's half-open code range ++ ``codes[offsets[c] + lo : offsets[c] + hi]`` of the cell-contiguous ++ ``(M, m)`` uint8 code array (fully coalesced), ++ 2. accumulates ``dist = sum_s LUT[query, probe, s, code[s]]`` -- one ++ HBM gather per sub-quantizer ``s`` into the compact per-(query, ++ list) table built by :mod:`...ivf_pq.triton.lut` -- and ++ 3. maintains an on-chip register top-K with the flash_knn ++ argmin/argmax insert loop. ++ ++Hard contract (identical to IVF-Flat): **no ``(nq x candidates)`` ++distance matrix is ever materialised in HBM** -- the only intermediates ++are the compact ``(nq, P, m, 256)`` LUT and the standard flash-decode ++partial buffer ``(nq, nprobe * n_splits, TOPK_PAD)``, merged host-side ++with a single ``torch.topk``. ++ ++One ``(query, probe, split)`` per program. ``n_splits`` chops long lists ++into wave-filling sub-ranges for the small-batch / online regime -- the ++same wave-count idea as the IVF-Flat / flash_knn decode split. ++""" ++from __future__ import annotations ++ ++from typing import Optional ++ ++import torch ++import triton ++import triton.language as tl ++ ++from ivfpqlib._kernels._common import _next_pow2 ++ ++ ++@triton.jit ++def _ivf_pq_fine_scan_kernel( ++ codes_ptr, probed_ptr, offsets_ptr, lut_ptr, ++ pv_ptr, pi_ptr, ++ stride_codes_m, stride_codes_s, ++ stride_pr_n, stride_pr_p, ++ stride_lut_n, stride_lut_p, stride_lut_s, stride_lut_j, ++ stride_pv_n, stride_pv_p, stride_pv_k, ++ stride_pi_n, stride_pi_p, stride_pi_k, ++ M, ++ MSUB: tl.constexpr, K: tl.constexpr, ++ N_SPLITS: tl.constexpr, ++ BM: tl.constexpr, ++ TOPK_PAD: tl.constexpr, MAX_STEPS: tl.constexpr, MAX_CHUNKS: tl.constexpr, ++): ++ """Grid: ``(nq, nprobe, N_SPLITS)``. One query, one probed list, one split. ++ ++ Writes ``(TOPK_PAD,)`` partial vals/idxs to slot ++ ``pid_p * N_SPLITS + pid_s`` for query ``pid_n``. ``idx`` is the ++ *stored-row* position into ``codes`` (mapped to original ids host-side). ++ """ ++ pid_n = tl.program_id(0).to(tl.int64) ++ pid_p = tl.program_id(1) ++ pid_s = tl.program_id(2) ++ ++ # Which inverted list this (query, probe-slot) scans. ++ c = tl.load(probed_ptr + pid_n * stride_pr_n + pid_p * stride_pr_p).to(tl.int64) ++ start = tl.load(offsets_ptr + c) ++ end = tl.load(offsets_ptr + c + 1) ++ list_len = end - start ++ ++ # Per-split sub-range within the list. ++ split_len = (list_len + N_SPLITS - 1) // N_SPLITS ++ lo = pid_s.to(tl.int64) * split_len ++ hi = tl.minimum(list_len, lo + split_len) ++ ++ # Base of this (query, probe-slot) LUT (stride_lut_p == 0 for non-residual). ++ lut_qp = lut_ptr + pid_n * stride_lut_n + pid_p * stride_lut_p ++ ++ bm_range = tl.arange(0, BM) ++ k_range = tl.arange(0, TOPK_PAD) ++ topk_vals = tl.full([TOPK_PAD], float("inf"), dtype=tl.float32) ++ topk_idx = tl.full([TOPK_PAD], -1, dtype=tl.int32) ++ topk_max = tl.max(topk_vals) ++ ++ for ci in range(MAX_CHUNKS): ++ within = lo + ci * BM + bm_range.to(tl.int64) # (BM,) within-list idx ++ valid = within < hi # (BM,) ++ pos = start + within # (BM,) stored-row pos ++ pos_safe = tl.minimum(tl.maximum(pos, 0), M - 1) ++ ++ # ADC: sum over sub-quantizers of LUT[s, code[s]] (one gather each). ++ acc = tl.zeros([BM], dtype=tl.float32) ++ for s in range(MSUB): ++ code_s = tl.load( ++ codes_ptr + pos_safe * stride_codes_m + s * stride_codes_s, ++ mask=valid, other=0, ++ ).to(tl.int64) # (BM,) in [0, 256) ++ lut_s = tl.load( ++ lut_qp + s * stride_lut_s + code_s * stride_lut_j, ++ mask=valid, other=0.0, ++ ) # (BM,) ++ acc += lut_s ++ ++ score = tl.where(valid, acc, float("inf")) # (BM,) ++ base = start + lo + ci * BM # scalar stored-row base ++ ++ # Iterative argmin-insert into the register top-K (flash_knn body). ++ chunk_best = tl.min(score) ++ if chunk_best < topk_max: ++ _active = tl.full([1], 1, dtype=tl.int32) ++ for _step in range(MAX_STEPS): ++ if tl.max(_active) > 0: ++ row_min = tl.min(score) ++ row_arg = tl.argmin(score, axis=0) ++ if row_min < topk_max: ++ worst = tl.argmax(topk_vals, axis=0) ++ repl = k_range == worst ++ topk_vals = tl.where(repl, row_min, topk_vals) ++ topk_idx = tl.where( ++ repl, (base + row_arg.to(tl.int64)).to(tl.int32), topk_idx ++ ) ++ topk_max = tl.max(topk_vals) ++ score = tl.where(bm_range == row_arg, float("inf"), score) ++ else: ++ _active = tl.full([1], 0, dtype=tl.int32) ++ ++ pslot = pid_p * N_SPLITS + pid_s ++ tl.store( ++ pv_ptr + pid_n * stride_pv_n + pslot * stride_pv_p + k_range * stride_pv_k, ++ topk_vals, ++ ) ++ tl.store( ++ pi_ptr + pid_n * stride_pi_n + pslot * stride_pi_p + k_range * stride_pi_k, ++ topk_idx, ++ ) ++ ++ ++def _pick_n_splits(nq: int, nprobe: int, max_list_len: int, sm_count: int) -> int: ++ """Chop lists into enough splits to roughly fill ~2 waves of SMs. ++ ++ For large query batches ``nq * nprobe`` already saturates the GPU and ++ ``n_splits == 1`` (each program owns a whole list). For tiny batches ++ (online / single-query search) we raise ``n_splits`` so the SMs are ++ not left idle -- the flash-decode wave-targeting idea, on the list axis. ++ """ ++ base = max(1, nq * nprobe) ++ target = 2 * max(sm_count, 1) ++ n_splits = (target + base - 1) // base ++ n_splits = max(1, min(n_splits, 64, max(1, max_list_len))) ++ return int(n_splits) ++ ++ ++def ivf_pq_fine_scan( ++ codes: torch.Tensor, ++ probed: torch.Tensor, ++ list_offsets: torch.Tensor, ++ lut: torch.Tensor, ++ k: int, ++ *, ++ by_residual: bool, ++ max_list_len: int, ++ BM: int = 128, ++ n_splits: Optional[int] = None, ++): ++ """Launch the fused ADC fine-scan + host-side merge. ++ ++ Args: ++ codes: ``(M, m)`` uint8 cell-contiguous PQ codes. ++ probed: ``(nq, nprobe)`` int32 inverted-list ids (from coarse search). ++ list_offsets: ``(nlist + 1,)`` int64 CSR offsets. ++ lut: ``(nq, P, m, ksub)`` fp32 ADC tables (``P = nprobe`` if ++ ``by_residual`` else ``1``). ++ k: neighbours per query. ++ by_residual: selects the LUT probe stride (0 when the LUT is ++ query-only, i.e. non-residual). ++ max_list_len: max inverted-list length (bounds the static chunk loop). ++ ++ Returns: ++ ``(vals, pos)`` -- ``vals`` ``(nq, k)`` ADC squared-L2 (fp32), ++ ``pos`` ``(nq, k)`` int64 stored-row positions into ``codes`` ++ (``-1`` where fewer than ``k`` candidates were available). ++ """ ++ assert codes.is_cuda and probed.is_cuda and list_offsets.is_cuda and lut.is_cuda ++ nq, nprobe = probed.shape ++ M, m = codes.shape ++ ++ codes = codes.contiguous() ++ probed = probed.contiguous().to(torch.int32) ++ list_offsets = list_offsets.contiguous().to(torch.int64) ++ lut = lut.contiguous() ++ ++ if n_splits is None: ++ sm_count = torch.cuda.get_device_properties(codes.device).multi_processor_count ++ n_splits = _pick_n_splits(nq, nprobe, max_list_len, sm_count) ++ n_splits = max(1, min(int(n_splits), max(1, max_list_len))) ++ ++ TOPK_PAD = _next_pow2(k) ++ MAX_STEPS = min(k, BM) ++ per_split_max = (max_list_len + n_splits - 1) // n_splits ++ MAX_CHUNKS = max(1, (per_split_max + BM - 1) // BM) ++ ++ # stride 0 on the probe axis when the LUT depends only on the query. ++ stride_lut_p = lut.stride(1) if by_residual else 0 ++ ++ P = nprobe * n_splits ++ partial_vals = torch.full((nq, P, TOPK_PAD), float("inf"), ++ device=codes.device, dtype=torch.float32) ++ partial_idx = torch.full((nq, P, TOPK_PAD), -1, ++ device=codes.device, dtype=torch.int32) ++ ++ grid = (nq, nprobe, n_splits) ++ _ivf_pq_fine_scan_kernel[grid]( ++ codes, probed, list_offsets, lut, ++ partial_vals, partial_idx, ++ codes.stride(0), codes.stride(1), ++ probed.stride(0), probed.stride(1), ++ lut.stride(0), stride_lut_p, lut.stride(2), lut.stride(3), ++ partial_vals.stride(0), partial_vals.stride(1), partial_vals.stride(2), ++ partial_idx.stride(0), partial_idx.stride(1), partial_idx.stride(2), ++ M, ++ MSUB=m, K=k, ++ N_SPLITS=n_splits, ++ BM=BM, ++ TOPK_PAD=TOPK_PAD, MAX_STEPS=MAX_STEPS, MAX_CHUNKS=MAX_CHUNKS, ++ num_warps=4, ++ ) ++ ++ # Stage-2: merge the P partial top-Ks per query (no HBM cross matrix). ++ pv = partial_vals.view(nq, -1) ++ pi = partial_idx.view(nq, -1) ++ vals, sel = pv.topk(k, dim=-1, largest=False, sorted=True) ++ pos = pi.gather(-1, sel).to(torch.int64) ++ pos = torch.where(vals.isinf(), torch.full_like(pos, -1), pos) ++ return vals, pos ++ ++ ++__all__ = ["ivf_pq_fine_scan", "_pick_n_splits"] +diff --git a/ivfpqlib/_kernels/fine_scan_batch.py b/ivfpqlib/_kernels/fine_scan_batch.py +new file mode 100644 +index 0000000..38b2925 +--- /dev/null ++++ b/ivfpqlib/_kernels/fine_scan_batch.py +@@ -0,0 +1,224 @@ ++"""IVF-PQ fine-scan, throughput variant: group-by-list code-sharing ADC. ++ ++The online :mod:`...ivf_pq.triton.fine_scan` kernel owns one ++``(query, list)`` pair per program and re-reads a list's PQ codes once ++per probing query. For batched search that leaves reuse on the table: ++many queries probe the same list, and the codes (and their layout) can ++be shared across a whole query tile. ++ ++This kernel does what the IVF-Flat GEMM variant does, adapted to ADC: ++ ++ 1. **Group queries by the list they probe** (host-side argsort of the ++ ``(query, list)`` pairs) so every query probing list ``c`` forms a ++ contiguous run. ++ 2. For each ``(list, query-tile)`` it streams the list's ``(BM, m)`` ++ codes **once** and, for the ``BN`` queries in the tile, accumulates ++ the ``(BN, BM)`` ADC score block as ``m`` lookup-table gathers -- ++ each query indexing its *own* precomputed LUT row (the right ++ probe-slot, recovered from the sorted-pair id). A per-query on-chip ++ top-K (the flash_knn 2-D insert body) is reduced in place. ++ ++Unlike the IVF-Flat GEMM kernel, ADC is computed **exactly** here (a sum ++of table lookups, not an x²-free cross term), so there is no oversampled ++re-rank: the returned distances *are* the ADC distances and match the ++online kernel bit-for-bit. Still no ``(nq x candidates)`` HBM matrix -- ++the score block lives in registers and is consumed into the top-K. ++""" ++from __future__ import annotations ++ ++from typing import Tuple ++ ++import torch ++import triton ++import triton.language as tl ++ ++from ivfpqlib._kernels._common import _next_pow2 ++ ++ ++@triton.jit ++def _ivf_pq_fine_batch_kernel( ++ codes_ptr, sorted_qid_ptr, sorted_pslot_ptr, lut_ptr, ++ q_offsets_ptr, list_offsets_ptr, ++ pv_ptr, pi_ptr, ++ stride_codes_m, stride_codes_s, ++ stride_lut_n, stride_lut_p, stride_lut_s, stride_lut_j, ++ stride_pv_p, stride_pv_k, ++ stride_pi_p, stride_pi_k, ++ MSUB: tl.constexpr, K: tl.constexpr, ++ BN: tl.constexpr, BM: tl.constexpr, ++ TOPK_PAD: tl.constexpr, MAX_STEPS: tl.constexpr, MAX_M_CHUNKS: tl.constexpr, ++): ++ """Grid: ``(nlist, MAX_QTILES)``. One ``(list, query-tile)`` per program.""" ++ pid_c = tl.program_id(0) ++ pid_qt = tl.program_id(1) ++ ++ qstart = tl.load(q_offsets_ptr + pid_c) ++ qend = tl.load(q_offsets_ptr + pid_c + 1) ++ qcount = qend - qstart ++ if pid_qt * BN >= qcount: ++ return ++ ++ i_range = tl.arange(0, BN) ++ q_local = pid_qt * BN + i_range ++ q_mask = q_local < qcount # (BN,) ++ pair_pos = (qstart + q_local).to(tl.int64) # (BN,) sorted-pair rows ++ qid = tl.load(sorted_qid_ptr + pair_pos, mask=q_mask, other=0).to(tl.int64) ++ pslot = tl.load(sorted_pslot_ptr + pair_pos, mask=q_mask, other=0).to(tl.int64) ++ # Per-query LUT base offset (probe-slot stride is 0 for non-residual). ++ lut_base = qid * stride_lut_n + pslot * stride_lut_p # (BN,) ++ ++ c_start = tl.load(list_offsets_ptr + pid_c) ++ c_end = tl.load(list_offsets_ptr + pid_c + 1) ++ ++ topk_vals = tl.full([BN, TOPK_PAD], float("inf"), dtype=tl.float32) ++ topk_idxs = tl.full([BN, TOPK_PAD], -1, dtype=tl.int32) ++ topk_max = tl.full([BN], float("inf"), dtype=tl.float32) ++ k_range = tl.arange(0, TOPK_PAD) ++ bm_range = tl.arange(0, BM) ++ ++ for ci in range(MAX_M_CHUNKS): ++ m_start = c_start + ci * BM ++ m_offs = m_start + bm_range.to(tl.int64) # (BM,) stored rows ++ m_mask = m_offs < c_end ++ ++ # ADC score block: shared code stream, per-query LUT gather. ++ acc = tl.zeros([BN, BM], dtype=tl.float32) ++ for s in range(MSUB): ++ code_s = tl.load( ++ codes_ptr + m_offs * stride_codes_m + s * stride_codes_s, ++ mask=m_mask, other=0, ++ ).to(tl.int64) # (BM,) shared across BN ++ off = ( ++ lut_base[:, None] ++ + s * stride_lut_s ++ + code_s[None, :] * stride_lut_j ++ ) # (BN, BM) ++ acc += tl.load( ++ lut_ptr + off, mask=q_mask[:, None] & m_mask[None, :], other=0.0, ++ ) ++ ++ score = tl.where(q_mask[:, None] & m_mask[None, :], acc, float("inf")) ++ ++ chunk_best = tl.min(score) ++ threshold_worst = tl.max(topk_max) ++ if chunk_best < threshold_worst: ++ _active = tl.full([1], 1, dtype=tl.int32) ++ for _step in range(MAX_STEPS): ++ if tl.max(_active) > 0: ++ row_min = tl.min(score, axis=1) ++ row_argmin = tl.argmin(score, axis=1) ++ do_insert = row_min < topk_max ++ n_inserts = tl.sum(do_insert.to(tl.int32)) ++ if n_inserts > 0: ++ topk_argmax = tl.argmax(topk_vals, axis=1) ++ replace_mask = k_range[None, :] == topk_argmax[:, None] ++ insert_mask = do_insert[:, None] & replace_mask ++ topk_vals = tl.where(insert_mask, row_min[:, None], topk_vals) ++ topk_idxs = tl.where( ++ insert_mask, ++ (m_start + row_argmin.to(tl.int64))[:, None].to(tl.int32), ++ topk_idxs, ++ ) ++ topk_max = tl.max(topk_vals, axis=1) ++ used_mask = bm_range[None, :] == row_argmin[:, None] ++ score = tl.where(used_mask & do_insert[:, None], float("inf"), score) ++ _active = tl.where( ++ n_inserts > 0, ++ tl.full([1], 1, dtype=tl.int32), ++ tl.full([1], 0, dtype=tl.int32), ++ ) ++ ++ write_mask = q_mask[:, None] & (k_range[None, :] < TOPK_PAD) ++ tl.store( ++ pv_ptr + pair_pos[:, None] * stride_pv_p + k_range[None, :] * stride_pv_k, ++ topk_vals, mask=write_mask, ++ ) ++ tl.store( ++ pi_ptr + pair_pos[:, None] * stride_pi_p + k_range[None, :] * stride_pi_k, ++ topk_idxs, mask=write_mask, ++ ) ++ ++ ++def _avg_group_size(nq: int, nprobe: int, nlist: int) -> float: ++ """Average number of queries probing a list -- the code-reuse factor.""" ++ return (nq * nprobe) / max(nlist, 1) ++ ++ ++def ivf_pq_fine_scan_batch( ++ codes: torch.Tensor, ++ probed: torch.Tensor, ++ list_offsets: torch.Tensor, ++ lut: torch.Tensor, ++ k: int, ++ *, ++ by_residual: bool, ++ max_list_len: int, ++ BN: int = 64, ++ BM: int = 64, ++) -> Tuple[torch.Tensor, torch.Tensor]: ++ """Group-by-list code-sharing ADC fine scan + host merge. ++ ++ Args mirror :func:`...fine_scan.ivf_pq_fine_scan`. Returns ``(vals, ++ pos)`` with ``vals`` ``(nq, k)`` ADC squared-L2 (fp32, identical to ++ the online kernel) and ``pos`` ``(nq, k)`` int64 stored-row positions ++ into ``codes`` (``-1`` where unavailable). ++ """ ++ assert codes.is_cuda and probed.is_cuda and lut.is_cuda ++ nq, nprobe = probed.shape ++ nlist = list_offsets.shape[0] - 1 ++ device = codes.device ++ ++ codes = codes.contiguous() ++ list_offsets = list_offsets.contiguous().to(torch.int64) ++ lut = lut.contiguous() ++ ++ # ── group (query, list) pairs by list id ─────────────────────────── ++ flat = probed.reshape(-1).contiguous().to(torch.int64) # (P,) pair -> list id ++ P = flat.numel() ++ perm = torch.argsort(flat, stable=True) # sorted-pair -> orig-pair ++ sorted_qid = (perm // nprobe).to(torch.int32) # query id per sorted pair ++ sorted_pslot = (perm % nprobe).to(torch.int32) # probe-slot (LUT row) ++ qcounts = torch.bincount(flat, minlength=nlist) # (nlist,) ++ q_offsets = torch.zeros(nlist + 1, dtype=torch.int64, device=device) ++ q_offsets[1:] = qcounts.cumsum(0) ++ max_qcount = int(qcounts.max().item()) ++ MAX_QTILES = max(1, (max_qcount + BN - 1) // BN) ++ ++ TOPK_PAD = _next_pow2(k) ++ MAX_STEPS = min(k, BM) ++ MAX_M_CHUNKS = max(1, (max_list_len + BM - 1) // BM) ++ stride_lut_p = lut.stride(1) if by_residual else 0 ++ ++ pv_sorted = torch.full((P, TOPK_PAD), float("inf"), device=device, dtype=torch.float32) ++ pi_sorted = torch.full((P, TOPK_PAD), -1, device=device, dtype=torch.int32) ++ ++ grid = (nlist, MAX_QTILES) ++ _ivf_pq_fine_batch_kernel[grid]( ++ codes, sorted_qid, sorted_pslot, lut, ++ q_offsets, list_offsets, ++ pv_sorted, pi_sorted, ++ codes.stride(0), codes.stride(1), ++ lut.stride(0), stride_lut_p, lut.stride(2), lut.stride(3), ++ pv_sorted.stride(0), pv_sorted.stride(1), ++ pi_sorted.stride(0), pi_sorted.stride(1), ++ MSUB=codes.shape[1], K=k, ++ BN=BN, BM=BM, ++ TOPK_PAD=TOPK_PAD, MAX_STEPS=MAX_STEPS, MAX_M_CHUNKS=MAX_M_CHUNKS, ++ num_warps=4, ++ ) ++ ++ # ── scatter partials back to per-query order, then merge ─────────── ++ pv = torch.empty_like(pv_sorted) ++ pi = torch.empty_like(pi_sorted) ++ pv[perm] = pv_sorted ++ pi[perm] = pi_sorted ++ pv = pv.view(nq, nprobe * TOPK_PAD) ++ pi = pi.view(nq, nprobe * TOPK_PAD) ++ ++ vals, sel = pv.topk(k, dim=-1, largest=False, sorted=True) ++ pos = pi.gather(-1, sel).to(torch.int64) ++ pos = torch.where(vals.isinf(), torch.full_like(pos, -1), pos) ++ return vals, pos ++ ++ ++__all__ = ["ivf_pq_fine_scan_batch", "_avg_group_size"] +diff --git a/ivfpqlib/_kernels/fine_scan_gemm.py b/ivfpqlib/_kernels/fine_scan_gemm.py +new file mode 100644 +index 0000000..f6a5ad4 +--- /dev/null ++++ b/ivfpqlib/_kernels/fine_scan_gemm.py +@@ -0,0 +1,416 @@ ++"""IVF-PQ fine-scan, throughput variant: cluster-centric **decode + GEMM**. ++ ++The online :mod:`...ivf_pq.triton.fine_scan` (and group-by-list ++:mod:`...ivf_pq.triton.fine_scan_batch`) kernels score candidates by ++*gathering* from an asymmetric-distance lookup table (ADC LUT): ``m`` ++table reads per ``(query, code)``. That is **gather-throughput bound** -- ++exactly the wall hand-written CUDA (cuVS) climbs with shared-memory LUTs ++that Triton cannot express -- and it forces an ``(nq, nprobe, m, 256)`` ++LUT that blows up with ``nprobe`` (tens of GB). ++ ++This kernel takes the other road the user asked for -- **no LUT** -- and ++turns ADC into a tensor-core GEMM, the same trick IVF-Flat uses: ++ ++ 1. **Coarse + inverse map.** Each query picks ``nprobe`` lists; we ++ argsort the ``(query, list)`` pairs so all queries probing list ++ ``c`` form a contiguous run (host side, in the driver). ++ 2. **Cluster sweep (this kernel).** Grid ``(nlist, query-tile)``. For ++ one ``(list c, query-tile)`` the kernel ++ a. forms the residual query tile ``rq = q - centroid_c`` (or ++ ``rq = q`` when not ``by_residual``), ``(BN, Dp)``; ++ b. streams the list's codes ``(BM, m)`` **once** and *decodes* ++ them to reconstructed sub-vectors ``xhat`` ``(BM, Dp)`` by ++ gathering the (tiny, L2-resident) PQ codebook -- shared across ++ all ``BN`` queries in the tile; ++ c. computes the cross term ``⟨rq, xhat⟩`` as a **WGMMA ``tl.dot``** ++ and the ADC distance ``‖rq‖² + ‖xhat‖² - 2⟨rq, xhat⟩`` (note: ++ ``‖rq‖²`` is kept -- with residual encoding it differs per list, ++ so dropping it would make cross-list partials incomparable); ++ d. reduces a per-query on-chip top-k (flash_knn insert body) and ++ writes one ``(BN, TOPK_PAD)`` partial at the pair's sorted row. ++ 3. **Reduce (driver).** Scatter partials to per-query order, merge the ++ ``nprobe`` partials, and **exact-re-rank** an oversampled pool with ++ :func:`_pq_rerank_kernel` (direct ``‖rq - xhat‖²`` decode) so the ++ returned distances are ADC-exact despite the tf32 GEMM ranking. ++ ++Why this wins: the per-``(query, code)`` LUT gathers become one decode ++gather per *code* (amortised over the whole query tile) plus a tensor-core ++GEMM, so large batches run **3-12x** faster than the LUT path with ++identical recall and ADC-exact distances -- and there is **no LUT**, so ++the ``nprobe``-scaling memory blow-up disappears (partials are only ++``nq*nprobe*k`` floats). ++""" ++from __future__ import annotations ++ ++from typing import Optional, Tuple ++ ++import torch ++import triton ++import triton.language as tl ++ ++from ivfpqlib._kernels._common import _next_pow2 ++ ++ ++@triton.jit ++def _ivf_pq_decode_gemm_kernel( ++ q_ptr, cent_ptr, sorted_qid_ptr, ++ q_offsets_ptr, list_offsets_ptr, ++ codes_ptr, cb_ptr, ++ pv_ptr, pi_ptr, ++ stride_q_n, stride_q_d, ++ stride_cent_c, stride_cent_d, ++ stride_codes_m, stride_codes_s, ++ stride_cb_m, stride_cb_j, stride_cb_d, ++ stride_pv_p, stride_pv_k, ++ stride_pi_p, stride_pi_k, ++ BY_RESIDUAL: tl.constexpr, ++ MSUB: tl.constexpr, DSUB: tl.constexpr, DP: tl.constexpr, D_INNER: tl.constexpr, ++ BN: tl.constexpr, BM: tl.constexpr, ++ TOPK_PAD: tl.constexpr, MAX_STEPS: tl.constexpr, ++): ++ """Grid ``(nlist, MAX_QTILES)``; one ``(list, query-tile)`` per program. ++ ++ Decodes the list's PQ codes to reconstructed sub-vectors and scores the ++ query tile against them with a tensor-core cross term -- no ADC LUT. ++ Writes ``(BN, TOPK_PAD)`` partials (approximate-ADC ranked; the driver ++ re-ranks the pool exactly) to the contiguous sorted-pair rows. ++ """ ++ pid_c = tl.program_id(0) ++ pid_qt = tl.program_id(1) ++ ++ qstart = tl.load(q_offsets_ptr + pid_c) ++ qend = tl.load(q_offsets_ptr + pid_c + 1) ++ qcount = qend - qstart ++ if pid_qt * BN >= qcount: ++ return ++ ++ i_range = tl.arange(0, BN) ++ q_local = pid_qt * BN + i_range ++ q_mask = q_local < qcount # (BN,) ++ pair_pos = (qstart + q_local).to(tl.int64) # (BN,) sorted-pair rows ++ qid = tl.load(sorted_qid_ptr + pair_pos, mask=q_mask, other=0).to(tl.int64) ++ ++ c_start = tl.load(list_offsets_ptr + pid_c) ++ c_end = tl.load(list_offsets_ptr + pid_c + 1) ++ ++ DREAL = MSUB * DSUB ++ ++ # Persistent residual query tile while padded Dp fits one D_INNER tile; ++ # above that the corpus loop re-forms rq in D_INNER-wide chunks. ++ if D_INNER >= DP: ++ d_offs = tl.arange(0, D_INNER) ++ d_mask = d_offs < DREAL ++ s_of_d = d_offs // DSUB ++ o_of_d = d_offs % DSUB ++ q_tile = tl.load( ++ q_ptr + qid[:, None] * stride_q_n + d_offs[None, :] * stride_q_d, ++ mask=q_mask[:, None] & d_mask[None, :], other=0.0, ++ ) ++ if BY_RESIDUAL: ++ cent = tl.load( ++ cent_ptr + pid_c * stride_cent_c + d_offs * stride_cent_d, ++ mask=d_mask, other=0.0, ++ ) ++ rq = q_tile - cent[None, :] ++ else: ++ rq = q_tile ++ rq = tl.where(q_mask[:, None] & d_mask[None, :], rq, 0.0) # (BN, D_INNER) ++ rq_sq = tl.sum(rq * rq, axis=1) # (BN,) ++ ++ topk_vals = tl.full([BN, TOPK_PAD], float("inf"), dtype=tl.float32) ++ topk_idxs = tl.full([BN, TOPK_PAD], -1, dtype=tl.int32) ++ topk_max = tl.full([BN], float("inf"), dtype=tl.float32) ++ k_range = tl.arange(0, TOPK_PAD) ++ bm_range = tl.arange(0, BM) ++ ++ # Per-list chunk count (data-dependent), NOT a global constexpr: lists are ++ # very uneven (SIFT: max 3.5k vs avg ~1k), so looping the longest list's ++ # chunk count for every program would run ~3x masked-empty chunks that ++ # still execute the full decode + tl.dot. Bounding by this list's own ++ # length cuts ~25% wall time (measured). ++ n_chunks = tl.cdiv(c_end - c_start, BM) ++ for ci in range(n_chunks): ++ m_start = c_start + ci * BM ++ m_offs = m_start + bm_range.to(tl.int64) # (BM,) stored rows ++ m_mask = m_offs < c_end ++ ++ if D_INNER >= DP: ++ # decode xhat (BM, D_INNER): xhat[bm,d] = cb[s_of_d, codes[bm, s_of_d], o_of_d] ++ code_col = tl.load( ++ codes_ptr + m_offs[:, None] * stride_codes_m ++ + s_of_d[None, :] * stride_codes_s, ++ mask=m_mask[:, None] & d_mask[None, :], other=0, ++ ).to(tl.int64) ++ xhat = tl.load( ++ cb_ptr + s_of_d[None, :] * stride_cb_m + code_col * stride_cb_j ++ + o_of_d[None, :] * stride_cb_d, ++ mask=m_mask[:, None] & d_mask[None, :], other=0.0, ++ ) # (BM, D_INNER) ++ xhat_sq = tl.sum(xhat * xhat, axis=1) # (BM,) ++ cross = tl.dot(rq, tl.trans(xhat), input_precision="tf32") # (BN, BM) ++ dist = rq_sq[:, None] + xhat_sq[None, :] - 2.0 * cross ++ else: ++ cross = tl.zeros([BN, BM], dtype=tl.float32) ++ xhat_sq = tl.zeros([BM], dtype=tl.float32) ++ rq_sq = tl.zeros([BN], dtype=tl.float32) ++ for d_start in range(0, DP, D_INNER): ++ d_offs = d_start + tl.arange(0, D_INNER) ++ d_mask = d_offs < DREAL ++ s_of_d = d_offs // DSUB ++ o_of_d = d_offs % DSUB ++ q_sub = tl.load( ++ q_ptr + qid[:, None] * stride_q_n + d_offs[None, :] * stride_q_d, ++ mask=q_mask[:, None] & d_mask[None, :], other=0.0, ++ ) ++ if BY_RESIDUAL: ++ cent = tl.load( ++ cent_ptr + pid_c * stride_cent_c + d_offs * stride_cent_d, ++ mask=d_mask, other=0.0, ++ ) ++ rq_sub = q_sub - cent[None, :] ++ else: ++ rq_sub = q_sub ++ rq_sub = tl.where(q_mask[:, None] & d_mask[None, :], rq_sub, 0.0) ++ rq_sq += tl.sum(rq_sub * rq_sub, axis=1) ++ code_col = tl.load( ++ codes_ptr + m_offs[:, None] * stride_codes_m ++ + s_of_d[None, :] * stride_codes_s, ++ mask=m_mask[:, None] & d_mask[None, :], other=0, ++ ).to(tl.int64) ++ xhat_sub = tl.load( ++ cb_ptr + s_of_d[None, :] * stride_cb_m + code_col * stride_cb_j ++ + o_of_d[None, :] * stride_cb_d, ++ mask=m_mask[:, None] & d_mask[None, :], other=0.0, ++ ) ++ xhat_sq += tl.sum(xhat_sub * xhat_sub, axis=1) ++ cross += tl.dot(rq_sub, tl.trans(xhat_sub), input_precision="tf32") ++ dist = rq_sq[:, None] + xhat_sq[None, :] - 2.0 * cross ++ ++ score = tl.where(q_mask[:, None] & m_mask[None, :], dist, float("inf")) ++ ++ chunk_best = tl.min(score) ++ threshold_worst = tl.max(topk_max) ++ if chunk_best < threshold_worst: ++ _active = tl.full([1], 1, dtype=tl.int32) ++ for _step in range(MAX_STEPS): ++ if tl.max(_active) > 0: ++ row_min = tl.min(score, axis=1) ++ row_argmin = tl.argmin(score, axis=1) ++ do_insert = row_min < topk_max ++ n_inserts = tl.sum(do_insert.to(tl.int32)) ++ if n_inserts > 0: ++ topk_argmax = tl.argmax(topk_vals, axis=1) ++ replace_mask = k_range[None, :] == topk_argmax[:, None] ++ insert_mask = do_insert[:, None] & replace_mask ++ topk_vals = tl.where(insert_mask, row_min[:, None], topk_vals) ++ topk_idxs = tl.where( ++ insert_mask, ++ (m_start + row_argmin.to(tl.int64))[:, None].to(tl.int32), ++ topk_idxs, ++ ) ++ topk_max = tl.max(topk_vals, axis=1) ++ used_mask = bm_range[None, :] == row_argmin[:, None] ++ score = tl.where(used_mask & do_insert[:, None], float("inf"), score) ++ _active = tl.where( ++ n_inserts > 0, ++ tl.full([1], 1, dtype=tl.int32), ++ tl.full([1], 0, dtype=tl.int32), ++ ) ++ ++ write_mask = q_mask[:, None] & (k_range[None, :] < TOPK_PAD) ++ tl.store( ++ pv_ptr + pair_pos[:, None] * stride_pv_p + k_range[None, :] * stride_pv_k, ++ topk_vals, mask=write_mask, ++ ) ++ tl.store( ++ pi_ptr + pair_pos[:, None] * stride_pi_p + k_range[None, :] * stride_pi_k, ++ topk_idxs, mask=write_mask, ++ ) ++ ++ ++@triton.jit ++def _pq_rerank_kernel( ++ q_ptr, cent_ptr, codes_ptr, cb_ptr, pos_ptr, clist_ptr, out_ptr, ++ stride_q_n, stride_q_d, ++ stride_cent_c, stride_cent_d, ++ stride_codes_m, stride_codes_s, ++ stride_cb_m, stride_cb_j, stride_cb_d, ++ stride_pos_n, stride_pos_k, ++ stride_out_n, stride_out_k, ++ BY_RESIDUAL: tl.constexpr, MSUB: tl.constexpr, DSUB: tl.constexpr, ++ DP: tl.constexpr, KK: tl.constexpr, ++): ++ """One query per program: exact ADC ``‖rq - xhat‖²`` for its ``KK`` ++ candidate stored-rows (``rq = q - centroid[list_of_candidate]``). ++ ++ Decode is done with the same codebook gather as the GEMM kernel, but ++ distances are accumulated directly (no x²-free cross term) so the ++ result is ADC-exact and immune to the tf32 GEMM rounding used for ++ candidate *selection*. ++ """ ++ i = tl.program_id(0) ++ kk = tl.arange(0, KK) ++ pos = tl.load(pos_ptr + i * stride_pos_n + kk * stride_pos_k).to(tl.int64) # (KK,) ++ clist = tl.load(clist_ptr + i * stride_pos_n + kk * stride_pos_k).to(tl.int64) ++ valid = pos >= 0 ++ ++ d_range = tl.arange(0, DP) ++ s_of_d = d_range // DSUB ++ o_of_d = d_range % DSUB ++ d_mask = d_range < (MSUB * DSUB) ++ ++ q = tl.load(q_ptr + i * stride_q_n + d_range * stride_q_d, mask=d_mask, other=0.0) ++ if BY_RESIDUAL: ++ cent = tl.load( ++ cent_ptr + clist[:, None] * stride_cent_c + d_range[None, :] * stride_cent_d, ++ mask=valid[:, None] & d_mask[None, :], other=0.0, ++ ) ++ rq = q[None, :] - cent ++ else: ++ rq = tl.broadcast_to(q[None, :], (KK, DP)) ++ code_col = tl.load( ++ codes_ptr + pos[:, None] * stride_codes_m + s_of_d[None, :] * stride_codes_s, ++ mask=valid[:, None] & d_mask[None, :], other=0, ++ ).to(tl.int64) ++ xhat = tl.load( ++ cb_ptr + s_of_d[None, :] * stride_cb_m + code_col * stride_cb_j ++ + o_of_d[None, :] * stride_cb_d, ++ mask=valid[:, None] & d_mask[None, :], other=0.0, ++ ) ++ diff = tl.where(d_mask[None, :], rq - xhat, 0.0) ++ dist = tl.sum(diff * diff, axis=1) ++ dist = tl.where(valid, dist, float("inf")) ++ tl.store(out_ptr + i * stride_out_n + kk * stride_out_k, dist) ++ ++ ++def _avg_group_size(nq: int, nprobe: int, nlist: int) -> float: ++ """Average number of queries probing a list -- the GEMM's reuse factor.""" ++ return (nq * nprobe) / max(nlist, 1) ++ ++ ++def ivf_pq_fine_scan_gemm( ++ Qp: torch.Tensor, ++ centroids: torch.Tensor, ++ codebooks: torch.Tensor, ++ codes: torch.Tensor, ++ probed: torch.Tensor, ++ list_offsets: torch.Tensor, ++ k: int, ++ *, ++ by_residual: bool, ++ over: int = 2, ++ BN: int = 64, ++ BM: int = 64, ++ num_stages: int = 2, ++) -> Tuple[torch.Tensor, torch.Tensor]: ++ """Cluster-centric decode+GEMM fine scan, then exact ADC re-rank. ++ ++ Args: ++ Qp: ``(nq, Dp)`` padded fp32 queries. ++ centroids: ``(nlist, Dp)`` coarse centroids (fp32). ++ codebooks: ``(m, ksub, dsub)`` PQ sub-centroids (fp32). ++ codes: ``(M, m)`` uint8 codes, cell-contiguous. ++ probed: ``(nq, nprobe)`` int32 probed list ids. ++ list_offsets: ``(nlist + 1,)`` int64 CSR offsets. ++ k: neighbours per query. ++ by_residual: residual vs direct PQ encoding. ++ over: candidate-pool oversample factor for the exact re-rank. ++ ++ Returns ``(vals, pos)`` with ``vals`` ``(nq, k)`` ADC-exact squared-L2 ++ (fp32) and ``pos`` ``(nq, k)`` int64 stored-row positions into ++ ``codes`` (``-1`` where unavailable). ++ """ ++ assert Qp.is_cuda and codes.is_cuda and centroids.is_cuda ++ nq, Dp = Qp.shape ++ nprobe = probed.shape[1] ++ nlist = list_offsets.shape[0] - 1 ++ m = codes.shape[1] ++ dsub = codebooks.shape[2] ++ device = Qp.device ++ ++ Qp = Qp.contiguous() ++ centroids = centroids.contiguous() ++ codebooks = codebooks.contiguous() ++ codes = codes.contiguous() ++ list_offsets = list_offsets.contiguous().to(torch.int64) ++ ++ # ── inverse map: group (query, list) pairs by list id ────────────── ++ flat = probed.reshape(-1).contiguous().to(torch.int64) # (P,) pair -> list id ++ P = flat.numel() ++ perm = torch.argsort(flat, stable=True) # sorted-pair -> orig-pair ++ sorted_qid = (perm // nprobe).to(torch.int32) ++ qcounts = torch.bincount(flat, minlength=nlist) ++ q_offsets = torch.zeros(nlist + 1, dtype=torch.int64, device=device) ++ q_offsets[1:] = qcounts.cumsum(0) ++ max_qcount = int(qcounts.max().item()) ++ MAX_QTILES = max(1, (max_qcount + BN - 1) // BN) ++ ++ TOPK_PAD = _next_pow2(k) ++ D_INNER = _next_pow2(Dp) if Dp <= 256 else 128 ++ MAX_STEPS = min(k, BM) ++ ++ pv_sorted = torch.full((P, TOPK_PAD), float("inf"), device=device, dtype=torch.float32) ++ pi_sorted = torch.full((P, TOPK_PAD), -1, device=device, dtype=torch.int32) ++ ++ _ivf_pq_decode_gemm_kernel[(nlist, MAX_QTILES)]( ++ Qp, centroids, sorted_qid, ++ q_offsets, list_offsets, ++ codes, codebooks, ++ pv_sorted, pi_sorted, ++ Qp.stride(0), Qp.stride(1), ++ centroids.stride(0), centroids.stride(1), ++ codes.stride(0), codes.stride(1), ++ codebooks.stride(0), codebooks.stride(1), codebooks.stride(2), ++ pv_sorted.stride(0), pv_sorted.stride(1), ++ pi_sorted.stride(0), pi_sorted.stride(1), ++ BY_RESIDUAL=by_residual, ++ MSUB=m, DSUB=dsub, DP=Dp, D_INNER=D_INNER, ++ BN=BN, BM=BM, ++ TOPK_PAD=TOPK_PAD, MAX_STEPS=MAX_STEPS, ++ num_warps=4, num_stages=num_stages, ++ ) ++ ++ # ── scatter partials to per-query order, merge nprobe partials ───── ++ pv = torch.empty_like(pv_sorted) ++ pi = torch.empty_like(pi_sorted) ++ pv[perm] = pv_sorted ++ pi[perm] = pi_sorted ++ pv = pv.view(nq, nprobe * TOPK_PAD) ++ pi = pi.view(nq, nprobe * TOPK_PAD) ++ ++ # Candidate pool by the (tf32-ranked) ADC score, then EXACT re-rank. ++ KK = min(_next_pow2(k * over), pv.shape[1]) ++ _, sel = pv.topk(KK, dim=-1, largest=False, sorted=False) ++ cand = pi.gather(-1, sel).to(torch.int64) # (nq, KK) stored rows ++ if cand.shape[1] < _next_pow2(k * over): ++ pad = torch.full((nq, _next_pow2(k * over) - cand.shape[1]), -1, ++ device=device, dtype=torch.int64) ++ cand = torch.cat([cand, pad], dim=1) ++ KK = cand.shape[1] ++ ++ # List of each candidate (for the residual rq) via CSR searchsorted. ++ clist = (torch.searchsorted(list_offsets, cand, right=True) - 1).clamp_(0, nlist - 1) ++ clist = torch.where(cand >= 0, clist, torch.zeros_like(clist)).to(torch.int32) ++ cand_i32 = cand.to(torch.int32) ++ ++ true_d = torch.empty((nq, KK), device=device, dtype=torch.float32) ++ _pq_rerank_kernel[(nq,)]( ++ Qp, centroids, codes, codebooks, cand_i32, clist, true_d, ++ Qp.stride(0), Qp.stride(1), ++ centroids.stride(0), centroids.stride(1), ++ codes.stride(0), codes.stride(1), ++ codebooks.stride(0), codebooks.stride(1), codebooks.stride(2), ++ cand_i32.stride(0), cand_i32.stride(1), ++ true_d.stride(0), true_d.stride(1), ++ BY_RESIDUAL=by_residual, MSUB=m, DSUB=dsub, ++ DP=_next_pow2(Dp), KK=KK, ++ num_warps=4, ++ ) ++ ++ vals, fsel = true_d.topk(k, dim=-1, largest=False, sorted=True) ++ pos = cand.gather(-1, fsel) ++ pos = torch.where(vals.isinf(), torch.full_like(pos, -1), pos) ++ return vals, pos ++ ++ ++__all__ = ["ivf_pq_fine_scan_gemm", "_avg_group_size"] +diff --git a/ivfpqlib/_kernels/lut.py b/ivfpqlib/_kernels/lut.py +new file mode 100644 +index 0000000..cebd5c5 +--- /dev/null ++++ b/ivfpqlib/_kernels/lut.py +@@ -0,0 +1,147 @@ ++"""IVF-PQ asymmetric distance lookup-table (LUT) builder. ++ ++For a query ``q`` probing list ``c`` (centroid ``cc``) the residual query ++is ``rq = q - cc`` (``by_residual``) or ``rq = q``. The LUT is ++ ++ LUT[s, j] = || rq_s - codebook[s, j] ||^2 (m, ksub) ++ ++i.e. the squared-L2 distance from each of the ``m`` residual-query ++sub-vectors to every one of the ``ksub = 256`` sub-centroids. A ++candidate with codes ``code`` then scores ``sum_s LUT[s, code[s]]`` -- ++this is the asymmetric distance computation (ADC) the fine-scan consumes. ++ ++One Triton program owns one ``(query, probe-slot, sub-quantizer)`` and ++writes that table's ``ksub`` row. The reduction over ``dsub`` is done in ++registers, so the only thing written to HBM is the compact ++``(nq, P, m, ksub)`` LUT itself -- never an ``(nq x candidates)`` matrix. ++``P == nprobe`` for residual encoding (the table depends on the probed ++list's centroid) and ``P == 1`` otherwise (the table depends only on the ++query, so the fine-scan reads it with a probe stride of 0). ++ ++The caller (:mod:`...ivf_pq.triton.search`) invokes this per **query ++tile**, so ``nq`` here is a tile of queries (``q_tile``), not the whole ++batch: the residual LUT grows with ``nprobe`` and would be enormous for a ++big batch (e.g. 42 GB at ``nq=10k, nprobe=64, m=64``), so search tiles ++over queries and only ever materialises one tile's table at a time. ++""" ++from __future__ import annotations ++ ++import torch ++import triton ++import triton.language as tl ++ ++from ivfpqlib._kernels._common import _next_pow2 ++ ++ ++@triton.jit ++def _pq_lut_kernel( ++ q_ptr, centroids_ptr, probed_ptr, codebook_ptr, ++ lut_ptr, ++ stride_q_n, stride_q_d, ++ stride_cent_c, stride_cent_d, ++ stride_pr_n, stride_pr_p, ++ stride_cb_m, stride_cb_j, stride_cb_d, ++ stride_lut_n, stride_lut_p, stride_lut_s, stride_lut_j, ++ BY_RESIDUAL: tl.constexpr, ++ DSUB: tl.constexpr, KSUB: tl.constexpr, ++ BJ: tl.constexpr, BD: tl.constexpr, ++): ++ """Grid: ``(nq, P, m)``. Writes ``lut[pid_n, pid_p, pid_s, 0:KSUB]``.""" ++ pid_n = tl.program_id(0).to(tl.int64) ++ pid_p = tl.program_id(1) ++ pid_s = tl.program_id(2) ++ ++ if BY_RESIDUAL: ++ c = tl.load(probed_ptr + pid_n * stride_pr_n + pid_p * stride_pr_p).to(tl.int64) ++ ++ lut_row = ( ++ lut_ptr ++ + pid_n * stride_lut_n ++ + pid_p * stride_lut_p ++ + pid_s * stride_lut_s ++ ) ++ ++ for j0 in range(0, KSUB, BJ): ++ j_off = j0 + tl.arange(0, BJ) ++ j_mask = j_off < KSUB ++ dist = tl.zeros([BJ], dtype=tl.float32) ++ for d0 in range(0, DSUB, BD): ++ d_off = d0 + tl.arange(0, BD) ++ d_mask = d_off < DSUB ++ d_global = (pid_s * DSUB + d_off).to(tl.int64) # into the Dp-wide row ++ qs = tl.load( ++ q_ptr + pid_n * stride_q_n + d_global * stride_q_d, ++ mask=d_mask, other=0.0, ++ ).to(tl.float32) # (BD,) ++ if BY_RESIDUAL: ++ cs = tl.load( ++ centroids_ptr + c * stride_cent_c + d_global * stride_cent_d, ++ mask=d_mask, other=0.0, ++ ).to(tl.float32) ++ rq = qs - cs ++ else: ++ rq = qs ++ cb = tl.load( ++ codebook_ptr ++ + pid_s * stride_cb_m ++ + j_off[:, None].to(tl.int64) * stride_cb_j ++ + d_off[None, :].to(tl.int64) * stride_cb_d, ++ mask=j_mask[:, None] & d_mask[None, :], other=0.0, ++ ).to(tl.float32) # (BJ, BD) ++ diff = rq[None, :] - cb # (BJ, BD) ++ dist += tl.sum(diff * diff, axis=1) # (BJ,) ++ tl.store(lut_row + j_off * stride_lut_j, dist, mask=j_mask) ++ ++ ++def pq_build_lut( ++ Qp: torch.Tensor, ++ centroids: torch.Tensor, ++ probed: torch.Tensor, ++ codebooks: torch.Tensor, ++ *, ++ by_residual: bool, ++ BJ: int = 64, ++) -> torch.Tensor: ++ """Build the ADC lookup tables. ++ ++ Args: ++ Qp: ``(nq, Dp)`` queries (fp32, padded to ``Dp``). ++ centroids: ``(nlist, Dp)`` coarse centroids (fp32). ++ probed: ``(nq, nprobe)`` int32 probed-list ids (coarse search). ++ codebooks: ``(m, ksub, dsub)`` PQ sub-centroids (fp32). ++ ++ Returns: ++ ``lut`` ``(nq, P, m, ksub)`` fp32 where ``P = nprobe`` if ++ ``by_residual`` else ``1``. ++ """ ++ nq = Qp.shape[0] ++ nprobe = probed.shape[1] ++ m, ksub, dsub = codebooks.shape ++ P = nprobe if by_residual else 1 ++ ++ Qp = Qp.contiguous() ++ centroids = centroids.contiguous() ++ codebooks = codebooks.contiguous() ++ probed = probed.contiguous().to(torch.int32) ++ ++ lut = torch.empty((nq, P, m, ksub), device=Qp.device, dtype=torch.float32) ++ BD = min(_next_pow2(dsub), 64) ++ ++ grid = (nq, P, m) ++ _pq_lut_kernel[grid]( ++ Qp, centroids, probed, codebooks, ++ lut, ++ Qp.stride(0), Qp.stride(1), ++ centroids.stride(0), centroids.stride(1), ++ probed.stride(0), probed.stride(1), ++ codebooks.stride(0), codebooks.stride(1), codebooks.stride(2), ++ lut.stride(0), lut.stride(1), lut.stride(2), lut.stride(3), ++ BY_RESIDUAL=bool(by_residual), ++ DSUB=dsub, KSUB=ksub, ++ BJ=BJ, BD=BD, ++ num_warps=4, ++ ) ++ return lut ++ ++ ++__all__ = ["pq_build_lut"] +diff --git a/ivfpqlib/_kernels/search.py b/ivfpqlib/_kernels/search.py +new file mode 100644 +index 0000000..3680708 +--- /dev/null ++++ b/ivfpqlib/_kernels/search.py +@@ -0,0 +1,382 @@ ++"""IVF-PQ search (Triton/GPU path). Two fine-scan strategies, one API. ++ ++Every stage avoids the ``(nq x candidates)`` HBM matrix; the **coarse** ++step is shared -- :func:`klib.primitives.knn.flash_knn` over the ++``nlist`` centroids picks each query's ``nprobe`` nearest lists. The fine ++scan then takes one of two roads: ++ ++**1. No-LUT decode + GEMM** (``"gemm"``, the default for batched search) ++ The cluster-centric path the database asks for: group queries by the ++ list they probe (inverse map), then per ``(list, query-tile)`` *decode* ++ the list's PQ codes back to sub-vectors (gathering the tiny codebook, ++ shared across the tile) and score them with a tensor-core cross term -- ++ ADC as a **GEMM**, no lookup table at all. Distances are made ADC-exact ++ by an oversampled re-rank. This sidesteps the gather-throughput wall of ++ the LUT scan (3-12x faster on Hopper) *and* removes the LUT entirely, so ++ nothing scales with ``nprobe`` in memory. See ++ :mod:`...ivf_pq.triton.fine_scan_gemm`. ++ ++**2. ADC LUT scan** (``"online"`` / ``"batch"``, best for tiny batches) ++ Build the compact ``(BQ, P, m, 256)`` asymmetric-distance tables ++ (:func:`...ivf_pq.triton.lut.pq_build_lut`) and stream the probed codes, ++ ADC-scoring each candidate against the LUT with an on-chip top-k ++ (:mod:`...ivf_pq.triton.fine_scan`). The only structure that can blow up ++ is this LUT (residual: ``(nq, nprobe, m, 256)``, e.g. 42 GB at ++ ``nq=10k, nprobe=64, m=64``), so -- flash-attention style -- queries are ++ processed in ``q_tile`` blocks, each building and consuming only a ++ ``(q_tile, P, m, 256)`` LUT (bounded by :data:`_LUT_BUDGET_BYTES`) before ++ the next tile starts. The full LUT is never materialised and results are ++ identical to the untiled computation. ++ ++``"auto"`` (default) first picks the road -- the LUT scan for long PQ ++sub-vectors at modest batch, decode+GEMM for short sub-vectors *or* large ++batches (where it amortises each list's decode across the many queries ++probing it) -- then the implementation tier: the hand-written ++CuTe DSL kernels on Hopper (``cute_lut`` / ``cute_gemm``, ++:mod:`...ivf_pq.cutedsl`), the portable Triton kernels elsewhere ++(``online`` / ``gemm``). See :func:`_pick_variant`. At a fixed ++``(nlist, nprobe)`` and codebooks all variants return the same ADC ++ranking/distances (to fp tolerance) as a reference IVF-PQ; only the ++kernel implementation differs. ++""" ++from __future__ import annotations ++ ++from functools import lru_cache ++from typing import Optional ++ ++import torch ++ ++def is_cutedsl_available(): ++ # CuTe DSL Hopper tier disabled in the vendored path; the portable ++ # Triton LUT-scan / decode+GEMM kernels below are used everywhere. ++ return False ++from ivfpqlib.index import IvfPqIndex ++from ivfpqlib.build import _pad_features ++from ivfpqlib._kernels.fine_scan import ivf_pq_fine_scan ++from ivfpqlib._kernels.fine_scan_batch import ivf_pq_fine_scan_batch ++from ivfpqlib._kernels.fine_scan_gemm import ivf_pq_fine_scan_gemm ++from ivfpqlib._kernels.lut import pq_build_lut ++ ++ ++# Cap on the *live* ADC LUT (the only thing that scales with nq*nprobe). ++# Queries are tiled so a tile's (q_tile, P, m, 256) fp32 table stays under ++# this; 2 GiB keeps the table HBM/L2-friendly (and bounds the pathological ++# 42 GB residual LUT) while large enough that the tiling costs ~1% vs the ++# untiled path on typical batches. (Only the LUT variants tile; the default ++# "gemm" path builds no LUT and never needs it.) ++ ++def _coarse_probe(Qp, centroids, nprobe): ++ """Exact ``nprobe`` nearest coarse centroids per (padded) query. ++ ++ IVF-PQ's coarse step over the ``nlist`` centroids. Exact and a negligible ++ fraction of total search time, so a torch ``cdist`` + ``topk`` returns the ++ same probed lists as any brute-force KNN kernel. ++ """ ++ d2 = torch.cdist(Qp, centroids) ** 2 # (nq, nlist) ++ return d2.topk(nprobe, dim=1, largest=False).indices.to(torch.int32) ++ ++_LUT_BUDGET_BYTES = 1 << 31 # 2 GiB ++ ++# Decode+GEMM has a higher fixed floor than the LUT scan (~0.9 ms vs ++# ~0.45 ms: a host argsort of the nq*nprobe pairs, extra launches, an ++# exact re-rank), so for *tiny* batches / little total work the LUT scan ++# wins outright regardless of geometry: nq must clear _GEMM_MIN_NQ and the ++# candidate comparisons -- nq * nprobe * (M / nlist) -- must clear ++# _GEMM_MIN_WORK before the GEMM floor can be repaid. Calibrated on Hopper. ++_GEMM_MIN_NQ = 256 ++_GEMM_MIN_WORK = 2_000_000 ++ ++# Past the floor, routing is a crossover calibrated on Hopper (D=128/256/960, ++# re-swept after the coalesced ``cute_lut`` redesign roughly doubled its ++# crossover). Per probed candidate the LUT does ``m`` gathers; decode+GEMM ++# reconstructs ``D = m*dsub`` dims on the tensor cores, decoding each list ++# once and amortising it over the queries-per-list ``qpl = nq*nprobe/nlist``. ++# The LUT wins when both hold: ++# * sub-vectors aren't too short -- ``dsub >= _DSUB_LUT_MIN`` (tiny dsub ++# decodes cheaply on the tensor cores; e.g. SIFT dsub<=8 -> GEMM), and ++# * the batch per list is small enough -- ``qpl <= _QPL_LUT_SLOPE * dsub`` ++# (decode+GEMM's win region grows ~linearly in qpl). ++# Two geometries pick the LUT regardless of qpl (measured LUT-win out to the ++# swept qpl=512): ++# * many sub-quantizers ``m >= _M_LUT_MIN`` -- decode+GEMM's reconstruction ++# loop is unrolled over m and becomes the bottleneck (GIST m=64: cute_lut ++# ~2.6-6x faster than cute_gemm), and ++# * long sub-vectors ``dsub >= _DSUB_LUT_ALWAYS`` -- so few gathers per ++# candidate that the LUT stays ahead even at large batch. ++# vs the old ``qpl<=2*dsub`` this cuts mis-routes 13->3 / 48 swept points and ++# worst-case regret 294%->31% (the 294% case was GIST m=64 routed to GEMM). ++_DSUB_LUT_MIN = 9 ++_QPL_LUT_SLOPE = 4.0 ++_M_LUT_MIN = 48 ++_DSUB_LUT_ALWAYS = 48 ++ ++ ++def _auto_q_tile(nq: int, nprobe: int, m: int, by_residual: bool) -> int: ++ """Largest query tile whose LUT fits the budget (>= 256, <= nq).""" ++ P = nprobe if by_residual else 1 ++ per_query = P * m * 256 * 4 # fp32 LUT bytes for one query ++ bq = _LUT_BUDGET_BYTES // max(per_query, 1) ++ return int(max(256, min(nq, bq))) ++ ++ ++@lru_cache(maxsize=None) ++def _cutedsl_hopper() -> bool: ++ """True iff the CuTe DSL fine-scan kernels can run on this machine. ++ ++ They are hand-written for Hopper (SM90 WGMMA / shared-memory gathers) ++ and need the CUTLASS Python DSL; otherwise the router falls back to the ++ portable Triton kernels. Device arch is fixed per process, so cache it. ++ """ ++ if not is_cutedsl_available(): ++ return False ++ try: ++ return ( ++ torch.cuda.is_available() ++ and torch.cuda.get_device_properties(0).major >= 9 ++ ) ++ except Exception: ++ return False ++ ++ ++def _pick_regime( ++ nq: int, nprobe: int, avg_list_len: float, dsub: int, m: int, nlist: int, ++) -> str: ++ """Pick the fine-scan road: ``"lut"`` (ADC gather) or ``"gemm"``. ++ ++ Tiny batches / low total work don't amortise the GEMM floor -> LUT. ++ Short sub-vectors (``dsub < _DSUB_LUT_MIN``) decode cheaply on the tensor ++ cores -> GEMM. Many sub-quantizers (``m >= _M_LUT_MIN``) or long ++ sub-vectors (``dsub >= _DSUB_LUT_ALWAYS``) keep the LUT ahead at any ++ batch. Otherwise the LUT wins while the batch per list is small enough ++ (``qpl <= _QPL_LUT_SLOPE * dsub``). ++ """ ++ work = nq * nprobe * max(avg_list_len, 1.0) ++ if nq < _GEMM_MIN_NQ or work < _GEMM_MIN_WORK: ++ return "lut" ++ if dsub < _DSUB_LUT_MIN: ++ return "gemm" ++ if m >= _M_LUT_MIN or dsub >= _DSUB_LUT_ALWAYS: ++ return "lut" ++ qpl = nq * nprobe / max(nlist, 1) ++ if qpl <= _QPL_LUT_SLOPE * dsub: ++ return "lut" ++ return "gemm" ++ ++ ++def _pick_variant( ++ variant: str, nq: int, nprobe: int, avg_list_len: float, dsub: int, ++ m: int, nlist: int, ++) -> str: ++ """Resolve ``variant`` to a concrete fine-scan kernel. ++ ++ Explicit names pass through; ``"auto"`` first chooses the road ++ (:func:`_pick_regime`) then the implementation tier -- the fast CuTe ++ DSL kernels on Hopper, the portable Triton kernels elsewhere. ++ """ ++ if variant in ("gemm", "batch", "online", "cute_lut", "cute_gemm"): ++ return variant ++ if variant != "auto": ++ raise ValueError( ++ f"unknown variant {variant!r} " ++ "(auto|gemm|batch|online|cute_lut|cute_gemm)" ++ ) ++ regime = _pick_regime(nq, nprobe, avg_list_len, dsub, m, nlist) ++ if _cutedsl_hopper(): ++ return "cute_lut" if regime == "lut" else "cute_gemm" ++ return "online" if regime == "lut" else "gemm" ++ ++ ++def _search_gemm( ++ index: IvfPqIndex, ++ Qp: torch.Tensor, ++ centroids: torch.Tensor, ++ codebooks: torch.Tensor, ++ k: int, ++ nprobe: int, ++): ++ """No-LUT cluster-centric decode+GEMM search over the whole batch. ++ ++ Builds no ADC LUT, so there is nothing that scales with ``nprobe`` to ++ tile -- the only intermediate is the ``(nq*nprobe, k)`` partial table. ++ Returns ``(vals, ids)``. ++ """ ++ probed = _coarse_probe(Qp, centroids, nprobe) # (nq, nprobe) ++ vals, pos = ivf_pq_fine_scan_gemm( ++ Qp, centroids, codebooks, index.codes, probed, index.list_offsets, k, ++ by_residual=index.by_residual, ++ ) # (nq, k) ++ valid = pos >= 0 ++ pos_safe = pos.clamp_min(0) ++ ids = torch.where(valid, index.ids[pos_safe], torch.full_like(pos, -1)) ++ return vals, ids ++ ++ ++def _search_cute( ++ index: IvfPqIndex, ++ Qp: torch.Tensor, ++ centroids: torch.Tensor, ++ codebooks: torch.Tensor, ++ k: int, ++ nprobe: int, ++ method: str, ++): ++ """CuTe DSL fine scan: shared-memory ADC LUT (``cute_lut``) or ++ decode+WGMMA GEMM (``cute_gemm``). Coarse + reduce mirror the Triton ++ ``"gemm"`` path; only the fine-scan kernel differs. Returns ``(vals, ids)``. ++ """ ++ # Lazy import so non-CUTLASS environments still load the Triton path. ++ if method == "cute_lut": ++ from klib.primitives.ivf_pq.cutedsl.shared_lut import ( ++ ivf_pq_fine_scan_shared_lut as _fine, ++ ) ++ else: ++ from klib.primitives.ivf_pq.cutedsl.decode_gemm import ( ++ ivf_pq_fine_scan_decode_gemm as _fine, ++ ) ++ probed = _coarse_probe(Qp, centroids, nprobe) # (nq, nprobe) ++ vals, pos = _fine( ++ Qp, centroids, codebooks, index.codes, probed, index.list_offsets, k, ++ by_residual=index.by_residual, ++ ) ++ valid = pos >= 0 ++ pos_safe = pos.clamp_min(0) ++ ids = torch.where(valid, index.ids[pos_safe], torch.full_like(pos, -1)) ++ return vals, ids ++ ++ ++def _search_tile( ++ index: IvfPqIndex, ++ Qp: torch.Tensor, ++ centroids: torch.Tensor, ++ codebooks: torch.Tensor, ++ k: int, ++ nprobe: int, ++ variant: str, ++ max_list_len: int, ++): ++ """Coarse + LUT + fine-scan for one (already padded) query tile. ++ ++ Builds and consumes a single ``(BQ, P, m, 256)`` LUT, so the live ++ table is bounded by the tile size. Returns ``(vals, ids)``. ++ """ ++ # ── coarse: nprobe nearest centroids (lists) per query ───────────── ++ probed = _coarse_probe(Qp, centroids, nprobe) # (BQ, nprobe) ++ ++ # ── ADC lookup tables (compact, per-tile, no candidate matrix) ───── ++ lut = pq_build_lut( ++ Qp, centroids, probed, codebooks, by_residual=index.by_residual, ++ ) # (BQ, P, m, ksub) ++ ++ # ── fine: fused ragged-code scan + on-chip top-k ─────────────────── ++ # ``variant`` is already resolved to "online"/"batch" by the driver. ++ chosen = "batch" if variant == "batch" else "online" ++ if chosen == "batch": ++ vals, pos = ivf_pq_fine_scan_batch( ++ index.codes, probed, index.list_offsets, lut, k, ++ by_residual=index.by_residual, max_list_len=max_list_len, ++ ) ++ else: ++ vals, pos = ivf_pq_fine_scan( ++ index.codes, probed, index.list_offsets, lut, k, ++ by_residual=index.by_residual, max_list_len=max_list_len, ++ ) # (BQ, k) ++ ++ # Map stored-row positions back to original ids (guard -1 padding). ++ valid = pos >= 0 ++ pos_safe = pos.clamp_min(0) ++ ids = torch.where(valid, index.ids[pos_safe], torch.full_like(pos, -1)) ++ return vals, ids ++ ++ ++def ivf_pq_search_triton( ++ index: IvfPqIndex, ++ Q: torch.Tensor, ++ k: int, ++ *, ++ nprobe: Optional[int] = None, ++ variant: str = "auto", ++ q_tile: Optional[int] = None, ++): ++ """Search a built IVF-PQ index. Returns ``(vals, ids)``. ++ ++ Args: ++ index: a built :class:`IvfPqIndex`. ++ Q: ``(nq, D)`` query tensor on CUDA. ++ k: neighbours per query. ++ nprobe: lists to probe (defaults to ``index.nprobe``). ++ variant: fine-scan kernel. ``"auto"`` (default) routes by ++ sub-vector length and batch size to the best available kernel ++ (CuTe DSL on Hopper, Triton elsewhere; see ++ :func:`_pick_variant`). Explicit names: ``"cute_lut"`` / ++ ``"cute_gemm"`` are the Hopper CuTe DSL LUT / decode+GEMM ++ kernels; ``"gemm"`` is the portable Triton no-LUT decode+GEMM; ++ ``"online"`` / ``"batch"`` are the portable Triton ADC-LUT ++ gather kernels (best for tiny batches). All variants return the ++ same ADC distances to fp tolerance. ++ q_tile: queries per LUT tile (LUT variants only -- the ``"gemm"`` ++ path builds no LUT and ignores it). ``None`` picks the largest ++ tile whose ``(q_tile, P, m, 256)`` LUT fits the internal budget ++ (so the full LUT is never materialised); pass an int to override. ++ ++ Returns: ++ ``vals`` ``(nq, k)`` ADC squared-L2 (fp32) and ``ids`` ``(nq, k)`` ++ int64 original row ids (``-1`` padded where unavailable). ++ """ ++ if not Q.is_cuda or Q.ndim != 2: ++ raise ValueError("ivf_pq_search_triton requires a 2D CUDA tensor") ++ nprobe = int(nprobe or index.nprobe) ++ nprobe = max(1, min(nprobe, index.nlist)) ++ if not (1 <= k <= index.M): ++ raise ValueError(f"k must be in [1, M={index.M}] (got {k})") ++ ++ nq = Q.shape[0] ++ Qp_all = _pad_features(Q.to(torch.float32), index.Dp).contiguous() # (nq, Dp) ++ centroids = index.centroids.to(torch.float32) ++ codebooks = index.pq_codebooks.to(torch.float32) ++ max_list_len = index.max_list_len or int(index.list_lengths().max().item()) ++ avg_list_len = index.M / max(index.nlist, 1) ++ ++ # No-LUT decode+GEMM path: no LUT to materialise, so no query tiling. ++ chosen = _pick_variant( ++ variant, nq, nprobe, avg_list_len, index.dsub, index.m, index.nlist, ++ ) ++ # Clustered real datasets (SIFT/GIST) produce heavily imbalanced inverted ++ # lists. The decode+GEMM fine-scan paths ("gemm"/"cute_gemm"/"cute_lut") ++ # size a per-(query, probe) candidate buffer for near-balanced lists and ++ # read out of bounds on a long-tail list; the streaming ADC-LUT "online" ++ # path is length-safe. When the longest list far exceeds the average, ++ # route auto there. ++ if variant == "auto" and chosen not in ("online", "batch") \ ++ and max_list_len > 4.0 * avg_list_len: ++ chosen = "online" ++ if chosen == "gemm": ++ return _search_gemm(index, Qp_all, centroids, codebooks, k, nprobe) ++ if chosen in ("cute_lut", "cute_gemm"): ++ return _search_cute(index, Qp_all, centroids, codebooks, k, nprobe, chosen) ++ variant = chosen # "online" / "batch" -- passed straight to _search_tile ++ ++ if q_tile is None: ++ q_tile = _auto_q_tile(nq, nprobe, index.m, index.by_residual) ++ q_tile = max(1, min(int(q_tile), nq)) ++ ++ # Single tile: no extra allocations / copies (identical to untiled). ++ if q_tile >= nq: ++ return _search_tile( ++ index, Qp_all, centroids, codebooks, k, nprobe, variant, max_list_len, ++ ) ++ ++ # Flash-style query tiling: build + consume one LUT tile at a time. ++ out_vals = torch.empty((nq, k), device=Q.device, dtype=torch.float32) ++ out_ids = torch.empty((nq, k), device=Q.device, dtype=torch.int64) ++ for lo in range(0, nq, q_tile): ++ hi = min(lo + q_tile, nq) ++ vals, ids = _search_tile( ++ index, Qp_all[lo:hi].contiguous(), centroids, codebooks, ++ k, nprobe, variant, max_list_len, ++ ) ++ out_vals[lo:hi] = vals ++ out_ids[lo:hi] = ids ++ return out_vals, out_ids ++ ++ ++__all__ = ["ivf_pq_search_triton"] +diff --git a/ivfpqlib/search.py b/ivfpqlib/search.py +index 56b559b..8b32e5d 100644 +--- a/ivfpqlib/search.py ++++ b/ivfpqlib/search.py +@@ -1,78 +1,19 @@ +-"""Naive IVF-PQ search (pure torch reference -- the baseline you optimize). ++"""IVF-PQ search -- vendored best-in-repo Triton fine-scan path. + +-Coarse step: probe the ``nprobe`` nearest lists per query. Fine step: for each +-query, build the per-list residual ADC lookup table and score every probed +-candidate as ``sum_s LUT[s, code[s]]``, then keep the global top-``k``. The +-Python per-query loop is what makes this slow; a fast GPU implementation fuses +-the coarse search, LUT build, and ragged candidate scan into kernels. ++Delegates to the fused Triton LUT-scan / decode+GEMM kernels vendored under ++``ivfpqlib._kernels``. The coarse nprobe-nearest-list step is an exact torch ++cdist+topk. Public contract is unchanged: + + ivf_pq_search(index, Q, k, *, nprobe=None) -> (vals (nq,k) fp32, + ids (nq,k) int64, -1 padded) + """ + from __future__ import annotations + +-from typing import Optional ++from ivfpqlib._kernels.search import ivf_pq_search_triton + +-import torch + +-from ivfpqlib.build import _pad_features +- +-def ivf_pq_search(index, Q: torch.Tensor, k: int, *, nprobe: Optional[int] = None): +- """Reference coarse + ADC IVF-PQ search over a built index. +- +- Returns ``(vals, ids)`` where ``vals[i, j]`` is the (approximate) +- squared-L2 distance to the ``j``-th nearest PQ reconstruction and +- ``ids`` are original row ids. Mirrors the Triton kernel exactly: +- probe the ``nprobe`` nearest lists, build the per-(query, list) +- residual LUT, score each member as the sum of ``m`` table lookups, +- keep the global top-``k``. +- """ +- if Q.ndim != 2: +- raise ValueError("ivf_pq search expects a 2D (nq, D) query tensor") +- nprobe = int(nprobe or index.nprobe) +- nprobe = max(1, min(nprobe, index.nlist)) +- nq = Q.shape[0] +- Dp, m, dsub = index.Dp, index.m, index.dsub +- +- Qp = _pad_features(Q.to(torch.float32), Dp) +- centroids = index.centroids.to(torch.float32) # (nlist, Dp) +- codebooks = index.pq_codebooks.to(torch.float32) # (m, ksub, dsub) +- codes = index.codes # (M, m) uint8 +- offsets = index.list_offsets +- +- coarse_d2 = torch.cdist(Qp, centroids) ** 2 # (nq, nlist) +- probed = coarse_d2.topk(nprobe, dim=1, largest=False).indices # (nq, nprobe) +- +- out_vals = torch.full((nq, k), float("inf"), device=Q.device, dtype=torch.float32) +- out_ids = torch.full((nq, k), -1, device=Q.device, dtype=torch.int64) +- +- for i in range(nq): +- cand_dists = [] +- cand_ids = [] +- for p in range(nprobe): +- c = int(probed[i, p].item()) +- s0, e0 = int(offsets[c].item()), int(offsets[c + 1].item()) +- if e0 <= s0: +- continue +- rq = (Qp[i] - centroids[c]) if index.by_residual else Qp[i] # (Dp,) +- rq_sub = rq.reshape(m, dsub) # (m, dsub) +- # LUT[s, j] = ||rq_s - codebook[s, j]||^2 +- lut = ((rq_sub[:, None, :] - codebooks) ** 2).sum(-1) # (m, ksub) +- cc = codes[s0:e0].to(torch.int64) # (L, m) +- # dist[l] = sum_s LUT[s, cc[l, s]] +- dist = lut.gather(1, cc.t().contiguous()).sum(0) # (L,) +- cand_dists.append(dist) +- cand_ids.append(index.ids[s0:e0]) +- if not cand_dists: +- continue +- dist = torch.cat(cand_dists) +- ids = torch.cat(cand_ids) +- kk = min(k, dist.shape[0]) +- vals, sel = dist.topk(kk, largest=False) +- out_vals[i, :kk] = vals +- out_ids[i, :kk] = ids[sel] +- +- return out_vals, out_ids ++def ivf_pq_search(index, Q, k, *, nprobe=None): ++ return ivf_pq_search_triton(index, Q, k, nprobe=nprobe, variant="gemm") # pin fast WGMMA gemm path (auto under-picks it) + + + __all__ = ["ivf_pq_search"] diff --git a/2.0/problems/kmeans_gpu_kernel_optimization/DESIGN.md b/2.0/problems/kmeans_gpu_kernel_optimization/DESIGN.md new file mode 100644 index 000000000..db5f9fd8a --- /dev/null +++ b/2.0/problems/kmeans_gpu_kernel_optimization/DESIGN.md @@ -0,0 +1,148 @@ +# Design notes — kmeans_gpu_kernel_optimization + +Operator-facing. Not copied into the agent workspace by the adapter. + +## What this task measures + +Can an agent turn a naive GPU K-Means into a fast one? The agent patches +`kmeanslib`, whose public entry point is **one Lloyd iteration** +`step(x, centroids) -> (labels, new_centroids)`. **The judge owns the loop**: it +fixes the data + initial centroids and calls `step` a fixed number of times +(`max_iters`, feeding each output into the next), times the patched `step` +against a frozen naive baseline on hidden `(N, D, K)` workloads, and scores the +geometric-mean speedup, gated on clustering quality (inertia of the final +centroids). First of a four-task **flashlib kernel-optimization family** (KMeans, +KNN, TruncatedSVD, PCA) that all share one evaluator + one Modal GPU harness. + +## Execution: Modal GPU offload (mirrors the vllm task) + +The judge and the agent public test both **offload GPU work to Modal**; the +containers are light (ubuntu + `modal` + git, no torch/triton). `flash_gpu.py` +(baked at `/opt/flash_gpu.py` in both images) builds an ephemeral Modal app on a +GPU (`evaluation.gpu`, default H100), ships the frozen baseline + the patched +package **as data** (no per-submission image rebuild), runs the worker, and +returns per-workload speedups + verdicts. Needs `MODAL_TOKEN_ID` / +`MODAL_TOKEN_SECRET`. torch/triton/the vendored kernels run on the Modal image +(`evaluation.pip`). + +The evaluator is generic and identical across the four tasks; only three +constants differ (`PRIMITIVE`/`PKG`/`REF_MODULE`), and all workload/threshold/ +Modal values come from `config.yaml` (delivered judge-only as +`/judge/task_config.json`, never present in the agent image). + +## Reference solution = best in the repo + +`reference.patch` vendors flashlib's own optimized **Triton** K-Means path +(`primitives/kmeans/triton/{kmeans,assign,update}.py`) under +`kmeanslib/_kernels/…`, with imports rewritten and the string "flashlib" +scrubbed, plus a thin adapter mapping our `step(x, centroids)` contract onto +flashlib's single-iteration path (`_euclid_iter`: `euclid_assign_triton` + +`triton_centroid_update_sorted_euclid`, run once with the batch dim as `B=1`). +It is triton-only (runs on any sm≥80), imports cleanly on CPU, and applies + +passes the patch policy. Calibrate `speedup_target` from its measured geomean on +the first GPU trial. (The reference does not fuse assign+update into a single +kernel, so an agent that does can legitimately score >100.) + +## Anti-reward-hack + +The **primary, load-bearing** defense is the **`step` contract itself**: the +agent owns only one Lloyd iteration `step(x, centroids)`; the judge owns the +loop, the data, the initial centroids, and the iteration count. This structurally +closes the reward-hacks earlier trials found when the agent owned the whole +`kmeans(...)` call: + +* **iteration-skip** — the agent used to run fewer than `max_iters` iterations + and self-report the count; now the judge does the counting (`for _ in + range(max_iters): lab, c = mod.step(x, c)`), so the count is not the agent's to + fake; +* **fake / subsampled labels** — the judge re-derives inertia from the *final + centroids after its own loop*; a `step` that returns garbage labels or + subsamples rows/dims produces bad centroids that compound over the loop and + spike inertia. `max_iters = 2` (not 1) makes it a genuine loop where a bad step + actually propagates, while still being too few for a converged step to degrade + into a no-op; +* **precision** — data is bf16 (locked in `gen`), so fp16/tf32/fp32 buy nothing. + +Additional structural defenses inside the GPU worker (the only place the +untrusted submission runs): + +* timing primitives (`perf_counter`, `torch.cuda.synchronize`) are captured to + locals **before** the submission is imported → monkey-patching torch cannot + affect measurement; +* quality is **re-verified by the judge from the returned tensors** → no agent + number is trusted; +* an ephemeral Modal container per submission → no global state persists. + +The patch policy is surgical defense-in-depth (so it does not reject legitimate +vendored/optimized kernel source): only `/**` `.py` files may change; it +bans external-optimized-library *imports* (flashlib/cuml/cupy/faiss/sklearn/ +cutlass — none installed on the GPU image anyway) and process/network/ +measurement-tamper patterns (`subprocess`, `socket`, `os.system`, +`torch.cuda.synchronize =`, …). + +The agent-facing `readme` describes the contract, gate, scoring, and policy but +**not** how the reference is implemented. + +## Correctness gate — centroid inertia (the step contract does the rest) + +Under the `step` contract the earlier subsample / fake-label hacks are closed +**structurally** (the judge owns the loop and re-derives quality from the final +centroids — see Anti-reward-hack), so the separate tight *label* gate is no +longer needed. One gate remains: + +* **Planted clusters** (not random): `gen` builds `x` = well-separated blob + centres (`centers * 6`) + unit noise. This makes the assignment *matter* — a + step that assigns points to the wrong blob lands the final centroids in the + wrong place and spikes inertia. (Random data leaves inertia nearly + assignment-invariant, which is why it was hackable.) +* **Centroid inertia gate** (`inertia_tolerance`, recalibrate): after the judge's + loop, the agent's final centroids' nearest-centroid inertia must be + `<= (1+tol) * baseline`. Computed in fp32 (upcast) on identical data. A step + that does less real work than a true assign+update lands worse centroids after + two iterations and fails this. + +The gate is permutation/convention-independent and cheat-proof. `init_centroids` +are always supplied by the judge, so the baseline is deterministic in +`(x, init_centroids, max_iters)`. + +## Precision lock — bf16 (task-root, not prompt) + +To stop an agent from buying speed by dropping matmul precision (the earlier +codex solution picked fp16 over tf32 where the 2% inertia tolerance allowed it), +the **data is bf16**: `flash_gpu`'s `gen` returns `x` and `init_centroids` as +`bfloat16`. Because the inputs carry only bf16 precision, no solution — baseline, +reference, or agent — gains anything from fp16/tf32/fp32 (they are only slower on +bf16 data), so the only remaining lever is the kernel structure. The naive +baseline's `_assign` was switched from `torch.cdist` (which upcasts) to an +explicit **bf16 matmul** so the baseline is genuinely bf16-native too — otherwise +a trivial "use bf16 instead of the fp32 baseline" would win without kernel work. +The inertia metric itself is still computed in fp32 (upcast) so the gate is a +clean quality comparison. + +## Timing unit — the judge-owned loop + +The timed unit is the **judge's loop of `max_iters` `step` calls** (warmup +`warmup_iters`, then `timed_iters = 1` timed run) — baseline and agent time the +identical loop on the same seeded data + init. `max_iters = 2`: a genuine loop +(so a bad step propagates and is caught by the inertia gate) but too few for a +converged step to become a redundant no-op. The score isolates per-`step` kernel +efficiency, since loop count is identical on both sides. + +## Calibration (H100 Modal trial — done, step / max_iters=2) + +Final setup (bf16 + `step` contract + judge-owned loop of `max_iters=2` + planted +clusters + centroid inertia gate): the flashlib `reference.patch` compiles + runs +on bf16 and **passes the centroid gate on all 6 workloads**, running +**2.1 / 2.4 / 6.7 / 3.5 / 13.2 / 9.3x** over the bf16 naive baseline (geomean +**4.95x**). `speedup_target = 5.0`. The worst reference bf16 inertia drift is +**+1.47%** (w3, D=512), so `inertia_tolerance = 0.05` clears it with margin while +still rejecting genuinely bad centroids. This speedup is a *pure kernel-structure* +win — both sides bf16 (precision is not a lever), and the judge owns the loop (so +iteration-skip / fake-labels / subsample are not levers either). + +(Prior `max_iters=1` calibration, for the record: bf16 + planted clusters + +labelled/centroid gates gave the flashlib reference geomean **4.54x**, all 6 +passing. That setting let the agent own the whole `kmeans(...)` call, which is why +iteration-skip was still reachable and motivated this `step` restructure. The +geomean is essentially unchanged (4.54 → 4.95) because the reference does the same +per-step work; what changed is that the hack surface is now closed structurally.) diff --git a/2.0/problems/kmeans_gpu_kernel_optimization/config.yaml b/2.0/problems/kmeans_gpu_kernel_optimization/config.yaml new file mode 100644 index 000000000..6c818851b --- /dev/null +++ b/2.0/problems/kmeans_gpu_kernel_optimization/config.yaml @@ -0,0 +1,71 @@ +tag: systems +runtime: + language: python + timeout_seconds: 10800 + environment: "GPU K-Means kernel optimization. The agent patches the kmeanslib package; the judge offloads timing to a Modal GPU, comparing the patched kmeans against a frozen naive baseline on hidden (N, D, K) workloads with an inertia (clustering-quality) gate." + apt_packages: + - bash + - ca-certificates + - git + - python3 + - python3-pip + judge_apt_packages: + - bash + - ca-certificates + - git + - python3 + - python3-pip + judge_pip_packages: + - modal + docker: + # Experimental local images. Build with docker/build_images.sh before a local + # Harbor trial. Both are light (ubuntu + modal + git); torch/triton/flashlib + # run on the Modal GPU image defined in flash_gpu.py. The agent image bakes + # /app/kmeanslib (git) + /opt/kmeans_ref + /opt/flash_gpu.py; the judge image + # additionally bakes /opt/kmeanslib-clean. + image: frontiercs/kmeans-gpu-kernel-optimization-agent:experimental-v0.4.0 + judge_image: frontiercs/kmeans-gpu-kernel-optimization-judge:experimental-v0.4.0 +environment: + cpus: 4 + memory_mb: 8192 + storage_mb: 16384 + build_timeout_seconds: 3600 +evaluation: + # --- primitive wiring (consumed by the generic evaluator + flash_gpu) --- + primitive: kmeans + pkg: kmeanslib + ref_module: refkmeans + clean_source: /opt/kmeanslib-clean + baseline_source: /opt/kmeans_ref + # --- Modal GPU --- + gpu: "H100" # Modal GPU spec (H100 / A100 / L40S) + cuda_image: "nvidia/cuda:12.4.1-devel-ubuntu22.04" + pip: ["torch", "numpy"] + app_name: "kmeans-kernel-opt-eval" + modal_timeout_seconds: 1800 + # --- timing + quality gate --- + warmup_iters: 2 + timed_iters: 1 # one timed run per workload; the timed unit is the judge-owned loop of max_iters `step` calls + inertia_tolerance: 0.05 # clustering gate: inertia of the agent's RETURNED (labels, centroids) <= (1+tol) x the baseline's. Catches fake/garbage labels and badly-placed centroids. Inertia is dominated by within-cluster noise, so it is deliberately backed up by label_mismatch_tolerance below. + label_mismatch_tolerance: 0.002 # assignment gate: fraction of points whose returned label is not the exact (fp32) nearest centroid of the centroids that step was handed. Inertia barely moves when ~1% of points are misassigned (~+0.1%), so this is what actually rejects approximate distances -- e.g. assigning on only the leading feature dims. Set above the bf16-vs-fp32 tie rate the honest reference shows. + speedup_target: 20.0 # score cap = 2x the FUSED best-known geomean (9.97x on H100, all 6 pass). The vendored flashlib reference.patch runs assign and update as two separate kernels (two passes over x) and only reaches 7.92x (scores ~69/100); a kernel that fuses assign+scatter-update into one pass over x -- which is exactly what this task asks for -- reaches 9.97x (~77/100), and you must be 2x faster than THAT to max out. Calibrating on the unfused reference alone made the cap too easy (a merely-fused solution scored 94). + # High-K (K=512-2048, arithmetic intensity >= H100 bf16 roofline knee ~295 FLOP/byte) keeps the assign GEMM compute-bound: ~138 TFLOPS geomean (14% of bf16 peak) vs ~40 TFLOPS on the old low-K workloads. + base_seed: 20260701 + agent_workload_count: 6 # agent's submission feedback covers ALL graded workloads (no blind pass-3/fail-6 cliff); params are agent-visible anyway + expose_per_workload_metrics: true # show per-workload pass/fail so the agent can fix the specific workload it regresses on + # The agent provides one Lloyd `step(x, centroids)`; the JUDGE owns the loop and + # calls it exactly max_iters times (=2). The agent cannot skip iterations, change + # the data/init, or fake the count -- it can only make a single step faster + # (e.g. fuse assign+update). Data is bf16 (locked in gen) so precision is not a + # lever; planted well-separated blobs. max_iters=2 (not 1) is enough to be a real + # loop while too few for a converged step to become a redundant no-op. + workloads: + - {id: w0, N: 200000, D: 64, K: 512, max_iters: 2} + - {id: w1, N: 300000, D: 128, K: 1024, max_iters: 2} + - {id: w2, N: 500000, D: 128, K: 1024, max_iters: 2} + - {id: w3, N: 150000, D: 512, K: 512, max_iters: 2} + - {id: w4, N: 1000000, D: 64, K: 1024, max_iters: 2} + - {id: w5, N: 300000, D: 256, K: 2048, max_iters: 2} +submission: + kind: file + path: /app/solution.patch diff --git a/2.0/problems/kmeans_gpu_kernel_optimization/docker/README.md b/2.0/problems/kmeans_gpu_kernel_optimization/docker/README.md new file mode 100644 index 000000000..476074363 --- /dev/null +++ b/2.0/problems/kmeans_gpu_kernel_optimization/docker/README.md @@ -0,0 +1,56 @@ +# Experimental K-Means Kernel-Optimization Images + +Two **light** images (ubuntu + `modal` + git). torch, triton, and the vendored +flashlib kernels all run on the **Modal GPU image** defined in `flash_gpu.py`; +the containers themselves are CPU-only and offload GPU work to Modal. + +```bash +bash 2.0/problems/kmeans_gpu_kernel_optimization/docker/build_images.sh +``` + +Defaults: + +```text +AGENT_TAG=frontiercs/kmeans-gpu-kernel-optimization-agent:experimental-v0.2.0 +JUDGE_TAG=frontiercs/kmeans-gpu-kernel-optimization-judge:experimental-v0.2.0 +``` + +Agent image: + +```text +/app/kmeanslib # clean, git-tracked package (the agent edits this) +/opt/flash_gpu.py # shared Modal GPU harness (public test uses it) +/opt/kmeans_ref/refkmeans.py # frozen naive baseline (public-test speed denominator) +``` + +Judge image: + +```text +/opt/kmeanslib-clean/kmeanslib # pristine tree; the patch is applied to a copy +/opt/kmeans_ref/refkmeans.py # frozen naive baseline (speed denominator + quality oracle) +/opt/flash_gpu.py # shared Modal GPU harness (the evaluator uses it) +``` + +## Runtime requirements + +Both the judge and the agent public test **offload timing to a Modal GPU** and +therefore need Modal credentials in the environment: + +```text +MODAL_TOKEN_ID / MODAL_TOKEN_SECRET +``` + +`flash_gpu.py` builds an ephemeral Modal app on a GPU (`evaluation.gpu`, default +`H100`), ships the frozen baseline + the patched package as data, times both on +fresh per-iteration data, verifies quality each iteration, and returns the +speedups. No persistent deployment is used, so a fresh GPU container is spun up +per evaluation. Without Modal credentials the evaluator returns a patch-policy +smoke pass (which repo CI exercises). + +## Smoke test + +```bash +bash 2.0/problems/kmeans_gpu_kernel_optimization/docker/smoke_images.sh +``` + +Import-only (modal + flash_gpu + the baked packages); does not touch a GPU or Modal. diff --git a/2.0/problems/kmeans_gpu_kernel_optimization/docker/agent/Dockerfile b/2.0/problems/kmeans_gpu_kernel_optimization/docker/agent/Dockerfile new file mode 100644 index 000000000..c19a77dc1 --- /dev/null +++ b/2.0/problems/kmeans_gpu_kernel_optimization/docker/agent/Dockerfile @@ -0,0 +1,34 @@ +# Agent image for kmeans_gpu_kernel_optimization. +# +# Light image (no torch/triton): the agent edits /app/kmeanslib and runs the +# public test, which offloads timing to a Modal GPU (torch/triton live on the +# Modal image defined in flash_gpu.py). Needs MODAL_TOKEN_ID / MODAL_TOKEN_SECRET +# in the environment to run the GPU public test. +# +# Build context = the task directory: +# docker build -f docker/agent/Dockerfile -t . +FROM ubuntu:24.04 + +ARG DEBIAN_FRONTEND=noninteractive +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + bash ca-certificates git python3 python3-pip ripgrep && \ + rm -rf /var/lib/apt/lists/* + +RUN pip3 install --break-system-packages --no-cache-dir modal + +# The package the agent edits (git-tracked so make_submission.sh can diff it). +WORKDIR /app +COPY kmeanslib /app/kmeanslib +RUN printf '%s\n' '__pycache__/' '*.pyc' > /app/.gitignore && \ + git -C /app init -q && \ + git -C /app config user.email task@frontier-cs && \ + git -C /app config user.name frontier-cs && \ + git -C /app add -A -- kmeanslib .gitignore && \ + git -C /app -c commit.gpgsign=false commit -qm "pristine kmeanslib" + +# NOTE: the GPU harness (flash_gpu.py) and its data GENERATOR are deliberately +# NOT baked into the agent image -- they live only in the judge. The agent gets +# feedback by submitting to the judge (public_test == the graded submission), +# so it cannot read how the workloads are generated and hardcode a +# data-specific shortcut. See harbor_app/public_test.py. diff --git a/2.0/problems/kmeans_gpu_kernel_optimization/docker/build_images.sh b/2.0/problems/kmeans_gpu_kernel_optimization/docker/build_images.sh new file mode 100644 index 000000000..6d2f2ad51 --- /dev/null +++ b/2.0/problems/kmeans_gpu_kernel_optimization/docker/build_images.sh @@ -0,0 +1,17 @@ +#!/usr/bin/env bash +# Build the experimental agent + judge images for kmeans_gpu_kernel_optimization. +set -euo pipefail +HERE=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd) # the problem directory (build context) + +AGENT_TAG=${AGENT_TAG:-frontiercs/kmeans-gpu-kernel-optimization-agent:experimental-v0.4.0} +JUDGE_TAG=${JUDGE_TAG:-frontiercs/kmeans-gpu-kernel-optimization-judge:experimental-v0.4.0} + +echo "Building agent image: $AGENT_TAG" +docker build -f "$HERE/docker/agent/Dockerfile" -t "$AGENT_TAG" "$HERE" + +echo "Building judge image: $JUDGE_TAG" +docker build -f "$HERE/docker/judge/Dockerfile" -t "$JUDGE_TAG" "$HERE" + +echo "Built:" +echo " $AGENT_TAG" +echo " $JUDGE_TAG" diff --git a/2.0/problems/kmeans_gpu_kernel_optimization/docker/judge/Dockerfile b/2.0/problems/kmeans_gpu_kernel_optimization/docker/judge/Dockerfile new file mode 100644 index 000000000..71ec8b1d6 --- /dev/null +++ b/2.0/problems/kmeans_gpu_kernel_optimization/docker/judge/Dockerfile @@ -0,0 +1,26 @@ +# Judge image for kmeans_gpu_kernel_optimization. +# +# Light image (no torch/triton): the judge validates + applies the patch and +# offloads timing to a Modal GPU (torch/triton live on the Modal image defined +# in flash_gpu.py). Needs MODAL_TOKEN_ID / MODAL_TOKEN_SECRET in the environment. +# The Frontier-CS 2.0 adapter builds the final judge image ON TOP of this one. +# +# Build context = the task directory: +# docker build -f docker/judge/Dockerfile -t . +FROM ubuntu:24.04 + +ARG DEBIAN_FRONTEND=noninteractive +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + bash ca-certificates git python3 python3-pip && \ + rm -rf /var/lib/apt/lists/* + +RUN pip3 install --break-system-packages --no-cache-dir modal + +# Pristine tree the patch is applied to, the frozen naive baseline (speed +# denominator + quality oracle), and the shared Modal GPU harness. +COPY kmeanslib /opt/kmeanslib-clean/kmeanslib +COPY judge/refkmeans.py /opt/kmeans_ref/refkmeans.py +COPY flash_gpu.py /opt/flash_gpu.py + +WORKDIR /judge diff --git a/2.0/problems/kmeans_gpu_kernel_optimization/docker/smoke_images.sh b/2.0/problems/kmeans_gpu_kernel_optimization/docker/smoke_images.sh new file mode 100644 index 000000000..f82eb4fee --- /dev/null +++ b/2.0/problems/kmeans_gpu_kernel_optimization/docker/smoke_images.sh @@ -0,0 +1,22 @@ +#!/usr/bin/env bash +# Smoke test for the built (intentionally torch-less) images: confirm the host-side +# harness imports (modal + flash_gpu) and that the frozen baseline + package parse. +# torch/triton live on the Modal image, so the packages are py_compiled here (a +# torch-less parse check), not imported. No GPU / Modal tokens required. +set -euo pipefail + +AGENT_TAG=${AGENT_TAG:-frontiercs/kmeans-gpu-kernel-optimization-agent:experimental-v0.2.0} +JUDGE_TAG=${JUDGE_TAG:-frontiercs/kmeans-gpu-kernel-optimization-judge:experimental-v0.2.0} + +SMOKE='import sys, modal, py_compile, pathlib +sys.path.insert(0, "/opt"); import flash_gpu +for root in sys.argv[1:]: + for p in pathlib.Path(root).rglob("*.py"): + py_compile.compile(str(p), doraise=True) +print("ok: modal + flash_gpu import; kmeans baseline + kmeanslib parse")' + +echo "== agent image ==" +docker run --rm -w /app "$AGENT_TAG" python3 -c "$SMOKE" /opt/kmeans_ref /app/kmeanslib + +echo "== judge image ==" +docker run --rm "$JUDGE_TAG" python3 -c "$SMOKE" /opt/kmeans_ref /opt/kmeanslib-clean/kmeanslib diff --git a/2.0/problems/kmeans_gpu_kernel_optimization/evaluate.sh b/2.0/problems/kmeans_gpu_kernel_optimization/evaluate.sh new file mode 100644 index 000000000..49d06ffc0 --- /dev/null +++ b/2.0/problems/kmeans_gpu_kernel_optimization/evaluate.sh @@ -0,0 +1,12 @@ +#!/usr/bin/env bash +# Local CLI evaluation for kmeans_gpu_kernel_optimization. +# +# A full run needs a GPU plus the baked judge sources at /opt (the pristine +# /opt/kmeanslib-clean tree and the frozen /opt/kmeans_ref baseline; see the +# judge image). Without them the evaluator still validates the patch policy and +# returns a smoke score, which is what repository CI exercises (the reference +# patch passes the policy). Pass a patch path to score it directly. +set -euo pipefail +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +SOLUTION="${1:-$SCRIPT_DIR/reference.patch}" +exec python3 "$SCRIPT_DIR/evaluator.py" "$SOLUTION" diff --git a/2.0/problems/kmeans_gpu_kernel_optimization/evaluator.py b/2.0/problems/kmeans_gpu_kernel_optimization/evaluator.py new file mode 100644 index 000000000..eebe5d7c9 --- /dev/null +++ b/2.0/problems/kmeans_gpu_kernel_optimization/evaluator.py @@ -0,0 +1,380 @@ +"""Generic evaluator for the flashlib GPU kernel-optimization task family. + +Identical across all four tasks (kmeans / knn / pca / truncated_svd); every +primitive-specific value comes from the ``evaluation`` block of the task's +config (delivered to the judge as ``/judge/task_config.json``), which is not +present in the agent workspace. + +Flow: statically validate the patch (only ``/**`` may change; no external +optimized libs / env / process / network access) -> apply it to the pristine +package baked at ``clean_source`` -> hand the frozen baseline + patched package +sources to a Modal GPU worker (``flash_gpu.run_remote``) that times both on +fresh per-iteration data and verifies clustering/retrieval/subspace quality on +every iteration -> gate on the worker's per-workload verdict -> score by the +geometric-mean speedup. + +Without Modal credentials or the baked judge sources (e.g. repo CI on a box with +no GPU/Modal), the evaluator still validates the patch policy and returns a +smoke pass, so ``python3 evaluator.py reference.patch`` works anywhere. +""" +from __future__ import annotations + +import fnmatch +import hashlib +import json +import math +import os +import re +import shutil +import subprocess +import sys +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +TASK_CONFIG_PATH = Path("/judge/task_config.json") + + +def _load_task_config() -> dict[str, Any]: + try: + payload = json.loads(TASK_CONFIG_PATH.read_text(encoding="utf-8")) + except Exception: + return {} + return payload if isinstance(payload, dict) else {} + + +TASK_CONFIG = _load_task_config() +EVAL = TASK_CONFIG.get("evaluation", {}) if isinstance(TASK_CONFIG.get("evaluation"), dict) else {} + + +def _get(name: str, default): + return EVAL.get(name, default) + + +def _get_int(name: str, default: int) -> int: + try: + return int(EVAL.get(name, default)) + except Exception: + return default + + +def _get_float(name: str, default: float) -> float: + try: + return float(EVAL.get(name, default)) + except Exception: + return default + + +def _get_bool(name: str, default: bool) -> bool: + raw = EVAL.get(name, default) + if isinstance(raw, bool): + return raw + if isinstance(raw, str): + return raw.strip().lower() in {"1", "true", "yes", "on"} + return bool(raw) + + +# --- Per-task identity: the ONLY lines that differ across the four tasks. --- +# Kept as constants (not config) so the patch policy is enforceable locally / in +# CI without /judge/task_config.json. Everything else is config-driven and only +# needed on the GPU path. +PRIMITIVE = "kmeans" +PKG = "kmeanslib" +REF_MODULE = "refkmeans" + +MAX_PATCH_BYTES = _get_int("max_patch_bytes", 2_000_000) +MAX_CHANGED_FILES = _get_int("max_changed_files", 80) +CLEAN_SOURCE = Path(f"/opt/{PKG}-clean") +BASELINE_SOURCE = Path(f"/opt/{PRIMITIVE}_ref") +SPEEDUP_TARGET = _get_float("speedup_target", 8.0) + +ALLOWED_PATTERNS = (f"{PKG}/**",) +DENIED_PATTERNS = ( + "evaluator.py", + "flash_gpu.py", + "reference.patch", + f"{REF_MODULE}.py", + f"**/{REF_MODULE}.py", + "**/conftest.py", + "**/test_*.py", + "pyproject.toml", + "setup.py", + "setup.cfg", +) +# Defense in depth (the load-bearing anti-tamper defense is structural, inside +# the GPU worker: it captures its perf_counter / synchronize references before +# importing the submission, regenerates data every iteration, and re-verifies +# quality each iteration). These patterns are matched surgically so that +# legitimate vendored/optimized kernel source is not rejected: +# * external optimized libraries, matched as IMPORT statements (not bare words, +# so a comment mentioning a name is fine); none are installed on the GPU +# image anyway; +# * process / network / measurement-tamper patterns that never appear in a +# real GPU kernel. +FORBIDDEN_IMPORT_RE = re.compile( + r"(?:^|\n)\s*(?:import|from)\s+(flashlib|cuml|cudf|cupy|faiss|sklearn|scikit|cutlass|cuvs)\b" +) +FORBIDDEN_PATTERN_RE = re.compile( + r"\bsubprocess\b|\bsocket\b|\burllib\b|\brequests\b|os\.system|" + r"torch\.cuda\.synchronize\s*=|time\.perf_counter\s*=|setattr\(\s*torch\.cuda" +) + + +@dataclass(frozen=True) +class PatchFile: + old_path: str + new_path: str + added_lines: tuple[str, ...] + + @property + def path(self) -> str: + return self.new_path if self.new_path != "/dev/null" else self.old_path + + +def _match(path: str, patterns: tuple[str, ...]) -> bool: + return any(fnmatch.fnmatch(path, p) for p in patterns) + + +def _invalid(message: str, metrics: dict[str, Any] | None = None): + payload = metrics or {} + payload.setdefault("valid_patch", 0) + return 0.0, 0.0, message, payload + + +def _parse_patch(text: str) -> list[PatchFile]: + files: list[PatchFile] = [] + old = new = "" + added: list[str] = [] + in_file = False + for line in text.splitlines(): + if line.startswith("diff --git "): + if in_file: + files.append(PatchFile(old, new, tuple(added))) + in_file, old, new, added = True, "", "", [] + continue + if not in_file: + continue + if line.startswith("--- "): + old = line[4:].strip() + old = old[2:] if old.startswith("a/") else old + elif line.startswith("+++ "): + new = line[4:].strip() + new = new[2:] if new.startswith("b/") else new + elif line.startswith("+") and not line.startswith("+++ "): + added.append(line[1:]) + if in_file: + files.append(PatchFile(old, new, tuple(added))) + return files + + +def _validate_path(path: str) -> tuple[bool, str]: + if not path or path == "/dev/null": + return True, "" + if path.startswith("/") or ".." in Path(path).parts: + return False, f"unsafe patch path: {path}" + if _match(path, DENIED_PATTERNS): + return False, f"changed file is outside task boundary: {path}" + if not _match(path, ALLOWED_PATTERNS): + return False, f"changed file is not allowlisted (only {PKG}/** may change): {path}" + if not path.endswith(".py"): + return False, f"only Python files may change: {path}" + return True, "" + + +def validate_patch(patch_path: Path) -> tuple[bool, str, dict[str, Any]]: + if not patch_path.exists(): + return False, "solution patch does not exist", {} + size = patch_path.stat().st_size + if size > MAX_PATCH_BYTES: + return False, f"patch is too large ({size} bytes > {MAX_PATCH_BYTES})", {} + text = patch_path.read_text(encoding="utf-8", errors="replace") + files = _parse_patch(text) + metrics: dict[str, Any] = { + "patch_bytes": size, + "patch_sha256": hashlib.sha256(text.encode("utf-8", "replace")).hexdigest(), + "changed_files": len(files), + } + if len(files) > MAX_CHANGED_FILES: + return False, f"too many changed files ({len(files)} > {MAX_CHANGED_FILES})", metrics + for pf in files: + path = pf.path + if pf.new_path == "/dev/null": + return False, f"deleting files is outside task boundary: {pf.old_path}", metrics + if pf.old_path != "/dev/null" and pf.old_path != pf.new_path: + ok, err = _validate_path(pf.old_path) + if not ok: + return False, f"rename source is outside task boundary: {err}", metrics + ok, err = _validate_path(path) + if not ok: + return False, err, metrics + added = "\n".join(pf.added_lines) + m = FORBIDDEN_IMPORT_RE.search(added) + if m: + return False, f"{path}: forbidden import of an external optimized library ({m.group(1)})", metrics + m = FORBIDDEN_PATTERN_RE.search(added) + if m: + return False, f"{path}: forbidden pattern in added code ({m.group(0).strip()[:40]})", metrics + metrics["valid_patch"] = 1 + return True, "patch accepted by static policy", metrics + + +def sanitize(text: str) -> str: + text = re.sub(r"/tmp/[A-Za-z0-9_./-]+", "", text or "") + text = re.sub(r"/opt/[A-Za-z0-9_./-]+", "", text) + text = re.sub(r"\bN=\d+|\bseed[=:]?\s*\d+", "", text, flags=re.IGNORECASE) + return text[-600:] + + +def geometric_mean(values: list[float]) -> float: + if not values: + return 0.0 + return math.exp(sum(math.log(max(v, 1e-9)) for v in values) / len(values)) + + +def score_from_speedup(gm: float) -> float: + if gm <= 0: + return 0.0 + raw = 100.0 * math.log(gm) / math.log(max(SPEEDUP_TARGET, 1.0000001)) + return max(0.0, min(100.0, raw)) + + +def is_final_role() -> bool: + return os.environ.get("FRONTIER_SUBMISSION_ROLE", "agent") == "final" + + +def _workloads(final_role: bool) -> list[dict[str, Any]]: + workloads = list(_get("workloads", []) or []) + if not final_role: + n = _get_int("agent_workload_count", 3) + workloads = workloads[:n] + base = _get_int("base_seed", 20260701) + for i, w in enumerate(workloads): + w.setdefault("seed", base + 1000 * (i + 1)) + return workloads + + +def _build_cfg() -> dict[str, Any]: + return { + "primitive": PRIMITIVE, + "pkg": PKG, + "ref_module": REF_MODULE, + "gpu": str(_get("gpu", "H100")), + "cuda_image": str(_get("cuda_image", "nvidia/cuda:12.4.1-devel-ubuntu22.04")), + "pip": list(_get("pip", ["torch==2.5.1", "triton==3.1.0", "numpy"])), + "app_name": str(_get("app_name", "flash-kernel-eval")), + "modal_timeout_seconds": _get_int("modal_timeout_seconds", 1800), + "warmup": _get_int("warmup_iters", 3), + "iters": _get_int("timed_iters", 7), + "inertia_tolerance": _get_float("inertia_tolerance", 0.02), + "label_mismatch_tolerance": _get_float("label_mismatch_tolerance", 0.002), + "recall_threshold": _get_float("recall_threshold", 0.99), + "captured_tolerance": _get_float("captured_tolerance", 0.02), + "ortho_tolerance": _get_float("ortho_tolerance", 0.02), + "ari_threshold": _get_float("ari_threshold", 0.99), + } + + +def _read_dir(root: Path, rel_to: Path) -> dict[str, str]: + out: dict[str, str] = {} + for p in root.rglob("*.py"): + out[str(p.relative_to(rel_to))] = p.read_text(encoding="utf-8", errors="replace") + return out + + +def full_evaluation(patch_path: Path, metrics: dict[str, Any]): + final_role = is_final_role() + metrics["submission_role"] = "final" if final_role else "agent" + workloads = _workloads(final_role) + + # Import the Modal harness baked into the judge image; missing harness or + # missing credentials/sources -> policy smoke pass. + sys.path.insert(0, "/opt") + try: + import flash_gpu # type: ignore + except Exception: + metrics["full_benchmark"] = 0 + return 1.0, 1.0, "patch policy smoke passed; GPU harness not available in this environment", metrics + if not flash_gpu.modal_available() or not CLEAN_SOURCE.exists() or not BASELINE_SOURCE.exists(): + metrics["full_benchmark"] = 0 + return 1.0, 1.0, "patch policy smoke passed; Modal/judge sources not configured in this environment", metrics + + with tempfile.TemporaryDirectory(prefix="flash_kernel_opt_") as tmp: + tmp_root = Path(tmp) + patched = tmp_root / "patched" + shutil.copytree(CLEAN_SOURCE, patched) + env = {"PATH": os.environ.get("PATH", "/usr/bin:/bin"), "HOME": str(tmp_root), + "LC_ALL": "C", "LANG": "C"} + if int(metrics.get("changed_files", 0)) > 0: + try: + subprocess.run(["git", "apply", "--check", str(patch_path)], cwd=patched, env=env, + check=True, capture_output=True, text=True, timeout=60) + subprocess.run(["git", "apply", str(patch_path)], cwd=patched, env=env, + check=True, capture_output=True, text=True, timeout=60) + except subprocess.CalledProcessError as exc: + metrics["stderr_tail"] = sanitize(exc.stderr or "") + return _invalid("patch does not apply to the clean package tree", metrics) + metrics["applied_patch"] = 1 + else: + metrics["used_empty_patch"] = 1 + + payload = { + "baseline_files": _read_dir(BASELINE_SOURCE, BASELINE_SOURCE), + "patched_files": _read_dir(patched, patched), + "workloads": workloads, + "cfg": _build_cfg(), + } + try: + result = flash_gpu.run_remote(payload) + except Exception as exc: + metrics["error_detail"] = sanitize(str(exc)) + return _invalid("GPU evaluation failed", metrics) + + if not result.get("ok"): + metrics["worker_error"] = sanitize(str(result.get("error", "unknown"))) + return _invalid("GPU benchmark worker failed", metrics) + + speedups: list[float] = [] + per_workload: dict[str, Any] = {} + for row in result.get("rows", []): + if not row.get("ok"): + metrics["failed_workload"] = {"reason": row.get("reason", "error")} + return _invalid("submission failed the quality gate or crashed on a hidden workload", metrics) + speedups.append(max(float(row["speedup"]), 0.01)) + per_workload[row["id"]] = {"speedup": row["speedup"]} + if not speedups: + return _invalid("no workloads were evaluated", metrics) + + gm = geometric_mean(speedups) + bounded = score_from_speedup(gm) + metrics.update({"full_benchmark": 1, "workload_count": len(speedups), "geomean_speedup": gm}) + if _get_bool("expose_per_workload_metrics", False): + metrics["per_workload"] = per_workload + return bounded, bounded, f"{PRIMITIVE} geomean speedup {gm:.3f}x over the naive baseline", metrics + + +def evaluate(solution_path: str) -> tuple[float, float, str, dict[str, Any]]: + patch_path = Path(solution_path) + ok, message, metrics = validate_patch(patch_path) + if not ok: + return _invalid(message, metrics) + try: + return full_evaluation(patch_path, metrics) + except Exception as exc: # noqa: BLE001 + metrics["error_type"] = type(exc).__name__ + metrics["error_detail"] = sanitize(str(exc)) + return _invalid("evaluation failed", metrics) + + +def main(argv: list[str]) -> int: + if len(argv) != 2: + print("Usage: evaluator.py SOLUTION_PATCH", file=sys.stderr) + return 2 + score, unbounded, message, metrics = evaluate(argv[1]) + print(json.dumps({"score": score, "score_unbounded": unbounded, + "message": message, "metrics": metrics}, indent=2, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv)) diff --git a/2.0/problems/kmeans_gpu_kernel_optimization/flash_gpu.py b/2.0/problems/kmeans_gpu_kernel_optimization/flash_gpu.py new file mode 100644 index 000000000..520912dd7 --- /dev/null +++ b/2.0/problems/kmeans_gpu_kernel_optimization/flash_gpu.py @@ -0,0 +1,446 @@ +"""Modal GPU evaluation harness for the flashlib kernel-optimization task family. + +Shared by the judge (evaluator.py) and the agent public test. Baked into both +the agent and judge images at ``/opt/flash_gpu.py``. It runs a frozen naive +baseline and the (patched or agent-edited) package on a Modal GPU and returns +per-workload timings + quality verdicts. + +Anti-reward-hack properties (all enforced inside the remote worker, which is the +only place the untrusted code runs): + +* **Timing primitives are captured before any agent code is imported** — the + worker binds ``time.perf_counter`` and ``torch.cuda.synchronize`` to locals up + front, so a submission that monkey-patches ``torch.cuda.synchronize`` (or any + torch attribute) cannot affect measurement. +* **Fresh data every timed iteration** — each measured iteration generates a new + input from a fresh seed, so a submission cannot memoize on a repeated input. +* **Quality is verified on every timed iteration** — the judge recomputes the + quality metric from the *returned tensors* on that iteration's data and gates + it, so returning a stale cached result for a different input is caught. +* **The judge computes all metrics** — no agent-provided number is trusted. +* **Fresh container per submission** — an ephemeral Modal app is used, so no + global state survives across evaluations. + +The remote worker is a single self-contained function serialized to Modal +(``serialized=True``), so the remote does not import this module and there is no +module-mounting to get wrong. The GPU image ships torch + triton (+ any extra +pip packages a task needs for its reference build). +""" +from __future__ import annotations + +import os + + +# --------------------------------------------------------------------------- # +# Remote worker (runs on the Modal GPU). Fully self-contained: every helper is +# nested so cloudpickle ships the whole thing. Takes and returns plain data. +# --------------------------------------------------------------------------- # +def _gpu_worker(payload: dict) -> dict: + import importlib + import os as _os + import sys as _sys + import tempfile + import time + import traceback + + import torch + + # Capture timing + sync references BEFORE importing any submission code, so a + # monkey-patch of torch.cuda.synchronize cannot affect measurement. + _perf = time.perf_counter + _sync = torch.cuda.synchronize + if not torch.cuda.is_available(): + return {"ok": False, "error": "cuda_unavailable"} + dev = "cuda" + + cfg = payload["cfg"] + prim = cfg["primitive"] + warmup = int(cfg.get("warmup", 3)) + iters = int(cfg.get("iters", 7)) + + def _materialize(files: dict, tag: str) -> str: + root = tempfile.mkdtemp(prefix=f"flash_{tag}_") + for rel, content in files.items(): + path = _os.path.join(root, rel) + parent = _os.path.dirname(path) + if parent: + _os.makedirs(parent, exist_ok=True) + with open(path, "w", encoding="utf-8") as fh: + fh.write(content) + _sys.path.insert(0, root) + return root + + _materialize(payload["baseline_files"], "baseline") + _materialize(payload["patched_files"], "patched") + ref = importlib.import_module(cfg["ref_module"]) + pkg = importlib.import_module(cfg["pkg"]) + + # ---- per-primitive: fixed context, data gen, call, quality metric ----- # + def setup(w, seed): + # Built ONCE per workload (not per timed iter). IVF-PQ builds the fixed + # index (the "database") here; queries are generated fresh per iteration. + if prim == "ivfpq": + g = torch.Generator(device=dev).manual_seed(int(seed)) + X = torch.randn(int(w["M"]), int(w["D"]), generator=g, device=dev, dtype=torch.float32) + index = ref.ivf_pq_build(X, int(w["nlist"]), m=int(w["m"]), + nprobe=int(w["nprobe"]), seed=int(seed)) + return {"index": index, "D": int(w["D"])} + return {} + + def gen(w, seed, ctx): + g = torch.Generator(device=dev).manual_seed(int(seed)) + if prim == "ivfpq": + q = torch.randn(int(w["Q"]), int(ctx["D"]), generator=g, device=dev, dtype=torch.float32) + return {"queries": q} + if prim == "knn": + db = torch.randn(w["M"], w["D"], generator=g, device=dev, dtype=torch.float32) + q = torch.randn(w["Q"], w["D"], generator=g, device=dev, dtype=torch.float32) + return {"queries": q, "database": db} + if prim == "dbscan": + nc, D, N = int(w["n_centers"]), int(w["D"]), int(w["N"]) + centers = torch.randn(nc, D, generator=g, device=dev, dtype=torch.float32) * 10.0 + asg = torch.randint(0, nc, (N,), generator=g, device=dev) + xb = centers.index_select(0, asg) + torch.randn(N, D, generator=g, device=dev, dtype=torch.float32) * 0.5 + nz = int(float(w.get("noise_frac", 0.0)) * N) + if nz > 0: + lo = xb.min(0).values; hi = xb.max(0).values + xb[:nz] = lo + (hi - lo) * torch.rand(nz, D, generator=g, device=dev) + return {"x": xb, "eps": float(w["eps"]), "min_samples": int(w["min_samples"])} + x = torch.randn(w["N"], w["D"], generator=g, device=dev, dtype=torch.float32) + if prim == "kmeans": + # PRECISION LOCK (task-root, not prompt): kmeans data is bf16, so the + # inputs carry only bf16 precision. No solution -- baseline, reference, + # or agent -- can trade accuracy for speed by picking a different matmul + # dtype (fp16/tf32/fp32 buy nothing on bf16 data and are only slower). + # + # PLANTED CLUSTERS: x = blob centers + unit noise (NOT random). This + # makes the assignment matter -- inertia is sensitive to which centroid + # each point takes -- so a `step` that subsamples rows, fakes the labels, + # or drops feature dims lands points on the wrong centroid and regresses + # the inertia gate. Random data would leave inertia nearly + # assignment-invariant and thus hackable by doing less work. + # + # SEPARATION IS SPREAD OVER ALL D DIMS. The nearest-competitor margin is + # ~ scale * sqrt(D/2) sigma of the unit noise. With a FIXED scale the + # margin grows as sqrt(D): at scale 6.0 / D=512 it is ~96 sigma, so a + # solution can assign using only the leading 128 dims and still land the + # exact same labels -- free work-skipping no output gate can see. Scaling + # by sqrt(2/D) pins the full-D margin at ~4 sigma, so using only a + # fraction f of the dims cuts it to 4*sqrt(f) (D=512 -> 128 dims: 2 sigma) + # and misassigns boundary points, which the inertia gate does see. Every + # dimension must be read. + K = int(w["K"]) + D = int(w["D"]) + centers = torch.randn(K, D, generator=g, device=dev) * (4.0 * (2.0 / D) ** 0.5) + asg = torch.randint(0, K, (int(w["N"]),), generator=g, device=dev) + x = (centers.index_select(0, asg) + x).to(torch.bfloat16) # top randn = unit noise + perm = torch.randperm(w["N"], generator=g, device=dev)[: K] + return {"x": x, "init": x.index_select(0, perm).clone()} + return {"x": x} + + def call(mod, w, data, ctx): + if prim == "ivfpq": + return mod.ivf_pq_search(ctx["index"], data["queries"], int(w["k"]), nprobe=int(w["nprobe"])) + if prim == "kmeans": + # The JUDGE owns the Lloyd loop: call the solution's one-iteration + # `step(x, centroids)` exactly max_iters times. The solution cannot + # change the iteration count, the data, or the init centroids -- so it + # cannot skip iterations or fake the number of steps; it can only make + # a single step faster (e.g. fuse assign+update). This whole loop is + # what gets timed. + c = data["init"].clone() + lab = None + c_in = None + for _ in range(int(w["max_iters"])): + c_in = c # centroids the LAST step was handed + lab, c = mod.step(data["x"], c) + # c_in is returned so the gate can re-derive the exact assignment the + # final `step` was supposed to produce and compare it to `lab`. + return lab, c, int(w["max_iters"]), c_in + if prim == "knn": + return mod.knn(data["queries"], data["database"], w["k"]) + if prim == "dbscan": + return mod.dbscan(data["x"], data["eps"], data["min_samples"]) + if prim == "pca": + return mod.pca(data["x"], w["k"]) + if prim == "tsvd": + return mod.truncated_svd(data["x"], w["k"]) + raise ValueError(prim) + + def check_shape(w, out): + if prim == "kmeans": + lab, cen = out[0], out[1] + if tuple(cen.shape) != (w["K"], w["D"]) or not torch.isfinite(cen.float()).all(): + raise ValueError("bad kmeans centroids") + if lab.ndim != 1 or int(lab.shape[0]) != int(w["N"]): + raise ValueError("bad kmeans labels shape") # step must return a full (N,) assignment + elif prim == "knn": + d, i = out + if tuple(d.shape) != (w["Q"], w["k"]) or tuple(i.shape) != (w["Q"], w["k"]): + raise ValueError("bad knn output") + if not torch.isfinite(d).all(): + raise ValueError("non-finite distances") + elif prim == "ivfpq": + vals, ids = out + if tuple(ids.shape) != (int(w["Q"]), int(w["k"])) or \ + tuple(vals.shape) != (int(w["Q"]), int(w["k"])): + raise ValueError("bad ivfpq output shape") + elif prim == "dbscan": + if out.ndim != 1 or int(out.shape[0]) != int(w["N"]): + raise ValueError("bad dbscan labels shape") + else: + a, b = out + comps = a if prim == "pca" else b + if tuple(comps.shape) != (w["k"], w["D"]) or not torch.isfinite(comps).all(): + raise ValueError("bad decomposition output") + + def inertia(x, centroids): + c = centroids.to(torch.float32) + cn = (c * c).sum(1) + total = 0.0 + for j in range(0, x.shape[0], 16384): + xb = x[j:j + 16384].to(torch.float32) # metric always fp32 (x may be bf16) + d = (xb * xb).sum(1, keepdim=True) - 2.0 * (xb @ c.t()) + cn[None, :] + total += float(d.min(1).values.clamp_min(0).sum().item()) + return total + + def inertia_labeled(x, centroids, labels): + # sum_i ||x_i - centroids[labels_i]||^2 : ties the returned LABELS to the + # returned centroids. Fake/garbage labels or a wrong (subsampled) assignment + # inflate this even if the centroids alone look fine -- closes the "return + # good centroids + junk labels" and "subsample the assignment" hacks. + c = centroids.to(torch.float32) + lab = labels.long().clamp_(0, c.shape[0] - 1) + total = 0.0 + for j in range(0, x.shape[0], 16384): + xb = x[j:j + 16384].to(torch.float32) + cl = c.index_select(0, lab[j:j + 16384]) + total += float(((xb - cl) ** 2).sum().item()) + return total + + def label_mismatch(x, centroids, labels): + # Fraction of points whose returned label is NOT the exact nearest centroid + # (fp32 metric) of the centroids that `step` was given. Inertia alone is a + # blunt detector: it is dominated by within-cluster noise, so misassigning + # 1% of the points moves it by only ~0.1% and hides under any usable + # tolerance. The assignment itself is what the contract specifies, so check + # it directly -- this is what catches assigning on a subset of the feature + # dims (or any other approximate distance) while the inertia stays flat. + c = centroids.to(torch.float32) + cn = (c * c).sum(1) + lab = labels.long() + bad = 0 + for j in range(0, x.shape[0], 16384): + xb = x[j:j + 16384].to(torch.float32) + d = (xb * xb).sum(1, keepdim=True) - 2.0 * (xb @ c.t()) + cn[None, :] + exact = d.argmin(1) + got = lab[j:j + 16384] + # a point sitting on a true tie may legitimately go either way: only + # count it when the returned centroid is strictly farther than the best + best = d.gather(1, exact[:, None]).squeeze(1) + mine = d.gather(1, got.clamp(0, c.shape[0] - 1)[:, None]).squeeze(1) + bad += int((mine > best * (1.0 + 1e-6) + 1e-6).sum().item()) + return bad / max(1, x.shape[0]) + + def recall(agent_idx, ref_idx): + a = agent_idx.long(); b = ref_idx.long() + hit = (a.unsqueeze(2) == b.unsqueeze(1)).any(2) + return float(hit.sum().item()) / (b.shape[0] * b.shape[1]) + + def ari(a, b): + # Adjusted Rand Index between two integer label vectors (N,); noise (-1) + # is treated as its own label (dense-remapped first). + a = torch.unique(a.long(), return_inverse=True)[1] + b = torch.unique(b.long(), return_inverse=True)[1] + na = int(a.max().item()) + 1; nb = int(b.max().item()) + 1 + cont = torch.bincount(a * nb + b, minlength=na * nb).reshape(na, nb).double() + ai = cont.sum(1); bj = cont.sum(0); n = cont.sum() + c2 = lambda z: z * (z - 1) / 2.0 + sij = c2(cont).sum(); sa = c2(ai).sum(); sb = c2(bj).sum() + exp = (sa * sb / c2(n)) if float(n) > 1 else 0.0 + den = 0.5 * (sa + sb) - exp + return float(((sij - exp) / den).item()) if float(den) != 0.0 else 1.0 + + def ortho_err(comps): + c = comps.to(torch.float32); k = c.shape[0] + return float((c @ c.t() - torch.eye(k, device=c.device)).abs().max().item()) + + def captured(x, comps, center): + c = comps.to(torch.float32) + mean = x.mean(0) if center else None + total = 0.0 + for j in range(0, x.shape[0], 16384): + xb = x[j:j + 16384] + if center: + xb = xb - mean + p = xb @ c.t() + total += float((p * p).sum().item()) + return total / (x.shape[0] - 1) if center else total + + def verdict(w, data, ctx, ref_out, agent_out): + """Return (ok, reason, ref_val, agent_val) gating the agent output.""" + check_shape(w, agent_out) + if prim == "ivfpq": + # iso-result recall: agent ids vs frozen-baseline ids on the SAME index. + rec = recall(agent_out[1], ref_out[1]) + ok = rec >= float(cfg["recall_threshold"]) + return ok, ("recall_regression" if not ok else ""), 1.0, rec + if prim == "kmeans": + # Clustering-quality gate on the inertia of the RETURNED (labels, + # centroids) pair -- NOT the centroids alone. Using the agent's own + # labels closes the "return good centroids + junk/zeroed labels" hack + # (e.g. labels = torch.empty(N)) and the "subsample the assignment" + # hack: a fake or partial assignment inflates the agent's inertia even + # when its centroids look fine. The judge owns the Lloyd loop, so the + # solution can only make one honest step faster. + rv = inertia_labeled(data["x"], ref_out[1], ref_out[0]) + av = inertia_labeled(data["x"], agent_out[1], agent_out[0]) + tol = float(cfg["inertia_tolerance"]) + ok = av <= (1.0 + tol) * rv + 1e-6 + if not ok: + return False, "inertia_regression", rv, av + # Assignment gate: the returned labels must BE the exact nearest-centroid + # assignment for the centroids the final step was handed. Catches + # approximate distances (e.g. assigning on a subset of the feature dims) + # that barely move inertia but are not the step the contract asks for. + mis = label_mismatch(data["x"], agent_out[3], agent_out[0]) + lt = float(cfg.get("label_mismatch_tolerance", 0.002)) + if mis > lt: + return False, "assignment_not_nearest_centroid", rv, av + return True, "", rv, av + if prim == "knn": + rec = recall(agent_out[1], ref_out[1]) + ok = rec >= float(cfg["recall_threshold"]) + return ok, ("recall_regression" if not ok else ""), 1.0, rec + if prim == "dbscan": + val = ari(agent_out, ref_out) + ok = val >= float(cfg["ari_threshold"]) + return ok, ("ari_regression" if not ok else ""), 1.0, val + center = (prim == "pca") + comps = agent_out[0] if prim == "pca" else agent_out[1] + ref_comps = ref_out[0] if prim == "pca" else ref_out[1] + oe = ortho_err(comps) + if oe > float(cfg["ortho_tolerance"]): + return False, "not_orthonormal", 0.0, oe + rv = captured(data["x"], ref_comps, center) + av = captured(data["x"], comps, center) + ok = av >= (1.0 - float(cfg["captured_tolerance"])) * rv - 1e-6 + return ok, ("captured_regression" if not ok else ""), rv, av + + def time_call(mod, w, data, ctx): + _sync(); t0 = _perf(); out = call(mod, w, data, ctx); _sync() + return (_perf() - t0) * 1000.0, out + + rows = [] + try: + for w in payload["workloads"]: + base_seed = int(w["seed"]) + ctx = setup(w, base_seed) # fixed context (e.g. IVF-PQ index) + # warmup on fresh data (kernels / autotune) — not measured + for i in range(warmup): + d = gen(w, base_seed + 100 + i, ctx) + call(ref, w, d, ctx); call(pkg, w, d, ctx); _sync() + del d + torch.cuda.empty_cache() + ratios, ref_val, agent_val = [], None, None + bad = None + for i in range(iters): + d = gen(w, base_seed + 10000 + i, ctx) # fresh every iteration + rt, rout = time_call(ref, w, d, ctx) + at, aout = time_call(pkg, w, d, ctx) + ok, reason, rv, av = verdict(w, d, ctx, rout, aout) # verify THIS iter + if not ok: + bad = reason; ref_val, agent_val = rv, av + break + ratios.append(rt / at if at > 0 else 0.01) + ref_val, agent_val = rv, av + del d, rout, aout + torch.cuda.empty_cache() + if bad is not None: + rows.append({"id": w["id"], "ok": False, "reason": bad, + "ref_val": ref_val, "agent_val": agent_val}) + continue + # aggregate the per-run speedups: "mean" (average of the N runs) or + # "median" (default, noise-robust). config sets this per task. + if str(cfg.get("aggregate", "median")).lower() == "mean": + sp = sum(ratios) / len(ratios) + else: + ratios.sort(); sp = ratios[len(ratios) // 2] + rows.append({"id": w["id"], "ok": True, "speedup": sp, + "ref_val": ref_val, "agent_val": agent_val}) + except Exception as exc: + return {"ok": False, "error": f"{type(exc).__name__}: {str(exc)[:200]}", + "trace": traceback.format_exc()[-500:]} + return {"ok": True, "rows": rows} + + +# --------------------------------------------------------------------------- # +# Host-side wrappers (run in the judge / agent container; need MODAL tokens). +# --------------------------------------------------------------------------- # +def _build_image(cfg: dict): + import modal + pip = list(cfg.get("pip", ["torch==2.5.1", "triton==3.1.0", "numpy"])) + base = cfg.get("cuda_image", "nvidia/cuda:12.4.1-devel-ubuntu22.04") + # add_python must match the judge/agent container Python (ubuntu:24.04 -> 3.12) + # because the worker ships via serialized=True (cloudpickle is version-sensitive). + # add_local_python_source mounts THIS module into the image so the remote can + # import it when deserializing the (module-level) worker function. + return (modal.Image.from_registry(base, add_python=str(cfg.get("python", "3.12"))) + .entrypoint([]) + .pip_install(*pip) + .add_local_python_source("flash_gpu")) + + +# Substrings that mark a transient Modal control-plane / image-build failure +# (evicted build, gateway hiccup, app stopped mid-run) rather than a real error +# in the submission — safe to retry. +_TRANSIENT_MARKERS = ( + "external shut-down", "terminated due to external", "please try again", + "app_state_stopped", "conflicterror", "deadline exceeded", "connection reset", + "502 bad gateway", "503 service", "temporarily unavailable", "timed out", + "gateway", "eviction", "internalfailure", "internalerror", + "failed to get new inputs", "runner failed", "task exited", +) + + +def _is_transient(text: str) -> bool: + low = (text or "").lower() + return any(m in low for m in _TRANSIENT_MARKERS) + + +def run_remote(payload: dict) -> dict: + """Run the GPU worker on Modal and return its result dict. + + Retries transient Modal control-plane failures; a real error in the + submission is non-transient and surfaces immediately. Raises RuntimeError + with a sanitized message on unrecoverable failure. + """ + import time as _time + + import modal + + cfg = payload["cfg"] + attempts = max(1, int(cfg.get("modal_retries", 3))) + last = "modal_gpu_eval_failed" + for attempt in range(1, attempts + 1): + app = modal.App(cfg.get("app_name", "flash-kernel-eval")) + remote = app.function( + gpu=cfg.get("gpu", "H100"), + image=_build_image(cfg), + timeout=int(cfg.get("modal_timeout_seconds", 1800)), + serialized=True, + max_containers=1, + )(_gpu_worker) + try: + with modal.enable_output(), app.run(): + return remote.remote(payload) + except Exception as exc: # noqa: BLE001 + last = str(exc)[-400:] + if not _is_transient(last) or attempt == attempts: + raise RuntimeError(f"modal_gpu_eval_failed: {last}") from exc + _time.sleep(min(45, 10 * attempt)) + raise RuntimeError(f"modal_gpu_eval_failed: {last}") + + +def modal_available() -> bool: + return bool(os.environ.get("MODAL_TOKEN_ID") and os.environ.get("MODAL_TOKEN_SECRET")) diff --git a/2.0/problems/kmeans_gpu_kernel_optimization/harbor/app/README.md b/2.0/problems/kmeans_gpu_kernel_optimization/harbor/app/README.md new file mode 100644 index 000000000..10536efd3 --- /dev/null +++ b/2.0/problems/kmeans_gpu_kernel_optimization/harbor/app/README.md @@ -0,0 +1,34 @@ +# K-Means kernel optimization — submission workflow + +You are optimizing the `kmeanslib` package at `/app/kmeanslib`. Edit the package +(rewrite the internals of `kmeans`, add Triton kernel modules under +`kmeanslib/`), then submit a patch. + +## Iterate locally + +Run the public self-test to check correctness and get a rough speed signal on +the two public shapes (needs a GPU in the agent container): + +```bash +bash /app/public_test.sh +``` + +## Submit + +```bash +bash /app/make_submission.sh # writes /app/solution.patch (kmeanslib diff) +bash /app/submit.sh # enqueues it for the black-box judge +``` + +Submissions are asynchronous. Submit early and keep improving; use +`bash /app/submissions.sh` and `bash /app/wait_submission.sh ` to inspect +results. + +## Rules + +- Only files under `kmeanslib/` may change. +- Do not import external optimized libraries (write the kernels yourself), and + do not access the environment, spawn processes, or use the network. +- Keep the public `kmeans(...)` signature and return contract unchanged. +- Clustering quality is gated (inertia vs the naive baseline); do not sacrifice + correctness for speed beyond the allowed tolerance. diff --git a/2.0/problems/kmeans_gpu_kernel_optimization/harbor/app/make_submission.sh b/2.0/problems/kmeans_gpu_kernel_optimization/harbor/app/make_submission.sh new file mode 100644 index 000000000..8a6473ab6 --- /dev/null +++ b/2.0/problems/kmeans_gpu_kernel_optimization/harbor/app/make_submission.sh @@ -0,0 +1,24 @@ +#!/usr/bin/env bash +# Package the current kmeanslib edits into a unified diff for submission. +set -euo pipefail + +APP_DIR="${APP_DIR:-/app}" +OUT="${1:-/app/solution.patch}" + +if [[ ! -d "$APP_DIR/kmeanslib" ]]; then + echo "kmeanslib package not found at $APP_DIR/kmeanslib" >&2 + exit 2 +fi +if [[ ! -d "$APP_DIR/.git" ]]; then + echo "git repo not found at $APP_DIR (expected in the agent image)" >&2 + exit 2 +fi + +# Stage and diff only the package, so submission helper scripts under /app are +# never included. New kernel files under kmeanslib/ are captured via `git add`. +git -C "$APP_DIR" add -A -- kmeanslib +git -C "$APP_DIR" diff --cached -- kmeanslib > "$OUT" +git -C "$APP_DIR" reset -q + +bytes=$(wc -c < "$OUT" | tr -d ' ') +echo "Wrote $OUT ($bytes bytes)" diff --git a/2.0/problems/kmeans_gpu_kernel_optimization/harbor/app/public_test.py b/2.0/problems/kmeans_gpu_kernel_optimization/harbor/app/public_test.py new file mode 100644 index 000000000..ea3bf13db --- /dev/null +++ b/2.0/problems/kmeans_gpu_kernel_optimization/harbor/app/public_test.py @@ -0,0 +1,62 @@ +"""Public self-test == the graded evaluation, run by the judge. + +The data generator is judge-only and is NOT present in this (agent) image, so you +cannot read how the workloads are produced and hardcode a data-specific shortcut. +Instead, this packages your current `/app/kmeanslib`, submits it to the judge -- +which runs the EXACT graded workloads on a GPU and grades your result against the +frozen reference -- and prints your per-workload result + score. It is +byte-for-byte the same evaluation used for your final grade. + +(Submission is asynchronous; this waits for the result. You can also submit +directly with `bash /app/make_submission.sh && bash /app/submit.sh` and poll with +`bash /app/submissions.sh`.) +""" +from __future__ import annotations + +import json +import re +import subprocess +import sys +from pathlib import Path + +SUBMISSIONS_LOG = Path("/logs/agent/submissions.jsonl") +_UUID = re.compile(r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}") + + +def _latest_uuid_from_log() -> str | None: + if not SUBMISSIONS_LOG.exists(): + return None + lines = [l for l in SUBMISSIONS_LOG.read_text(encoding="utf-8").splitlines() if l.strip()] + for l in reversed(lines): + try: + u = json.loads(l).get("submission_uuid") + if u: + return u + except Exception: # noqa: BLE001 + continue + return None + + +def main() -> int: + # 1) package the current package into /app/solution.patch + subprocess.run(["bash", "/app/make_submission.sh"], check=True) + # 2) submit to the judge (async); it runs the exact graded eval on GPU + print("Submitting to the judge (runs the exact graded workloads + returns your score)...", + flush=True) + r = subprocess.run(["python3", "/app/submit.py", *sys.argv[1:]], + capture_output=True, text=True) + sys.stdout.write(r.stdout) + sys.stderr.write(r.stderr) + if r.returncode != 0: + return r.returncode + # 3) find the submission uuid and wait for the result + m = _UUID.findall(r.stdout + r.stderr) + uuid = m[-1] if m else _latest_uuid_from_log() + if not uuid: + print("Submitted. Poll for the result with: bash /app/submissions.sh") + return 0 + return subprocess.call(["python3", "/app/wait_submission.py", uuid]) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/2.0/problems/kmeans_gpu_kernel_optimization/harbor/app/public_test.sh b/2.0/problems/kmeans_gpu_kernel_optimization/harbor/app/public_test.sh new file mode 100644 index 000000000..ca70fc405 --- /dev/null +++ b/2.0/problems/kmeans_gpu_kernel_optimization/harbor/app/public_test.sh @@ -0,0 +1,7 @@ +#!/usr/bin/env bash +# Public self-test: run the patched kmeanslib on the public shapes and print a +# correctness + rough speedup summary. Requires a GPU in the agent container. +set -euo pipefail +APP_DIR="${APP_DIR:-/app}" +cd "$APP_DIR" +exec python3 "$APP_DIR/public_test.py" diff --git a/2.0/problems/kmeans_gpu_kernel_optimization/harbor/app/solution.patch b/2.0/problems/kmeans_gpu_kernel_optimization/harbor/app/solution.patch new file mode 100644 index 000000000..e69de29bb diff --git a/2.0/problems/kmeans_gpu_kernel_optimization/judge/refkmeans.py b/2.0/problems/kmeans_gpu_kernel_optimization/judge/refkmeans.py new file mode 100644 index 000000000..98994eaec --- /dev/null +++ b/2.0/problems/kmeans_gpu_kernel_optimization/judge/refkmeans.py @@ -0,0 +1,44 @@ +"""Frozen naive one-step K-Means baseline used by the judge as the speed +denominator. Behaviourally identical to the shipped ``kmeanslib.kmeans`` +(``step`` + ``_assign`` + ``_update``); imported under its own module name so the +judge worker can load the frozen baseline and the patched ``kmeanslib`` package +in one process. The judge owns the Lloyd loop and calls ``step`` a fixed number +of times. +""" +from __future__ import annotations + +import torch + + +def _assign(x: torch.Tensor, centroids: torch.Tensor) -> torch.Tensor: + c = centroids.to(x.dtype) + c_sq = (c.float() * c.float()).sum(1) + labels = torch.empty(x.shape[0], device=x.device, dtype=torch.long) + for lo in range(0, x.shape[0], 16384): + xb = x[lo:lo + 16384] + dist = c_sq[None, :] - 2.0 * (xb @ c.t()).float() # bf16 matmul, fp32 accum + labels[lo:lo + 16384] = torch.argmin(dist, dim=1) + return labels + + +def _update(x: torch.Tensor, labels: torch.Tensor, n_clusters: int, + old_centroids: torch.Tensor) -> torch.Tensor: + N, D = x.shape + sums = torch.zeros((n_clusters, D), device=x.device, dtype=torch.float32) + counts = torch.zeros((n_clusters,), device=x.device, dtype=torch.float32) + sums.index_add_(0, labels, x.float()) + counts.index_add_(0, labels, torch.ones(N, device=x.device, dtype=torch.float32)) + empty = counts == 0 + counts = counts.clamp_min(1.0) + new = sums / counts[:, None] + if empty.any(): + new[empty] = old_centroids[empty].float() + return new.to(x.dtype) + + +def step(x: torch.Tensor, centroids: torch.Tensor): + if x.ndim != 2: + raise ValueError(f"x must be 2-D (N, D); got shape {tuple(x.shape)}") + labels = _assign(x, centroids) + new_centroids = _update(x, labels, centroids.shape[0], centroids) + return labels, new_centroids diff --git a/2.0/problems/kmeans_gpu_kernel_optimization/kmeanslib/__init__.py b/2.0/problems/kmeans_gpu_kernel_optimization/kmeanslib/__init__.py new file mode 100644 index 000000000..fed5e9319 --- /dev/null +++ b/2.0/problems/kmeans_gpu_kernel_optimization/kmeanslib/__init__.py @@ -0,0 +1,16 @@ +"""kmeanslib -- a tiny GPU K-Means library you are asked to make fast. + +The public entry point is :func:`step`, one Lloyd iteration +``step(x, centroids) -> (labels, new_centroids)``. The judge owns the iteration +loop and calls it a fixed number of times. The shipped implementation is correct +but deliberately unoptimised. Rewrite the internals (new Triton kernels, a fused +assign+update pass, better memory traffic, ...) so that ``step`` runs as fast as +possible while producing the same clustering. The public function signature and +return contract must not change. +""" +from __future__ import annotations + +from kmeanslib.kmeans import step + +__all__ = ["step"] +__version__ = "0.1.0" diff --git a/2.0/problems/kmeans_gpu_kernel_optimization/kmeanslib/kmeans.py b/2.0/problems/kmeans_gpu_kernel_optimization/kmeanslib/kmeans.py new file mode 100644 index 000000000..2eadfdbff --- /dev/null +++ b/2.0/problems/kmeans_gpu_kernel_optimization/kmeanslib/kmeans.py @@ -0,0 +1,80 @@ +"""One Lloyd K-Means step -- the function you optimise. + +The judge owns the iteration loop, the data, the initial centroids, and the +iteration count; you provide a single Lloyd step and it is called repeatedly. +This shipped step is intentionally naive (materialise a (chunk, K) distance +matrix with a bf16 matmul + argmin, then a PyTorch scatter update). + +Contract (do NOT change): + + step(x, centroids) -> (labels, new_centroids) + + x : (N, D) bfloat16 CUDA tensor of points. + centroids : (K, D) bfloat16 tensor -- the current centroids. + labels : (N,) int64 -- nearest-centroid assignment of every point to + `centroids` (a full assignment of all N points). + new_centroids: (K, D) -- centroids recomputed as the mean of each cluster + (empty clusters keep their previous centroid). + +This is exactly one Lloyd iteration: assign to `centroids`, then update. You may +add modules/kernels under ``kmeanslib`` and rewrite the body of ``step``, +``_assign`` and ``_update`` freely -- including **fusing the assign + update into +a single kernel** -- as long as ``step(x, centroids) -> (labels, new_centroids)`` +is preserved. You cannot change how many times it runs, the data, or the initial +centroids; the judge owns the loop and calls ``step`` a fixed number of times. +""" +from __future__ import annotations + +import torch + + +def _assign(x: torch.Tensor, centroids: torch.Tensor) -> torch.Tensor: + """Nearest-centroid assignment by squared-L2 distance. + + Naive: for each chunk of points, materialise the (chunk, K) distance matrix + with a bf16 matmul, then argmin. ``argmin ||x-c||^2 == argmin (||c||^2 - 2 + x.c^T)`` since ``||x||^2`` is constant per row. bf16 in, fp32 accumulation. + """ + c = centroids.to(x.dtype) + c_sq = (c.float() * c.float()).sum(1) # (K,) fp32 + labels = torch.empty(x.shape[0], device=x.device, dtype=torch.long) + for lo in range(0, x.shape[0], 16384): + xb = x[lo:lo + 16384] + dist = c_sq[None, :] - 2.0 * (xb @ c.t()).float() # (chunk, K) fp32 + labels[lo:lo + 16384] = torch.argmin(dist, dim=1) + return labels + + +def _update( + x: torch.Tensor, + labels: torch.Tensor, + n_clusters: int, + old_centroids: torch.Tensor, +) -> torch.Tensor: + """Recompute each centroid as the mean of its assigned points. + + Empty clusters keep their previous centroid. + """ + N, D = x.shape + sums = torch.zeros((n_clusters, D), device=x.device, dtype=torch.float32) + counts = torch.zeros((n_clusters,), device=x.device, dtype=torch.float32) + sums.index_add_(0, labels, x.float()) + counts.index_add_(0, labels, torch.ones(N, device=x.device, dtype=torch.float32)) + empty = counts == 0 + counts = counts.clamp_min(1.0) + new = sums / counts[:, None] + if empty.any(): + new[empty] = old_centroids[empty].float() + return new.to(x.dtype) + + +def step(x: torch.Tensor, centroids: torch.Tensor): + """One Lloyd iteration: assign to `centroids`, then recompute them. + + Returns ``(labels, new_centroids)``. See the module docstring for the contract. + """ + if x.ndim != 2: + raise ValueError(f"x must be 2-D (N, D); got shape {tuple(x.shape)}") + labels = _assign(x, centroids) + new_centroids = _update(x, labels, centroids.shape[0], centroids) + return labels, new_centroids diff --git a/2.0/problems/kmeans_gpu_kernel_optimization/readme b/2.0/problems/kmeans_gpu_kernel_optimization/readme new file mode 100644 index 000000000..048eb2c57 --- /dev/null +++ b/2.0/problems/kmeans_gpu_kernel_optimization/readme @@ -0,0 +1,133 @@ +# GPU K-Means Kernel Optimization + +## Problem + +You are given a small GPU K-Means library, `kmeanslib`, in the Harbor workspace +at `/app/kmeanslib`. Its public entry point is a **single Lloyd iteration**: + +```python +kmeanslib.step(x, centroids) -> (labels, new_centroids) +``` + +`x` is an `(N, D)` **bfloat16** CUDA tensor of points and `centroids` is the +current `(K, D)` bfloat16 tensor. One call performs exactly one Euclidean +(squared-L2) Lloyd step: assign every point to its nearest centroid +(`labels`, an `(N,)` int64 full assignment), then recompute each centroid as the +mean of its assigned points (`new_centroids`, `(K, D)`; empty clusters keep their +previous centroid). All data is bfloat16 — treat it as the fixed working +precision. The shipped implementation is a straightforward, correct PyTorch +version (a bf16 matmul + argmin, then a scatter update). + +**The judge owns the iteration loop.** It fixes the data and the initial +centroids, then calls your `step` a fixed number of times per workload, feeding +each call's `new_centroids` into the next. You do not control the data, the +initial centroids, or how many iterations run — you control only how fast a +single `step` executes. + +Your goal is to make `kmeanslib.step` **as fast as possible** on the GPU while +producing the same clustering. You may rewrite the internals of the package +however you like and add new modules (including Triton kernels) under +`kmeanslib/` — in particular you may **fuse the assign and update into a single +kernel**. The public function signature and return contract above must not +change, and each result must remain a deterministic function of its inputs. + +## Workload + +The graded workloads are a family of held-out dense clustering problems that +vary `(N, D, K)`, spanning small/medium/large point counts and both wide-feature +(large `D`) and many-cluster (large `K`) regimes, in **bfloat16**, with a fixed +number of judge-owned `step` calls per workload and caller-supplied initial +centroids. How the point sets are generated is not disclosed — implement a +**general** exact Lloyd step (correct assignment to the given centroids + exact +cluster-mean update) and do not special-case the data. + +## Iterate on a GPU + +The agent workspace has no GPU, and the data generator is judge-only (it is not in +your image). To check your current code, **submit it to the judge** — it runs the +exact graded workloads on a GPU and returns your per-workload result + score. This +one command packages, submits, and waits for the result: + +```bash +bash /app/public_test.sh +``` + +It reports, per graded workload, whether your result passes the quality gate and +your speedup, plus the geometric-mean speedup and your score (0-100) — the +identical evaluation used for your final grade. (Equivalently: +`bash /app/make_submission.sh && bash /app/submit.sh`, then poll with +`bash /app/submissions.sh`.) Submissions are asynchronous; submit early and iterate. + +## Submission + +The submitted artifact is a patch over the `kmeanslib` package: + +```text +/app/solution.patch +``` + +After editing `/app/kmeanslib`, generate and submit: + +```bash +bash /app/make_submission.sh +bash /app/submit.sh +``` + +Submissions are asynchronous; submit early and keep iterating. The judge applies +your patch to a clean copy of `kmeanslib` and times it against the original +baseline on a GPU, on the same seeded data and initial centroids. + +## Correctness + +Correctness is a gate. After running the fixed loop of `step` calls, the judge +computes the **inertia of your returned clustering** — the sum of squared +distances of every point to `new_centroids[labels]`, i.e. using **the `labels` +and `new_centroids` your final `step` returned** — and compares it to the +baseline's, requiring your inertia to stay within a small relative tolerance. +Because the gate reads your returned `labels`, they must be the real full `(N,)` +nearest-centroid assignment to the `centroids` each `step` was given: returning +fake/empty labels, a partial or sub-sampled assignment, or centroids computed +from a subset of points all inflate this inertia and fail the gate. Each `step` +returns `labels` `(N,)` and `new_centroids` `(K, D)`. Crashes, non-finite output, +wrong shapes/dtypes, timeouts, and clustering that regresses beyond the tolerance +are penalized before speed is considered. The working precision is fixed at bfloat16 for every solution (the +inputs are bfloat16), so it is not a tuning knob — speed comes from the kernel, +not from the arithmetic precision. + +## Scoring + +Valid submissions are scored by speedup relative to the baseline on the same +hardware and workloads. For each workload: + +```text +speedup = baseline_time / your_time +``` + +The objective is the geometric mean of per-workload speedups, so broad speedups +are preferred over a single large outlier. A result no faster than the baseline +earns 0; regressions earn 0. The raw geometric-mean speedup is reported in the +evaluator metrics. + +## Patch Policy + +The evaluator validates the patch before running it. Only Python files under the +package may change: + +```text +kmeanslib/** +``` + +New Python modules inside `kmeanslib/` are allowed. Patches may not: modify +anything outside `kmeanslib/`; import or call an external optimized ML/kernel +library (write the kernels yourself); read or write environment variables, spawn +processes, or access the network; or otherwise tamper with the measurement +framework. The judge owns the loop, the data, and the initial centroids and +re-verifies quality from your final centroids, so faking labels, skipping the +real assign/update work, or caching results across calls does not help. + +## Resource Budget + +```text +GPU: single Modal GPU (H100 reference; the Triton paths also run on L40S / A100) +Agent container: CPU-only (GPU work is offloaded to Modal) +``` diff --git a/2.0/problems/kmeans_gpu_kernel_optimization/reference.patch b/2.0/problems/kmeans_gpu_kernel_optimization/reference.patch new file mode 100644 index 000000000..1f87526dd --- /dev/null +++ b/2.0/problems/kmeans_gpu_kernel_optimization/reference.patch @@ -0,0 +1,2530 @@ +diff --git a/kmeanslib/_kernels/__init__.py b/kmeanslib/_kernels/__init__.py +new file mode 100644 +index 0000000..e69de29 +diff --git a/kmeanslib/_kernels/primitives/__init__.py b/kmeanslib/_kernels/primitives/__init__.py +new file mode 100644 +index 0000000..e69de29 +diff --git a/kmeanslib/_kernels/primitives/kmeans/__init__.py b/kmeanslib/_kernels/primitives/kmeans/__init__.py +new file mode 100644 +index 0000000..e69de29 +diff --git a/kmeanslib/_kernels/primitives/kmeans/triton/__init__.py b/kmeanslib/_kernels/primitives/kmeans/triton/__init__.py +new file mode 100644 +index 0000000..043b4b7 +--- /dev/null ++++ b/kmeanslib/_kernels/primitives/kmeans/triton/__init__.py +@@ -0,0 +1,43 @@ ++"""kmeans triton backend. ++ ++Re-exports the public Python wrappers from each component file. ++``@triton.jit`` kernels stay private to their file. ++""" ++from kmeanslib._kernels.primitives.kmeans.triton.assign import ( ++ _ceil_div, ++ euclid_assign_triton, ++ cosine_assign_triton, ++) ++from kmeanslib._kernels.primitives.kmeans.triton.kmeans import ( ++ COMPILE_FLAG, ++ batch_kmeans_Euclid, ++ batch_kmeans_Cosine, ++ batch_kmeans_Dot, ++) ++from kmeanslib._kernels.primitives.kmeans.triton.update import ( ++ triton_centroid_update_cosine, ++ torch_loop_centroid_update_cosine, ++ triton_centroid_update_euclid, ++ triton_centroid_update_sorted_cosine, ++ triton_centroid_update_sorted_euclid, ++ triton_centroid_finalize, ++ triton_lloyd_centroid_step_euclid, ++ main, ++) ++ ++__all__ = [ ++ "euclid_assign_triton", ++ "cosine_assign_triton", ++ "COMPILE_FLAG", ++ "batch_kmeans_Euclid", ++ "batch_kmeans_Cosine", ++ "batch_kmeans_Dot", ++ "triton_centroid_update_cosine", ++ "torch_loop_centroid_update_cosine", ++ "triton_centroid_update_euclid", ++ "triton_centroid_update_sorted_cosine", ++ "triton_centroid_update_sorted_euclid", ++ "triton_centroid_finalize", ++ "triton_lloyd_centroid_step_euclid", ++ "main", ++] +diff --git a/kmeanslib/_kernels/primitives/kmeans/triton/assign.py b/kmeanslib/_kernels/primitives/kmeans/triton/assign.py +new file mode 100644 +index 0000000..5014a67 +--- /dev/null ++++ b/kmeanslib/_kernels/primitives/kmeans/triton/assign.py +@@ -0,0 +1,1476 @@ ++from typing import Optional ++import torch ++import triton ++import triton.language as tl ++ ++# =============================================================== ++# Triton kernel: compute nearest-centroid IDs (Euclidean distance). ++# ++# x²-free score: the kernel ranks ``s(n, k) = c_sq[k] − 2·`` ++# which is ``||x − c||² − ||x||²`` and so produces the same argmin ++# without loading any ``x_sq`` tensor. See assign_euclid in ++# klib/primitives/kmeans/torch_fallback.py for the matching CPU ++# reference and the gather kernel in klib/kernels/distance for ++# recovering the true distances post-hoc. ++# ++# Inputs: ++# x : (B, N, D) float16 / float32 ++# centroids : (B, K, D) same dtype as x ++# Output: ++# cluster_ids : (B, N) int32 -- nearest centroid index per point ++# =============================================================== ++ ++ ++def _ceil_div(a: int, b: int) -> int: ++ return (a + b - 1) // b ++ ++ ++def _next_power_of_2(n: int) -> int: ++ p = 1 ++ while p < n: ++ p <<= 1 ++ return p ++ ++ ++# ----------------------------------------------------------------------------- ++# Auto-tuning setup – explore various tile sizes / warp counts ++# ----------------------------------------------------------------------------- ++ ++_TUNE_CONFIGS = [ ++ triton.Config({"BLOCK_N": BN, "BLOCK_K": BK}, num_stages=num_stages, num_warps=wp) ++ for BN in [32, 64, 128] ++ for BK in [32, 64, 128] ++ for wp in [4, 8] ++ for num_stages in [1, 2, 4] ++] ++ ++ ++def _cfg_keep(conf): ++ """Basic heuristic to prune unbalanced configs.""" ++ BN = conf.kwargs["BLOCK_N"] ++ BK = conf.kwargs["BLOCK_K"] ++ # Avoid tiny tiles on many warps ++ if BN * BK < 32 * 32 and conf.num_warps > 4: ++ return False ++ return True ++ ++_TUNE_CONFIGS = list(filter(_cfg_keep, _TUNE_CONFIGS)) ++ ++ ++# Tuning grid for the split-D kernel. Adds a third tile dim ``BLOCK_D`` ++# that controls how much of the feature dimension is materialised inside ++# each program at a time. Pruned to keep peak SMEM and register pressure ++# under control on conservative GPUs (GB10). ++_TUNE_CONFIGS_SPLIT_D = [ ++ triton.Config({"BLOCK_N": BN, "BLOCK_K": BK, "BLOCK_D": BD}, ++ num_stages=num_stages, num_warps=wp) ++ for BN in [32, 64, 128] ++ for BK in [32, 64, 128] ++ for BD in [32, 64, 128] ++ for wp in [4, 8] ++ for num_stages in [1, 2, 4] ++] ++ ++ ++def _cfg_keep_split_d(conf): ++ BN = conf.kwargs["BLOCK_N"] ++ BK = conf.kwargs["BLOCK_K"] ++ BD = conf.kwargs["BLOCK_D"] ++ # Tiny tiles do not need many warps. ++ if BN * BK < 32 * 32 and conf.num_warps > 4: ++ return False ++ # Cap (BN, BK) tile size for register/SMEM safety. The (BN, BK) fp32 ++ # cross accumulator is the same size as the small-D kernel, so keeping ++ # BN*BK <= 128*128 prevents new spill regressions. ++ if BN * BK > 128 * 128: ++ return False ++ # Prune the largest combined work tiles to keep tuning wall-clock down ++ # without losing useful configs. ++ if BN * BK * BD > 128 * 128 * 128: ++ return False ++ return True ++ ++ ++_TUNE_CONFIGS_SPLIT_D = list(filter(_cfg_keep_split_d, _TUNE_CONFIGS_SPLIT_D)) ++ ++_HALF_DTYPES = (torch.float16, torch.bfloat16) ++ ++ ++def _dtype_bytes(dtype) -> int: ++ """Element size in bytes for a torch / numpy-ish dtype. ++ ++ Falls back to 2 (fp16) when the dtype is unknown to keep prior behaviour ++ (the heuristic was originally tuned with fp16 in mind). ++ """ ++ if dtype is None: ++ return 2 ++ if isinstance(dtype, torch.dtype): ++ return torch.tensor([], dtype=dtype).element_size() ++ # Allow callers to pass a raw byte size. ++ if isinstance(dtype, int): ++ return dtype ++ return 2 ++ ++ ++def _is_half_dtype(dtype) -> bool: ++ """True for fp16/bf16 (the original tuning regime). ++ ++ For these dtypes we skip the SMEM-fitting fallback entirely so heuristic ++ selection on already-validated GPUs (H100/H200/A100) is byte-for-byte ++ identical to the previous behaviour. ++ """ ++ if dtype is None: ++ return True ++ if isinstance(dtype, torch.dtype): ++ return dtype in _HALF_DTYPES ++ return False ++ ++ ++def _smem_bytes(D: int, BN: int, BK: int, num_stages: int, dtype_bytes: int) -> int: ++ """Approximate dynamic shared-memory usage of `_euclid_assign_kernel`. ++ ++ The kernel materialises: ++ - one ``x_tile`` of shape (BN, D) outside the K loop, and ++ - ``num_stages`` copies of ``c_tile`` of shape (D, BK) for the software ++ pipelined K loop. ++ ++ Other buffers (c_sq, masks, accumulators) are negligible compared ++ to these and are ignored. The x²-free kernel does not load ``x_sq``. ++ """ ++ return D * dtype_bytes * (BN + num_stages * BK) ++ ++ ++def _smem_limit(device) -> int: ++ """Per-block dynamic shared-memory budget for ``device``. ++ ++ Triton uses opt-in dynamic shared memory; prefer that attribute when ++ available, fall back to the static limit, and finally to a conservative ++ 48 KiB for very old PyTorch builds. ++ """ ++ props = torch.cuda.get_device_properties(device) ++ for attr in ( ++ "shared_memory_per_block_optin", ++ "max_shared_memory_per_block_optin", ++ "shared_memory_per_block", ++ "max_shared_memory_per_block", ++ ): ++ v = getattr(props, attr, None) ++ if v: ++ return int(v) ++ return 48 * 1024 ++ ++ ++def _fit_config_to_smem( ++ cfg: dict, ++ D: int, ++ dtype_bytes: int, ++ smem_limit: int, ++) -> dict: ++ """Return a config that fits ``smem_limit`` and is closest to ``cfg``. ++ ++ The original config is returned unchanged whenever it already fits. If ++ not, we enumerate all power-of-two ``(BLOCK_N, BLOCK_K, num_stages)`` ++ that are no larger than the original and pick the one that maximises ++ work-per-program tile (``BLOCK_N * BLOCK_K * num_stages``), breaking ++ ties towards the original aspect ratio. This avoids the pitfall of a ++ pure greedy halving (e.g. shrinking BLOCK_K all the way to 16 when only ++ a single halving was needed). ++ ++ Raises ``RuntimeError`` if even ``(BN=16, BK=16, S=1)`` does not fit – ++ this only happens for absurdly large D combined with fp32 on tiny-SMEM ++ GPUs. ++ """ ++ BN0 = int(cfg["BLOCK_N"]) ++ BK0 = int(cfg["BLOCK_K"]) ++ W0 = int(cfg["num_warps"]) ++ S0 = int(cfg["num_stages"]) ++ ++ if _smem_bytes(D, BN0, BK0, S0, dtype_bytes) <= smem_limit: ++ return {"BLOCK_N": BN0, "BLOCK_K": BK0, "num_warps": W0, "num_stages": S0} ++ ++ def _pow2_down_to_16(v): ++ out = [] ++ x = v ++ while x >= 16: ++ out.append(x) ++ x //= 2 ++ return out ++ ++ best = None ++ best_key = None ++ for BN in _pow2_down_to_16(BN0): ++ for BK in _pow2_down_to_16(BK0): ++ for S in range(S0, 0, -1): ++ if _smem_bytes(D, BN, BK, S, dtype_bytes) > smem_limit: ++ continue ++ # Prefer larger total tile work, then closer aspect ratio ++ # to the original, then larger BLOCK_N (more parallelism ++ # along N), then larger num_stages (better pipelining). ++ aspect_penalty = abs( ++ (BN / max(BK, 1)) - (BN0 / max(BK0, 1)) ++ ) ++ key = (BN * BK * S, -aspect_penalty, BN, S) ++ if best_key is None or key > best_key: ++ best_key = key ++ best = (BN, BK, S) ++ ++ if best is None: ++ raise RuntimeError( ++ f"euclid_assign_triton: cannot fit kernel into shared memory " ++ f"(D={D}, dtype_bytes={dtype_bytes}, smem_limit={smem_limit}). " ++ f"Even BLOCK_N=16, BLOCK_K=16, num_stages=1 needs " ++ f"{_smem_bytes(D, 16, 16, 1, dtype_bytes)} bytes." ++ ) ++ ++ BN, BK, S = best ++ W = W0 ++ # Tiny tiles do not benefit from many warps and may even fail to compile ++ # for some Triton versions; cap to 4. ++ if BN * BK <= 32 * 32 and W > 4: ++ W = 4 ++ ++ return {"BLOCK_N": BN, "BLOCK_K": BK, "num_warps": W, "num_stages": S} ++ ++ ++def _smem_bytes_split_d(BD: int, BN: int, BK: int, num_stages: int, dtype_bytes: int) -> int: ++ """SMEM estimate for ``_euclid_assign_kernel_split_d`` per program. ++ ++ The split-D kernel materialises: ++ - one ``x_chunk`` of shape (BN, BD) per D-tile, and ++ - ``num_stages`` copies of ``c_chunk`` of shape (BD, BK) for the software ++ pipelined inner D loop. ++ """ ++ return BD * dtype_bytes * (BN + num_stages * BK) ++ ++ ++def _smallD_kernel_fits_smem(D: int, dtype_bytes: int, smem_limit: int) -> bool: ++ """Return True if even the tiniest small-D kernel config fits SMEM. ++ ++ Used by the wrapper to fall back to split-D when the small-D kernel can ++ not run at all (e.g., GB10 + fp32 + D=448). ++ """ ++ return _smem_bytes(D, 16, 16, 1, dtype_bytes) <= smem_limit ++ ++ ++def _fit_config_to_smem_split_d( ++ cfg: dict, ++ D: int, ++ dtype_bytes: int, ++ smem_limit: int, ++) -> dict: ++ """Shrink a split-D config until it fits ``smem_limit``. ++ ++ Mirrors ``_fit_config_to_smem`` with the additional ``BLOCK_D`` axis. ++ Returns the largest-work config that fits, breaking ties towards the ++ original aspect ratio. ``BLOCK_D`` is also clamped to D when D < BD0 ++ (no point materialising more dims than exist). ++ """ ++ BN0 = int(cfg["BLOCK_N"]) ++ BK0 = int(cfg["BLOCK_K"]) ++ W0 = int(cfg["num_warps"]) ++ S0 = int(cfg["num_stages"]) ++ BD0 = int(cfg["BLOCK_D"]) ++ # No point letting BD exceed next_pow2(D) (the loop would still run ++ # once). We round D up to the next power-of-2 because BLOCK_D feeds ++ # into ``tl.arange(0, BLOCK_D)`` which Triton requires to be pow2. ++ D_ceil = _next_power_of_2(max(D, 16)) ++ BD0 = min(BD0, D_ceil) ++ ++ if _smem_bytes_split_d(BD0, BN0, BK0, S0, dtype_bytes) <= smem_limit: ++ return { ++ "BLOCK_N": BN0, "BLOCK_K": BK0, "BLOCK_D": BD0, ++ "num_warps": W0, "num_stages": S0, ++ } ++ ++ def _pow2_down_to_16(v): ++ out = [] ++ x = v ++ while x >= 16: ++ out.append(x) ++ x //= 2 ++ return out ++ ++ best = None ++ best_key = None ++ for BN in _pow2_down_to_16(BN0): ++ for BK in _pow2_down_to_16(BK0): ++ for BD in _pow2_down_to_16(BD0): ++ for S in range(S0, 0, -1): ++ if _smem_bytes_split_d(BD, BN, BK, S, dtype_bytes) > smem_limit: ++ continue ++ aspect_penalty = abs( ++ (BN / max(BK, 1)) - (BN0 / max(BK0, 1)) ++ ) ++ key = (BN * BK * BD * S, -aspect_penalty, BN, BD, S) ++ if best_key is None or key > best_key: ++ best_key = key ++ best = (BN, BK, BD, S) ++ ++ if best is None: ++ raise RuntimeError( ++ f"euclid_assign_triton (split-D): cannot fit kernel into shared " ++ f"memory (D={D}, dtype_bytes={dtype_bytes}, smem_limit={smem_limit})." ++ ) ++ ++ BN, BK, BD, S = best ++ W = W0 ++ if BN * BK <= 32 * 32 and W > 4: ++ W = 4 ++ return { ++ "BLOCK_N": BN, "BLOCK_K": BK, "BLOCK_D": BD, ++ "num_warps": W, "num_stages": S, ++ } ++ ++ ++# ----------------------------------------------------------------------------- ++# Per-arch small-D heuristic functions. These bodies are the original ++# hand-tuned tables, moved verbatim here so the top-level dispatcher stays ++# small and so the split-D path can live alongside without touching them. ++# Any change here must be guarded by examples/regression_fp16_smalld.py. ++# ----------------------------------------------------------------------------- ++ ++ ++def _heuristic_euclid_config_h200_smallD(N: int, K: int, D: int, dtype) -> dict: ++ """H200 small-D heuristic for the x²-free Euclidean kernel. ++ ++ Derived from a fresh grid sweep (``scripts/tune_euclid_h200.py`` in ++ flash-kmeans) after dropping the ``||x||²`` load and the underflow ++ clamp from ``_euclid_assign_kernel``: N ∈ {65536, 1048576}, ++ K ∈ {256, 4096, 65536, 200000}, D ∈ {64, 128, 256, 512}, B=1, ++ fp16 + fp32. ++ ++ fp32: ++ - D=64 K<=256 : BN=128 BK=64 W=4 S=1 ++ - D=64 25665K : BN=128 BK=64 W=4 S=1 ++ - D=128 K<=4K : BN=128 BK=64 W=4 S=1 ++ - D=128 K>4K : BN=128 BK=64 W=8 S=1 ++ - D=256 : BN=128 BK=64 W=8 S=1 ++ - D=512 (and 320/384/448) : BN=64 BK=32 W=4 S=1 ++ (fp32 + wide-D only fits H200's 226 KiB SMEM with the smallest tile) ++ ++ fp16 / bf16: ++ - D=64 K<=256 : BN=128 BK=128 W=4 S=2 ++ - D=64 256=200K : BN=128 BK=64 W=4 S=4 ++ - D=128 K<=256 : BN=128 BK=64 W=4 S=1 ++ - D=128 25665K : BN=128 BK=64 W=4 S=1 ++ - D=256 K<=4K : BN=128 BK=64 W=4 S=1 ++ - D=256 K>4K : BN=128 BK=64 W=8 S=1 ++ - D>=320 (i.e. D=512) : BN=128 BK=64 W=8 S=1 ++ ++ Tiny N (<65536) shrinks BN to 64 to avoid wasted work -- kept as a ++ guard since the new sweep only covered N>=65536. ++ """ ++ half = _is_half_dtype(dtype) ++ ++ if not half: ++ if D <= 64: ++ if K <= 256: ++ return {"BLOCK_N": 128, "BLOCK_K": 64, "num_warps": 4, "num_stages": 1} ++ if K <= 65536: ++ return {"BLOCK_N": 128, "BLOCK_K": 128, "num_warps": 4, "num_stages": 1} ++ return {"BLOCK_N": 128, "BLOCK_K": 64, "num_warps": 4, "num_stages": 1} ++ if D <= 128: ++ if K <= 4096: ++ return {"BLOCK_N": 128, "BLOCK_K": 64, "num_warps": 4, "num_stages": 1} ++ return {"BLOCK_N": 128, "BLOCK_K": 64, "num_warps": 8, "num_stages": 1} ++ if D <= 256: ++ return {"BLOCK_N": 128, "BLOCK_K": 64, "num_warps": 8, "num_stages": 1} ++ cfg = {"BLOCK_N": 64, "BLOCK_K": 32, "num_warps": 4, "num_stages": 1} ++ if N < 65536: ++ cfg = dict(cfg) ++ cfg["BLOCK_N"] = 32 ++ return cfg ++ ++ if D <= 64: ++ if K <= 256: ++ cfg = {"BLOCK_N": 128, "BLOCK_K": 128, "num_warps": 4, "num_stages": 2} ++ elif K < 200_000: ++ cfg = {"BLOCK_N": 64, "BLOCK_K": 128, "num_warps": 4, "num_stages": 4} ++ else: ++ cfg = {"BLOCK_N": 128, "BLOCK_K": 64, "num_warps": 4, "num_stages": 4} ++ elif D <= 128: ++ if K <= 256: ++ cfg = {"BLOCK_N": 128, "BLOCK_K": 64, "num_warps": 4, "num_stages": 1} ++ elif K <= 65536: ++ cfg = {"BLOCK_N": 128, "BLOCK_K": 128, "num_warps": 8, "num_stages": 2} ++ else: ++ cfg = {"BLOCK_N": 128, "BLOCK_K": 64, "num_warps": 4, "num_stages": 1} ++ elif D <= 256: ++ if K <= 4096: ++ cfg = {"BLOCK_N": 128, "BLOCK_K": 64, "num_warps": 4, "num_stages": 1} ++ else: ++ cfg = {"BLOCK_N": 128, "BLOCK_K": 64, "num_warps": 8, "num_stages": 1} ++ else: ++ cfg = {"BLOCK_N": 128, "BLOCK_K": 64, "num_warps": 8, "num_stages": 1} ++ ++ if N < 65536: ++ cfg = dict(cfg) ++ cfg["BLOCK_N"] = 64 ++ ++ return cfg ++ ++ ++def _heuristic_euclid_config_h100_smallD(N: int, K: int, D: int, dtype) -> dict: ++ # H100 tuned heuristic (more conservative on D=64 mid-K vs H200). ++ block_n = 128 ++ block_k = 64 ++ num_warps = 4 ++ num_stages = 1 ++ ++ if D >= 512: ++ block_n = 128 ++ block_k = 64 ++ num_warps = 8 ++ num_stages = 1 ++ elif D >= 256: ++ block_n = 128 ++ block_k = 64 ++ if K <= 1024: ++ num_warps = 8 ++ num_stages = 1 ++ elif K <= 16384: ++ num_warps = 4 ++ num_stages = 1 ++ else: ++ num_warps = 8 ++ num_stages = 1 ++ else: ++ # D <= 128 ++ if D <= 64: ++ if K <= 1024: ++ block_k = 64 ++ num_warps = 4 ++ num_stages = 2 ++ elif K <= 16384: ++ block_k = 64 ++ num_warps = 4 ++ num_stages = 2 ++ elif K <= 65536: ++ block_k = 128 ++ num_warps = 4 ++ num_stages = 4 ++ else: ++ block_k = 64 ++ num_warps = 4 ++ num_stages = 4 ++ else: ++ # D == 128 ++ if K <= 1024: ++ block_k = 64 ++ num_warps = 4 ++ num_stages = 1 ++ elif K <= 65536: ++ block_k = 128 ++ num_warps = 8 ++ num_stages = 2 ++ else: ++ block_k = 64 ++ num_warps = 4 ++ num_stages = 4 ++ ++ if N < 65536: ++ block_n = 64 ++ ++ return { ++ "BLOCK_N": block_n, ++ "BLOCK_K": block_k, ++ "num_warps": num_warps, ++ "num_stages": num_stages, ++ } ++ ++ ++def _heuristic_euclid_config_a100_smallD(N: int, K: int, D: int, dtype) -> dict: ++ # Robust default on A100 across tuned grid. ++ block_n = 128 ++ block_k = 32 ++ num_warps = 4 ++ num_stages = 2 ++ ++ if D == 128: ++ # Small-N cases tend to prefer a larger K tile. ++ if N <= 65536: ++ block_k = 64 ++ elif D == 256: ++ # D=256 benefits from deeper pipeline at larger K. ++ if K >= 65536: ++ block_k = 32 ++ num_stages = 4 ++ elif K >= 1024 and N <= 262144: ++ block_k = 64 ++ num_stages = 4 ++ ++ return { ++ "BLOCK_N": block_n, ++ "BLOCK_K": block_k, ++ "num_warps": num_warps, ++ "num_stages": num_stages, ++ } ++ ++ ++def _heuristic_euclid_config_gb10_smallD(N: int, K: int, D: int, dtype) -> dict: ++ # GB10 (Grace Blackwell, ~80 SMs, ~99 KiB SMEM/SM) tuned heuristic. ++ # Derived from a grid sweep over N in {65536, 262144, 1048576}, ++ # K in {256..200000}, D in {64,128,256,512}, B in {1, 32}, fp16. ++ if D >= 512: ++ if K <= 256: ++ return {"BLOCK_N": 64, "BLOCK_K": 32, "num_warps": 4, "num_stages": 1} ++ return {"BLOCK_N": 64, "BLOCK_K": 32, "num_warps": 8, "num_stages": 1} ++ ++ if D >= 256: ++ if K <= 256: ++ return {"BLOCK_N": 64, "BLOCK_K": 32, "num_warps": 4, "num_stages": 1} ++ return {"BLOCK_N": 128, "BLOCK_K": 64, "num_warps": 4, "num_stages": 2} ++ ++ if D >= 128: ++ if K <= 256: ++ return {"BLOCK_N": 64, "BLOCK_K": 64, "num_warps": 4, "num_stages": 1} ++ if K <= 1024: ++ if N <= 65536: ++ return {"BLOCK_N": 64, "BLOCK_K": 64, "num_warps": 4, "num_stages": 1} ++ return {"BLOCK_N": 128, "BLOCK_K": 32, "num_warps": 4, "num_stages": 2} ++ if K <= 65536: ++ return {"BLOCK_N": 128, "BLOCK_K": 32, "num_warps": 4, "num_stages": 1} ++ return {"BLOCK_N": 128, "BLOCK_K": 64, "num_warps": 4, "num_stages": 1} ++ ++ # D <= 64 ++ if K <= 256 and N <= 65536: ++ return {"BLOCK_N": 64, "BLOCK_K": 64, "num_warps": 4, "num_stages": 1} ++ return {"BLOCK_N": 128, "BLOCK_K": 32, "num_warps": 4, "num_stages": 1} ++ ++ ++def _heuristic_euclid_config_fallback_smallD(N: int, K: int, D: int, dtype) -> dict: ++ # Conservative default for unknown architectures (prioritize avoiding OOR). ++ return { ++ "BLOCK_N": 64, ++ "BLOCK_K": 32, ++ "num_warps": 4, ++ "num_stages": 1, ++ } ++ ++ ++_KNOWN_ARCHS = ("H200", "H100", "A100", "GB10") ++ ++ ++def _is_known_arch(gpu_name: str) -> bool: ++ return any(tag in gpu_name for tag in _KNOWN_ARCHS) ++ ++ ++def _arch_smallD_picker(gpu_name: str): ++ if "H200" in gpu_name: ++ return _heuristic_euclid_config_h200_smallD ++ if "H100" in gpu_name: ++ return _heuristic_euclid_config_h100_smallD ++ if "A100" in gpu_name: ++ return _heuristic_euclid_config_a100_smallD ++ if "GB10" in gpu_name: ++ return _heuristic_euclid_config_gb10_smallD ++ return _heuristic_euclid_config_fallback_smallD ++ ++ ++def _heuristic_euclid_config( ++ N: int, ++ K: int, ++ D: int, ++ *, ++ device: Optional[torch.device] = None, ++ dtype=None, ++): ++ """Architecture-aware heuristic config selection without autotune. ++ ++ Per-GPU sub-functions own the actual lookup tables. This function only ++ routes to the right one and (for non-half dtypes) post-processes the ++ config through ``_fit_config_to_smem`` so fp32 / large-D inputs do not ++ OOR on small-SMEM GPUs. ++ ++ For fp16/bf16 the picked config is returned **byte-for-byte** as the ++ original tables produced (``examples/regression_fp16_smalld.py`` ++ enforces this). ++ """ ++ if device is None: ++ device = torch.device("cuda") ++ gpu_name = torch.cuda.get_device_properties(device).name.upper() ++ ++ cfg = _arch_smallD_picker(gpu_name)(N, K, D, dtype) ++ ++ if _is_half_dtype(dtype): ++ return cfg ++ ++ dtype_bytes = _dtype_bytes(dtype) ++ smem_limit = _smem_limit(device) ++ return _fit_config_to_smem(cfg, D, dtype_bytes, smem_limit) ++ ++ ++# ----------------------------------------------------------------------------- ++# Split-D heuristic. Per-arch tables stay conservative for now and rely on ++# ``_fit_config_to_smem_split_d`` for SMEM safety. H200 and H100 have freshly ++# tuned tables; A100/GB10 still use a single default until tuning data lands. ++# ----------------------------------------------------------------------------- ++ ++ ++def _heuristic_euclid_config_h200_largeD(N: int, K: int, D: int, dtype) -> dict: ++ """H200 split-D heuristic. ++ ++ Derived from a focused grid sweep over D ∈ {1024, 2048, 4096}, ++ K ∈ {256..65536}, N ∈ {65536, 262144, 1048576}, B=1, fp16+fp32. ++ Patterns: ++ - fp16/bf16 (2 bytes): wide tile (BN=128, BK=128) with BD=64 is the ++ clear winner across D=1024 and D=4096 (3/3 N votes per K bucket). ++ D=2048 with K ≥ 4096 prefers the deeper-D tile (BN=64, BK=128, BD=128) ++ because the centroid stream dominates and a fatter D chunk amortises ++ loads better. ++ - fp32 (4 bytes): same shape but BD shrinks to 32 to keep SMEM under ++ budget. D=1024 with K ≥ 4096 splits BN→64 and grows BD→64 (more ++ D-axis amortisation when the centroid set is large). ++ """ ++ half = _is_half_dtype(dtype) ++ ++ if half: ++ if D >= 4096: ++ return {"BLOCK_N": 128, "BLOCK_K": 128, "BLOCK_D": 64, ++ "num_warps": 8, "num_stages": 4} ++ if D >= 2048: ++ if K <= 1024: ++ return {"BLOCK_N": 128, "BLOCK_K": 128, "BLOCK_D": 64, ++ "num_warps": 8, "num_stages": 4} ++ return {"BLOCK_N": 64, "BLOCK_K": 128, "BLOCK_D": 128, ++ "num_warps": 4, "num_stages": 4} ++ # D ≈ 1024 (also covers D in (512, 1024) when small-D kernel doesn't fit) ++ return {"BLOCK_N": 128, "BLOCK_K": 128, "BLOCK_D": 64, ++ "num_warps": 8, "num_stages": 4} ++ ++ # fp32 / wider ++ if D >= 2048: ++ return {"BLOCK_N": 128, "BLOCK_K": 128, "BLOCK_D": 32, ++ "num_warps": 8, "num_stages": 4} ++ # D ≈ 1024 fp32 ++ if K <= 1024: ++ return {"BLOCK_N": 128, "BLOCK_K": 128, "BLOCK_D": 32, ++ "num_warps": 8, "num_stages": 4} ++ return {"BLOCK_N": 64, "BLOCK_K": 128, "BLOCK_D": 64, ++ "num_warps": 4, "num_stages": 4} ++ ++ ++def _heuristic_euclid_config_h100_largeD(N: int, K: int, D: int, dtype) -> dict: ++ """H100 split-D heuristic. ++ ++ Derived from a focused grid sweep over D ∈ {1024, 2048, 4096}, ++ K ∈ {256, 1024, 4096, 16384, 65536}, N ∈ {65536, 262144, 1048576}, ++ B=1, fp16+fp32. Two cells (K=65536 D=4096 N=1048576 fp16/fp32) were ++ skipped due to multi-hour cost; their config is extrapolated from the ++ matching K=65536 K=16384/N=1048576 winners (same dominant shape). ++ ++ Patterns: ++ - fp16/bf16: BN=128 BK=128 BD=64 W=8 S=4 dominates D=1024 (K≥1024) ++ and D=4096. D=2048 with K ∈ [1024, 16384] prefers BN=64 BK=128 ++ BD=128 W=4 S=4 — a deeper D tile amortises centroid loads when ++ the centroid set is moderate; the wider N tile only pays off ++ once K=65536 makes per-program K-streaming dominate. K=256 ++ uniformly prefers the smaller N tile (BN=64, W=4) — too few K ++ tiles to justify the wider N-axis program. ++ - fp32: BN=128 BK=128 BD=32 W=8 S=4 dominates D ≥ 2048. D=1024 ++ prefers BN=64 BK=128 BD=64 W=4 S=4 across all K — distinct from ++ H200 which keeps BD=32 at small K, because H100's narrower L2 ++ benefits from the wider D chunk amortising the centroid stream. ++ """ ++ half = _is_half_dtype(dtype) ++ ++ if half: ++ # K=256 has too few centroid tiles to justify a wide N program. ++ if K <= 256: ++ return {"BLOCK_N": 64, "BLOCK_K": 128, "BLOCK_D": 64, ++ "num_warps": 4, "num_stages": 4} ++ if D >= 4096: ++ return {"BLOCK_N": 128, "BLOCK_K": 128, "BLOCK_D": 64, ++ "num_warps": 8, "num_stages": 4} ++ if D >= 2048: ++ # K ∈ [1024, 16384]: deeper D tile (BD=128) amortises centroid ++ # loads better. K=65536: the long K stream dominates and the ++ # wider N tile (BN=128, BD=64) wins. ++ if K <= 16384: ++ return {"BLOCK_N": 64, "BLOCK_K": 128, "BLOCK_D": 128, ++ "num_warps": 4, "num_stages": 4} ++ return {"BLOCK_N": 128, "BLOCK_K": 128, "BLOCK_D": 64, ++ "num_warps": 8, "num_stages": 4} ++ # D ≈ 1024 (also covers D ∈ (512, 1024) when small-D doesn't fit). ++ return {"BLOCK_N": 128, "BLOCK_K": 128, "BLOCK_D": 64, ++ "num_warps": 8, "num_stages": 4} ++ ++ # fp32 / wider ++ if D >= 2048: ++ return {"BLOCK_N": 128, "BLOCK_K": 128, "BLOCK_D": 32, ++ "num_warps": 8, "num_stages": 4} ++ # D ≈ 1024 fp32 — wider BD=64 wins consistently across K on H100. ++ return {"BLOCK_N": 64, "BLOCK_K": 128, "BLOCK_D": 64, ++ "num_warps": 4, "num_stages": 4} ++ ++ ++def _heuristic_euclid_config_a100_largeD(N: int, K: int, D: int, dtype) -> dict: ++ return { ++ "BLOCK_N": 64, ++ "BLOCK_K": 32, ++ "BLOCK_D": 32, ++ "num_warps": 4, ++ "num_stages": 2, ++ } ++ ++ ++def _heuristic_euclid_config_gb10_largeD(N: int, K: int, D: int, dtype) -> dict: ++ return { ++ "BLOCK_N": 64, ++ "BLOCK_K": 32, ++ "BLOCK_D": 32, ++ "num_warps": 4, ++ "num_stages": 1, ++ } ++ ++ ++def _heuristic_euclid_config_fallback_largeD(N: int, K: int, D: int, dtype) -> dict: ++ return { ++ "BLOCK_N": 32, ++ "BLOCK_K": 32, ++ "BLOCK_D": 32, ++ "num_warps": 4, ++ "num_stages": 1, ++ } ++ ++ ++def _arch_largeD_picker(gpu_name: str): ++ if "H200" in gpu_name: ++ return _heuristic_euclid_config_h200_largeD ++ if "H100" in gpu_name: ++ return _heuristic_euclid_config_h100_largeD ++ if "A100" in gpu_name: ++ return _heuristic_euclid_config_a100_largeD ++ if "GB10" in gpu_name: ++ return _heuristic_euclid_config_gb10_largeD ++ return _heuristic_euclid_config_fallback_largeD ++ ++ ++def _heuristic_euclid_config_split_d( ++ N: int, ++ K: int, ++ D: int, ++ *, ++ device: Optional[torch.device] = None, ++ dtype=None, ++): ++ """Heuristic config picker for the split-D Euclid assign kernel. ++ ++ Always post-processes through ``_fit_config_to_smem_split_d`` so the ++ selected config is guaranteed to fit SMEM regardless of dtype/D. ++ """ ++ if device is None: ++ device = torch.device("cuda") ++ gpu_name = torch.cuda.get_device_properties(device).name.upper() ++ cfg = _arch_largeD_picker(gpu_name)(N, K, D, dtype) ++ dtype_bytes = _dtype_bytes(dtype) ++ smem_limit = _smem_limit(device) ++ return _fit_config_to_smem_split_d(cfg, D, dtype_bytes, smem_limit) ++ ++ ++# Hard threshold: D > this triggers split-D dispatch even when SMEM would ++# nominally fit, so the fp16/bf16/fp32 + D ≤ 512 regime stays on the ++# original kernel path (matches the regime the existing tables were tuned ++# in). Larger D goes through the split-D path. ++_SMALL_D_MAX = 512 ++ ++ ++def _is_power_of_2(n: int) -> bool: ++ return n > 0 and (n & (n - 1)) == 0 ++ ++ ++def _need_split_d(D: int, dtype, device) -> bool: ++ """Decide whether to dispatch to the split-D kernel. ++ ++ Split-D is needed when: ++ 1. D exceeds the small-D regime (the original kernel can't tile D). ++ 2. D is **not a power of 2** — the small-D kernel uses ++ ``tl.arange(0, D)`` which Triton requires to be a power of 2. ++ The split-D kernel uses ``tl.arange(0, BLOCK_D)`` (always a ++ power of 2) with a ``d_mask = d_offsets < D`` guard, so any D ++ value works. ++ 3. The small-D kernel cannot fit even at minimum tile (BN=16, BK=16, ++ S=1) — SMEM safety net for awkward dtype/D/GPU triples. ++ 4. The GPU is unknown to the heuristic. The small-D path relies on ++ per-arch tuning tables; on unfamiliar architectures we have no ++ data to trust those configs and the SMEM probe may also be ++ unreliable. Split-D with the conservative fallback (small BN/BK/BD, ++ num_stages=1) is the safer choice — ``_fit_config_to_smem_split_d`` ++ then guarantees the launch fits regardless of how small the SMEM ++ budget actually turns out to be. ++ """ ++ if D > _SMALL_D_MAX: ++ return True ++ if not _is_power_of_2(D): ++ return True ++ # tl.dot requires the contraction axis (D) to be >= 16. ++ if D < 16: ++ return True ++ if device is None: ++ device = torch.device("cuda") ++ gpu_name = torch.cuda.get_device_properties(device).name.upper() ++ if not _is_known_arch(gpu_name): ++ return True ++ dtype_bytes = _dtype_bytes(dtype) ++ smem_limit = _smem_limit(device) ++ return not _smallD_kernel_fits_smem(D, dtype_bytes, smem_limit) ++ ++ ++@triton.jit ++def _euclid_assign_kernel( ++ x_ptr, # *f16 / *f32 [B, N, D] ++ c_ptr, # *f16 / *f32 [B, K, D] ++ c_sq_ptr, # *f32 [B, K] ++ out_ptr, # *i32 [B, N] ++ B: tl.constexpr, ++ N: tl.constexpr, ++ K: tl.constexpr, ++ D: tl.constexpr, ++ stride_x_b: tl.constexpr, ++ stride_x_n: tl.constexpr, ++ stride_x_d: tl.constexpr, ++ stride_c_b: tl.constexpr, ++ stride_c_k: tl.constexpr, ++ stride_c_d: tl.constexpr, ++ stride_csq_b: tl.constexpr, ++ stride_csq_k: tl.constexpr, ++ stride_out_b: tl.constexpr, ++ stride_out_n: tl.constexpr, ++ BLOCK_N: tl.constexpr, ++ BLOCK_K: tl.constexpr, ++): ++ """Each program handles a tile of BLOCK_N points for a given batch element. ++ ++ The kernel iterates over the centroid dimension K in chunks of BLOCK_K and ++ maintains the running minimum distance as well as the corresponding index ++ for every point in the tile. ++ """ ++ pid_n = tl.program_id(0) # tile index along N dimension ++ pid_b = tl.program_id(1) # batch index ++ pid_b = pid_b.to(tl.int64) ++ ++ n_start = pid_n * BLOCK_N ++ n_offsets = n_start + tl.arange(0, BLOCK_N) ++ n_offsets = n_offsets.to(tl.int64) ++ n_mask = n_offsets < N ++ ++ # ------------------------------------------------------------------ ++ # Load x tile (BLOCK_N, D) ++ # ------------------------------------------------------------------ ++ offs_d = tl.arange(0, D).to(tl.int64) ++ # Compute pointer for x block: base + b*stride_x_b + n*stride_x_n + d*stride_x_d ++ x_ptrs = ( ++ x_ptr ++ + pid_b * stride_x_b ++ + n_offsets[:, None] * stride_x_n ++ + offs_d[None, :] * stride_x_d ++ ) ++ x_tile = tl.load(x_ptrs, mask=n_mask[:, None], other=0.0) ++ x_tile = x_tile # compute in f32 ++ ++ # Init best distance / index ++ best_dist = tl.full((BLOCK_N,), 3.4e38, tl.float32) # large number ++ best_idx = tl.zeros((BLOCK_N,), tl.int32) ++ ++ # ------------------------------------------------------------------ ++ # Iterate over centroids in chunks of BLOCK_K ++ # ------------------------------------------------------------------ ++ for k_start in range(0, K, BLOCK_K): ++ k_offsets = k_start + tl.arange(0, BLOCK_K) ++ k_offsets = k_offsets.to(tl.int64) ++ k_mask = k_offsets < K ++ ++ # Load centroid tile (D, BLOCK_K) ++ c_ptrs = ( ++ c_ptr ++ + pid_b * stride_c_b ++ + k_offsets[None, :] * stride_c_k ++ + offs_d[:, None] * stride_c_d ++ ) ++ c_tile = tl.load(c_ptrs, mask=k_mask[None, :], other=0.0) ++ c_tile = c_tile ++ ++ # load c_sq for the tile (BLOCK_K,) ++ csq_ptrs = c_sq_ptr + pid_b * stride_csq_b + k_offsets * stride_csq_k ++ cent_sq = tl.load(csq_ptrs, mask=k_mask, other=0.0).to(tl.float32) ++ ++ # # Compute centroid squared norms (BLOCK_K,) ++ # cent_sq = tl.sum(c_tile * c_tile, axis=0).to(tl.float32) ++ ++ # Compute cross term (BLOCK_N, BLOCK_K) = x_tile @ c_tile ++ cross = tl.dot(x_tile, c_tile).to(tl.float32) ++ ++ # Squared-distance rank score (same argmin as ||x-c||^2); ||x||^2 omitted. ++ dist = cent_sq[None, :] - 2.0 * cross ++ ++ # Mask out invalid centroid columns before reduction ++ dist = tl.where(k_mask[None, :], dist, 3.4e38) ++ ++ curr_min = tl.min(dist, axis=1) ++ curr_idx = tl.argmin(dist, axis=1) ++ ++ update = curr_min < best_dist ++ best_dist = tl.where(update, curr_min, best_dist) ++ best_idx = tl.where(update, k_start + curr_idx, best_idx) ++ ++ # ------------------------------------------------------------------ ++ # Write results ++ # ------------------------------------------------------------------ ++ out_ptrs = out_ptr + pid_b * stride_out_b + n_offsets * stride_out_n ++ tl.store(out_ptrs, best_idx, mask=n_mask) ++ ++_euclid_assign_kernel_autotuned = triton.autotune(_TUNE_CONFIGS, key=["N", "K"])(_euclid_assign_kernel) ++ ++ ++# =============================================================== ++# Split-D Euclid assign kernel. ++# ++# This kernel mirrors ``_euclid_assign_kernel`` but tiles the feature ++# dimension D into chunks of size ``BLOCK_D`` so the per-program SMEM ++# footprint is bounded by BLOCK_D rather than D. The K-streaming property ++# (no full distance matrix materialised) is preserved: outer loop is K, ++# inner D loop accumulates ``cross (BN, BK)`` in registers across D chunks, ++# distance and best-index are computed at the end of each K iteration. ++# =============================================================== ++@triton.jit ++def _euclid_assign_kernel_split_d( ++ x_ptr, # *f16 / *f32 [B, N, D] ++ c_ptr, # *f16 / *f32 [B, K, D] ++ c_sq_ptr, # *f32 [B, K] ++ out_ptr, # *i32 [B, N] ++ B: tl.constexpr, ++ N: tl.constexpr, ++ K: tl.constexpr, ++ D: tl.constexpr, ++ stride_x_b: tl.constexpr, ++ stride_x_n: tl.constexpr, ++ stride_x_d: tl.constexpr, ++ stride_c_b: tl.constexpr, ++ stride_c_k: tl.constexpr, ++ stride_c_d: tl.constexpr, ++ stride_csq_b: tl.constexpr, ++ stride_csq_k: tl.constexpr, ++ stride_out_b: tl.constexpr, ++ stride_out_n: tl.constexpr, ++ BLOCK_N: tl.constexpr, ++ BLOCK_K: tl.constexpr, ++ BLOCK_D: tl.constexpr, ++): ++ pid_n = tl.program_id(0) ++ pid_b = tl.program_id(1) ++ pid_b = pid_b.to(tl.int64) ++ ++ n_start = pid_n * BLOCK_N ++ n_offsets = n_start + tl.arange(0, BLOCK_N) ++ n_offsets = n_offsets.to(tl.int64) ++ n_mask = n_offsets < N ++ ++ offs_d = tl.arange(0, BLOCK_D).to(tl.int64) ++ ++ best_dist = tl.full((BLOCK_N,), 3.4e38, tl.float32) ++ best_idx = tl.zeros((BLOCK_N,), tl.int32) ++ ++ for k_start in range(0, K, BLOCK_K): ++ k_offsets = k_start + tl.arange(0, BLOCK_K) ++ k_offsets = k_offsets.to(tl.int64) ++ k_mask = k_offsets < K ++ ++ csq_ptrs = c_sq_ptr + pid_b * stride_csq_b + k_offsets * stride_csq_k ++ cent_sq = tl.load(csq_ptrs, mask=k_mask, other=0.0).to(tl.float32) ++ ++ # cross accumulator lives in registers across the D loop. ++ cross = tl.zeros((BLOCK_N, BLOCK_K), dtype=tl.float32) ++ ++ for d_start in range(0, D, BLOCK_D): ++ d_offsets = d_start + offs_d ++ d_mask = d_offsets < D ++ ++ x_ptrs = ( ++ x_ptr ++ + pid_b * stride_x_b ++ + n_offsets[:, None] * stride_x_n ++ + d_offsets[None, :] * stride_x_d ++ ) ++ x_chunk = tl.load( ++ x_ptrs, ++ mask=n_mask[:, None] & d_mask[None, :], ++ other=0.0, ++ ) ++ ++ c_ptrs = ( ++ c_ptr ++ + pid_b * stride_c_b ++ + k_offsets[None, :] * stride_c_k ++ + d_offsets[:, None] * stride_c_d ++ ) ++ c_chunk = tl.load( ++ c_ptrs, ++ mask=k_mask[None, :] & d_mask[:, None], ++ other=0.0, ++ ) ++ ++ cross += tl.dot(x_chunk, c_chunk).to(tl.float32) ++ ++ dist = cent_sq[None, :] - 2.0 * cross ++ dist = tl.where(k_mask[None, :], dist, 3.4e38) ++ ++ curr_min = tl.min(dist, axis=1) ++ curr_idx = tl.argmin(dist, axis=1) ++ ++ update = curr_min < best_dist ++ best_dist = tl.where(update, curr_min, best_dist) ++ best_idx = tl.where(update, k_start + curr_idx, best_idx) ++ ++ out_ptrs = out_ptr + pid_b * stride_out_b + n_offsets * stride_out_n ++ tl.store(out_ptrs, best_idx, mask=n_mask) ++ ++ ++_euclid_assign_kernel_split_d_autotuned = triton.autotune( ++ _TUNE_CONFIGS_SPLIT_D, key=["N", "K", "D"] ++)(_euclid_assign_kernel_split_d) ++ ++ ++@triton.jit ++def _cosine_assign_kernel( ++ x_ptr, # *f16 / *f32 [B, N, D] ++ c_ptr, # *f16 / *f32 [B, K, D] ++ out_ptr, # *i32 [B, N] ++ B: tl.constexpr, ++ N: tl.constexpr, ++ K: tl.constexpr, ++ D: tl.constexpr, ++ stride_x_b: tl.constexpr, ++ stride_x_n: tl.constexpr, ++ stride_x_d: tl.constexpr, ++ stride_c_b: tl.constexpr, ++ stride_c_k: tl.constexpr, ++ stride_c_d: tl.constexpr, ++ stride_out_b: tl.constexpr, ++ stride_out_n: tl.constexpr, ++ BLOCK_N: tl.constexpr, ++ BLOCK_K: tl.constexpr, ++): ++ """Each program handles a tile of BLOCK_N points for a given batch element. ++ ++ The kernel iterates over the centroid dimension K in chunks of BLOCK_K and ++ maintains the running minimum distance as well as the corresponding index ++ for every point in the tile. ++ """ ++ pid_n = tl.program_id(0) # tile index along N dimension ++ pid_b = tl.program_id(1) # batch index ++ pid_b = pid_b.to(tl.int64) ++ ++ n_start = pid_n * BLOCK_N ++ n_offsets = n_start + tl.arange(0, BLOCK_N) ++ n_offsets = n_offsets.to(tl.int64) ++ n_mask = n_offsets < N ++ ++ # ------------------------------------------------------------------ ++ # Load x tile (BLOCK_N, D) ++ # ------------------------------------------------------------------ ++ offs_d = tl.arange(0, D).to(tl.int64) ++ # Compute pointer for x block: base + b*stride_x_b + n*stride_x_n + d*stride_x_d ++ x_ptrs = ( ++ x_ptr ++ + pid_b * stride_x_b ++ + n_offsets[:, None] * stride_x_n ++ + offs_d[None, :] * stride_x_d ++ ) ++ x_tile = tl.load(x_ptrs, mask=n_mask[:, None], other=0.0) ++ x_tile = x_tile # compute in f32 ++ ++ # Init best distance / index ++ best_dist = tl.full((BLOCK_N,), -3.4e38, tl.float32) # less is worse ++ best_idx = tl.zeros((BLOCK_N,), tl.int32) ++ ++ # ------------------------------------------------------------------ ++ # Iterate over centroids in chunks of BLOCK_K ++ # ------------------------------------------------------------------ ++ for k_start in range(0, K, BLOCK_K): ++ k_offsets = k_start + tl.arange(0, BLOCK_K) ++ k_offsets = k_offsets.to(tl.int64) ++ k_mask = k_offsets < K ++ ++ # Load centroid tile (D, BLOCK_K) ++ c_ptrs = ( ++ c_ptr ++ + pid_b * stride_c_b ++ + k_offsets[None, :] * stride_c_k ++ + offs_d[:, None] * stride_c_d ++ ) ++ c_tile = tl.load(c_ptrs, mask=k_mask[None, :], other=0.0) ++ c_tile = c_tile ++ ++ # Compute cosine distance (BLOCK_N, BLOCK_K) = x_tile @ c_tile ++ cross = tl.dot(x_tile, c_tile).to(tl.float32) ++ ++ # Mask out invalid centroid columns before reduction. ++ # Use a sentinel below any real cosine/dot score so masked tail lanes ++ # never win the argmax. (0.0 is a real cosine value: any input row whose ++ # cosine to all real centroids is < 0 would otherwise route to a masked ++ # lane, returning an out-of-range cluster id.) ++ dist = tl.where(k_mask[None, :], cross, -torch.finfo(torch.float32).max) ++ ++ curr_max = tl.max(dist, axis=1) ++ curr_idx = tl.argmax(dist, axis=1) ++ ++ update = curr_max > best_dist ++ best_dist = tl.where(update, curr_max, best_dist) ++ best_idx = tl.where(update, k_start + curr_idx, best_idx) ++ ++ # ------------------------------------------------------------------ ++ # Write results ++ # ------------------------------------------------------------------ ++ out_ptrs = out_ptr + pid_b * stride_out_b + n_offsets * stride_out_n ++ tl.store(out_ptrs, best_idx, mask=n_mask) ++ ++_cosine_assign_kernel_autotuned = triton.autotune(_TUNE_CONFIGS, key=["N", "K"])(_cosine_assign_kernel) ++ ++ ++# =============================================================== ++# Split-D Cosine assign kernel. Same loop structure as Euclid split-D ++# but tracks running argmax over the dot-product (cosine score with ++# normalized inputs). ++# =============================================================== ++@triton.jit ++def _cosine_assign_kernel_split_d( ++ x_ptr, ++ c_ptr, ++ out_ptr, ++ B: tl.constexpr, ++ N: tl.constexpr, ++ K: tl.constexpr, ++ D: tl.constexpr, ++ stride_x_b: tl.constexpr, ++ stride_x_n: tl.constexpr, ++ stride_x_d: tl.constexpr, ++ stride_c_b: tl.constexpr, ++ stride_c_k: tl.constexpr, ++ stride_c_d: tl.constexpr, ++ stride_out_b: tl.constexpr, ++ stride_out_n: tl.constexpr, ++ BLOCK_N: tl.constexpr, ++ BLOCK_K: tl.constexpr, ++ BLOCK_D: tl.constexpr, ++): ++ pid_n = tl.program_id(0) ++ pid_b = tl.program_id(1) ++ pid_b = pid_b.to(tl.int64) ++ ++ n_start = pid_n * BLOCK_N ++ n_offsets = n_start + tl.arange(0, BLOCK_N) ++ n_offsets = n_offsets.to(tl.int64) ++ n_mask = n_offsets < N ++ ++ offs_d = tl.arange(0, BLOCK_D).to(tl.int64) ++ ++ best_dist = tl.full((BLOCK_N,), -3.4e38, tl.float32) ++ best_idx = tl.zeros((BLOCK_N,), tl.int32) ++ ++ for k_start in range(0, K, BLOCK_K): ++ k_offsets = k_start + tl.arange(0, BLOCK_K) ++ k_offsets = k_offsets.to(tl.int64) ++ k_mask = k_offsets < K ++ ++ cross = tl.zeros((BLOCK_N, BLOCK_K), dtype=tl.float32) ++ ++ for d_start in range(0, D, BLOCK_D): ++ d_offsets = d_start + offs_d ++ d_mask = d_offsets < D ++ ++ x_ptrs = ( ++ x_ptr ++ + pid_b * stride_x_b ++ + n_offsets[:, None] * stride_x_n ++ + d_offsets[None, :] * stride_x_d ++ ) ++ x_chunk = tl.load( ++ x_ptrs, ++ mask=n_mask[:, None] & d_mask[None, :], ++ other=0.0, ++ ) ++ ++ c_ptrs = ( ++ c_ptr ++ + pid_b * stride_c_b ++ + k_offsets[None, :] * stride_c_k ++ + d_offsets[:, None] * stride_c_d ++ ) ++ c_chunk = tl.load( ++ c_ptrs, ++ mask=k_mask[None, :] & d_mask[:, None], ++ other=0.0, ++ ) ++ ++ cross += tl.dot(x_chunk, c_chunk).to(tl.float32) ++ ++ # Mask invalid centroid columns to a sentinel below any real score. ++ dist = tl.where(k_mask[None, :], cross, -torch.finfo(torch.float32).max) ++ ++ curr_max = tl.max(dist, axis=1) ++ curr_idx = tl.argmax(dist, axis=1) ++ ++ update = curr_max > best_dist ++ best_dist = tl.where(update, curr_max, best_dist) ++ best_idx = tl.where(update, k_start + curr_idx, best_idx) ++ ++ out_ptrs = out_ptr + pid_b * stride_out_b + n_offsets * stride_out_n ++ tl.store(out_ptrs, best_idx, mask=n_mask) ++ ++ ++_cosine_assign_kernel_split_d_autotuned = triton.autotune( ++ _TUNE_CONFIGS_SPLIT_D, key=["N", "K", "D"] ++)(_cosine_assign_kernel_split_d) ++ ++ ++# --------------------------------------------------------------- ++# Python wrapper ++# --------------------------------------------------------------- ++ ++def euclid_assign_triton( ++ x: torch.Tensor, ++ centroids: torch.Tensor, ++ out: torch.Tensor = None, ++ c_sq: torch.Tensor = None, ++ *, ++ BLOCK_N: int = 128, ++ BLOCK_K: int = 128, ++ num_warps: Optional[int] = None, ++ num_stages: Optional[int] = None, ++ config: Optional[dict] = None, ++ use_heuristic: bool = True, ++) -> torch.Tensor: ++ """Return nearest-centroid indices using Triton kernel. ++ ++ Args: ++ x : (B, N, D) float16 / float32 (on CUDA) ++ centroids : (B, K, D) same dtype/device as x ++ out : (B, N) int32 – (option) pre-allocated output tensor (on CUDA) ++ c_sq : (B, K) float32 – (option) ||centroids||^2 per centroid (on CUDA) ++ ++ Returns: ++ cluster_ids (B, N) int32 (callers can cast to int64 if desired) ++ Extra: ++ config : {"BLOCK_N","BLOCK_K","num_warps","num_stages"} to force a config ++ use_heuristic : use a fixed heuristic config instead of autotune ++ """ ++ assert x.is_cuda and centroids.is_cuda, "All tensors must be on CUDA" ++ # assert x.dtype in (torch.float16, torch.float32), "x must be fp16/fp32" ++ assert centroids.dtype == x.dtype, "centroids dtype mismatch" ++ ++ B, N, D = x.shape ++ K = centroids.shape[1] ++ assert centroids.shape == (B, K, D), "centroids shape mismatch" ++ ++ # x = x.contiguous() ++ # centroids = centroids.contiguous() ++ ++ if out is None: ++ out = torch.empty((B, N), device=x.device, dtype=torch.int32) ++ if c_sq is None: ++ c_sq = (centroids.to(torch.float32) ** 2).sum(-1) ++ ++ # Strides (in elements) ++ stride_x_b, stride_x_n, stride_x_d = x.stride() ++ stride_c_b, stride_c_k, stride_c_d = centroids.stride() ++ stride_csq_b, stride_csq_k = c_sq.stride() ++ stride_out_b, stride_out_n = out.stride() ++ ++ grid = lambda META: (triton.cdiv(N, META["BLOCK_N"]), B) ++ ++ use_split_d = _need_split_d(D, x.dtype, x.device) ++ ++ selected_config = None ++ if config is not None: ++ selected_config = config ++ elif num_warps is not None or num_stages is not None: ++ if num_warps is None or num_stages is None: ++ raise ValueError("num_warps and num_stages must be set together") ++ selected_config = { ++ "BLOCK_N": BLOCK_N, ++ "BLOCK_K": BLOCK_K, ++ "num_warps": num_warps, ++ "num_stages": num_stages, ++ } ++ if use_split_d and "BLOCK_D" not in selected_config: ++ # Caller supplied a small-D-shaped config but D demands split-D. ++ # Use a sane default for BLOCK_D and let SMEM fitter shrink. ++ selected_config = dict(selected_config) ++ selected_config["BLOCK_D"] = 64 ++ selected_config = _fit_config_to_smem_split_d( ++ selected_config, ++ D, ++ _dtype_bytes(x.dtype), ++ _smem_limit(x.device), ++ ) ++ elif use_heuristic: ++ if use_split_d: ++ selected_config = _heuristic_euclid_config_split_d( ++ N, K, D, device=x.device, dtype=x.dtype ++ ) ++ else: ++ selected_config = _heuristic_euclid_config( ++ N, K, D, device=x.device, dtype=x.dtype ++ ) ++ ++ if use_split_d: ++ if selected_config is None: ++ _euclid_assign_kernel_split_d_autotuned[grid]( ++ x, centroids, c_sq, out, ++ B, N, K, D, ++ stride_x_b, stride_x_n, stride_x_d, ++ stride_c_b, stride_c_k, stride_c_d, ++ stride_csq_b, stride_csq_k, ++ stride_out_b, stride_out_n, ++ ) ++ else: ++ _euclid_assign_kernel_split_d[grid]( ++ x, centroids, c_sq, out, ++ B, N, K, D, ++ stride_x_b, stride_x_n, stride_x_d, ++ stride_c_b, stride_c_k, stride_c_d, ++ stride_csq_b, stride_csq_k, ++ stride_out_b, stride_out_n, ++ BLOCK_N=selected_config["BLOCK_N"], ++ BLOCK_K=selected_config["BLOCK_K"], ++ BLOCK_D=selected_config["BLOCK_D"], ++ num_warps=selected_config["num_warps"], ++ num_stages=selected_config["num_stages"], ++ ) ++ return out ++ ++ if selected_config is not None: ++ _euclid_assign_kernel[grid]( ++ x, ++ centroids, ++ c_sq, ++ out, ++ B, ++ N, ++ K, ++ D, ++ stride_x_b, ++ stride_x_n, ++ stride_x_d, ++ stride_c_b, ++ stride_c_k, ++ stride_c_d, ++ stride_csq_b, ++ stride_csq_k, ++ stride_out_b, ++ stride_out_n, ++ BLOCK_N=selected_config["BLOCK_N"], ++ BLOCK_K=selected_config["BLOCK_K"], ++ num_warps=selected_config["num_warps"], ++ num_stages=selected_config["num_stages"], ++ ) ++ else: ++ _euclid_assign_kernel_autotuned[grid]( ++ x, ++ centroids, ++ c_sq, ++ out, ++ B, ++ N, ++ K, ++ D, ++ stride_x_b, ++ stride_x_n, ++ stride_x_d, ++ stride_c_b, ++ stride_c_k, ++ stride_c_d, ++ stride_csq_b, ++ stride_csq_k, ++ stride_out_b, ++ stride_out_n, ++ ) ++ return out ++ ++ ++def cosine_assign_triton(x: torch.Tensor, centroids: torch.Tensor, out: torch.Tensor = None, ++ *, BLOCK_N: int = 128, BLOCK_K: int = 128) -> torch.Tensor: ++ """Return nearest(cosine similarity)-centroid indices using Triton kernel. ++ ++ Args: ++ x : (B, N, D) float16 / float32 (on CUDA) ++ centroids : (B, K, D) same dtype/device as x ++ ++ Returns: ++ cluster_ids (B, N) int32 (callers can cast to int64 if desired) ++ """ ++ assert x.is_cuda and centroids.is_cuda, "All tensors must be on CUDA" ++ # assert x.dtype in (torch.float16, torch.float32), "x must be fp16/fp32" ++ assert centroids.dtype == x.dtype, "centroids dtype mismatch" ++ ++ B, N, D = x.shape ++ K = centroids.shape[1] ++ assert centroids.shape == (B, K, D), "centroids shape mismatch" ++ ++ # x = x.contiguous() ++ # centroids = centroids.contiguous() ++ ++ if out is None: ++ out = torch.empty((B, N), device=x.device, dtype=torch.int32) ++ ++ # Strides (in elements) ++ stride_x_b, stride_x_n, stride_x_d = x.stride() ++ stride_c_b, stride_c_k, stride_c_d = centroids.stride() ++ stride_out_b, stride_out_n = out.stride() ++ ++ grid = lambda META: (triton.cdiv(N, META["BLOCK_N"]), B) ++ ++ if _need_split_d(D, x.dtype, x.device): ++ _cosine_assign_kernel_split_d_autotuned[grid]( ++ x, ++ centroids, ++ out, ++ B, ++ N, ++ K, ++ D, ++ stride_x_b, ++ stride_x_n, ++ stride_x_d, ++ stride_c_b, ++ stride_c_k, ++ stride_c_d, ++ stride_out_b, ++ stride_out_n, ++ ) ++ return out ++ ++ # Small-D path. The existing autotune sweep includes configs that may ++ # not fit SMEM (e.g. fp16 D=512 BN=64 BK=64 S=4 → 320 KiB > H200 limit; ++ # fp32 D=512 BN=64 BK=32 S=4 → 393 KiB). The cosine kernel has the ++ # same tile shapes as euclid, so reuse the SMEM-aware euclid heuristic ++ # for config selection regardless of dtype. This keeps fp16/bf16 small-D ++ # behaviour byte-identical to the euclid path's heuristic (verified by ++ # examples/regression_fp16_smalld.py). ++ cfg = _heuristic_euclid_config(N, K, D, device=x.device, dtype=x.dtype) ++ _cosine_assign_kernel[grid]( ++ x, ++ centroids, ++ out, ++ B, ++ N, ++ K, ++ D, ++ stride_x_b, ++ stride_x_n, ++ stride_x_d, ++ stride_c_b, ++ stride_c_k, ++ stride_c_d, ++ stride_out_b, ++ stride_out_n, ++ BLOCK_N=cfg["BLOCK_N"], ++ BLOCK_K=cfg["BLOCK_K"], ++ num_warps=cfg["num_warps"], ++ num_stages=cfg["num_stages"], ++ ) ++ return out +diff --git a/kmeanslib/_kernels/primitives/kmeans/triton/kmeans.py b/kmeanslib/_kernels/primitives/kmeans/triton/kmeans.py +new file mode 100644 +index 0000000..db5fb79 +--- /dev/null ++++ b/kmeanslib/_kernels/primitives/kmeans/triton/kmeans.py +@@ -0,0 +1,234 @@ ++import torch ++import torch.nn.functional as F ++from torch.cuda import nvtx ++from kmeanslib._kernels.primitives.kmeans.triton.assign import euclid_assign_triton, cosine_assign_triton ++from kmeanslib._kernels.primitives.kmeans.triton.update import ( ++ triton_centroid_update_cosine, ++ triton_centroid_update_euclid, ++ triton_centroid_update_sorted_euclid, ++ triton_centroid_update_sorted_cosine, ++ triton_lloyd_centroid_step_euclid, ++) ++try: ++ from tqdm import trange ++except Exception: ++ trange = range ++ ++# -------------------- Compiled single-iteration kernels -------------------- ++ ++# 1. Euclidean ++def _euclid_iter(x, centroids, use_heuristic=True): ++ ++ cluster_ids = euclid_assign_triton(x, centroids, use_heuristic=use_heuristic) ++ centroids_new = triton_centroid_update_sorted_euclid(x, cluster_ids, centroids) ++ ++ shift = (centroids_new - centroids).norm(dim=-1).max() ++ return centroids_new, shift, cluster_ids ++ ++# 2. Cosine ++def _cosine_iter(x_norm, centroids): ++ # cos_sim = torch.einsum('bnd,bkd->bnk', x_norm, centroids) ++ # cluster_ids = cos_sim.argmax(dim=-1) ++ cluster_ids = cosine_assign_triton(x_norm, centroids) ++ centroids_new = triton_centroid_update_sorted_cosine(x_norm, cluster_ids, centroids) ++ # centroids_new = centroids_new.clone() ++ shift = (centroids_new - centroids).norm(dim=-1).max() ++ return centroids_new, shift, cluster_ids ++ ++# 3. Dot-product ++def _dot_iter(x, centroids): ++ # sim = torch.einsum('bnd,bkd->bnk', x, centroids) ++ # cluster_ids = sim.argmax(dim=-1) ++ cluster_ids = cosine_assign_triton(x, centroids) ++ centroids_new = triton_centroid_update_sorted_cosine(x, cluster_ids, centroids) ++ # centroids_new = centroids_new.clone() ++ shift = (centroids_new - centroids).norm(dim=-1).max() ++ return centroids_new, shift, cluster_ids ++ ++COMPILE_FLAG = False ++ ++try: ++ if COMPILE_FLAG: ++ _euclid_iter_compiled = torch.compile(_euclid_iter, dynamic=True, mode="reduce-overhead") ++ _cosine_iter_compiled = torch.compile(_cosine_iter, dynamic=True, mode="reduce-overhead") ++ _dot_iter_compiled = torch.compile(_dot_iter, dynamic=True, mode="reduce-overhead") ++ else: ++ _euclid_iter_compiled = _euclid_iter ++ _cosine_iter_compiled = _cosine_iter ++ _dot_iter_compiled = _dot_iter ++except Exception: # pragma: no cover ++ _euclid_iter_compiled = _euclid_iter ++ _cosine_iter_compiled = _cosine_iter ++ _dot_iter_compiled = _dot_iter ++ ++def batch_kmeans_Euclid( ++ x, ++ n_clusters, ++ max_iters=100, ++ tol=0.0, ++ init_centroids=None, ++ verbose=False, ++ *, ++ use_heuristic=True, ++ fused=True, ++): ++ """ ++ Batched KMeans clustering in PyTorch using Euclidean distance. ++ ++ Args: ++ x: Tensor of shape (B, N, D), batch_size B, N points per batch, D dims. ++ n_clusters: Number of clusters. ++ max_iters: Max number of iterations. ++ tol: Relative tolerance for center movement. ++ verbose: Print loss for each iter. ++ use_heuristic: Use heuristic Triton config (skip autotune). ++ fused: If True (default), use the fused Lloyd path with preallocated ++ sums/cnts/new/shift buffers and ping-pong centroid swap (no ++ .clone() per iter). Falls back to per-iter alloc when False. ++ Returns: ++ cluster_ids: (B, N) LongTensor, cluster assignment for each point. ++ centroids: (B, n_clusters, D) final cluster centers. ++ """ ++ B, N, D = x.shape ++ K = n_clusters ++ ++ if init_centroids is None: ++ # Randomly select initial centers from x ++ indices = torch.randint(0, N, (B, K), device=x.device) ++ centroids = torch.gather( ++ x, ++ dim=1, ++ index=indices[..., None].expand(-1, -1, D) ++ ) # (B, K, D) ++ else: ++ centroids = init_centroids ++ ++ centroids = centroids.view(B, K, D).contiguous() ++ ++ if not fused: ++ # ----- per-iter alloc + .clone() path ----- ++ for it in range(max_iters): ++ centroids_new, center_shift, cluster_ids = _euclid_iter_compiled( ++ x, centroids, use_heuristic ++ ) ++ if verbose: ++ print(f"Iter {it}, center shift: {center_shift.item():.6f}") ++ if center_shift < tol: ++ break ++ centroids = centroids_new.clone() ++ return cluster_ids, centroids, it + 1 ++ ++ # ----- fused path: preallocated buffers + ping-pong centroid swap ----- ++ # Two centroid buffers swapped each iter so we never .clone(). ++ cent_a = centroids ++ cent_b = torch.empty_like(centroids) ++ ++ sums_buf = torch.zeros((B, K, D), device=x.device, dtype=torch.float32) ++ cnts_buf = torch.zeros((B, K), device=x.device, dtype=torch.int32) ++ shift_buf = torch.empty((B, K), device=x.device, dtype=torch.float32) ++ ++ cur, nxt = cent_a, cent_b ++ cluster_ids = None ++ it = 0 ++ for it in range(max_iters): ++ cluster_ids = euclid_assign_triton(x, cur, use_heuristic=use_heuristic) ++ # writes new centroids into `nxt`, returns scalar GPU tensor for shift ++ new_cent, _, max_shift = triton_lloyd_centroid_step_euclid( ++ x, cluster_ids, cur, ++ sums_buf=sums_buf, ++ cnts_buf=cnts_buf, ++ new_buf=nxt, ++ shift_buf=shift_buf, ++ ) ++ if verbose: ++ print(f"Iter {it}, center shift: {max_shift.item():.6f}") ++ # swap before convergence check so `cur` always points to the latest ++ cur, nxt = nxt, cur ++ # Convergence check: `max_shift` is a 0-D GPU tensor, so ++ # `if max_shift < tol` triggers `tensor.__bool__()` which forces ++ # a per-iter cuda sync (~2.4 ms/iter on H200, drains the kernel ++ # pipeline). The short-circuit on `tol > 0.0` keeps the default ++ # `tol=0.0` path sync-free; users opting into early-exit accept ++ # the sync as part of that contract. ++ if tol > 0.0 and max_shift < tol: ++ break ++ ++ return cluster_ids, cur, it + 1 ++ ++ ++def batch_kmeans_Cosine(x, n_clusters, max_iters=100, tol=0.0, init_centroids=None, verbose=False): ++ """ ++ Batched KMeans clustering in PyTorch using Cosine similarity. ++ ++ Args: ++ x: Tensor of shape (B, N, D), batch_size B, N points per batch, D dims. ++ n_clusters: Number of clusters. ++ max_iters: Max number of iterations. ++ tol: Relative tolerance for center movement. ++ verbose: Print loss for each iter. ++ Returns: ++ cluster_ids: (B, N) LongTensor, cluster assignment for each point. ++ centroids: (B, n_clusters, D) final cluster centers. ++ """ ++ B, N, D = x.shape ++ ++ # Normalize input vectors for cosine similarity ++ x_norm = F.normalize(x, p=2, dim=-1) # (B, N, D) ++ ++ if init_centroids is None: ++ # Randomly select initial centers from x_norm ++ indices = torch.randint(0, N, (B, n_clusters), device=x.device) ++ centroids = torch.gather( ++ x_norm, ++ dim=1, ++ index=indices[..., None].expand(-1, -1, D) ++ ) # (B, n_clusters, D) ++ else: ++ centroids = init_centroids ++ ++ centroids = centroids.view(B, n_clusters, D) ++ centroids = F.normalize(centroids, p=2, dim=-1) # Ensure centroids are normalized ++ ++ for it in range(max_iters): ++ # ---- compiled single iteration ---- ++ centroids_new, center_shift, cluster_ids = _cosine_iter_compiled(x_norm, centroids) ++ ++ # 4. Check for convergence ++ if verbose: ++ print(f"Iter {it}, center shift: {center_shift.item():.6f}") ++ if center_shift < tol: ++ break ++ centroids = centroids_new.clone() ++ ++ return cluster_ids, centroids, it + 1 ++ ++ ++def batch_kmeans_Dot(x, n_clusters, max_iters=100, tol=0.0, init_centroids=None, verbose=False): ++ """ ++ Batched KMeans clustering in PyTorch using raw dot-product as similarity. ++ ++ """ ++ B, N, D = x.shape ++ ++ if init_centroids is None: ++ indices = torch.randint(0, N, (B, n_clusters), device=x.device) ++ centroids = torch.gather( ++ x, ++ dim=1, ++ index=indices[..., None].expand(-1, -1, D) ++ ) ++ else: ++ centroids = init_centroids ++ ++ centroids = centroids.view(B, n_clusters, D) ++ ++ for it in range(max_iters): ++ centroids_new, center_shift, cluster_ids = _dot_iter_compiled(x, centroids) ++ ++ if verbose: ++ print(f"Iter {it} (dot), center shift: {center_shift.item():.6f}") ++ if center_shift < tol: ++ break ++ centroids = centroids_new.clone() ++ ++ return cluster_ids, centroids, it + 1 +diff --git a/kmeanslib/_kernels/primitives/kmeans/triton/update.py b/kmeanslib/_kernels/primitives/kmeans/triton/update.py +new file mode 100644 +index 0000000..88ff697 +--- /dev/null ++++ b/kmeanslib/_kernels/primitives/kmeans/triton/update.py +@@ -0,0 +1,640 @@ ++import torch ++import torch.nn.functional as F ++import triton ++import triton.language as tl ++try: ++ from tqdm import trange ++except Exception: ++ trange = range ++ ++ ++def _ceil_div(a: int, b: int) -> int: ++ return (a + b - 1) // b ++ ++ ++@triton.jit ++def _centroid_update_kernel( ++ x_ptr, # *f16 / *f32 [B, N, D] ++ cluster_ptr, # *i32 [B, N] ++ sum_ptr, # *f32 [B, K, D] ++ count_ptr, # *i32 [B, K] ++ # --- strides (elements) --- ++ stride_x_b, stride_x_n, stride_x_d, ++ stride_sum_b, stride_sum_k, stride_sum_d, ++ stride_count_b, stride_count_k, ++ B: tl.constexpr, ++ N: tl.constexpr, ++ D: tl.constexpr, ++ K: tl.constexpr, ++ BLOCK_D: tl.constexpr, # number of dims processed per program ++): ++ """Each program processes 1 token across BLOCK_D dims using atomics with general strides.""" ++ pid = tl.program_id(axis=0) ++ token_idx = pid # range: [0, B*N) ++ ++ # Derive (b, n) ++ b = (token_idx // N).to(tl.int64) ++ n = (token_idx % N).to(tl.int64) ++ ++ # pointer to this token's feature vector ++ x_offset = b * stride_x_b + n * stride_x_n ++ x_tok_ptr = x_ptr + x_offset ++ ++ cluster_idx = tl.load(cluster_ptr + b * N + n) ++ cluster_idx = tl.where(cluster_idx < K, cluster_idx, 0) ++ cluster_idx = cluster_idx.to(tl.int64) ++ ++ # base ptr for centroid accum array ++ centroid_base = b * stride_sum_b + cluster_idx * stride_sum_k ++ ++ offs = tl.arange(0, BLOCK_D).to(tl.int64) ++ for d_start in range(0, D, BLOCK_D): ++ mask = offs + d_start < D ++ feats = tl.load(x_tok_ptr + (d_start + offs) * stride_x_d, mask=mask, other=0.0) ++ feats = feats.to(tl.float32) ++ dest_ptr = sum_ptr + centroid_base + (d_start + offs) * stride_sum_d ++ tl.atomic_add(dest_ptr, feats, mask=mask) ++ ++ tl.atomic_add(count_ptr + b * stride_count_b + cluster_idx * stride_count_k, 1) ++ ++ ++def triton_centroid_update_cosine(x_norm: torch.Tensor, cluster_ids: torch.Tensor, old_centroids: torch.Tensor): ++ """Compute centroids using custom Triton kernel. ++ ++ Args: ++ x_norm (Tensor): (B, N, D) normalized input vectors (float16/float32) ++ cluster_ids (LongTensor): (B, N) cluster assignment per point ++ old_centroids (Tensor): (B, K, D) previous centroids (same dtype as x_norm) ++ ++ Returns: ++ Tensor: (B, K, D) updated and L2-normalized centroids (dtype == x_norm.dtype) ++ """ ++ assert x_norm.is_cuda and cluster_ids.is_cuda, "Input tensors must be on CUDA device" ++ B, N, D = x_norm.shape ++ K = old_centroids.shape[1] ++ assert cluster_ids.shape == (B, N) ++ ++ # Allocate accumulation buffers ++ centroid_sums = torch.zeros((B, K, D), device=x_norm.device, dtype=torch.float32) ++ centroid_counts = torch.zeros((B, K), device=x_norm.device, dtype=torch.int32) ++ ++ # Launch Triton kernel – one program per token ++ total_tokens = B * N ++ BLOCK_D = 128 # tuneable ++ ++ grid = (total_tokens,) ++ _centroid_update_kernel[grid]( ++ x_norm, ++ cluster_ids.to(torch.int32), ++ centroid_sums, ++ centroid_counts, ++ x_norm.stride(0), x_norm.stride(1), x_norm.stride(2), ++ centroid_sums.stride(0), centroid_sums.stride(1), centroid_sums.stride(2), ++ centroid_counts.stride(0), centroid_counts.stride(1), ++ B, N, D, K, ++ BLOCK_D=BLOCK_D, ++ ) ++ ++ # Compute means; keep old centroid if empty cluster ++ counts_f = centroid_counts.to(torch.float32).unsqueeze(-1).clamp(min=1.0) ++ centroids = centroid_sums / counts_f ++ ++ # For clusters with zero count, revert to old centroids ++ zero_mask = (centroid_counts == 0).unsqueeze(-1) ++ centroids = torch.where(zero_mask, old_centroids.to(torch.float32), centroids) ++ ++ centroids = centroids.to(x_norm.dtype) ++ centroids = F.normalize(centroids, p=2, dim=-1) ++ return centroids ++ ++ ++def torch_loop_centroid_update_cosine(x_norm: torch.Tensor, cluster_ids: torch.Tensor, old_centroids: torch.Tensor): ++ """Reference Python implementation (double for-loop)""" ++ B, N, D = x_norm.shape ++ K = old_centroids.shape[1] ++ new_centroids = torch.zeros_like(old_centroids) ++ for b in range(B): ++ for k in range(K): ++ mask = cluster_ids[b] == k ++ if mask.any(): ++ new_centroids[b, k] = F.normalize(x_norm[b][mask].mean(dim=0, dtype=x_norm.dtype), p=2, dim=0) ++ else: ++ new_centroids[b, k] = old_centroids[b, k] ++ return new_centroids ++ ++ ++def triton_centroid_update_euclid(x: torch.Tensor, cluster_ids: torch.Tensor, old_centroids: torch.Tensor): ++ """Compute centroids for Euclidean KMeans using Triton. ++ ++ Args: ++ x (Tensor): (B, N, D) input vectors (float16/float32) ++ cluster_ids (LongTensor): (B, N) cluster assignment per point ++ old_centroids (Tensor): (B, K, D) previous centroids (same dtype as x) ++ ++ Returns: ++ Tensor: (B, K, D) updated centroids (dtype == x.dtype) ++ """ ++ assert x.is_cuda and cluster_ids.is_cuda, "Input tensors must be on CUDA device" ++ B, N, D = x.shape ++ K = old_centroids.shape[1] ++ assert cluster_ids.shape == (B, N) ++ ++ # Allocate accumulation buffers ++ centroid_sums = torch.zeros((B, K, D), device=x.device, dtype=torch.float32) ++ centroid_counts = torch.zeros((B, K), device=x.device, dtype=torch.int32) ++ ++ total_tokens = B * N ++ BLOCK_D = 128 # tuneable ++ grid = (total_tokens,) ++ ++ _centroid_update_kernel[grid]( ++ x, ++ cluster_ids.to(torch.int32), ++ centroid_sums, ++ centroid_counts, ++ x.stride(0), x.stride(1), x.stride(2), ++ centroid_sums.stride(0), centroid_sums.stride(1), centroid_sums.stride(2), ++ centroid_counts.stride(0), centroid_counts.stride(1), ++ B, N, D, K, ++ BLOCK_D=BLOCK_D, ++ ) ++ ++ # Compute means; keep old centroid if empty cluster ++ counts_f = centroid_counts.to(torch.float32).unsqueeze(-1).clamp(min=1.0) ++ centroids = centroid_sums / counts_f ++ ++ # For clusters with zero count, revert to old centroids ++ zero_mask = (centroid_counts == 0).unsqueeze(-1) ++ centroids = torch.where(zero_mask, old_centroids.to(torch.float32), centroids) ++ ++ return centroids.to(x.dtype) ++ ++ ++# ------------------------------ NEW: chunk-wise centroid update (sorted ids) ------------------------------ ++ ++def _next_power_of_2(n: int) -> int: ++ p = 1 ++ while p < n: ++ p <<= 1 ++ return p ++ ++ ++@triton.jit ++def _centroid_update_chunk_kernel( ++ x_ptr, # *f16 / *f32 [B, N, D] – ORIGINAL ORDER ++ sorted_idx_ptr, # *i32 [B, N] – indices after sort ++ sorted_cluster_ptr, # *i32 [B, N] – cluster ids in sorted order ++ sum_ptr, # *f32 [B, K, D] ++ count_ptr, # *i32 [B, K] ++ # strides ++ stride_x_b, stride_x_n, stride_x_d, ++ stride_idx_b, stride_idx_n, stride_cluster_b, stride_cluster_n, ++ stride_sum_b, stride_sum_k, stride_sum_d, ++ stride_count_b, stride_count_k, ++ B: tl.constexpr, ++ N: tl.constexpr, ++ D: tl.constexpr, ++ K: tl.constexpr, ++ BLOCK_N: tl.constexpr, ++ BLOCK_D: tl.constexpr, ++): ++ """Each program processes **BLOCK_N consecutive, already-sorted tokens**. ++ ++ Because the tokens are sorted by cluster id, identical ids appear in ++ contiguous runs. We therefore accumulate a local sum/count for the ++ current run and perform **a single atomic update per run**, instead of ++ per-token. ++ ++ ``BLOCK_D`` tiles the feature dimension so that non-power-of-2 D ++ values work (Triton requires ``tl.arange`` ranges to be power-of-2). ++ When ``BLOCK_D >= D`` the D-loop executes once — no perf regression ++ for already-power-of-2 inputs. ++ """ ++ pid_chunk = tl.program_id(axis=0) ++ pid_b = tl.program_id(axis=1) ++ ++ b = pid_b.to(tl.int64) ++ chunk_start = (pid_chunk * BLOCK_N).to(tl.int64) ++ ++ if chunk_start >= N: ++ return ++ ++ idx_batch_base = sorted_idx_ptr + b * stride_idx_b ++ cid_batch_base = sorted_cluster_ptr + b * stride_cluster_b ++ x_batch_base = x_ptr + b * stride_x_b ++ ++ offs_token = tl.arange(0, BLOCK_N).to(tl.int64) ++ offs_dim = tl.arange(0, BLOCK_D).to(tl.int64) ++ ++ token_idx = chunk_start + offs_token ++ valid_tok = token_idx < N ++ first_token_idx = chunk_start ++ last_token_idx = tl.minimum(chunk_start + BLOCK_N, N) - 1 ++ ++ first_id = tl.load(cid_batch_base + first_token_idx) ++ last_id = tl.load(cid_batch_base + last_token_idx) ++ all_ids = tl.load(cid_batch_base + token_idx * stride_cluster_n, mask=valid_tok, other=-1) ++ ++ all_tokens_idxs = tl.load(idx_batch_base + token_idx * stride_idx_n, mask=valid_tok, other=-1) ++ all_tokens_idxs = all_tokens_idxs.to(tl.int64) ++ ++ for cid in range(first_id, last_id + 1): ++ cluster_mask = all_ids == cid ++ cluster_size = tl.sum(cluster_mask.to(tl.int32)) ++ if cluster_size != 0: ++ tl.atomic_add(count_ptr + b*stride_count_b + cid*stride_count_k, cluster_size) ++ for d_start in range(0, D, BLOCK_D): ++ d_offsets = d_start + offs_dim ++ d_mask = d_offsets < D ++ row_ptrs = (x_batch_base ++ + all_tokens_idxs[:, None] * stride_x_n ++ + d_offsets[None, :] * stride_x_d) ++ cluster_feats = tl.load( ++ row_ptrs, ++ mask=cluster_mask[:, None] & d_mask[None, :], ++ other=0.0, ++ ) ++ cluster_feats = cluster_feats.to(tl.float32) ++ sum_feats = tl.sum(cluster_feats, axis=0) ++ dest_ptr = (sum_ptr + b * stride_sum_b ++ + cid * stride_sum_k ++ + d_offsets * stride_sum_d) ++ tl.atomic_add(dest_ptr, sum_feats, mask=d_mask) ++ ++ ++# --------------------------------------------------------------------------------------------- ++ ++def triton_centroid_update_sorted_cosine(x_norm: torch.Tensor, cluster_ids: torch.Tensor, old_centroids: torch.Tensor, ++ *, BLOCK_N: int = 256): ++ """Fast centroid update assuming **cluster_ids are sorted along N**. ++ ++ This helper will sort the assignments (together with `x_norm`) and launch the ++ chunk kernel above. Compared to the naive per-token kernel it performs *one ++ atomic add per run of identical ids* instead of per token, providing large ++ speed-ups when clusters are reasonably sized. ++ """ ++ assert x_norm.is_cuda and cluster_ids.is_cuda, "Inputs must be on CUDA" ++ B, N, D = x_norm.shape ++ K = old_centroids.shape[1] ++ assert cluster_ids.shape == (B, N) ++ ++ # -------- sort per-batch -------- ++ sorted_cluster_ids, sorted_idx = torch.sort(cluster_ids, dim=-1) ++ sorted_idx_int = sorted_idx.to(torch.int32) ++ ++ # accumulation buffers ++ centroid_sums = torch.zeros((B, K, D), device=x_norm.device, dtype=torch.float32) ++ centroid_cnts = torch.zeros((B, K), device=x_norm.device, dtype=torch.int32) ++ ++ BLOCK_D = _next_power_of_2(D) ++ grid = (triton.cdiv(N, BLOCK_N), B) ++ _centroid_update_chunk_kernel[grid]( ++ x_norm, ++ sorted_idx_int, ++ sorted_cluster_ids.to(torch.int32), ++ centroid_sums, ++ centroid_cnts, ++ x_norm.stride(0), x_norm.stride(1), x_norm.stride(2), ++ sorted_idx_int.stride(0), sorted_idx_int.stride(1), ++ sorted_cluster_ids.stride(0), sorted_cluster_ids.stride(1), ++ centroid_sums.stride(0), centroid_sums.stride(1), centroid_sums.stride(2), ++ centroid_cnts.stride(0), centroid_cnts.stride(1), ++ B, N, D, K, ++ BLOCK_N=BLOCK_N, ++ BLOCK_D=BLOCK_D, ++ ) ++ ++ # finalise – convert to means, handle empty clusters, renormalise ++ counts_f = centroid_cnts.to(torch.float32).unsqueeze(-1).clamp(min=1.0) ++ centroids = centroid_sums / counts_f ++ empty_mask = (centroid_cnts == 0).unsqueeze(-1) ++ centroids = torch.where(empty_mask, old_centroids.to(torch.float32), centroids) ++ centroids = centroids.to(x_norm.dtype) ++ centroids = F.normalize(centroids, p=2, dim=-1) ++ return centroids ++ ++def triton_centroid_update_sorted_euclid(x: torch.Tensor, cluster_ids: torch.Tensor, old_centroids: torch.Tensor, ++ *, BLOCK_N: int = 256, centroid_sums: torch.Tensor = None, centroid_cnts: torch.Tensor = None, calculate_new: bool = True): ++ """Fast centroid update for *Euclidean* KMeans assuming cluster IDs are pre-sorted. ++ ++ Parameters ++ ---------- ++ x : Tensor [B, N, D] ++ Input feature vectors (no normalization assumed). ++ cluster_ids : LongTensor [B, N] ++ Cluster assignment for each point. ++ old_centroids : Tensor [B, K, D] ++ Previous centroids (used to fill empty clusters). ++ BLOCK_N : int, optional ++ Tokens per Triton program (affects occupancy/perf). ++ centroid_sums : Tensor [B, K, D], optional ++ Pre-allocated accumulation buffer for sums. If None, a new buffer is created. ++ centroid_cnts : Tensor [B, K], optional ++ Pre-allocated accumulation buffer for counts. If None, a new buffer is created. ++ calculate_new : bool, default=True ++ If True, compute and return the new centroids. If False, only update the ++ accumulation buffers. ++ ++ Returns ++ _________ ++ centroids_new : Tensor [B, K, D] or None ++ Updated centroids if `calculate_new` is True; otherwise None. ++ """ ++ assert x.is_cuda and cluster_ids.is_cuda, "Inputs must be on CUDA device" ++ B, N, D = x.shape ++ K = old_centroids.shape[1] ++ ++ # Batch-wise sort of cluster assignments ++ sorted_cluster_ids, sorted_idx = torch.sort(cluster_ids, dim=-1) ++ sorted_idx_int = sorted_idx.to(torch.int32) ++ ++ if centroid_sums is None: ++ centroid_sums = torch.zeros((B, K, D), device=x.device, dtype=torch.float32) ++ else: ++ assert centroid_sums.shape == (B, K, D) ++ ++ if centroid_cnts is None: ++ centroid_cnts = torch.zeros((B, K), device=x.device, dtype=torch.int32) ++ else: ++ assert centroid_cnts.shape == (B, K) ++ ++ BLOCK_D = _next_power_of_2(D) ++ grid = (triton.cdiv(N, BLOCK_N), B) ++ _centroid_update_chunk_kernel[grid]( ++ x, # original features ++ sorted_idx_int, # gather indices ++ sorted_cluster_ids.to(torch.int32), ++ centroid_sums, ++ centroid_cnts, ++ x.stride(0), x.stride(1), x.stride(2), ++ sorted_idx_int.stride(0), sorted_idx_int.stride(1), ++ sorted_cluster_ids.stride(0), sorted_cluster_ids.stride(1), ++ centroid_sums.stride(0), centroid_sums.stride(1), centroid_sums.stride(2), ++ centroid_cnts.stride(0), centroid_cnts.stride(1), ++ B, N, D, K, ++ BLOCK_N=BLOCK_N, ++ BLOCK_D=BLOCK_D, ++ ) ++ ++ if calculate_new: ++ # Convert sums to means; replace empty clusters with old centroids ++ counts_f = centroid_cnts.to(torch.float32).unsqueeze(-1).clamp(min=1.0) ++ centroids = centroid_sums / counts_f ++ empty_mask = (centroid_cnts == 0).unsqueeze(-1) ++ centroids = torch.where(empty_mask, old_centroids.to(torch.float32), centroids) ++ return centroids.to(x.dtype) ++ else: ++ return None ++# ------------------------------ END new implementation ------------------------------ ++ ++ ++# ============================================================================= ++# Fused centroid finalize + per-iter Lloyd helper ++# ============================================================================= ++# ++# `_centroid_finalize_kernel` collapses 6 host-side ops into one Triton kernel: ++# 1. cnts.float() count cast f32 ++# 2. clamp(cnts, min=1) safe-divide guard ++# 3. sums / cnts per-element divide ++# 4. empty-mask where fall back to old centroid where cnts==0 ++# 5. cast back to original dtype ++# 6. (new - old).norm(dim=-1) per-cluster shift (host then takes max) ++# ++# Outputs (B, K) per-cluster shift so the host max reduction is over K, not B*K*D. ++# ++# `triton_lloyd_centroid_step_euclid` glues sort + chunk-update + finalize using ++# preallocated sums / cnts / new / shift buffers reused across Lloyd iterations, ++# so per-iter allocation cost is zero. ++ ++@triton.jit ++def _centroid_finalize_kernel( ++ sums_ptr, # *f32 [B, K, D] ++ cnts_ptr, # *i32 [B, K] ++ old_ptr, # *T [B, K, D] (T = output dtype) ++ new_ptr, # *T [B, K, D] (output) ++ shift_ptr, # *f32 [B, K] (output: per-cluster ‖new − old‖₂) ++ stride_sums_b, stride_sums_k, stride_sums_d, ++ stride_cnts_b, stride_cnts_k, ++ stride_old_b, stride_old_k, stride_old_d, ++ stride_new_b, stride_new_k, stride_new_d, ++ stride_shift_b, stride_shift_k, ++ K: tl.constexpr, ++ D: tl.constexpr, ++ BLOCK_D: tl.constexpr, ++): ++ pid_k = tl.program_id(0) ++ pid_b = tl.program_id(1) ++ ++ cnt = tl.load(cnts_ptr + pid_b * stride_cnts_b + pid_k * stride_cnts_k).to(tl.float32) ++ inv = 1.0 / tl.maximum(cnt, 1.0) ++ is_empty = cnt == 0.0 ++ ++ sq_acc = tl.zeros([], dtype=tl.float32) ++ offs_d = tl.arange(0, BLOCK_D) ++ n_blocks = (D + BLOCK_D - 1) // BLOCK_D ++ ++ for blk in range(n_blocks): ++ d_idx = blk * BLOCK_D + offs_d ++ d_mask = d_idx < D ++ ++ sum_off = pid_b * stride_sums_b + pid_k * stride_sums_k + d_idx * stride_sums_d ++ old_off = pid_b * stride_old_b + pid_k * stride_old_k + d_idx * stride_old_d ++ new_off = pid_b * stride_new_b + pid_k * stride_new_k + d_idx * stride_new_d ++ ++ s = tl.load(sums_ptr + sum_off, mask=d_mask, other=0.0).to(tl.float32) ++ old_v = tl.load(old_ptr + old_off, mask=d_mask, other=0.0).to(tl.float32) ++ ++ # divide; if empty, fall back to old centroid (shift contribution = 0) ++ new_v = tl.where(is_empty, old_v, s * inv) ++ ++ # accumulate squared shift ++ diff = new_v - old_v ++ sq_acc += tl.sum(tl.where(d_mask, diff * diff, 0.0)) ++ ++ tl.store(new_ptr + new_off, new_v, mask=d_mask) ++ ++ shift = tl.sqrt(sq_acc) ++ tl.store(shift_ptr + pid_b * stride_shift_b + pid_k * stride_shift_k, shift) ++ ++ ++def triton_centroid_finalize( ++ sums: torch.Tensor, # (B, K, D) fp32 ++ cnts: torch.Tensor, # (B, K) int32 ++ old_centroids: torch.Tensor, # (B, K, D) original dtype ++ *, ++ out: torch.Tensor = None, ++ shift: torch.Tensor = None, ++ BLOCK_D: int = 128, ++): ++ """Fused finalize: sums/cnts → new centroids + per-cluster shifts. ++ ++ Replaces the host pipeline: ++ cnts_f = cnts.float().unsqueeze(-1).clamp(min=1) ++ new = sums / cnts_f ++ new = where(cnts==0, old, new) ++ new = new.to(old.dtype) ++ shift = (new - old).norm(dim=-1) # (B, K) ++ ++ Returns (new_centroids, shift) where shift has shape (B, K) — caller takes ++ `.max()` over the K axis for the convergence criterion. ++ """ ++ assert sums.is_cuda and cnts.is_cuda and old_centroids.is_cuda ++ B, K, D = sums.shape ++ assert cnts.shape == (B, K) ++ assert old_centroids.shape == (B, K, D) ++ ++ if out is None: ++ out = torch.empty_like(old_centroids) ++ if shift is None: ++ shift = torch.empty((B, K), device=sums.device, dtype=torch.float32) ++ ++ grid = (K, B) ++ _centroid_finalize_kernel[grid]( ++ sums, cnts, old_centroids, out, shift, ++ sums.stride(0), sums.stride(1), sums.stride(2), ++ cnts.stride(0), cnts.stride(1), ++ old_centroids.stride(0), old_centroids.stride(1), old_centroids.stride(2), ++ out.stride(0), out.stride(1), out.stride(2), ++ shift.stride(0), shift.stride(1), ++ K=K, D=D, BLOCK_D=BLOCK_D, ++ ) ++ return out, shift ++ ++ ++def triton_lloyd_centroid_step_euclid( ++ x: torch.Tensor, ++ cluster_ids: torch.Tensor, ++ old_centroids: torch.Tensor, ++ *, ++ BLOCK_N: int = 256, ++ sums_buf: torch.Tensor = None, ++ cnts_buf: torch.Tensor = None, ++ new_buf: torch.Tensor = None, ++ shift_buf: torch.Tensor = None, ++): ++ """Single Lloyd centroid step: sort → chunk-update → fused finalize. ++ ++ All accumulator + output buffers can be preallocated and reused across ++ iterations for zero per-iter allocation cost. ++ ++ Returns ++ ------- ++ new_centroids : (B, K, D), original dtype (== `new_buf` if provided) ++ cluster_ids : (B, N) int (echoed back; unchanged) ++ max_shift : scalar fp32 GPU tensor — `(new - old).norm(-1).max()` ++ Caller can `.item()` to get host-side scalar. ++ """ ++ B, N, D = x.shape ++ K = old_centroids.shape[1] ++ ++ if sums_buf is None: ++ sums_buf = torch.zeros((B, K, D), device=x.device, dtype=torch.float32) ++ else: ++ sums_buf.zero_() ++ if cnts_buf is None: ++ cnts_buf = torch.zeros((B, K), device=x.device, dtype=torch.int32) ++ else: ++ cnts_buf.zero_() ++ ++ # Reuse the existing sorted-update kernel to fill sums_buf / cnts_buf (no ++ # final divide — we delegate that to the fused finalize below). ++ triton_centroid_update_sorted_euclid( ++ x, cluster_ids, old_centroids, ++ BLOCK_N=BLOCK_N, ++ centroid_sums=sums_buf, ++ centroid_cnts=cnts_buf, ++ calculate_new=False, ++ ) ++ ++ new_centroids, shift_per_k = triton_centroid_finalize( ++ sums_buf, cnts_buf, old_centroids, ++ out=new_buf, ++ shift=shift_buf, ++ ) ++ # max over K → (B,), then max over B → scalar (matches existing ++ # `(cnew - cold).norm(dim=-1).max()` semantics) ++ max_shift = shift_per_k.amax(dim=-1).amax() ++ return new_centroids, cluster_ids, max_shift ++ ++ ++def main(): ++ torch.manual_seed(0) ++ ++ B, N, D = 32, 74256, 128 # modest sizes for quick correctness test ++ K = 1000 ++ dtype = torch.float16 ++ ++ x = torch.randn(B, N, D, device="cuda", dtype=dtype) ++ x_norm = F.normalize(x, p=2, dim=-1) ++ ++ cluster_ids = torch.randint(0, K, (B, N), device="cuda", dtype=torch.int32) ++ ++ # Random old centroids for handling empty clusters ++ old_centroids = F.normalize(torch.randn(B, K, D, device="cuda", dtype=dtype), p=2, dim=-1) ++ ++ # ---------------- Correctness check (compile Triton kernel) ---------------- ++ ref_centroids = torch_loop_centroid_update_cosine(x_norm, cluster_ids, old_centroids) ++ tri_centroids = triton_centroid_update_cosine(x_norm, cluster_ids, old_centroids) # this call triggers compilation ++ tri_sorted_centroids = triton_centroid_update_sorted_cosine(x_norm, cluster_ids, old_centroids) ++ ++ # Validate correctness (includes first-run compile) ++ if torch.allclose(ref_centroids, tri_centroids, atol=1e-3, rtol=1e-3): ++ print("Centroid update: PASS ✅") ++ else: ++ max_diff = (ref_centroids - tri_centroids).abs().max().item() ++ print(f"Centroid update: FAIL ❌ | max diff = {max_diff}") ++ ++ # Validate new sorted kernel ++ if torch.allclose(ref_centroids, tri_sorted_centroids, atol=1e-3, rtol=1e-3): ++ print("Sorted centroid update: PASS ✅") ++ else: ++ max_diff = (ref_centroids - tri_sorted_centroids).abs().max().item() ++ print(f"Sorted centroid update: FAIL ❌ | max diff = {max_diff}") ++ ++ ++ # show some examples ++ print(f"ref_centroids[0,0:5,0:5]: {ref_centroids[0,0:5,0:5]}") ++ print(f"tri_centroids[0,0:5,0:5]: {tri_centroids[0,0:5,0:5]}") ++ print(f"tri_sorted_centroids[0,0:5,0:5]: {tri_sorted_centroids[0,0:5,0:5]}") ++ ++ # ---------------- Efficiency benchmark (exclude compile) ---------------- ++ repeats = 20 ++ ++ # Torch loop timing ++ torch.cuda.synchronize() ++ start = torch.cuda.Event(enable_timing=True) ++ end = torch.cuda.Event(enable_timing=True) ++ start.record() ++ for _ in trange(repeats): ++ torch_loop_centroid_update_cosine(x_norm, cluster_ids, old_centroids) ++ end.record(); torch.cuda.synchronize() ++ torch_time = start.elapsed_time(end) / repeats # average per run (ms) ++ ++ # Triton timing (already compiled) ++ torch.cuda.synchronize() ++ start = torch.cuda.Event(enable_timing=True) ++ end = torch.cuda.Event(enable_timing=True) ++ start.record() ++ for _ in trange(repeats): ++ triton_centroid_update_cosine(x_norm, cluster_ids, old_centroids) ++ end.record(); torch.cuda.synchronize() ++ triton_time = start.elapsed_time(end) / repeats # average per run (ms) ++ ++ # Sorted Triton timing (already compiled) ++ torch.cuda.synchronize() ++ start = torch.cuda.Event(enable_timing=True) ++ end = torch.cuda.Event(enable_timing=True) ++ start.record() ++ for _ in trange(repeats): ++ triton_centroid_update_sorted_cosine(x_norm, cluster_ids, old_centroids) ++ end.record(); torch.cuda.synchronize() ++ triton_sorted_time = start.elapsed_time(end) / repeats # average per run (ms) ++ ++ print(f"\n=== Efficiency (average over {repeats} runs, exclude compile) ===") ++ print(f"Torch loop : {torch_time:.2f} ms") ++ print(f"Triton kernel: {triton_time:.2f} ms (speed-up x{torch_time / triton_time:.2f})") ++ print(f"Triton sorted: {triton_sorted_time:.2f} ms (speed-up x{torch_time / triton_sorted_time:.2f})") ++ ++ ++if __name__ == "__main__": ++ main() +diff --git a/kmeanslib/kmeans.py b/kmeanslib/kmeans.py +index 2eadfdb..9823476 100644 +--- a/kmeanslib/kmeans.py ++++ b/kmeanslib/kmeans.py +@@ -1,80 +1,31 @@ +-"""One Lloyd K-Means step -- the function you optimise. ++"""Euclidean (squared-L2) one Lloyd step -- vendored best-in-repo Triton path. + +-The judge owns the iteration loop, the data, the initial centroids, and the +-iteration count; you provide a single Lloyd step and it is called repeatedly. +-This shipped step is intentionally naive (materialise a (chunk, K) distance +-matrix with a bf16 matmul + argmin, then a PyTorch scatter update). +- +-Contract (do NOT change): ++Delegates to the Triton assign + sorted centroid-update kernels vendored under ++``kmeanslib._kernels`` (the fastest single-iteration Euclidean K-Means path in ++the source repo). Public contract is unchanged: + + step(x, centroids) -> (labels, new_centroids) + +- x : (N, D) bfloat16 CUDA tensor of points. +- centroids : (K, D) bfloat16 tensor -- the current centroids. +- labels : (N,) int64 -- nearest-centroid assignment of every point to +- `centroids` (a full assignment of all N points). +- new_centroids: (K, D) -- centroids recomputed as the mean of each cluster +- (empty clusters keep their previous centroid). +- +-This is exactly one Lloyd iteration: assign to `centroids`, then update. You may +-add modules/kernels under ``kmeanslib`` and rewrite the body of ``step``, +-``_assign`` and ``_update`` freely -- including **fusing the assign + update into +-a single kernel** -- as long as ``step(x, centroids) -> (labels, new_centroids)`` +-is preserved. You cannot change how many times it runs, the data, or the initial +-centroids; the judge owns the loop and calls ``step`` a fixed number of times. ++This is exactly one Lloyd iteration -- assign every point to its nearest ++``centroids`` (Triton distance kernel), then recompute the centroids as the mean ++of each cluster (Triton sorted scatter-update). The judge owns the loop. + """ + from __future__ import annotations + + import torch + ++from kmeanslib._kernels.primitives.kmeans.triton.assign import euclid_assign_triton ++from kmeanslib._kernels.primitives.kmeans.triton.update import ( ++ triton_centroid_update_sorted_euclid, ++) + +-def _assign(x: torch.Tensor, centroids: torch.Tensor) -> torch.Tensor: +- """Nearest-centroid assignment by squared-L2 distance. +- +- Naive: for each chunk of points, materialise the (chunk, K) distance matrix +- with a bf16 matmul, then argmin. ``argmin ||x-c||^2 == argmin (||c||^2 - 2 +- x.c^T)`` since ``||x||^2`` is constant per row. bf16 in, fp32 accumulation. +- """ +- c = centroids.to(x.dtype) +- c_sq = (c.float() * c.float()).sum(1) # (K,) fp32 +- labels = torch.empty(x.shape[0], device=x.device, dtype=torch.long) +- for lo in range(0, x.shape[0], 16384): +- xb = x[lo:lo + 16384] +- dist = c_sq[None, :] - 2.0 * (xb @ c.t()).float() # (chunk, K) fp32 +- labels[lo:lo + 16384] = torch.argmin(dist, dim=1) +- return labels +- +- +-def _update( +- x: torch.Tensor, +- labels: torch.Tensor, +- n_clusters: int, +- old_centroids: torch.Tensor, +-) -> torch.Tensor: +- """Recompute each centroid as the mean of its assigned points. +- +- Empty clusters keep their previous centroid. +- """ +- N, D = x.shape +- sums = torch.zeros((n_clusters, D), device=x.device, dtype=torch.float32) +- counts = torch.zeros((n_clusters,), device=x.device, dtype=torch.float32) +- sums.index_add_(0, labels, x.float()) +- counts.index_add_(0, labels, torch.ones(N, device=x.device, dtype=torch.float32)) +- empty = counts == 0 +- counts = counts.clamp_min(1.0) +- new = sums / counts[:, None] +- if empty.any(): +- new[empty] = old_centroids[empty].float() +- return new.to(x.dtype) +- +- +-def step(x: torch.Tensor, centroids: torch.Tensor): +- """One Lloyd iteration: assign to `centroids`, then recompute them. + +- Returns ``(labels, new_centroids)``. See the module docstring for the contract. +- """ ++def step(x, centroids): + if x.ndim != 2: +- raise ValueError(f"x must be 2-D (N, D); got shape {tuple(x.shape)}") +- labels = _assign(x, centroids) +- new_centroids = _update(x, labels, centroids.shape[0], centroids) +- return labels, new_centroids ++ raise ValueError(f"x must be 2-D (N, D); got {tuple(x.shape)}") ++ # The vendored kernels are batched (B, N, D); run a single instance as B=1. ++ xb = x.unsqueeze(0).contiguous() ++ cb = centroids.to(x.dtype).unsqueeze(0).contiguous() ++ labels = euclid_assign_triton(xb, cb, use_heuristic=True) # (1, N) ++ new_c = triton_centroid_update_sorted_euclid(xb, labels, cb) # (1, K, D) ++ return labels.squeeze(0).to(torch.long), new_c.squeeze(0).to(x.dtype) diff --git a/2.0/problems/knn_gpu_kernel_optimization/DESIGN.md b/2.0/problems/knn_gpu_kernel_optimization/DESIGN.md new file mode 100644 index 000000000..285c31089 --- /dev/null +++ b/2.0/problems/knn_gpu_kernel_optimization/DESIGN.md @@ -0,0 +1,84 @@ +# Design notes — knn_gpu_kernel_optimization + +Operator-facing. Not copied into the agent workspace by the adapter. + +## What this task measures + +Can an agent turn a naive GPU brute-force k-NN into a fast one? The agent +patches `knnlib`; the judge times the patched `knn` against a frozen naive +baseline on hidden `(Q, M, D, k)` workloads and scores the geometric-mean +speedup, gated on nearest-neighbor quality (recall@k). Second of a four-task +**flashlib kernel-optimization family** (KMeans, KNN, TruncatedSVD, PCA) that +all share one evaluator + one Modal GPU harness. + +## Execution: Modal GPU offload (mirrors the vllm task) + +The judge and the agent public test both **offload GPU work to Modal**; the +containers are light (ubuntu + `modal` + git, no torch/triton). `flash_gpu.py` +(baked at `/opt/flash_gpu.py` in both images) builds an ephemeral Modal app on a +GPU (`evaluation.gpu`, default H100), ships the frozen baseline + the patched +package **as data** (no per-submission image rebuild), runs the worker, and +returns per-workload speedups + verdicts. Needs `MODAL_TOKEN_ID` / +`MODAL_TOKEN_SECRET`. torch/triton/the vendored kernels run on the Modal image +(`evaluation.pip`). + +The evaluator is generic and identical across the four tasks; only three +constants differ (`PRIMITIVE`/`PKG`/`REF_MODULE`), and all workload/threshold/ +Modal values come from `config.yaml` (delivered judge-only as +`/judge/task_config.json`, never present in the agent image). + +## Reference solution = best in the repo + +`reference.patch` vendors flashlib's own optimized **Triton** k-NN path +(`primitives/knn/triton/{dispatch,insert,sortmerge,_common,_row_norm}.py` plus +the squared-L2 gather kernel `kernels/distance/triton/knn_gather_l2sq.py`) under +`knnlib/_kernels/…`, with imports rewritten and the string "flashlib" scrubbed, +plus a thin adapter mapping our `(Q, D)` / `(M, D)` contract onto flashlib's +batched entry (`flash_knn_triton` for the indices + `triton_knn_gather_sqdist` +for the true squared distances). It is triton-only (runs on any sm≥80), imports +cleanly on CPU, and applies + passes the patch policy. Calibrate `speedup_target` +from its measured geomean on the first GPU trial. + +## Anti-reward-hack + +The **load-bearing** defenses are structural, inside the GPU worker (the only +place the untrusted submission runs): + +* timing primitives (`perf_counter`, `torch.cuda.synchronize`) are captured to + locals **before** the submission is imported → monkey-patching torch cannot + affect measurement; +* **fresh data every timed iteration** → memoizing on a repeated input is + useless; +* quality is **re-verified from the returned tensors on every iteration** → + returning a stale cached result for a different input is caught; +* the judge computes all metrics (no agent number is trusted); +* an ephemeral Modal container per submission → no global state persists. + +The patch policy is surgical defense-in-depth (so it does not reject legitimate +vendored/optimized kernel source): only `/**` `.py` files may change; it +bans external-optimized-library *imports* (flashlib/cuml/cupy/faiss/sklearn/ +cutlass — none installed on the GPU image anyway) and process/network/ +measurement-tamper patterns (`subprocess`, `socket`, `os.system`, +`torch.cuda.synchronize =`, …). + +The agent-facing `readme` describes the contract, gate, scoring, and policy but +**not** how the reference is implemented. + +## Correctness gate + +recall@k = fraction of the true `k` nearest database indices the submission +recovers, averaged over all queries, judge-computed from the returned indices on +each iteration's data against the frozen exact baseline (`refknn`, a `cdist` + +`topk`). Tie-break- and convention-independent (exact-set based), robust to fp +drift, cheat-proof: high recall requires actually finding the neighbors. +`queries`/`database` are generated from a per-workload seed, so the baseline — +and thus the true top-k — is a deterministic function of `(Q, M, D, k, seed)`. + +## Calibration TODO (needs a Modal GPU trial) + +- Run `reference.patch` on Modal; confirm it passes the recall@k gate on all + workloads (bump `recall_threshold` slack only if flashlib's low-precision + cross-term drifts a hair) and set `speedup_target` from its geomean. +- Confirm the pinned `torch`/`triton` in `evaluation.pip` compile the vendored + kernels, and hidden shapes fit the GPU. +- Confirm Modal token wiring in the Harbor judge/agent containers. diff --git a/2.0/problems/knn_gpu_kernel_optimization/config.yaml b/2.0/problems/knn_gpu_kernel_optimization/config.yaml new file mode 100644 index 000000000..8c8ed35ef --- /dev/null +++ b/2.0/problems/knn_gpu_kernel_optimization/config.yaml @@ -0,0 +1,73 @@ +tag: systems +runtime: + language: python + timeout_seconds: 10800 + environment: "GPU brute-force k-NN kernel optimization. The agent patches the knnlib package; the judge offloads timing to a Modal GPU, comparing the patched nearest-neighbor search against a frozen naive baseline on hidden (Q, M, D, k) workloads with a recall@k (retrieval-quality) gate." + apt_packages: + - bash + - ca-certificates + - git + - python3 + - python3-pip + judge_apt_packages: + - bash + - ca-certificates + - git + - python3 + - python3-pip + judge_pip_packages: + - modal + docker: + # Experimental local images. Build with docker/build_images.sh before a local + # Harbor trial. Both are light (ubuntu + modal + git); torch/triton/flashlib + # run on the Modal GPU image defined in flash_gpu.py. The agent image bakes + # /app/knnlib (git) + /opt/knn_ref + /opt/flash_gpu.py; the judge image + # additionally bakes /opt/knnlib-clean. + image: frontiercs/knn-gpu-kernel-optimization-agent:experimental-v0.6.0 + judge_image: frontiercs/knn-gpu-kernel-optimization-judge:experimental-v0.6.0 +environment: + cpus: 4 + memory_mb: 8192 + storage_mb: 16384 + build_timeout_seconds: 3600 +evaluation: + # --- primitive wiring (consumed by the generic evaluator + flash_gpu) --- + primitive: knn + pkg: knnlib + ref_module: refknn + clean_source: /opt/knnlib-clean + baseline_source: /opt/knn_ref + # --- Modal GPU --- + gpu: "H100" + cuda_image: "nvidia/cuda:12.4.1-devel-ubuntu22.04" + pip: ["torch", "numpy", "h5py"] # match the original (validated) knn image; torch brings its matched triton. h5py loads SIFT/GIST HDF5 + app_name: "knn-kernel-opt-eval" + dataset_volume: "flashlib-ann-datasets" # Modal Volume holding sift.hdf5 / gist.hdf5 (ann-benchmarks); mounted read-only at /data + modal_timeout_seconds: 3600 # exact brute-force over a 1M-row real database is heavier than the synthetic shapes; give margin + warmup_iters: 2 + timed_iters: 7 + recall_threshold: 0.88 # ball-recall gate on bf16 data (see knn_ball_recall): a returned neighbour counts if its exact fp32 distance is within the true k-NN ball. bf16 distances tie heavily near rank k, so even a correct bf16 kernel only reaches ~0.94-1.0 (the flashlib reference measured 0.935-1.0). 0.88 clears that with margin while still rejecting a sub-bf16 (e.g. fp8) search, whose coarser distances drop far lower. Precision is pinned by the bf16 DATA, not this threshold. + distance_ball_tolerance: 0.0 # relative slack on the k-NN ball radius (the check already carries a 1e-6 absolute epsilon for fp32 reduction jitter) + speedup_target: 4.0 # score cap = 2x the flashlib reference geomean on the bf16 real workloads (1.94x on H100: SIFT k10/50/100 = 4.43/2.60/1.48x, GIST = 1.71/1.67/1.10x, ball-recall >=0.94 all 6). bf16 data + modest cuBLAS-cdist headroom keep the factor small. + base_seed: 20260701 + agent_workload_count: 6 # agent's submission feedback covers ALL graded workloads (no blind pass-3/fail-6 cliff); params are agent-visible anyway + expose_per_workload_metrics: true # show per-workload pass/fail so the agent can fix the specific workload it regresses on + # REAL-DATA WORKLOADS: exact brute-force k-NN over real descriptor vectors + # (SIFT 128d / GIST 960d, 1M-row bases). ANTI-CACHING: every iteration draws a + # fresh random Q-subset of the held-out real query set AND a fresh random + # M-subset of the real base as the database, so neither the queries nor the + # database tensor repeats across calls -- a solution cannot cache a prebuilt + # index/norms keyed on a fixed database and skip the timed search (a fixed + # database let a trial hit 80x by caching an index instead of writing a fast + # kernel). M = database rows drawn per iteration; k varies the neighbour count. + # Q <= #queries (SIFT 10000, GIST 1000); M <= base rows (1M). + workloads: + - {id: sift_k10, dataset: sift, M: 500000, Q: 1024, k: 10} + - {id: sift_k50, dataset: sift, M: 500000, Q: 1024, k: 50} + - {id: sift_k100, dataset: sift, M: 500000, Q: 1024, k: 100} + - {id: gist_k10, dataset: gist, M: 500000, Q: 512, k: 10} + - {id: gist_k50, dataset: gist, M: 500000, Q: 512, k: 50} + - {id: gist_k100, dataset: gist, M: 500000, Q: 512, k: 100} +submission: + kind: file + path: /app/solution.patch diff --git a/2.0/problems/knn_gpu_kernel_optimization/docker/README.md b/2.0/problems/knn_gpu_kernel_optimization/docker/README.md new file mode 100644 index 000000000..1071e3002 --- /dev/null +++ b/2.0/problems/knn_gpu_kernel_optimization/docker/README.md @@ -0,0 +1,56 @@ +# Experimental Brute-Force k-NN Kernel-Optimization Images + +Two **light** images (ubuntu + `modal` + git). torch, triton, and the vendored +flashlib kernels all run on the **Modal GPU image** defined in `flash_gpu.py`; +the containers themselves are CPU-only and offload GPU work to Modal. + +```bash +bash 2.0/problems/knn_gpu_kernel_optimization/docker/build_images.sh +``` + +Defaults: + +```text +AGENT_TAG=frontiercs/knn-gpu-kernel-optimization-agent:experimental-v0.2.0 +JUDGE_TAG=frontiercs/knn-gpu-kernel-optimization-judge:experimental-v0.2.0 +``` + +Agent image: + +```text +/app/knnlib # clean, git-tracked package (the agent edits this) +/opt/flash_gpu.py # shared Modal GPU harness (public test uses it) +/opt/knn_ref/refknn.py # frozen naive baseline (public-test speed denominator) +``` + +Judge image: + +```text +/opt/knnlib-clean/knnlib # pristine tree; the patch is applied to a copy +/opt/knn_ref/refknn.py # frozen naive baseline (speed denominator + quality oracle) +/opt/flash_gpu.py # shared Modal GPU harness (the evaluator uses it) +``` + +## Runtime requirements + +Both the judge and the agent public test **offload timing to a Modal GPU** and +therefore need Modal credentials in the environment: + +```text +MODAL_TOKEN_ID / MODAL_TOKEN_SECRET +``` + +`flash_gpu.py` builds an ephemeral Modal app on a GPU (`evaluation.gpu`, default +`H100`), ships the frozen baseline + the patched package as data, times both on +fresh per-iteration data, verifies quality each iteration, and returns the +speedups. No persistent deployment is used, so a fresh GPU container is spun up +per evaluation. Without Modal credentials the evaluator returns a patch-policy +smoke pass (which repo CI exercises). + +## Smoke test + +```bash +bash 2.0/problems/knn_gpu_kernel_optimization/docker/smoke_images.sh +``` + +Import-only (modal + flash_gpu + the baked packages); does not touch a GPU or Modal. diff --git a/2.0/problems/knn_gpu_kernel_optimization/docker/agent/Dockerfile b/2.0/problems/knn_gpu_kernel_optimization/docker/agent/Dockerfile new file mode 100644 index 000000000..e2e5acb1b --- /dev/null +++ b/2.0/problems/knn_gpu_kernel_optimization/docker/agent/Dockerfile @@ -0,0 +1,34 @@ +# Agent image for knn_gpu_kernel_optimization. +# +# Light image (no torch/triton): the agent edits /app/knnlib and runs the +# public test, which offloads timing to a Modal GPU (torch/triton live on the +# Modal image defined in flash_gpu.py). Needs MODAL_TOKEN_ID / MODAL_TOKEN_SECRET +# in the environment to run the GPU public test. +# +# Build context = the task directory: +# docker build -f docker/agent/Dockerfile -t . +FROM ubuntu:24.04 + +ARG DEBIAN_FRONTEND=noninteractive +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + bash ca-certificates git python3 python3-pip ripgrep && \ + rm -rf /var/lib/apt/lists/* + +RUN pip3 install --break-system-packages --no-cache-dir modal + +# The package the agent edits (git-tracked so make_submission.sh can diff it). +WORKDIR /app +COPY knnlib /app/knnlib +RUN printf '%s\n' '__pycache__/' '*.pyc' > /app/.gitignore && \ + git -C /app init -q && \ + git -C /app config user.email task@frontier-cs && \ + git -C /app config user.name frontier-cs && \ + git -C /app add -A -- knnlib .gitignore && \ + git -C /app -c commit.gpgsign=false commit -qm "pristine knnlib" + +# Shared Modal GPU harness + the frozen naive baseline (the public-test speed +# denominator). The baseline is exactly the code the agent starts from, so it is +# not secret. +COPY flash_gpu.py /opt/flash_gpu.py +COPY judge/refknn.py /opt/knn_ref/refknn.py diff --git a/2.0/problems/knn_gpu_kernel_optimization/docker/build_images.sh b/2.0/problems/knn_gpu_kernel_optimization/docker/build_images.sh new file mode 100755 index 000000000..1b91827ca --- /dev/null +++ b/2.0/problems/knn_gpu_kernel_optimization/docker/build_images.sh @@ -0,0 +1,17 @@ +#!/usr/bin/env bash +# Build the experimental agent + judge images for knn_gpu_kernel_optimization. +set -euo pipefail +HERE=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd) # the problem directory (build context) + +AGENT_TAG=${AGENT_TAG:-frontiercs/knn-gpu-kernel-optimization-agent:experimental-v0.5.0} +JUDGE_TAG=${JUDGE_TAG:-frontiercs/knn-gpu-kernel-optimization-judge:experimental-v0.5.0} + +echo "Building agent image: $AGENT_TAG" +docker build -f "$HERE/docker/agent/Dockerfile" -t "$AGENT_TAG" "$HERE" + +echo "Building judge image: $JUDGE_TAG" +docker build -f "$HERE/docker/judge/Dockerfile" -t "$JUDGE_TAG" "$HERE" + +echo "Built:" +echo " $AGENT_TAG" +echo " $JUDGE_TAG" diff --git a/2.0/problems/knn_gpu_kernel_optimization/docker/judge/Dockerfile b/2.0/problems/knn_gpu_kernel_optimization/docker/judge/Dockerfile new file mode 100644 index 000000000..92aaa9a26 --- /dev/null +++ b/2.0/problems/knn_gpu_kernel_optimization/docker/judge/Dockerfile @@ -0,0 +1,26 @@ +# Judge image for knn_gpu_kernel_optimization. +# +# Light image (no torch/triton): the judge validates + applies the patch and +# offloads timing to a Modal GPU (torch/triton live on the Modal image defined +# in flash_gpu.py). Needs MODAL_TOKEN_ID / MODAL_TOKEN_SECRET in the environment. +# The Frontier-CS 2.0 adapter builds the final judge image ON TOP of this one. +# +# Build context = the task directory: +# docker build -f docker/judge/Dockerfile -t . +FROM ubuntu:24.04 + +ARG DEBIAN_FRONTEND=noninteractive +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + bash ca-certificates git python3 python3-pip && \ + rm -rf /var/lib/apt/lists/* + +RUN pip3 install --break-system-packages --no-cache-dir modal + +# Pristine tree the patch is applied to, the frozen naive baseline (speed +# denominator + quality oracle), and the shared Modal GPU harness. +COPY knnlib /opt/knnlib-clean/knnlib +COPY judge/refknn.py /opt/knn_ref/refknn.py +COPY flash_gpu.py /opt/flash_gpu.py + +WORKDIR /judge diff --git a/2.0/problems/knn_gpu_kernel_optimization/docker/smoke_images.sh b/2.0/problems/knn_gpu_kernel_optimization/docker/smoke_images.sh new file mode 100755 index 000000000..0ef7f8b62 --- /dev/null +++ b/2.0/problems/knn_gpu_kernel_optimization/docker/smoke_images.sh @@ -0,0 +1,22 @@ +#!/usr/bin/env bash +# Smoke test for the built (intentionally torch-less) images: confirm the host-side +# harness imports (modal + flash_gpu) and that the frozen baseline + package parse. +# torch/triton live on the Modal image, so the packages are py_compiled here (a +# torch-less parse check), not imported. No GPU / Modal tokens required. +set -euo pipefail + +AGENT_TAG=${AGENT_TAG:-frontiercs/knn-gpu-kernel-optimization-agent:experimental-v0.2.0} +JUDGE_TAG=${JUDGE_TAG:-frontiercs/knn-gpu-kernel-optimization-judge:experimental-v0.2.0} + +SMOKE='import sys, modal, py_compile, pathlib +sys.path.insert(0, "/opt"); import flash_gpu +for root in sys.argv[1:]: + for p in pathlib.Path(root).rglob("*.py"): + py_compile.compile(str(p), doraise=True) +print("ok: modal + flash_gpu import; knn baseline + knnlib parse")' + +echo "== agent image ==" +docker run --rm -w /app "$AGENT_TAG" python3 -c "$SMOKE" /opt/knn_ref /app/knnlib + +echo "== judge image ==" +docker run --rm "$JUDGE_TAG" python3 -c "$SMOKE" /opt/knn_ref /opt/knnlib-clean/knnlib diff --git a/2.0/problems/knn_gpu_kernel_optimization/evaluate.sh b/2.0/problems/knn_gpu_kernel_optimization/evaluate.sh new file mode 100755 index 000000000..3982ca60b --- /dev/null +++ b/2.0/problems/knn_gpu_kernel_optimization/evaluate.sh @@ -0,0 +1,12 @@ +#!/usr/bin/env bash +# Local CLI evaluation for knn_gpu_kernel_optimization. +# +# A full run needs a GPU plus the baked judge sources at /opt (the pristine +# /opt/knnlib-clean tree and the frozen /opt/knn_ref baseline; see the judge +# image). Without them the evaluator still validates the patch policy and +# returns a smoke score, which is what repository CI exercises (the reference +# patch passes the policy). Pass a patch path to score it directly. +set -euo pipefail +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +SOLUTION="${1:-$SCRIPT_DIR/reference.patch}" +exec python3 "$SCRIPT_DIR/evaluator.py" "$SOLUTION" diff --git a/2.0/problems/knn_gpu_kernel_optimization/evaluator.py b/2.0/problems/knn_gpu_kernel_optimization/evaluator.py new file mode 100644 index 000000000..b0315dbc2 --- /dev/null +++ b/2.0/problems/knn_gpu_kernel_optimization/evaluator.py @@ -0,0 +1,381 @@ +"""Generic evaluator for the flashlib GPU kernel-optimization task family. + +Identical across all four tasks (kmeans / knn / pca / truncated_svd); every +primitive-specific value comes from the ``evaluation`` block of the task's +config (delivered to the judge as ``/judge/task_config.json``), which is not +present in the agent workspace. + +Flow: statically validate the patch (only ``/**`` may change; no external +optimized libs / env / process / network access) -> apply it to the pristine +package baked at ``clean_source`` -> hand the frozen baseline + patched package +sources to a Modal GPU worker (``flash_gpu.run_remote``) that times both on +fresh per-iteration data and verifies clustering/retrieval/subspace quality on +every iteration -> gate on the worker's per-workload verdict -> score by the +geometric-mean speedup. + +Without Modal credentials or the baked judge sources (e.g. repo CI on a box with +no GPU/Modal), the evaluator still validates the patch policy and returns a +smoke pass, so ``python3 evaluator.py reference.patch`` works anywhere. +""" +from __future__ import annotations + +import fnmatch +import hashlib +import json +import math +import os +import re +import shutil +import subprocess +import sys +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +TASK_CONFIG_PATH = Path("/judge/task_config.json") + + +def _load_task_config() -> dict[str, Any]: + try: + payload = json.loads(TASK_CONFIG_PATH.read_text(encoding="utf-8")) + except Exception: + return {} + return payload if isinstance(payload, dict) else {} + + +TASK_CONFIG = _load_task_config() +EVAL = TASK_CONFIG.get("evaluation", {}) if isinstance(TASK_CONFIG.get("evaluation"), dict) else {} + + +def _get(name: str, default): + return EVAL.get(name, default) + + +def _get_int(name: str, default: int) -> int: + try: + return int(EVAL.get(name, default)) + except Exception: + return default + + +def _get_float(name: str, default: float) -> float: + try: + return float(EVAL.get(name, default)) + except Exception: + return default + + +def _get_bool(name: str, default: bool) -> bool: + raw = EVAL.get(name, default) + if isinstance(raw, bool): + return raw + if isinstance(raw, str): + return raw.strip().lower() in {"1", "true", "yes", "on"} + return bool(raw) + + +# --- Per-task identity: the ONLY lines that differ across the four tasks. --- +# Kept as constants (not config) so the patch policy is enforceable locally / in +# CI without /judge/task_config.json. Everything else is config-driven and only +# needed on the GPU path. +PRIMITIVE = "knn" +PKG = "knnlib" +REF_MODULE = "refknn" + +MAX_PATCH_BYTES = _get_int("max_patch_bytes", 2_000_000) +MAX_CHANGED_FILES = _get_int("max_changed_files", 80) +CLEAN_SOURCE = Path(f"/opt/{PKG}-clean") +BASELINE_SOURCE = Path(f"/opt/{PRIMITIVE}_ref") +SPEEDUP_TARGET = _get_float("speedup_target", 8.0) + +ALLOWED_PATTERNS = (f"{PKG}/**",) +DENIED_PATTERNS = ( + "evaluator.py", + "flash_gpu.py", + "reference.patch", + f"{REF_MODULE}.py", + f"**/{REF_MODULE}.py", + "**/conftest.py", + "**/test_*.py", + "pyproject.toml", + "setup.py", + "setup.cfg", +) +# Defense in depth (the load-bearing anti-tamper defense is structural, inside +# the GPU worker: it captures its perf_counter / synchronize references before +# importing the submission, regenerates data every iteration, and re-verifies +# quality each iteration). These patterns are matched surgically so that +# legitimate vendored/optimized kernel source is not rejected: +# * external optimized libraries, matched as IMPORT statements (not bare words, +# so a comment mentioning a name is fine); none are installed on the GPU +# image anyway; +# * process / network / measurement-tamper patterns that never appear in a +# real GPU kernel. +FORBIDDEN_IMPORT_RE = re.compile( + r"(?:^|\n)\s*(?:import|from)\s+(flashlib|cuml|cudf|cupy|faiss|sklearn|scikit|cutlass|cuvs)\b" +) +FORBIDDEN_PATTERN_RE = re.compile( + r"\bsubprocess\b|\bsocket\b|\burllib\b|\brequests\b|os\.system|" + r"torch\.cuda\.synchronize\s*=|time\.perf_counter\s*=|setattr\(\s*torch\.cuda" +) + + +@dataclass(frozen=True) +class PatchFile: + old_path: str + new_path: str + added_lines: tuple[str, ...] + + @property + def path(self) -> str: + return self.new_path if self.new_path != "/dev/null" else self.old_path + + +def _match(path: str, patterns: tuple[str, ...]) -> bool: + return any(fnmatch.fnmatch(path, p) for p in patterns) + + +def _invalid(message: str, metrics: dict[str, Any] | None = None): + payload = metrics or {} + payload.setdefault("valid_patch", 0) + return 0.0, 0.0, message, payload + + +def _parse_patch(text: str) -> list[PatchFile]: + files: list[PatchFile] = [] + old = new = "" + added: list[str] = [] + in_file = False + for line in text.splitlines(): + if line.startswith("diff --git "): + if in_file: + files.append(PatchFile(old, new, tuple(added))) + in_file, old, new, added = True, "", "", [] + continue + if not in_file: + continue + if line.startswith("--- "): + old = line[4:].strip() + old = old[2:] if old.startswith("a/") else old + elif line.startswith("+++ "): + new = line[4:].strip() + new = new[2:] if new.startswith("b/") else new + elif line.startswith("+") and not line.startswith("+++ "): + added.append(line[1:]) + if in_file: + files.append(PatchFile(old, new, tuple(added))) + return files + + +def _validate_path(path: str) -> tuple[bool, str]: + if not path or path == "/dev/null": + return True, "" + if path.startswith("/") or ".." in Path(path).parts: + return False, f"unsafe patch path: {path}" + if _match(path, DENIED_PATTERNS): + return False, f"changed file is outside task boundary: {path}" + if not _match(path, ALLOWED_PATTERNS): + return False, f"changed file is not allowlisted (only {PKG}/** may change): {path}" + if not path.endswith(".py"): + return False, f"only Python files may change: {path}" + return True, "" + + +def validate_patch(patch_path: Path) -> tuple[bool, str, dict[str, Any]]: + if not patch_path.exists(): + return False, "solution patch does not exist", {} + size = patch_path.stat().st_size + if size > MAX_PATCH_BYTES: + return False, f"patch is too large ({size} bytes > {MAX_PATCH_BYTES})", {} + text = patch_path.read_text(encoding="utf-8", errors="replace") + files = _parse_patch(text) + metrics: dict[str, Any] = { + "patch_bytes": size, + "patch_sha256": hashlib.sha256(text.encode("utf-8", "replace")).hexdigest(), + "changed_files": len(files), + } + if len(files) > MAX_CHANGED_FILES: + return False, f"too many changed files ({len(files)} > {MAX_CHANGED_FILES})", metrics + for pf in files: + path = pf.path + if pf.new_path == "/dev/null": + return False, f"deleting files is outside task boundary: {pf.old_path}", metrics + if pf.old_path != "/dev/null" and pf.old_path != pf.new_path: + ok, err = _validate_path(pf.old_path) + if not ok: + return False, f"rename source is outside task boundary: {err}", metrics + ok, err = _validate_path(path) + if not ok: + return False, err, metrics + added = "\n".join(pf.added_lines) + m = FORBIDDEN_IMPORT_RE.search(added) + if m: + return False, f"{path}: forbidden import of an external optimized library ({m.group(1)})", metrics + m = FORBIDDEN_PATTERN_RE.search(added) + if m: + return False, f"{path}: forbidden pattern in added code ({m.group(0).strip()[:40]})", metrics + metrics["valid_patch"] = 1 + return True, "patch accepted by static policy", metrics + + +def sanitize(text: str) -> str: + text = re.sub(r"/tmp/[A-Za-z0-9_./-]+", "", text or "") + text = re.sub(r"/opt/[A-Za-z0-9_./-]+", "", text) + text = re.sub(r"\bN=\d+|\bseed[=:]?\s*\d+", "", text, flags=re.IGNORECASE) + return text[-600:] + + +def geometric_mean(values: list[float]) -> float: + if not values: + return 0.0 + return math.exp(sum(math.log(max(v, 1e-9)) for v in values) / len(values)) + + +def score_from_speedup(gm: float) -> float: + if gm <= 0: + return 0.0 + raw = 100.0 * math.log(gm) / math.log(max(SPEEDUP_TARGET, 1.0000001)) + return max(0.0, min(100.0, raw)) + + +def is_final_role() -> bool: + return os.environ.get("FRONTIER_SUBMISSION_ROLE", "agent") == "final" + + +def _workloads(final_role: bool) -> list[dict[str, Any]]: + workloads = list(_get("workloads", []) or []) + if not final_role: + n = _get_int("agent_workload_count", 3) + workloads = workloads[:n] + base = _get_int("base_seed", 20260701) + for i, w in enumerate(workloads): + w.setdefault("seed", base + 1000 * (i + 1)) + return workloads + + +def _build_cfg() -> dict[str, Any]: + return { + "primitive": PRIMITIVE, + "pkg": PKG, + "ref_module": REF_MODULE, + "gpu": str(_get("gpu", "H100")), + "cuda_image": str(_get("cuda_image", "nvidia/cuda:12.4.1-devel-ubuntu22.04")), + "pip": list(_get("pip", ["torch==2.5.1", "triton==3.1.0", "numpy"])), + "app_name": str(_get("app_name", "flash-kernel-eval")), + "dataset_volume": _get("dataset_volume", None), # Modal Volume with real SIFT/GIST datasets, mounted at /data + "modal_timeout_seconds": _get_int("modal_timeout_seconds", 1800), + "warmup": _get_int("warmup_iters", 3), + "iters": _get_int("timed_iters", 7), + "inertia_tolerance": _get_float("inertia_tolerance", 0.02), + "recall_threshold": _get_float("recall_threshold", 0.99), + "distance_ball_tolerance": _get_float("distance_ball_tolerance", 0.0), + "captured_tolerance": _get_float("captured_tolerance", 0.02), + "ortho_tolerance": _get_float("ortho_tolerance", 0.02), + "ari_threshold": _get_float("ari_threshold", 0.99), + } + + +def _read_dir(root: Path, rel_to: Path) -> dict[str, str]: + out: dict[str, str] = {} + for p in root.rglob("*.py"): + out[str(p.relative_to(rel_to))] = p.read_text(encoding="utf-8", errors="replace") + return out + + +def full_evaluation(patch_path: Path, metrics: dict[str, Any]): + final_role = is_final_role() + metrics["submission_role"] = "final" if final_role else "agent" + workloads = _workloads(final_role) + + # Import the Modal harness baked into the judge image; missing harness or + # missing credentials/sources -> policy smoke pass. + sys.path.insert(0, "/opt") + try: + import flash_gpu # type: ignore + except Exception: + metrics["full_benchmark"] = 0 + return 1.0, 1.0, "patch policy smoke passed; GPU harness not available in this environment", metrics + if not flash_gpu.modal_available() or not CLEAN_SOURCE.exists() or not BASELINE_SOURCE.exists(): + metrics["full_benchmark"] = 0 + return 1.0, 1.0, "patch policy smoke passed; Modal/judge sources not configured in this environment", metrics + + with tempfile.TemporaryDirectory(prefix="flash_kernel_opt_") as tmp: + tmp_root = Path(tmp) + patched = tmp_root / "patched" + shutil.copytree(CLEAN_SOURCE, patched) + env = {"PATH": os.environ.get("PATH", "/usr/bin:/bin"), "HOME": str(tmp_root), + "LC_ALL": "C", "LANG": "C"} + if int(metrics.get("changed_files", 0)) > 0: + try: + subprocess.run(["git", "apply", "--check", str(patch_path)], cwd=patched, env=env, + check=True, capture_output=True, text=True, timeout=60) + subprocess.run(["git", "apply", str(patch_path)], cwd=patched, env=env, + check=True, capture_output=True, text=True, timeout=60) + except subprocess.CalledProcessError as exc: + metrics["stderr_tail"] = sanitize(exc.stderr or "") + return _invalid("patch does not apply to the clean package tree", metrics) + metrics["applied_patch"] = 1 + else: + metrics["used_empty_patch"] = 1 + + payload = { + "baseline_files": _read_dir(BASELINE_SOURCE, BASELINE_SOURCE), + "patched_files": _read_dir(patched, patched), + "workloads": workloads, + "cfg": _build_cfg(), + } + try: + result = flash_gpu.run_remote(payload) + except Exception as exc: + metrics["error_detail"] = sanitize(str(exc)) + return _invalid("GPU evaluation failed", metrics) + + if not result.get("ok"): + metrics["worker_error"] = sanitize(str(result.get("error", "unknown"))) + return _invalid("GPU benchmark worker failed", metrics) + + speedups: list[float] = [] + per_workload: dict[str, Any] = {} + for row in result.get("rows", []): + if not row.get("ok"): + metrics["failed_workload"] = {"reason": row.get("reason", "error")} + return _invalid("submission failed the quality gate or crashed on a hidden workload", metrics) + speedups.append(max(float(row["speedup"]), 0.01)) + per_workload[row["id"]] = {"speedup": row["speedup"]} + if not speedups: + return _invalid("no workloads were evaluated", metrics) + + gm = geometric_mean(speedups) + bounded = score_from_speedup(gm) + metrics.update({"full_benchmark": 1, "workload_count": len(speedups), "geomean_speedup": gm}) + if _get_bool("expose_per_workload_metrics", False): + metrics["per_workload"] = per_workload + return bounded, bounded, f"{PRIMITIVE} geomean speedup {gm:.3f}x over the naive baseline", metrics + + +def evaluate(solution_path: str) -> tuple[float, float, str, dict[str, Any]]: + patch_path = Path(solution_path) + ok, message, metrics = validate_patch(patch_path) + if not ok: + return _invalid(message, metrics) + try: + return full_evaluation(patch_path, metrics) + except Exception as exc: # noqa: BLE001 + metrics["error_type"] = type(exc).__name__ + metrics["error_detail"] = sanitize(str(exc)) + return _invalid("evaluation failed", metrics) + + +def main(argv: list[str]) -> int: + if len(argv) != 2: + print("Usage: evaluator.py SOLUTION_PATCH", file=sys.stderr) + return 2 + score, unbounded, message, metrics = evaluate(argv[1]) + print(json.dumps({"score": score, "score_unbounded": unbounded, + "message": message, "metrics": metrics}, indent=2, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv)) diff --git a/2.0/problems/knn_gpu_kernel_optimization/flash_gpu.py b/2.0/problems/knn_gpu_kernel_optimization/flash_gpu.py new file mode 100644 index 000000000..4edd879ec --- /dev/null +++ b/2.0/problems/knn_gpu_kernel_optimization/flash_gpu.py @@ -0,0 +1,427 @@ +"""Modal GPU evaluation harness for the flashlib kernel-optimization task family. + +Shared by the judge (evaluator.py) and the agent public test. Baked into both +the agent and judge images at ``/opt/flash_gpu.py``. It runs a frozen naive +baseline and the (patched or agent-edited) package on a Modal GPU and returns +per-workload timings + quality verdicts. + +Anti-reward-hack properties (all enforced inside the remote worker, which is the +only place the untrusted code runs): + +* **Timing primitives are captured before any agent code is imported** — the + worker binds ``time.perf_counter`` and ``torch.cuda.synchronize`` to locals up + front, so a submission that monkey-patches ``torch.cuda.synchronize`` (or any + torch attribute) cannot affect measurement. +* **Fresh data every timed iteration** — each measured iteration generates a new + input from a fresh seed, so a submission cannot memoize on a repeated input. +* **Quality is verified on every timed iteration** — the judge recomputes the + quality metric from the *returned tensors* on that iteration's data and gates + it, so returning a stale cached result for a different input is caught. +* **The judge computes all metrics** — no agent-provided number is trusted. +* **Fresh container per submission** — an ephemeral Modal app is used, so no + global state survives across evaluations. + +The remote worker is a single self-contained function serialized to Modal +(``serialized=True``), so the remote does not import this module and there is no +module-mounting to get wrong. The GPU image ships torch + triton (+ any extra +pip packages a task needs for its reference build). +""" +from __future__ import annotations + +import os + + +# --------------------------------------------------------------------------- # +# Remote worker (runs on the Modal GPU). Fully self-contained: every helper is +# nested so cloudpickle ships the whole thing. Takes and returns plain data. +# --------------------------------------------------------------------------- # +def _gpu_worker(payload: dict) -> dict: + import importlib + import os as _os + import sys as _sys + import tempfile + import time + import traceback + + import torch + + # Capture timing + sync references BEFORE importing any submission code, so a + # monkey-patch of torch.cuda.synchronize cannot affect measurement. + _perf = time.perf_counter + _sync = torch.cuda.synchronize + if not torch.cuda.is_available(): + return {"ok": False, "error": "cuda_unavailable"} + dev = "cuda" + + cfg = payload["cfg"] + prim = cfg["primitive"] + warmup = int(cfg.get("warmup", 3)) + iters = int(cfg.get("iters", 7)) + + def _materialize(files: dict, tag: str) -> str: + root = tempfile.mkdtemp(prefix=f"flash_{tag}_") + for rel, content in files.items(): + path = _os.path.join(root, rel) + parent = _os.path.dirname(path) + if parent: + _os.makedirs(parent, exist_ok=True) + with open(path, "w", encoding="utf-8") as fh: + fh.write(content) + _sys.path.insert(0, root) + return root + + _materialize(payload["baseline_files"], "baseline") + _materialize(payload["patched_files"], "patched") + ref = importlib.import_module(cfg["ref_module"]) + pkg = importlib.import_module(cfg["pkg"]) + + # ---- per-primitive: fixed context, data gen, call, quality metric ----- # + _ann_cache = {} + def _load_ann(name): + # Load a real ANN dataset (ann-benchmarks HDF5) from the mounted Modal + # Volume at /data. Cached across workloads in one worker call. + if name not in _ann_cache: + import h5py + import numpy as _np + with h5py.File(f"/data/{name}.hdf5", "r") as f: + base = torch.from_numpy(_np.asarray(f["train"], dtype=_np.float32)).to(dev) + queries = torch.from_numpy(_np.asarray(f["test"], dtype=_np.float32)).to(dev) + _ann_cache[name] = (base.contiguous(), queries.contiguous()) + return _ann_cache[name] + + def setup(w, seed): + # Built ONCE per workload (not per timed iter). IVF-PQ builds the fixed + # index (the "database") here; queries are generated fresh per iteration. + if prim == "ivfpq": + g = torch.Generator(device=dev).manual_seed(int(seed)) + X = torch.randn(int(w["M"]), int(w["D"]), generator=g, device=dev, dtype=torch.float32) + index = ref.ivf_pq_build(X, int(w["nlist"]), m=int(w["m"]), + nprobe=int(w["nprobe"]), seed=int(seed)) + return {"index": index, "D": int(w["D"])} + if prim == "knn" and w.get("dataset"): + # Real dataset: the fixed database is the full real base; queries are + # drawn per-iteration from the held-out query set. + base, queries_all = _load_ann(w["dataset"]) + return {"base": base, "queries_all": queries_all} + return {} + + def gen(w, seed, ctx): + g = torch.Generator(device=dev).manual_seed(int(seed)) + if prim == "ivfpq": + q = torch.randn(int(w["Q"]), int(ctx["D"]), generator=g, device=dev, dtype=torch.float32) + return {"queries": q} + if prim == "knn": + if w.get("dataset"): + # Fresh random subset of the real query set AND of the real base + # every iteration. The database must differ per call: handing back + # one fixed tensor lets a solution cache a prebuilt index/norms + # keyed on it (data_ptr or content) and skip the search work the + # timer is supposed to measure -- exactly what the synthetic path + # avoids by regenerating the database each iteration. + # PRECISION LOCK (bf16): hand out bf16 queries + database, so the + # inputs carry only bf16 precision. No solution can gain from a + # different distance-matmul dtype -- fp32/TF32 buy nothing over bf16 + # tensor cores on bf16 inputs (only slower), and fp8 loses too much + # to clear the recall gate. Baseline and agent compute on the same + # bf16 data, so speed is purely about the kernel, not precision. + qa, base = ctx["queries_all"], ctx["base"] + qperm = torch.randperm(qa.shape[0], generator=g, device=dev)[: int(w["Q"])] + dperm = torch.randperm(base.shape[0], generator=g, device=dev)[: int(w["M"])] + return {"queries": qa.index_select(0, qperm).to(torch.bfloat16).contiguous(), + "database": base.index_select(0, dperm).to(torch.bfloat16).contiguous()} + db = torch.randn(w["M"], w["D"], generator=g, device=dev, dtype=torch.float32) + q = torch.randn(w["Q"], w["D"], generator=g, device=dev, dtype=torch.float32) + return {"queries": q, "database": db} + if prim == "dbscan": + nc, D, N = int(w["n_centers"]), int(w["D"]), int(w["N"]) + centers = torch.randn(nc, D, generator=g, device=dev, dtype=torch.float32) * 10.0 + asg = torch.randint(0, nc, (N,), generator=g, device=dev) + xb = centers.index_select(0, asg) + torch.randn(N, D, generator=g, device=dev, dtype=torch.float32) * 0.5 + nz = int(float(w.get("noise_frac", 0.0)) * N) + if nz > 0: + lo = xb.min(0).values; hi = xb.max(0).values + xb[:nz] = lo + (hi - lo) * torch.rand(nz, D, generator=g, device=dev) + return {"x": xb, "eps": float(w["eps"]), "min_samples": int(w["min_samples"])} + x = torch.randn(w["N"], w["D"], generator=g, device=dev, dtype=torch.float32) + if prim == "kmeans": + perm = torch.randperm(w["N"], generator=g, device=dev)[: w["K"]] + return {"x": x, "init": x.index_select(0, perm).clone()} + return {"x": x} + + def call(mod, w, data, ctx): + if prim == "ivfpq": + return mod.ivf_pq_search(ctx["index"], data["queries"], int(w["k"]), nprobe=int(w["nprobe"])) + if prim == "kmeans": + return mod.kmeans(data["x"], w["K"], max_iters=w["max_iters"], + init_centroids=data["init"], tol=0.0) + if prim == "knn": + return mod.knn(data["queries"], data["database"], w["k"]) + if prim == "dbscan": + return mod.dbscan(data["x"], data["eps"], data["min_samples"]) + if prim == "pca": + return mod.pca(data["x"], w["k"]) + if prim == "tsvd": + return mod.truncated_svd(data["x"], w["k"]) + raise ValueError(prim) + + def check_shape(w, out): + if prim == "kmeans": + _, cen, _ = out + if tuple(cen.shape) != (w["K"], w["D"]) or not torch.isfinite(cen).all(): + raise ValueError("bad kmeans output") + elif prim == "knn": + d, i = out + if tuple(d.shape) != (w["Q"], w["k"]) or tuple(i.shape) != (w["Q"], w["k"]): + raise ValueError("bad knn output") + if not torch.isfinite(d).all(): + raise ValueError("non-finite distances") + elif prim == "ivfpq": + vals, ids = out + if tuple(ids.shape) != (int(w["Q"]), int(w["k"])) or \ + tuple(vals.shape) != (int(w["Q"]), int(w["k"])): + raise ValueError("bad ivfpq output shape") + elif prim == "dbscan": + if out.ndim != 1 or int(out.shape[0]) != int(w["N"]): + raise ValueError("bad dbscan labels shape") + else: + a, b = out + comps = a if prim == "pca" else b + if tuple(comps.shape) != (w["k"], w["D"]) or not torch.isfinite(comps).all(): + raise ValueError("bad decomposition output") + + def inertia(x, centroids): + c = centroids.to(torch.float32) + cn = (c * c).sum(1) + total = 0.0 + for j in range(0, x.shape[0], 16384): + xb = x[j:j + 16384] + d = (xb * xb).sum(1, keepdim=True) - 2.0 * (xb @ c.t()) + cn[None, :] + total += float(d.min(1).values.clamp_min(0).sum().item()) + return total + + def recall(agent_idx, ref_idx): + a = agent_idx.long(); b = ref_idx.long() + hit = (a.unsqueeze(2) == b.unsqueeze(1)).any(2) + return float(hit.sum().item()) / (b.shape[0] * b.shape[1]) + + def knn_ball_recall(queries, database, agent_idx, ref_idx, tol): + # Tie-robust recall: a returned neighbour counts if its EXACT (fp32) + # squared-L2 distance is within the true k-NN ball, i.e. <= the baseline's + # k-th distance x (1+tol). Plain index-recall is unusable on bf16 data -- + # coarse distances tie heavily, so two *correct* searches that break ties + # differently disagree at ~0.9-0.97. This gate ignores which tie member is + # picked (any point at the k-th distance is correct) but still rejects a + # genuinely-farther point, so a lower-precision search that returns wrong + # neighbours fails while an honest bf16 kernel passes exactly. + # + # DISTINCT-index counting (anti-hack): a returned index counts at most ONCE + # per query. Without this, returning the single nearest neighbour k times + # scores 1.0 (the 1-NN is trivially inside every k-NN ball) -- i.e. a search + # that finds only argmin and pads with duplicates fakes full recall. Counting + # distinct in-ball indices makes that hack score ~1/k; an honest top-k returns + # k distinct in-ball ids and still scores ~1.0. + qf = queries.to(torch.float32); dbf = database.to(torch.float32) + M = dbf.shape[0] + # k-NN ball radius per query = exact distance to the baseline's k-th nbr + kth = dbf.index_select(0, ref_idx[:, -1].long()) # (Q, D) + d_k = ((qf - kth) ** 2).sum(1) # (Q,) + thr = d_k * (1.0 + tol) + 1e-6 + a = agent_idx.long(); Q, k = a.shape + sent = M + torch.arange(k, device=a.device) # distinct sentinels >= M + hits = 0 + for j in range(0, Q, 256): + qb = qf[j:j + 256]; ab = a[j:j + 256] # (q, D), (q, k) + cand = dbf.index_select(0, ab.reshape(-1)).reshape(ab.shape[0], k, -1) + ad = ((qb[:, None, :] - cand) ** 2).sum(2) # (q, k) exact dists + inball = ad <= thr[j:j + 256, None] # (q, k) bool + # out-of-ball slots -> per-slot distinct sentinels (>= M) so they never + # collide with a real id and are dropped by the (< M) filter below. + idx_ib = torch.where(inball, ab, sent[None, :]) # (q, k) + srt, _ = idx_ib.sort(dim=1) + first = torch.ones_like(srt, dtype=torch.bool) + first[:, 1:] = srt[:, 1:] != srt[:, :-1] # first occurrence of each value + hits += int((first & (srt < M)).sum().item()) # distinct real in-ball ids + return hits / (Q * k) + + def ari(a, b): + # Adjusted Rand Index between two integer label vectors (N,); noise (-1) + # is treated as its own label (dense-remapped first). + a = torch.unique(a.long(), return_inverse=True)[1] + b = torch.unique(b.long(), return_inverse=True)[1] + na = int(a.max().item()) + 1; nb = int(b.max().item()) + 1 + cont = torch.bincount(a * nb + b, minlength=na * nb).reshape(na, nb).double() + ai = cont.sum(1); bj = cont.sum(0); n = cont.sum() + c2 = lambda z: z * (z - 1) / 2.0 + sij = c2(cont).sum(); sa = c2(ai).sum(); sb = c2(bj).sum() + exp = (sa * sb / c2(n)) if float(n) > 1 else 0.0 + den = 0.5 * (sa + sb) - exp + return float(((sij - exp) / den).item()) if float(den) != 0.0 else 1.0 + + def ortho_err(comps): + c = comps.to(torch.float32); k = c.shape[0] + return float((c @ c.t() - torch.eye(k, device=c.device)).abs().max().item()) + + def captured(x, comps, center): + c = comps.to(torch.float32) + mean = x.mean(0) if center else None + total = 0.0 + for j in range(0, x.shape[0], 16384): + xb = x[j:j + 16384] + if center: + xb = xb - mean + p = xb @ c.t() + total += float((p * p).sum().item()) + return total / (x.shape[0] - 1) if center else total + + def verdict(w, data, ctx, ref_out, agent_out): + """Return (ok, reason, ref_val, agent_val) gating the agent output.""" + check_shape(w, agent_out) + if prim == "ivfpq": + # iso-result recall: agent ids vs frozen-baseline ids on the SAME index. + rec = recall(agent_out[1], ref_out[1]) + ok = rec >= float(cfg["recall_threshold"]) + return ok, ("recall_regression" if not ok else ""), 1.0, rec + if prim == "kmeans": + rv = inertia(data["x"], ref_out[1]); av = inertia(data["x"], agent_out[1]) + tol = float(cfg["inertia_tolerance"]) + ok = av <= (1.0 + tol) * rv + 1e-6 + return ok, ("inertia_regression" if not ok else ""), rv, av + if prim == "knn": + rec = knn_ball_recall(data["queries"], data["database"], + agent_out[1], ref_out[1], + float(cfg.get("distance_ball_tolerance", 0.0))) + ok = rec >= float(cfg["recall_threshold"]) + return ok, ("recall_regression" if not ok else ""), 1.0, rec + if prim == "dbscan": + val = ari(agent_out, ref_out) + ok = val >= float(cfg["ari_threshold"]) + return ok, ("ari_regression" if not ok else ""), 1.0, val + center = (prim == "pca") + comps = agent_out[0] if prim == "pca" else agent_out[1] + ref_comps = ref_out[0] if prim == "pca" else ref_out[1] + oe = ortho_err(comps) + if oe > float(cfg["ortho_tolerance"]): + return False, "not_orthonormal", 0.0, oe + rv = captured(data["x"], ref_comps, center) + av = captured(data["x"], comps, center) + ok = av >= (1.0 - float(cfg["captured_tolerance"])) * rv - 1e-6 + return ok, ("captured_regression" if not ok else ""), rv, av + + def time_call(mod, w, data, ctx): + _sync(); t0 = _perf(); out = call(mod, w, data, ctx); _sync() + return (_perf() - t0) * 1000.0, out + + rows = [] + try: + for w in payload["workloads"]: + base_seed = int(w["seed"]) + ctx = setup(w, base_seed) # fixed context (e.g. IVF-PQ index) + # warmup on fresh data (kernels / autotune) — not measured + for i in range(warmup): + d = gen(w, base_seed + 100 + i, ctx) + call(ref, w, d, ctx); call(pkg, w, d, ctx); _sync() + del d + torch.cuda.empty_cache() + ratios, ref_val, agent_val = [], None, None + bad = None + for i in range(iters): + d = gen(w, base_seed + 10000 + i, ctx) # fresh every iteration + rt, rout = time_call(ref, w, d, ctx) + at, aout = time_call(pkg, w, d, ctx) + ok, reason, rv, av = verdict(w, d, ctx, rout, aout) # verify THIS iter + if not ok: + bad = reason; ref_val, agent_val = rv, av + break + ratios.append(rt / at if at > 0 else 0.01) + ref_val, agent_val = rv, av + del d, rout, aout + torch.cuda.empty_cache() + if bad is not None: + rows.append({"id": w["id"], "ok": False, "reason": bad, + "ref_val": ref_val, "agent_val": agent_val}) + continue + ratios.sort() + rows.append({"id": w["id"], "ok": True, + "speedup": ratios[len(ratios) // 2], + "ref_val": ref_val, "agent_val": agent_val}) + except Exception as exc: + return {"ok": False, "error": f"{type(exc).__name__}: {str(exc)[:200]}", + "trace": traceback.format_exc()[-500:]} + return {"ok": True, "rows": rows} + + +# --------------------------------------------------------------------------- # +# Host-side wrappers (run in the judge / agent container; need MODAL tokens). +# --------------------------------------------------------------------------- # +def _build_image(cfg: dict): + import modal + pip = list(cfg.get("pip", ["torch==2.5.1", "triton==3.1.0", "numpy"])) + base = cfg.get("cuda_image", "nvidia/cuda:12.4.1-devel-ubuntu22.04") + # add_python must match the judge/agent container Python (ubuntu:24.04 -> 3.12) + # because the worker ships via serialized=True (cloudpickle is version-sensitive). + # add_local_python_source mounts THIS module into the image so the remote can + # import it when deserializing the (module-level) worker function. + return (modal.Image.from_registry(base, add_python=str(cfg.get("python", "3.12"))) + .entrypoint([]) + .pip_install(*pip) + .add_local_python_source("flash_gpu")) + + +# Substrings that mark a transient Modal control-plane / image-build failure +# (evicted build, gateway hiccup, app stopped mid-run) rather than a real error +# in the submission — safe to retry. +_TRANSIENT_MARKERS = ( + "external shut-down", "terminated due to external", "please try again", + "app_state_stopped", "conflicterror", "deadline exceeded", "connection reset", + "502 bad gateway", "503 service", "temporarily unavailable", "timed out", + "gateway", "eviction", "internalfailure", "internalerror", + "failed to get new inputs", "runner failed", "task exited", +) + + +def _is_transient(text: str) -> bool: + low = (text or "").lower() + return any(m in low for m in _TRANSIENT_MARKERS) + + +def run_remote(payload: dict) -> dict: + """Run the GPU worker on Modal and return its result dict. + + Retries transient Modal control-plane failures; a real error in the + submission is non-transient and surfaces immediately. Raises RuntimeError + with a sanitized message on unrecoverable failure. + """ + import time as _time + + import modal + + cfg = payload["cfg"] + attempts = max(1, int(cfg.get("modal_retries", 3))) + last = "modal_gpu_eval_failed" + for attempt in range(1, attempts + 1): + app = modal.App(cfg.get("app_name", "flash-kernel-eval")) + fn_kwargs = dict( + gpu=cfg.get("gpu", "H100"), + image=_build_image(cfg), + timeout=int(cfg.get("modal_timeout_seconds", 1800)), + serialized=True, + max_containers=1, + ) + # Mount the cached real-dataset Volume (SIFT/GIST HDF5) read-only at /data. + vol_name = cfg.get("dataset_volume") + if vol_name: + fn_kwargs["volumes"] = {"/data": modal.Volume.from_name(vol_name)} + remote = app.function(**fn_kwargs)(_gpu_worker) + try: + with modal.enable_output(), app.run(): + return remote.remote(payload) + except Exception as exc: # noqa: BLE001 + last = str(exc)[-400:] + if not _is_transient(last) or attempt == attempts: + raise RuntimeError(f"modal_gpu_eval_failed: {last}") from exc + _time.sleep(min(45, 10 * attempt)) + raise RuntimeError(f"modal_gpu_eval_failed: {last}") + + +def modal_available() -> bool: + return bool(os.environ.get("MODAL_TOKEN_ID") and os.environ.get("MODAL_TOKEN_SECRET")) diff --git a/2.0/problems/knn_gpu_kernel_optimization/harbor/app/README.md b/2.0/problems/knn_gpu_kernel_optimization/harbor/app/README.md new file mode 100644 index 000000000..c044f496a --- /dev/null +++ b/2.0/problems/knn_gpu_kernel_optimization/harbor/app/README.md @@ -0,0 +1,34 @@ +# Brute-force k-NN kernel optimization — submission workflow + +You are optimizing the `knnlib` package at `/app/knnlib`. Edit the package +(rewrite the internals of `knn`, add Triton kernel modules under `knnlib/`), +then submit a patch. + +## Iterate locally + +Run the public self-test to check correctness and get a rough speed signal on +the two public shapes (needs a GPU in the agent container): + +```bash +bash /app/public_test.sh +``` + +## Submit + +```bash +bash /app/make_submission.sh # writes /app/solution.patch (knnlib diff) +bash /app/submit.sh # enqueues it for the black-box judge +``` + +Submissions are asynchronous. Submit early and keep improving; use +`bash /app/submissions.sh` and `bash /app/wait_submission.sh ` to inspect +results. + +## Rules + +- Only files under `knnlib/` may change. +- Do not import external optimized libraries (write the kernels yourself), and + do not access the environment, spawn processes, or use the network. +- Keep the public `knn(...)` signature and return contract unchanged. +- Nearest-neighbor quality is gated (recall@k vs an exact baseline); do not + sacrifice correctness for speed beyond the allowed tolerance. diff --git a/2.0/problems/knn_gpu_kernel_optimization/harbor/app/make_submission.sh b/2.0/problems/knn_gpu_kernel_optimization/harbor/app/make_submission.sh new file mode 100755 index 000000000..b54576f18 --- /dev/null +++ b/2.0/problems/knn_gpu_kernel_optimization/harbor/app/make_submission.sh @@ -0,0 +1,24 @@ +#!/usr/bin/env bash +# Package the current knnlib edits into a unified diff for submission. +set -euo pipefail + +APP_DIR="${APP_DIR:-/app}" +OUT="${1:-/app/solution.patch}" + +if [[ ! -d "$APP_DIR/knnlib" ]]; then + echo "knnlib package not found at $APP_DIR/knnlib" >&2 + exit 2 +fi +if [[ ! -d "$APP_DIR/.git" ]]; then + echo "git repo not found at $APP_DIR (expected in the agent image)" >&2 + exit 2 +fi + +# Stage and diff only the package, so submission helper scripts under /app are +# never included. New kernel files under knnlib/ are captured via `git add`. +git -C "$APP_DIR" add -A -- knnlib +git -C "$APP_DIR" diff --cached -- knnlib > "$OUT" +git -C "$APP_DIR" reset -q + +bytes=$(wc -c < "$OUT" | tr -d ' ') +echo "Wrote $OUT ($bytes bytes)" diff --git a/2.0/problems/knn_gpu_kernel_optimization/harbor/app/public_test.py b/2.0/problems/knn_gpu_kernel_optimization/harbor/app/public_test.py new file mode 100644 index 000000000..e8de4a4bd --- /dev/null +++ b/2.0/problems/knn_gpu_kernel_optimization/harbor/app/public_test.py @@ -0,0 +1,171 @@ +"""Public GPU self-test -- IDENTICAL to the final graded evaluation. + +This runs the EXACT same workloads, thresholds, timing, and seeds that the hidden +judge uses to grade your submission (all read from ``/app/task_config.json``, +the same config the judge reads), on a Modal GPU, and reports the same +per-workload pass/fail + speedup + geomean + predicted score (0-100) you would +receive on submission. There is no longer a separate, smaller "public" workload +set -- what you see here is what you get graded on. + +Requires MODAL_TOKEN_ID / MODAL_TOKEN_SECRET in the environment. +""" +from __future__ import annotations + +import json +import math +import os +import sys +from pathlib import Path + +sys.path.insert(0, "/opt") # flash_gpu.py is baked at /opt + +APP_DIR = os.environ.get("APP_DIR", "/app") +TASK_CONFIG_PATH = os.environ.get("TASK_CONFIG_PATH", "/app/task_config.json") + + +def _load_eval() -> dict: + """The judge grades from task_config.json's `evaluation` block; read the same.""" + try: + doc = json.loads(Path(TASK_CONFIG_PATH).read_text(encoding="utf-8")) + except Exception as exc: # noqa: BLE001 + print(f"could not read {TASK_CONFIG_PATH}: {exc}") + return {} + return doc.get("evaluation", doc) + + +EV = _load_eval() + + +def g(key, default): + v = EV.get(key, default) + return default if v is None else v + + +PRIMITIVE = str(g("primitive", "")) +PKG = str(g("pkg", "")) +REF_MODULE = str(g("ref_module", "")) +SPEEDUP_TARGET = float(g("speedup_target", 8.0)) +BASELINE_DIR = str(g("baseline_source", "/opt/flash_ref")) + + +def _workloads() -> list: + """Exactly the judge's final-role workload set: ALL workloads, same seed + derivation ``base_seed + 1000*(i+1)`` -- identical data to the graded run.""" + wls = [dict(w) for w in g("workloads", [])] + base = int(g("base_seed", 20260701)) + for i, w in enumerate(wls): + w.setdefault("seed", base + 1000 * (i + 1)) + return wls + + +def _cfg() -> dict: + """Byte-for-byte the judge's _build_cfg().""" + return { + "primitive": PRIMITIVE, + "pkg": PKG, + "ref_module": REF_MODULE, + "gpu": str(g("gpu", "H100")), + "cuda_image": str(g("cuda_image", "nvidia/cuda:12.4.1-devel-ubuntu22.04")), + "pip": list(g("pip", ["torch", "numpy"])), + "app_name": str(g("app_name", "flash-kernel-public")), + "modal_timeout_seconds": int(g("modal_timeout_seconds", 1800)), + "warmup": int(g("warmup_iters", 3)), + "iters": int(g("timed_iters", 7)), + "inertia_tolerance": float(g("inertia_tolerance", 0.02)), + "recall_threshold": float(g("recall_threshold", 0.99)), + "captured_tolerance": float(g("captured_tolerance", 0.02)), + "ortho_tolerance": float(g("ortho_tolerance", 0.02)), + "ari_threshold": float(g("ari_threshold", 0.99)), + } + + +def geometric_mean(values: list) -> float: + if not values: + return 0.0 + return math.exp(sum(math.log(max(v, 1e-9)) for v in values) / len(values)) + + +def score_from_speedup(gm: float) -> float: + if gm <= 0: + return 0.0 + raw = 100.0 * math.log(gm) / math.log(max(SPEEDUP_TARGET, 1.0000001)) + return max(0.0, min(100.0, raw)) + + +def _read(root: str, sub: str = "") -> dict: + base = Path(root) + scan = base / sub if sub else base + return {str(p.relative_to(base)): p.read_text(encoding="utf-8", errors="replace") + for p in scan.rglob("*.py")} + + +def main() -> int: + import flash_gpu # baked at /opt/flash_gpu.py + if not flash_gpu.modal_available(): + print("Set MODAL_TOKEN_ID and MODAL_TOKEN_SECRET to run the GPU self-test.") + return 1 + workloads = _workloads() + cfg = _cfg() + if not workloads: + print("no workloads found in task_config.json; cannot run.") + return 1 + payload = { + "baseline_files": _read(BASELINE_DIR), + "patched_files": _read(APP_DIR, PKG), + "workloads": workloads, + "cfg": cfg, + } + try: + result = flash_gpu.run_remote(payload) + except Exception as exc: # noqa: BLE001 + print(f"GPU run failed: {exc}") + return 1 + if not result.get("ok"): + print(f"worker error: {result.get('error')}") + return 1 + + if PRIMITIVE in ("knn", "ivfpq"): + mlabel, gate = "recall@k", f">= {cfg['recall_threshold']}" + elif PRIMITIVE == "dbscan": + mlabel, gate = "ARI", f">= {cfg['ari_threshold']}" + elif PRIMITIVE == "kmeans": + mlabel, gate = "inertia", f"<= (1+{cfg['inertia_tolerance']}) x ref (lower is better)" + else: + mlabel, gate = "quality", "" + + print("=== FINAL-EQUIVALENT self-test: identical workloads / thresholds / seeds / " + "timing to the graded judge ===") + print(f"correctness gate: {mlabel} {gate} | speedup_target = {SPEEDUP_TARGET:g}") + print(f"{'workload':9s} {'status':22s} {'speedup':>9s} {mlabel}: agent / ref") + rows = result.get("rows", []) + speedups = [] + any_fail = False + for row in rows: + av, rv = row.get("agent_val"), row.get("ref_val") + q = (f"{av:.4f} / {rv:.4f}" + if isinstance(av, (int, float)) and isinstance(rv, (int, float)) else "n/a") + if row.get("ok"): + sp = f"{row['speedup']:.2f}x" + speedups.append(max(float(row["speedup"]), 0.01)) # same clamp as the judge + st = "OK" + else: + sp = "-" + st = "FAIL:" + str(row.get("reason", "")) + any_fail = True + print(f"{row['id']:9s} {st:22s} {sp:>9s} {q}") + + # Scoring is byte-for-byte the judge's full_evaluation: ANY gate failure -> 0. + print() + if any_fail or len(speedups) != len(rows) or not speedups: + print("RESULT: at least one workload FAILED the correctness gate -> a submission now " + "would be INVALID and score 0/100. Fix the failing workload(s) before submitting.") + return 0 + gm = geometric_mean(speedups) + score = score_from_speedup(gm) + print(f"geomean speedup = {gm:.3f}x over the naive baseline") + print(f"PREDICTED FINAL SCORE = {score:.2f} / 100 (100 == {SPEEDUP_TARGET:g}x geomean)") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/2.0/problems/knn_gpu_kernel_optimization/harbor/app/public_test.sh b/2.0/problems/knn_gpu_kernel_optimization/harbor/app/public_test.sh new file mode 100755 index 000000000..1157af76f --- /dev/null +++ b/2.0/problems/knn_gpu_kernel_optimization/harbor/app/public_test.sh @@ -0,0 +1,7 @@ +#!/usr/bin/env bash +# Public self-test: run the patched knnlib on the public shapes and print a +# correctness + rough speedup summary. Requires a GPU in the agent container. +set -euo pipefail +APP_DIR="${APP_DIR:-/app}" +cd "$APP_DIR" +exec python3 "$APP_DIR/public_test.py" diff --git a/2.0/problems/knn_gpu_kernel_optimization/harbor/app/solution.patch b/2.0/problems/knn_gpu_kernel_optimization/harbor/app/solution.patch new file mode 100644 index 000000000..e69de29bb diff --git a/2.0/problems/knn_gpu_kernel_optimization/judge/refknn.py b/2.0/problems/knn_gpu_kernel_optimization/judge/refknn.py new file mode 100644 index 000000000..0eb1b392e --- /dev/null +++ b/2.0/problems/knn_gpu_kernel_optimization/judge/refknn.py @@ -0,0 +1,18 @@ +"""Frozen naive brute-force k-NN baseline used by the judge. + +This is a standalone (non-package) copy of the ``knnlib.knn`` implementation the +agent starts from. It is imported under its own module name so the judge worker +can load the frozen baseline and the patched ``knnlib`` package in the same +process. The judge uses it both as the speed denominator and as the exact +oracle from which the true k nearest neighbours (for the recall@k gate) are +recomputed. Keep this behaviourally identical to the shipped ``knnlib/knn.py``. +""" +from __future__ import annotations + +import torch + + +def knn(queries, database, k): + d2 = torch.cdist(queries, database) ** 2 # (Q, M), materialized -- naive + dist, idx = torch.topk(d2, k, dim=1, largest=False) + return dist, idx.to(torch.long) diff --git a/2.0/problems/knn_gpu_kernel_optimization/knnlib/__init__.py b/2.0/problems/knn_gpu_kernel_optimization/knnlib/__init__.py new file mode 100644 index 000000000..399c81247 --- /dev/null +++ b/2.0/problems/knn_gpu_kernel_optimization/knnlib/__init__.py @@ -0,0 +1,16 @@ +"""knnlib -- a tiny GPU brute-force k-NN library you are asked to make fast. + +The public entry point is :func:`knn`. The shipped implementation is correct +but deliberately unoptimised: it materialises the full ``(Q, M)`` pairwise +distance matrix before selecting the ``k`` nearest neighbours. Your task is to +rewrite the internals (new Triton kernels, fused/streamed top-k passes, better +memory traffic, ...) so that :func:`knn` runs as fast as possible while +returning the same nearest neighbours. The public function signature and return +contract must not change. +""" +from __future__ import annotations + +from knnlib.knn import knn + +__all__ = ["knn"] +__version__ = "0.1.0" diff --git a/2.0/problems/knn_gpu_kernel_optimization/knnlib/knn.py b/2.0/problems/knn_gpu_kernel_optimization/knnlib/knn.py new file mode 100644 index 000000000..ad0c7b6c9 --- /dev/null +++ b/2.0/problems/knn_gpu_kernel_optimization/knnlib/knn.py @@ -0,0 +1,39 @@ +"""Brute-force squared-L2 k-nearest-neighbours -- the reference you must optimise. + +This implementation is intentionally naive. It computes the full ``(Q, M)`` +pairwise distance matrix with :func:`torch.cdist` and materialises it in HBM, +then runs :func:`torch.topk` over the whole matrix to pick the ``k`` nearest +database points for every query. It is correct and deterministic, but it moves +far more memory than necessary (the entire ``(Q, M)`` matrix) and cannot scale +to large ``M`` without exhausting device memory. + +Contract (do NOT change): + + knn(queries, database, k) -> (distances, indices) + + queries : (Q, D) float32 CUDA tensor of query points. + database : (M, D) float32 CUDA tensor of database points to search. + k : int, number of nearest neighbours to return per query. + + distances : (Q, k) float32 tensor of the SQUARED-L2 distances to the ``k`` + nearest database points, in ascending order (nearest first). + indices : (Q, k) int64 tensor of the corresponding database row indices. + +You may add modules/kernels inside the ``knnlib`` package and rewrite the body +of :func:`knn` freely (fused/streamed running top-k, tiled distance kernels, +Triton), as long as the public contract above is preserved. In particular you +should avoid ever materialising the full ``(Q, M)`` distance matrix. +""" +from __future__ import annotations + +import torch + + +def knn(queries, database, k): + """Return the (squared-L2) k nearest database points for each query. + + Naive: materialise the full (Q, M) distance matrix, then top-k. + """ + d2 = torch.cdist(queries, database) ** 2 # (Q, M), materialized -- naive + dist, idx = torch.topk(d2, k, dim=1, largest=False) + return dist, idx.to(torch.long) diff --git a/2.0/problems/knn_gpu_kernel_optimization/readme b/2.0/problems/knn_gpu_kernel_optimization/readme new file mode 100644 index 000000000..3ecba5ff3 --- /dev/null +++ b/2.0/problems/knn_gpu_kernel_optimization/readme @@ -0,0 +1,131 @@ +# GPU Brute-Force k-NN Kernel Optimization + +## Problem + +You are given a small GPU brute-force k-nearest-neighbor library, `knnlib`, in +the Harbor workspace at `/app/knnlib`. Its public entry point is: + +```python +knnlib.knn(queries, database, k) -> (distances, indices) +``` + +`queries` is a `(Q, D)` **bfloat16** CUDA tensor of query points, `database` is an +`(M, D)` **bfloat16** CUDA tensor of database points to search, and `k` is the +number of nearest neighbors to return per query. The inputs are bfloat16 by +design — treat it as the fixed working precision; upcasting to fp32/TF32 buys no +accuracy over bf16 tensor cores here (only latency), so precision is not a tuning +knob. The function returns `distances`, a `(Q, k)` float32 tensor of the +**squared-L2** distances to the `k` nearest database points (ascending, nearest +first), and `indices`, a `(Q, k)` int64 tensor of the corresponding database row +indices. The shipped implementation is a straightforward, correct PyTorch version. + +Your goal is to make `knnlib.knn` **as fast as possible** on the GPU while +returning the same nearest neighbors. You may rewrite the internals of the +package however you like and add new modules (including Triton kernels) under +`knnlib/`. The public function signature and return contract above must not +change, and every result must remain a deterministic function of its inputs. + +## Workload + +The graded workloads are a family of held-out dense search problems that vary +`(Q, M, D, k)`, spanning both wide-feature (large `D`) and many-neighbor (large +`k`) regimes. All use squared-L2 distance and float32 data. The points are real +image-descriptor vectors, not synthetic noise: expect clustered, non-uniform +data rather than i.i.d. gaussians. + +Every timed iteration draws a **fresh random `Q` queries and a fresh random `M` +database rows**, so neither the query set nor the `database` tensor repeats +across calls. Building an index or precomputing norms once and caching them for +reuse on a later call therefore buys nothing — each call sees a database it has +never seen. Optimize the per-call search itself. + +The public self-test now runs the **exact graded workloads** — the same shapes, +thresholds, seeds, and timing the judge uses to score you — so there are no +separate hidden shapes to guess at. Treat this as a general dense brute-force +k-NN kernel; the workloads are just where it is measured. + +## Iterate on a GPU + +The agent workspace has no GPU. Use the public test to check your current code's correctness and speed on +a GPU through Modal (needs `MODAL_TOKEN_ID` and `MODAL_TOKEN_SECRET` in your +environment): + +```bash +bash /app/public_test.sh +``` + +It runs the **identical evaluation the judge uses to grade you** and reports, per +graded workload, whether your result passes the quality gate, your speedup over +the baseline, the geometric-mean speedup, and your **predicted final score** +(0-100). Any workload that fails its gate makes the submission score 0. + +## Submission + +The submitted artifact is a patch over the `knnlib` package: + +```text +/app/solution.patch +``` + +After editing `/app/knnlib`, generate and submit: + +```bash +bash /app/make_submission.sh +bash /app/submit.sh +``` + +Submissions are asynchronous; submit early and keep iterating. The judge applies +your patch to a clean copy of `knnlib` and times it against the original +baseline on a GPU, on the same seeded queries and database. + +## Correctness + +Correctness is a gate. On every timed iteration the judge scores your neighbours +by **ball-recall**: each returned index counts if the exact (fp32) squared-L2 +distance from its query to that database row lies within the true k-nearest- +neighbour ball (the distance to the k-th true neighbour). This is tie-robust — +bf16 distances tie heavily near rank k, so which of several equidistant points +you return does not matter, but returning a genuinely farther point does. Your +ball-recall must stay at or above a high threshold. Because the inputs are bf16, +an honest bf16 kernel clears it comfortably; a search that drops below bf16 +precision (e.g. fp8) returns farther points and fails. Crashes, non-finite +output, wrong shapes/dtypes, and timeouts are penalized before speed is +considered. + +## Scoring + +Valid submissions are scored by speedup relative to the baseline on the same +hardware and workloads. For each workload: + +```text +speedup = baseline_time / your_time +``` + +The objective is the geometric mean of per-workload speedups, so broad speedups +are preferred over a single large outlier. A result no faster than the baseline +earns 0; regressions earn 0. The raw geometric-mean speedup is reported in the +evaluator metrics. + +## Patch Policy + +The evaluator validates the patch before running it. Only Python files under the +package may change: + +```text +knnlib/** +``` + +New Python modules inside `knnlib/` are allowed. Patches may not: modify +anything outside `knnlib/`; import or call an external optimized ML/kernel +library (write the kernels yourself); read or write environment variables, spawn +processes, or access the network; or otherwise tamper with the measurement +framework. The timing harness measures your code on freshly generated data every +iteration and re-verifies quality each time, so caching results across calls does +not help. + +## Resource Budget + +```text +GPU: single Modal GPU (H100 reference; the Triton paths also run on L40S / A100) +Agent container: CPU-only (GPU work is offloaded to Modal) +``` diff --git a/2.0/problems/knn_gpu_kernel_optimization/reference.patch b/2.0/problems/knn_gpu_kernel_optimization/reference.patch new file mode 100644 index 000000000..8a50b02b5 --- /dev/null +++ b/2.0/problems/knn_gpu_kernel_optimization/reference.patch @@ -0,0 +1,2160 @@ +diff --git a/knnlib/_kernels/__init__.py b/knnlib/_kernels/__init__.py +new file mode 100644 +index 0000000..e69de29 +diff --git a/knnlib/_kernels/kernels/__init__.py b/knnlib/_kernels/kernels/__init__.py +new file mode 100644 +index 0000000..e69de29 +diff --git a/knnlib/_kernels/kernels/distance/__init__.py b/knnlib/_kernels/kernels/distance/__init__.py +new file mode 100644 +index 0000000..e69de29 +diff --git a/knnlib/_kernels/kernels/distance/triton/__init__.py b/knnlib/_kernels/kernels/distance/triton/__init__.py +new file mode 100644 +index 0000000..bf669cf +--- /dev/null ++++ b/knnlib/_kernels/kernels/distance/triton/__init__.py +@@ -0,0 +1,11 @@ ++"""distance triton backend (kNN squared-L2 gather only). ++ ++Minimal re-export: only the tiled squared-L2 gather kernel that recovers ++true ``||x - c[idx]||^2`` per neighbour after the fused kNN pass. The ++other distance kernels from the upstream backend are not vendored here. ++""" ++from knnlib._kernels.kernels.distance.triton.knn_gather_l2sq import ( ++ triton_knn_gather_sqdist, ++) ++ ++__all__ = ["triton_knn_gather_sqdist"] +diff --git a/knnlib/_kernels/kernels/distance/triton/knn_gather_l2sq.py b/knnlib/_kernels/kernels/distance/triton/knn_gather_l2sq.py +new file mode 100644 +index 0000000..e2e3fbe +--- /dev/null ++++ b/knnlib/_kernels/kernels/distance/triton/knn_gather_l2sq.py +@@ -0,0 +1,232 @@ ++"""Tiled gather kernel: true squared L2 to selected kNN neighbours. ++ ++After a fused kNN pass that ranks candidates with the shift-invariant ++score ``c_sq - 2 * ``, this kernel writes the **true** ++``||x[b, n] - c[b, idx[b, n, k]]||^2`` for every neighbour. Computing ++the difference directly (``(x - y)^2`` per d) avoids the cancellation ++inherent to the ``x^2 + y^2 - 2*x*y`` expansion, so the result is ++accurate even for fp32 inputs where the fused kNN kernel may have used ++TF32 in the cross-term. ++ ++Complexity: ``O(B * N * K * D)`` -- cheap vs the ``O(B * N * M * D)`` ++GEMM that found the candidates. ++ ++Design ++------ ++One program owns a ``(BN, K_BLOCK)`` output tile. The query row tile ++``x[b, n_block, :]`` is loaded once and reused across the ``K_BLOCK`` ++neighbours, then streamed through ``D`` in ``BLOCK_D`` chunks. The ++corpus rows ``c[b, idx, :]`` are gather-loaded fresh per d-chunk into a ++3-D tile ``(BN, K_BLOCK, BLOCK_D)`` and subtracted directly from the ++broadcast x tile in fp32. The fp32 squared difference is summed along ++``D`` into the accumulator. ++ ++Three regime presets pick ``BN`` and ``K_BLOCK`` so the c-tile fits in ++SMEM/registers and amortises one x-load over many neighbours: ++ ++ * K <= 32: ``K_BLOCK = next_pow2(K)``, ``BN = 16``. ++ * 32 < K <= 512: ``K_BLOCK = 32``, ``BN = 8``; multiple K-tiles per row. ++ * K > 512: ``K_BLOCK = 32``, ``BN = 4``. ++ ++``BLOCK_D`` is chosen so the per-program SMEM footprint ++``BN * K_BLOCK * BLOCK_D * sizeof(dtype)`` stays under ~32 KB. The ++unmodified ``D`` (not padded) is passed as a constexpr so the Triton ++compiler unrolls the d-loop exactly. ++""" ++from __future__ import annotations ++ ++import torch ++import triton ++import triton.language as tl ++ ++ ++def _next_pow2(n: int) -> int: ++ if n <= 1: ++ return 1 ++ return 1 << (n - 1).bit_length() ++ ++ ++@triton.jit ++def _knn_gather_l2sq_kernel( ++ x_ptr, c_ptr, idx_ptr, out_ptr, ++ M, ++ stride_x_b, stride_x_n, stride_x_d, ++ stride_c_b, stride_c_m, stride_c_d, ++ stride_i_b, stride_i_n, stride_i_k, ++ stride_o_b, stride_o_n, stride_o_k, ++ N: tl.constexpr, D: tl.constexpr, K: tl.constexpr, ++ BN: tl.constexpr, K_BLOCK: tl.constexpr, ++ BLOCK_D: tl.constexpr, ++ SINGLE_D_TILE: tl.constexpr, ++): ++ """Tiled (BN, K_BLOCK, BLOCK_D) gather + sq-distance accumulate. ++ ++ Grid: ``(ceil(N / BN), ceil(K / K_BLOCK), B)``. ++ ++ Each program owns one ``(BN, K_BLOCK)`` output tile. The fp32 ++ accumulator stays in registers; the c-row tile ``(BN, K_BLOCK, ++ BLOCK_D)`` is reloaded per d-chunk to keep SMEM bounded. ++ """ ++ pid_n = tl.program_id(0) ++ pid_k = tl.program_id(1) ++ pid_b = tl.program_id(2).to(tl.int64) ++ ++ n_start = pid_n * BN ++ n_offs = (n_start + tl.arange(0, BN)).to(tl.int64) ++ n_mask = n_offs < N ++ ++ k_start = pid_k * K_BLOCK ++ k_offs = (k_start + tl.arange(0, K_BLOCK)).to(tl.int64) ++ k_mask = k_offs < K ++ ++ idx_tile = tl.load( ++ idx_ptr + pid_b * stride_i_b ++ + n_offs[:, None] * stride_i_n ++ + k_offs[None, :] * stride_i_k, ++ mask=n_mask[:, None] & k_mask[None, :], other=0, ++ ).to(tl.int64) ++ idx_tile = tl.maximum(idx_tile, 0) ++ idx_tile = tl.minimum(idx_tile, M - 1) ++ ++ acc = tl.zeros((BN, K_BLOCK), dtype=tl.float32) ++ ++ if SINGLE_D_TILE: ++ d_offs = tl.arange(0, BLOCK_D).to(tl.int64) ++ d_mask = d_offs < D ++ x_tile = tl.load( ++ x_ptr + pid_b * stride_x_b ++ + n_offs[:, None] * stride_x_n ++ + d_offs[None, :] * stride_x_d, ++ mask=n_mask[:, None] & d_mask[None, :], other=0.0, ++ ).to(tl.float32) ++ c_tile = tl.load( ++ c_ptr + pid_b * stride_c_b ++ + idx_tile[:, :, None] * stride_c_m ++ + d_offs[None, None, :] * stride_c_d, ++ mask=(n_mask[:, None] & k_mask[None, :])[:, :, None] ++ & d_mask[None, None, :], other=0.0, ++ ).to(tl.float32) ++ diff = x_tile[:, None, :] - c_tile ++ acc = tl.sum(diff * diff, axis=2) ++ else: ++ for d_start in tl.range(0, D, BLOCK_D): ++ d_offs = (d_start + tl.arange(0, BLOCK_D)).to(tl.int64) ++ d_mask = d_offs < D ++ x_tile = tl.load( ++ x_ptr + pid_b * stride_x_b ++ + n_offs[:, None] * stride_x_n ++ + d_offs[None, :] * stride_x_d, ++ mask=n_mask[:, None] & d_mask[None, :], other=0.0, ++ ).to(tl.float32) ++ c_tile = tl.load( ++ c_ptr + pid_b * stride_c_b ++ + idx_tile[:, :, None] * stride_c_m ++ + d_offs[None, None, :] * stride_c_d, ++ mask=(n_mask[:, None] & k_mask[None, :])[:, :, None] ++ & d_mask[None, None, :], other=0.0, ++ ).to(tl.float32) ++ diff = x_tile[:, None, :] - c_tile ++ acc += tl.sum(diff * diff, axis=2) ++ ++ tl.store( ++ out_ptr + pid_b * stride_o_b ++ + n_offs[:, None] * stride_o_n ++ + k_offs[None, :] * stride_o_k, ++ acc, mask=n_mask[:, None] & k_mask[None, :], ++ ) ++ ++ ++def _pick_tile(K: int, D: int, dtype_bytes: int) -> tuple[int, int, int, bool]: ++ """Pick ``(BN, K_BLOCK, BLOCK_D, single_d_tile)`` for the kernel. ++ ++ Regimes mirror the file docstring -- small ``K_BLOCK`` for large K, ++ larger ``BN`` for small K. ``BLOCK_D`` is sized so the 3-D c-tile ++ SMEM footprint stays around ~32 KB. ++ """ ++ K_pad = _next_pow2(K) ++ ++ if K_pad <= 32: ++ K_BLOCK = max(1, K_pad) ++ BN = 16 ++ elif K_pad <= 512: ++ K_BLOCK = 32 ++ BN = 8 ++ else: ++ K_BLOCK = 32 ++ BN = 4 ++ ++ target_bytes = 32 * 1024 ++ per_d = BN * K_BLOCK * dtype_bytes ++ block_d_cap = max(16, target_bytes // max(1, per_d)) ++ block_d_cap = min(block_d_cap, 256) ++ ++ if D <= block_d_cap: ++ BLOCK_D = max(16, _next_pow2(D)) ++ single_d_tile = True ++ else: ++ BLOCK_D = 64 if block_d_cap >= 64 else 32 ++ single_d_tile = False ++ ++ return BN, K_BLOCK, BLOCK_D, single_d_tile ++ ++ ++def triton_knn_gather_sqdist( ++ x: torch.Tensor, ++ c: torch.Tensor, ++ idx: torch.Tensor, ++ *, ++ out: torch.Tensor | None = None, ++) -> torch.Tensor: ++ """``out[b, n, k] = || x[b, n, :] - c[b, idx[b, n, k], :] ||^2`` (fp32). ++ ++ Args: ++ x: (B, N, D) query tensor, any float cuda dtype (bf16 / fp16 / fp32). ++ c: (B, M, D) corpus, same dtype and same batch as ``x``. ++ idx: (B, N, K) int32 / int64 column indices into ``c``. ++ out: optional pre-allocated (B, N, K) fp32 output buffer. ++ ++ Returns: ++ (B, N, K) fp32. The computation runs in fp32 throughout ++ (``diff = x.to(fp32) - c.to(fp32)``, ``sum(diff*diff)``) so the ++ result is bit-equivalent to the naive torch reference up to ++ accumulation order on the d-axis. ++ """ ++ assert x.is_cuda and c.is_cuda and idx.is_cuda ++ assert x.dim() == 3 and c.dim() == 3 and idx.dim() == 3 ++ B, N, D = x.shape ++ Bc, M, Dc = c.shape ++ Bi, Ni, K = idx.shape ++ assert B == Bc == Bi and N == Ni and D == Dc ++ ++ if not x.is_contiguous(): ++ x = x.contiguous() ++ if not c.is_contiguous(): ++ c = c.contiguous() ++ if not idx.is_contiguous(): ++ idx = idx.contiguous() ++ ++ if out is None: ++ out = torch.empty((B, N, K), device=x.device, dtype=torch.float32) ++ else: ++ assert out.shape == (B, N, K) and out.dtype == torch.float32 ++ ++ dtype_bytes = x.element_size() ++ BN, K_BLOCK, BLOCK_D, single_d_tile = _pick_tile(K, D, dtype_bytes) ++ ++ grid = (triton.cdiv(N, BN), triton.cdiv(K, K_BLOCK), B) ++ _knn_gather_l2sq_kernel[grid]( ++ x, c, idx, out, ++ M, ++ x.stride(0), x.stride(1), x.stride(2), ++ c.stride(0), c.stride(1), c.stride(2), ++ idx.stride(0), idx.stride(1), idx.stride(2), ++ out.stride(0), out.stride(1), out.stride(2), ++ N=N, D=D, K=K, ++ BN=BN, K_BLOCK=K_BLOCK, BLOCK_D=BLOCK_D, ++ SINGLE_D_TILE=single_d_tile, ++ num_warps=4, ++ ) ++ return out ++ ++ ++__all__ = ["triton_knn_gather_sqdist"] +diff --git a/knnlib/_kernels/primitives/__init__.py b/knnlib/_kernels/primitives/__init__.py +new file mode 100644 +index 0000000..e69de29 +diff --git a/knnlib/_kernels/primitives/knn/__init__.py b/knnlib/_kernels/primitives/knn/__init__.py +new file mode 100644 +index 0000000..e69de29 +diff --git a/knnlib/_kernels/primitives/knn/triton/__init__.py b/knnlib/_kernels/primitives/knn/triton/__init__.py +new file mode 100644 +index 0000000..c177f23 +--- /dev/null ++++ b/knnlib/_kernels/primitives/knn/triton/__init__.py +@@ -0,0 +1,41 @@ ++"""knn triton backend. ++ ++Re-exports the public Python wrappers from each component file. ++``@triton.jit`` kernels stay private to their file. ++ ++Kernel layout: ++ ++* :mod:`_common` -- shared helpers: ``_next_pow2``, micro-bench, ++ IEEE-sortable u32 transform, ``_INF_PACKED`` sentinel. ++* :mod:`sortmerge` -- packed-uint64 sort-merge top-K (x²-free). Routed ++ to only for the small-Q + medium-K Pattern-A corner. ++* :mod:`insert` -- iterative argmin-insert top-K (x²-free). The ++ general path -- BN in {8, 16, 32, 64, 128} covers everything from ++ small-Q search to large-Q build. ++* :mod:`_row_norm` -- ``_fast_row_sq`` / ``_get_or_compute_csq`` for ++ the CuteDSL FA3 wrapper (it needs ``c_sq`` as a kernel argument). ++* :mod:`dispatch` -- host-side router. Single public entry ++ :func:`flash_knn_triton` runs through :func:`_heuristic_config` ++ unconditionally -- the per-CTA-count check inside the heuristic ++ picks single-pass vs M-split. Per-shape config (BN/BM/mode/mps/ ++ ns_pipe) also falls out of :func:`_heuristic_config`. ++""" ++from knnlib._kernels.primitives.knn.triton._common import ( ++ _next_pow2, ++ _bench_quick, ++) ++from knnlib._kernels.primitives.knn.triton._row_norm import ( ++ _fast_row_sq, ++ _get_or_compute_csq, ++) ++from knnlib._kernels.primitives.knn.triton.dispatch import ( ++ flash_knn_triton, ++ flash_knn_triton_small_n, ++ flash_knn_triton_large_n, ++) ++ ++__all__ = [ ++ "flash_knn_triton", ++ "flash_knn_triton_small_n", ++ "flash_knn_triton_large_n", ++] +diff --git a/knnlib/_kernels/primitives/knn/triton/_common.py b/knnlib/_kernels/primitives/knn/triton/_common.py +new file mode 100644 +index 0000000..686e85b +--- /dev/null ++++ b/knnlib/_kernels/primitives/knn/triton/_common.py +@@ -0,0 +1,81 @@ ++"""Shared utilities for flash_knn Triton dispatch + kernels. ++ ++Lives in one place so :mod:`dispatch`, :mod:`sortmerge` and :mod:`insert` ++all import the same primitives: ++ ++ * :func:`_next_pow2` -- D / K padding helper. ++ * :func:`_bench_quick` -- micro-bench used by the autotuner. ++ * :func:`_fp32_to_sortable_u32` / :func:`_sortable_u32_to_fp32` -- ++ branch-free IEEE-sortable u32 transform. With the x²-free score ++ ``s = c² - 2⟨x, c⟩`` the quantity is signed, so the packed-uint64 ++ sort-merge / final-sort needs an ascending-u32 = ascending-fp32 ++ mapping. The transform is 3 ops, exact for non-NaN fp32. ++ * :data:`_INF_PACKED` -- (sortable +inf << 32) | -1 idx sentinel, ++ used as the initial fill of the sortmerge running top-K so the ++ early-exit ``chunk_best < sortable_inv(top)`` always fires until ++ real values arrive. ++""" ++from __future__ import annotations ++ ++import torch ++import triton ++import triton.language as tl ++ ++ ++def _next_pow2(n: int) -> int: ++ """Smallest power-of-two >= ``max(1, n)``. Used for D / K padding.""" ++ if n <= 1: ++ return 1 ++ return 1 << (n - 1).bit_length() ++ ++ ++def _bench_quick(fn, *, warmup: int = 3, reps: int = 5) -> float: ++ """Median-ish wall time in ms; cheap enough to call once per autotune cfg.""" ++ for _ in range(warmup): ++ fn() ++ torch.cuda.synchronize() ++ s = torch.cuda.Event(enable_timing=True) ++ e = torch.cuda.Event(enable_timing=True) ++ s.record() ++ for _ in range(reps): ++ fn() ++ e.record() ++ torch.cuda.synchronize() ++ return s.elapsed_time(e) / reps ++ ++ ++# Sortable upper-bits = sortable(+inf) = 0xFF800000; lower-bits = -1 (idx ++# sentinel) = 0xFFFFFFFF. ``_sortable_u32_to_fp32(0xFF800000) == +inf`` so ++# the sortmerge early-exit ``chunk_best < +inf`` always fires until the ++# running top-K starts to hold real values. ++# ++# Wrapped in ``tl.constexpr`` so ``@triton.jit`` kernels can reference it ++# directly as a module-level global (un-wrapped Python ints are not ++# accessible from inside a JIT'd kernel as of Triton 3.x). ++_INF_PACKED = tl.constexpr(0xFF800000_FFFFFFFF) ++ ++ ++@triton.jit ++def _fp32_to_sortable_u32(x): ++ """Map ascending fp32 -> ascending uint32 (branch-free, 3 ops). ++ ++ Standard IEEE-754 trick: arithmetic-shift the int32 view by 31 (so ++ MSB=1 / negative produces 0xFFFFFFFF, MSB=0 / positive produces 0), ++ OR with 0x80000000 to flip the sign bit of positives, then XOR with ++ the original bits. NaNs map into 0xFF800001..0xFFFFFFFF (unused in ++ practice). ++ """ ++ bits = x.to(tl.uint32, bitcast=True) ++ sign_mask = (bits.to(tl.int32, bitcast=True) >> 31).to(tl.uint32, bitcast=True) ++ flip = sign_mask | 0x80000000 ++ return bits ^ flip ++ ++ ++@triton.jit ++def _sortable_u32_to_fp32(s): ++ """Inverse of :func:`_fp32_to_sortable_u32`.""" ++ sign_mask = (s.to(tl.int32, bitcast=True) >> 31).to(tl.uint32, bitcast=True) ++ # In sortable space: MSB=1 -> was positive (flip back via 0x80000000); ++ # MSB=0 -> was negative (flip back via 0xFFFFFFFF). ++ flip = (sign_mask ^ 0xFFFFFFFF) | 0x80000000 ++ return (s ^ flip).to(tl.float32, bitcast=True) +diff --git a/knnlib/_kernels/primitives/knn/triton/_row_norm.py b/knnlib/_kernels/primitives/knn/triton/_row_norm.py +new file mode 100644 +index 0000000..60b0ec8 +--- /dev/null ++++ b/knnlib/_kernels/primitives/knn/triton/_row_norm.py +@@ -0,0 +1,93 @@ ++"""Fast row sum-of-squares (||x_i||^2) helper + per-tensor cache. ++ ++Tiny utility, but shared by the FA3 fused KNN path which needs the ++pre-computed query/corpus norms to fold into the squared-L2 distance ++formula ``||x - c||^2 = ||x||^2 + ||c||^2 - 2 ``. ++ ++This module deliberately stays narrow: ++ - ``_row_sq_kernel`` is a pure Triton row-norm kernel (no cross ++ matrix, no top-K) — it never materialises an N x M tensor. ++ - ``_fast_row_sq`` wraps the kernel. ++ - ``_get_or_compute_csq`` caches results per ``(data_ptr, shape, ++ dtype)`` so repeat queries on the same corpus skip the recompute. ++""" ++from __future__ import annotations ++ ++import math ++ ++import torch ++import triton ++import triton.language as tl ++ ++from knnlib._kernels.primitives.knn.triton._common import _next_pow2 ++ ++ ++@triton.jit ++def _row_sq_kernel( ++ x_ptr, out_ptr, ++ stride_x_b, stride_x_n, stride_x_d, ++ stride_o_b, stride_o_n, ++ B: tl.constexpr, N: tl.constexpr, D: tl.constexpr, ++ BN: tl.constexpr, BD: tl.constexpr, ++): ++ """Row sum-of-squares for ``(B, N, D) -> (B, N)`` fp32. ++ ++ Casts input to fp32 inside the kernel; reads input exactly once ++ and writes output exactly once. At BN=128 BD=128 nw=4 this hits ++ ~36% peak HBM BW on H200 (typical for tall-skinny reads). ++ """ ++ pid_n = tl.program_id(0) ++ pid_b = tl.program_id(1).to(tl.int64) ++ n_offs = (pid_n * BN + tl.arange(0, BN)).to(tl.int64) ++ n_mask = n_offs < N ++ acc = tl.zeros([BN], dtype=tl.float32) ++ for d_start in range(0, D, BD): ++ d_offs = (d_start + tl.arange(0, BD)).to(tl.int64) ++ d_mask = d_offs < D ++ x = tl.load( ++ x_ptr + pid_b * stride_x_b ++ + n_offs[:, None] * stride_x_n ++ + d_offs[None, :] * stride_x_d, ++ mask=n_mask[:, None] & d_mask[None, :], other=0.0, ++ ) ++ x_f = x.to(tl.float32) ++ acc += tl.sum(x_f * x_f, axis=1) ++ tl.store(out_ptr + pid_b * stride_o_b + n_offs * stride_o_n, ++ acc, mask=n_mask) ++ ++ ++def _fast_row_sq(x: torch.Tensor) -> torch.Tensor: ++ """Return ``(B, N)`` fp32 row sum-of-squares via the Triton kernel.""" ++ assert x.is_cuda and x.ndim == 3 ++ B, N, D = x.shape ++ out = torch.empty(B, N, device=x.device, dtype=torch.float32) ++ BN = 128 if N >= 128 else _next_pow2(max(N, 16)) ++ BD = 128 if D >= 128 else _next_pow2(max(D, 16)) ++ grid = (math.ceil(N / BN), B) ++ _row_sq_kernel[grid]( ++ x, out, ++ x.stride(0), x.stride(1), x.stride(2), ++ out.stride(0), out.stride(1), ++ B=B, N=N, D=D, BN=BN, BD=BD, ++ num_warps=4, ++ ) ++ return out ++ ++ ++# Per-data-ptr cache so repeated queries on the same corpus don't pay ++# the row-norm kernel cost. ++_csq_cache: dict = {} ++ ++ ++def _get_or_compute_csq(c: torch.Tensor) -> torch.Tensor: ++ """Cached ``||c_i||^2`` (fp32) keyed by ``(data_ptr, shape, dtype)``.""" ++ key = (c.data_ptr(), tuple(c.shape), c.dtype) ++ cached = _csq_cache.get(key) ++ if cached is not None and cached.device == c.device: ++ return cached ++ csq = _fast_row_sq(c) ++ _csq_cache[key] = csq ++ if len(_csq_cache) > 8: ++ for old_key in list(_csq_cache.keys())[:-8]: ++ del _csq_cache[old_key] ++ return csq +diff --git a/knnlib/_kernels/primitives/knn/triton/dispatch.py b/knnlib/_kernels/primitives/knn/triton/dispatch.py +new file mode 100644 +index 0000000..4c78f35 +--- /dev/null ++++ b/knnlib/_kernels/primitives/knn/triton/dispatch.py +@@ -0,0 +1,1178 @@ ++"""flash_knn -- host-side dispatcher: kernel/config selection + 2-stage launch. ++ ++Backs the public :func:`knnlib._kernels.primitives.knn.flash_knn` entry point. ++ ++Top-level wrappers ++------------------ ++ ++* :func:`flash_knn_triton` -- one-shot kNN, returns ``(B, N, k)`` int32 ++ indices. The shape-only heuristic now handles both build (large-N, ++ single-pass per CTA) and search (small-N, M-split flash-decode) ++ without any host-side SM-saturation gate -- ``ctas_no_split`` is ++ inspected after BN is picked, so build shapes (``BN=128`` -> ++ ``ctas_no_split`` already saturates 132 SMs) and search shapes ++ (small BN -> M-split for tail-fill) both fall out of the same ++ decision tree. ++* :func:`flash_knn_triton_small_n` / :func:`flash_knn_triton_large_n` ++ -- explicit overrides for callers (e.g. K-means assign) that know ++ their shape category at construction time. ++* :func:`_heuristic_config`, :func:`_autotune` -- shape-only heuristic ++ (default, fast first call) + opt-in brute-force autotune (slower ++ first call, cached steady-state). Both share the same kernel grid; ++ ``force_path="large_n"`` collapses M-splits down to one per CTA. ++ ++Two kernels live behind this dispatcher: ++ ++* :mod:`insert` -- iterative argmin-insert top-K. Picked by the ++ heuristic on virtually every shape. ++* :mod:`sortmerge` -- packed-uint64 sort-merge top-K with the IEEE- ++ sortable u32 transform. Routed to for the **Pattern-A small-Q** ++ corner (``B*N <= 8 or (B*N <= 16 and K in {32, 64})``) at medium-K ++ + ``M <= 200K`` where the wider sort per chunk beats insert's ++ argmin loop on this regime's autotune data. Otherwise insert wins; ++ sortmerge is kept in :func:`_gen_configs` so the offline tuner can ++ keep verifying that. ++ ++Pipelining depth ++---------------- ++Both kernels take a ``NUM_STAGES_PIPE`` constexpr for the M-loop's ++``tl.range`` pipelining factor. The dispatcher tunes it per shape from ++a 201-shape ns-sweep (upstream ``bench/sweep_ns_by_k.py``): K<=8 + ++small tile wants ``ns=3`` to hide HBM latency, K>=64 + big tile wants ++``ns=2`` (registers tight), wide-D wants ``ns=2`` uniformly, narrow ++tile + anything wants ``ns=1`` to avoid spills. See the comments ++inside :func:`_heuristic_config` for the empirical breakdown. ++ ++Distance recovery ++----------------- ++The kernels emit **signed shifted scores** ``s = c_sq - 2*``, not ++true squared L2. ``argmin-k`` over ``s`` matches ``argmin-k`` over true ++squared L2 (since ``-||x||^2`` is a per-row constant), but the returned ++score is not a valid distance. The :func:`flash_knn` public wrapper ++calls :func:`knnlib._kernels.kernels.distance.triton_knn_gather_sqdist` on the ++returned indices to write the true ``||x - c[idx]||^2`` per neighbour. ++""" ++from __future__ import annotations ++ ++import math ++from typing import Optional ++ ++import torch ++ ++from knnlib._kernels.primitives.knn.triton._common import _next_pow2, _bench_quick ++from knnlib._kernels.primitives.knn.triton.sortmerge import _flash_knn_sortmerge_kernel ++from knnlib._kernels.primitives.knn.triton.insert import _flash_knn_insert_kernel ++ ++ ++# ── SRAM estimation ──────────────────────────────────────────────────── ++ ++ ++def _estimate_sram(bn, bm, D, K, d_inner, num_stages, *, ++ sortmerge=False, dtype_bytes=2): ++ """Per-CTA SRAM estimate (bytes) for the unified kernel — coarse. ++ ++ Counts only what Triton **definitely** keeps in SMEM: the x tile ++ + c tile (per iteration of the inner D loop). Everything else ++ (cross fp32 accumulator, score tile, topK heap, sort scratch) is ++ register-resident at the tile sizes used by the dispatcher's ++ heuristic, and the NUM_STAGES_PIPE multi-buffering of the c tile ++ is sometimes elided by Triton's pipeliner. Including any of these ++ in the estimate produces 50-200 KB of over-count that **wrongly ++ shrinks** configs the kernel could actually launch — a bf16 ++ regression demonstrated by the ``D=256`` cells in ++ ``benchmarks/results/micro_knn_dtype_smem_fit.md``. ++ ++ This estimator is therefore **optimistic on purpose**. Its only ++ job is to bias the autotune candidate filter away from configs ++ that obviously can't fit (e.g. ``BN=128, BM=128, D=256, fp32`` ++ needs >232 KB by raw arithmetic). The source of truth for "does ++ this fit?" is ``_run``'s runtime ``OutOfResources`` catch + shrink ++ loop. ``dtype_bytes`` is still threaded through so fp32 inputs ++ correctly double the x/c contribution. ++ """ ++ x_sub = bn * d_inner * dtype_bytes ++ c_sub = bm * d_inner * dtype_bytes ++ c_sq = bm * 4 ++ base = x_sub + c_sub + c_sq ++ if sortmerge: ++ chunk = bn * bm * 8 ++ base += chunk ++ return base ++ ++ ++def _smem_limit(device) -> int: ++ """Per-block dynamic shared-memory budget for ``device``. ++ ++ Mirrors :func:`knnlib._kernels.primitives.kmeans.triton.assign._smem_limit` ++ — Triton uses opt-in dynamic SMEM; prefer that attribute when ++ available, fall back to static limit, and finally to a conservative ++ 48 KiB for old PyTorch builds. ++ """ ++ props = torch.cuda.get_device_properties(device) ++ for attr in ( ++ "shared_memory_per_block_optin", ++ "max_shared_memory_per_block_optin", ++ "shared_memory_per_block", ++ "max_shared_memory_per_block", ++ ): ++ v = getattr(props, attr, None) ++ if v: ++ return int(v) ++ return 48 * 1024 ++ ++ ++def _fit_config_to_smem(cfg: dict, D: int, K: int, ++ dtype_bytes: int, smem_limit: int) -> dict: ++ """Shrink ``(BN, BM, NUM_STAGES_PIPE)`` until the kernel fits SMEM. ++ ++ Returns ``cfg`` unchanged when it already fits. Otherwise enumerates ++ power-of-two reductions on the three SMEM-bearing axes and picks the ++ one that maximises ``BN * BM * num_stages`` (work-per-program tile) ++ among those that fit, breaking ties towards the original aspect ratio. ++ ++ Why only these three axes? ++ * ``D_INNER``: changing it flips between single-D-tile and D-split ++ paths, which would also change the kernel's perf profile. We ++ keep the heuristic's D-split decision and shrink the cheaper ++ axes first. ++ * ``TOPK_PAD``: determined by ``K``; cannot be reduced without ++ breaking correctness. ++ * ``kernel_mode``, ``num_warps``, ``M_PER_SPLIT``, ``NUM_SPLITS``: ++ these don't affect per-CTA SMEM. Stable. ++ ++ Pattern (and rationale) mirrors :func:`knnlib._kernels.primitives.kmeans.\ ++triton.assign._fit_config_to_smem`. ++ """ ++ BN0 = int(cfg["BN"]) ++ BM0 = int(cfg["BM"]) ++ D_INNER = int(cfg["D_INNER"]) ++ NS0 = int(cfg.get("NUM_STAGES_PIPE", 2)) ++ sortmerge = cfg.get("kernel_mode") == "sortmerge" ++ ++ cur = _estimate_sram(BN0, BM0, D, K, D_INNER, NS0, ++ sortmerge=sortmerge, dtype_bytes=dtype_bytes) ++ if cur <= smem_limit: ++ return cfg ++ ++ def _pow2_down_to(v, lo): ++ out = [] ++ x = v ++ while x >= lo: ++ out.append(x) ++ x //= 2 ++ return out ++ ++ # BN must stay >= 8 (matches _gen_configs lowest BN rung); ++ # BM must stay >= 32 for the inner loop to keep WGMMA shape; ++ # NUM_STAGES_PIPE must stay >= 1. ++ bn_cands = _pow2_down_to(BN0, 8) ++ bm_cands = _pow2_down_to(BM0, 32) ++ ns_cands = list(range(NS0, 0, -1)) ++ ++ best = None ++ best_key = None ++ for bn in bn_cands: ++ for bm in bm_cands: ++ # sortmerge requires BM == TOPK_PAD; do not shrink BM there. ++ if sortmerge and bm != BM0: ++ continue ++ for ns in ns_cands: ++ if _estimate_sram(bn, bm, D, K, D_INNER, ns, ++ sortmerge=sortmerge, ++ dtype_bytes=dtype_bytes) > smem_limit: ++ continue ++ aspect_penalty = abs((bn / max(bm, 1)) - (BN0 / max(BM0, 1))) ++ # Prefer: larger total tile work, closer aspect ratio, ++ # larger BN (more N-parallelism), larger NS (better pipelining). ++ key = (bn * bm * ns, -aspect_penalty, bn, ns) ++ if best_key is None or key > best_key: ++ best_key = key ++ best = (bn, bm, ns) ++ ++ if best is None: ++ raise RuntimeError( ++ f"flash_knn: cannot fit kernel into shared memory " ++ f"(D={D}, K={K}, D_INNER={D_INNER}, dtype_bytes={dtype_bytes}, " ++ f"smem_limit={smem_limit}). Original config " ++ f"BN={BN0}, BM={BM0}, NS={NS0} needs {cur} bytes." ++ ) ++ ++ bn, bm, ns = best ++ out = dict(cfg) ++ out["BN"] = bn ++ out["BM"] = bm ++ out["NUM_STAGES_PIPE"] = ns ++ return out ++ ++ ++def _pipe_stages_for(K, d_inner, D): ++ """``NUM_STAGES_PIPE`` candidates for the autotune M-loop sweep. ++ ++ The K-step inner loop adds register pressure that scales with K, so ++ deeper pipelining (which double-/triple-buffers the next C-tile) is ++ only profitable when K is small. Empirically: ++ ++ * D-split path (``d_inner < D``) -- kernel uses ns=1 by construction. ++ * K >= 64 -- only try {1, 2}; ns >= 3 spills registers and tanks perf ++ (probe showed ns=3 going from 750 -> 1220 us at K=128). ++ * K <= 32 -- try {1, 2, 3} so big pipelines can hide HBM latency ++ for huge-M shapes. ++ """ ++ if d_inner < D: ++ return [1] ++ if K >= 64: ++ return [1, 2] ++ return [1, 2, 3] ++ ++ ++def _gen_configs(D, K, *, dtype_bytes=2, smem_limit=None): ++ """Generate candidate kernel configs (sortmerge + insert), SRAM-validated. ++ ++ Covers ``BN ∈ {8, 16, 32, 64, 128}`` × ``num_warps ∈ {2, 4}`` × the ++ BM / D_INNER / NUM_STAGES_PIPE combinations that fit per-CTA SMEM ++ on H200. ++ ++ The ``BN = 8`` rung unlocks the small-N + medium-K Pattern-A regime ++ (``B*N <= 8`` + ``K ∈ {16, 32, 64, 128}``) where a small row tile + ++ sortmerge ``BM = max(32, K)`` wins over insert. ``tl.dot`` silently ++ pads the M dimension to 16 on Triton 3.x sm_90, so BN=8 still ++ compiles. ++ ++ ``dtype_bytes`` and ``smem_limit`` parameterise the SMEM check so ++ fp32 inputs don't slip through with a config that OOR's at launch. ++ Defaults preserve the pre-fix behaviour (bf16 / 220 KB) for any ++ caller that didn't yet plumb dtype awareness through. ++ """ ++ # 220 KB stays the default — slightly below the 227 KB H200 opt-in ++ # limit to leave room for static SMEM the compiler adds. ++ max_sram = 220_000 if smem_limit is None else smem_limit ++ d_pad = _next_pow2(D) ++ topk_pad_sm = max(32, _next_pow2(K)) ++ topk_pad_ins = _next_pow2(K) ++ d_inner_cands = [d_pad] if D <= 256 else [128] ++ ++ configs = [] ++ ++ # sort-merge: BM == TOPK_PAD; BN ∈ {8, 16, 32} (BN >= 64 never wins ++ # for sortmerge in autotune data). ++ bm_sm = topk_pad_sm ++ for bn in [8, 16, 32]: ++ for nw in [2, 4]: ++ for d_inner in d_inner_cands: ++ num_d_iters = math.ceil(D / d_inner) ++ ns = 2 if num_d_iters == 1 else 1 ++ sram = _estimate_sram( ++ bn, bm_sm, D, K, d_inner, ns, ++ sortmerge=True, dtype_bytes=dtype_bytes, ++ ) ++ if sram <= max_sram: ++ for ns_pipe in _pipe_stages_for(K, d_inner, D): ++ configs.append({ ++ "BN": bn, "BM": bm_sm, "D_INNER": d_inner, ++ "num_warps": nw, ++ "TOPK_PAD": topk_pad_sm, ++ "kernel_mode": "sortmerge", ++ "NUM_STAGES_PIPE": ns_pipe, ++ }) ++ ++ # iterative insert: BM decoupled from K, BN ∈ {8, 16, 32, 64, 128}. ++ for bn in [8, 16, 32, 64, 128]: ++ for nw in [2, 4]: ++ for bm in [64, 128, 256]: ++ for d_inner in d_inner_cands: ++ num_d_iters = math.ceil(D / d_inner) ++ ns = 2 if num_d_iters == 1 else 1 ++ sram = _estimate_sram( ++ bn, bm, D, K, d_inner, ns, ++ sortmerge=False, dtype_bytes=dtype_bytes, ++ ) ++ if sram <= max_sram: ++ for ns_pipe in _pipe_stages_for(K, d_inner, D): ++ configs.append({ ++ "BN": bn, "BM": bm, "D_INNER": d_inner, ++ "num_warps": nw, ++ "TOPK_PAD": topk_pad_ins, ++ "kernel_mode": "insert", ++ "NUM_STAGES_PIPE": ns_pipe, ++ }) ++ ++ return configs ++ ++ ++def _gen_m_splits(M, BM, *, B=1, N=1, BN=16): ++ """M_PER_SPLIT candidates targeting {1, 2, 4, 8, 16} waves on 132 SMs. ++ ++ The wave count is critical -- Stage-2 reduce scales linearly with ++ NUM_SPLITS, while too few splits leaves SMs idle. The autotune ++ sweep confirms that wave=2 is the winner on huge-M + medium-N ++ shapes (e.g. ``1×128×10M×64×10``), which the original {1, 4, 16} ++ grid missed. ++ ++ Capped at 4096×BM tiles (``MAX_MPS_TILES``) to bound the static ++ M-loop trip count for fast Triton compilation. Also includes the ++ single-pass case when M itself fits the cap (subsumes the large-N ++ kernel as ``num_splits = 1``). ++ """ ++ NUM_SMS = 132 ++ MAX_MPS_TILES = 4096 ++ num_n_tiles = max(1, (N + BN - 1) // BN) ++ max_mps = min(M, MAX_MPS_TILES * BM) ++ candidates = set() ++ for waves in [1, 2, 4, 8, 16]: ++ target_splits = max(2, (NUM_SMS * waves) // (num_n_tiles * B)) ++ target_splits = min(target_splits, max(2, M // BM)) ++ mps = (M + target_splits - 1) // target_splits ++ mps = ((mps + BM - 1) // BM) * BM ++ mps = max(mps, BM) ++ mps = min(mps, max_mps) ++ candidates.add(mps) ++ if M <= max_mps: ++ single = ((M + BM - 1) // BM) * BM ++ candidates.add(single) ++ return sorted(candidates) ++ ++ ++# ── shape-only heuristic ─────────────────────────────────────────────── ++ ++ ++def _heuristic_config(B, N, M, D, K, *, force_path=None, ++ dtype_bytes=2, smem_limit=None): ++ """Pick a kernel config from shape alone -- no autotune. ++ ++ When ``dtype_bytes`` and ``smem_limit`` are supplied (the dispatcher ++ plumbs them from ``x.element_size()`` and the device's opt-in SMEM ++ cap), the picked config is post-processed by :func:`_fit_config_to_smem` ++ so that fp32 inputs at large D never produce a config that OOR's at ++ launch. ++ ++ Derived from a 92-shape autotune sweep on H200/bf16. The data ++ showed clean decision boundaries on three axes: ++ ++ * **BN by NB-bucket + M-bump**: ++ NB <= 8 -> BN=8; ++ NB <= 32 -> BN=16 (32 at M>=5M); ++ NB <= 256 -> BN=64 (128 at M>=5M + narrow-D + small-K); ++ NB <= 2K -> BN=64 (128 at M>=5M); ++ NB >= 8K -> BN=128 (64 for D>=256 + K>=32). ++ ++ The M-bump rule fixes the ``(1, 128, 10M, 64, 10)`` regression ++ identified by klib -- at M=10M each extra N-tile costs ++ ~300us of c-replication HBM traffic. ++ ++ * **BM by K**: ++ K<=4 -> 128/256 (256 only at very-small M + small NB); ++ K=5..16 -> 128 default (256 at huge-M); ++ K=32 -> 128 default (256 at NB<=8 + M<=200K); ++ K>=64 -> 64 default (sortmerge at NB<=8). ++ ++ * **kernel_mode**: ++ ``sortmerge`` only when ``NB <= 8`` and ``K ∈ {32, 64, 128}`` ++ and ``D <= 256`` and ``M <= 200K`` (Pattern-A1), plus a ++ secondary ``NB <= 16 + K ∈ {32, 64} + M <= 200K`` corner ++ (Pattern-A2). All other shapes use ``insert``. ++ ++ * **num_warps**: 4 everywhere except tiny tiles (BN=8 + BM<=64 + ++ K<=4 + huge M, where nw=2 wins by reducing register pressure) ++ and the BN=16 + huge-M + small-K corner where nw=2 frees a ++ warp set for the topK epilogue. ++ ++ * **target waves**: 1 for build (NB>=8K) and very-large K; ++ 2 for medium/large M; 4 for small M + small K. ++ ++ * **NUM_STAGES_PIPE**: 2 by default; 3 for BN>=64 + BM<=128 + ++ K<=8 (small tile + tiny K hides HBM latency) and for ++ BN ∈ [16, 32] + BM=256 + K>=64 (big tile keeps K loop fed); ++ 1 for narrow-tile + any K (avoid register spills) and for the ++ D-split path; 2 for D_INNER>=256 and for BN>=128. Without ++ this rule, K>=64 + D=256 shapes hit catastrophic regressions ++ (e.g. ``(1,16,100K,256,128)`` 439us -> 2429us at ns=1, a 5.5x ++ slowdown); conversely K<=32 + huge-M shapes save 20-45% vs ++ ns=2 (e.g. ``(1,1,10M,128,4)`` 1426us -> 785us at ns=1). ++ ++ Args: ++ B, N, M, D, K: shape parameters. ++ force_path: ``"large_n"`` for single-pass (one M-split per CTA), ++ ``None`` for the wave-targeted M-split. ++ ++ Returns: ++ Config dict: ``BN, BM, D_INNER, TOPK_PAD, kernel_mode, ++ num_warps, M_PER_SPLIT, NUM_SPLITS, NUM_STAGES_PIPE``. ++ """ ++ NUM_SMS = 132 ++ NB = N * B ++ TOPK_PAD = _next_pow2(K) ++ D_INNER = _next_pow2(D) if D <= 256 else 128 ++ is_large_n = (force_path == "large_n") ++ # MAX_MPS_TILES bounds the per-CTA M-loop iteration count. The ++ # original cap of 512 was too tight for huge M with BM=128 (cap ++ # = 65K), which forced fewer splits than autotune wanted. Bump ++ # to 4096 so we never bottleneck on this cap. ++ MAX_MPS_TILES = 4096 ++ ++ # ─────────────────────────────────────────────────────────────── ++ # Step 1: kernel_mode, BN, BM, num_warps ++ # ─────────────────────────────────────────────────────────────── ++ if NB <= 8: ++ # Tiny query: BN=8 default. ++ BN = 8 ++ # Sortmerge fast-path: K ∈ {32, 64, 128}, M <= 200K. Now ++ # extended to ANY D (autotune ``(1,1,100K,1024,32)`` picks ++ # BN=8 BM=32 sort). ++ if (not is_large_n) and K in (32, 64, 128) and M <= 200_000: ++ BM = max(32, _next_pow2(K)) ++ sram = _estimate_sram(8, BM, D, K, D_INNER, ++ num_stages=(2 if D_INNER >= D else 1), ++ sortmerge=True) ++ if sram <= 220_000: ++ kernel_mode = "sortmerge" ++ num_warps = 2 if BM >= 128 else 4 ++ else: ++ kernel_mode = "insert" ++ BM = 256 if M <= 200_000 else 128 ++ num_warps = 4 ++ else: ++ kernel_mode = "insert" ++ if K <= 4: ++ # K=4 + D>=256 + huge M -> BM=64 nw=2 (autotune ++ # ``(1,1,1M,256,4)`` picks this). HBM-bound; BM=256 ++ # adds wasted SMEM. ++ if D >= 256 and M >= 500_000: ++ BM = 64 ++ else: ++ BM = 256 ++ elif K <= 16: ++ # D>=256 + K=16 -> BM=128 (autotune (1,1,100K,256,16)) ++ if D >= 256 and K >= 16: ++ BM = 128 ++ else: ++ BM = 256 if (M >= 500_000 or K >= 16) else 128 ++ elif K == 32: ++ # autotune (1,1,1M,256,32) prefers BM=128 at D>=256. ++ # At narrower D, BM=256 wins. ++ if D >= 256 and M >= 500_000: ++ BM = 128 ++ else: ++ BM = 256 ++ elif K <= 64: ++ BM = 64 ++ else: # K=128 ++ BM = 256 ++ # nw=2 for small/medium-D BM=64 cases ++ if BM == 64 and K <= 4 and D >= 512 and M >= 500_000: ++ # D >= 512 hits the D-split path (D_INNER capped at 128); ++ # nw=2 was calibrated for this regime at (1, 1, 1M, 256, 4) ++ # but later re-benching at D=256 showed nw=4 ns=1 wins by ++ # 1.48-1.81x there (single-D-iter, multibuffered C). Keep ++ # nw=2 only for the D-split path; D=256 falls to the ++ # default nw=4 below and gets the ns=1 override in Step 3. ++ num_warps = 2 ++ elif BM >= 128 and K <= 4 and D <= 64 and M >= 5_000_000: ++ # NB<=8 + K<=4 + narrow-D + huge-M: nw=2 wins by 20-30 %. ++ # 4 warps split the SM register file too thinly for the ++ # K=4 argmin-insert epilogue at BM>=128, so the compiler ++ # spills; halving the warp count doubles regs/warp and ++ # lets ILP recover. Verified at D=64 across M ∈ [5M, 60M] ++ # (D=128 strictly prefers nw=4 at this BM; do not bump ++ # the threshold without re-benching). ++ num_warps = 2 ++ else: ++ num_warps = 4 ++ ++ elif NB <= 32: ++ kernel_mode = "insert" ++ # Sortmerge also wins at NB <= 16 + K ∈ {32, 64} (autotune ++ # ``(1,16,100K,256,32)`` picks BN=8 BM=32 sortmerge). ++ if (not is_large_n) and NB <= 16 and K in (32, 64) and M <= 200_000: ++ BN = 8 ++ BM = max(32, _next_pow2(K)) ++ sram = _estimate_sram(8, BM, D, K, D_INNER, ++ num_stages=(2 if D_INNER >= D else 1), ++ sortmerge=True) ++ if sram <= 220_000: ++ kernel_mode = "sortmerge" ++ num_warps = 4 ++ else: ++ BN, BM = 16, 256 ++ if kernel_mode == "insert": ++ if M >= 5_000_000 and N >= 32: ++ BN = 32 ++ else: ++ BN = max(8, min(16, _next_pow2(N))) ++ if K <= 4: ++ BM = 128 ++ elif K <= 16: ++ BM = 256 if M <= 2_000_000 else 128 ++ else: # K >= 32 ++ BM = 256 ++ if K <= 4 and M >= 500_000: ++ num_warps = 2 ++ elif K <= 16 and M >= 5_000_000 and BN <= 16: ++ num_warps = 2 ++ elif K <= 16 and NB >= 17 and M >= 5_000_000: ++ # autotune (1,32,10M,64,8) BN=32 nw=2 wins ++ num_warps = 2 ++ else: ++ num_warps = 4 ++ ++ elif NB <= 256: ++ # Medium queries: BN must scale with M to control ++ # c-replication, but ONLY when K is very small (K<=4) or ++ # very large (K>=32). At K ∈ {5..16} the autotune consistently ++ # picks BN=64 even at M=10M, because K=8..16 hits a sweet ++ # spot where the topk-insert + Stage-2 cost of 1-wave ++ # BN=128 outweighs the c-HBM savings. ++ kernel_mode = "insert" ++ big_NB = (NB >= 192) ++ ++ # K=128 special-case: medium NB + K=128, autotune picks ++ # BN=8 BM=256 insert (NOT sortmerge!). The K-step argmin loop ++ # at BN=64 K=128 is too expensive -- better to chop N into ++ # many small tiles. ++ if K >= 128 and M <= 200_000: ++ BN = 8 ++ BM = 256 ++ elif M >= 5_000_000 and D <= 64 and (K <= 4 or K >= 32): ++ BN = 128 ++ elif big_NB and M >= 5_000_000: ++ BN = 128 ++ elif big_NB and D >= 256 and M <= 200_000 and K <= 4: ++ BN = 128 ++ elif D >= 256 and K <= 16 and (NB >= 64 and M >= 500_000): ++ # D=256 + K<=16 + (medium NB + medium M): BN=128 wins. ++ # autotune (1,128,1M,256,8) picks BN=128 BM=64. ++ BN = 128 ++ else: ++ BN = 64 ++ BN = min(BN, max(8, _next_pow2(N))) ++ # BM ++ if K >= 128: ++ BM = 256 # already covered above, keep consistent here ++ elif K <= 4 and D <= 128 and M <= 200_000: ++ BM = 256 ++ elif D >= 256 and K <= 16: ++ BM = 64 ++ elif K <= 32: ++ BM = 64 if (BN == 64 and K == 32 and D >= 256) else 128 ++ elif K <= 64: ++ BM = 128 ++ else: ++ BM = 128 ++ num_warps = 4 ++ ++ elif NB <= 2048: ++ kernel_mode = "insert" ++ if M >= 5_000_000: ++ BN = 128 ++ elif M >= 500_000 and D >= 256 and K <= 16: ++ BN = 128 ++ else: ++ BN = 64 ++ BN = min(BN, max(8, _next_pow2(N))) ++ if K <= 4 and M >= 500_000: ++ BM = 64 ++ elif K <= 16 and M >= 500_000 and D >= 128: ++ BM = 64 ++ elif K >= 64: ++ BM = 64 ++ else: ++ BM = 128 ++ num_warps = 4 ++ ++ else: # NB >= 8K (build / large-eval regime) ++ kernel_mode = "insert" ++ # autotune patterns: ++ # NB=50K-100K D<=128 K<=8 -> BN=128 (halves HBM) ++ # NB=10K K>=8 -> BN=64 (more topk work per row) ++ # NB>=10K D>=256 K<=16 -> BN=128 ++ # NB>=10K D>=256 K>=32 -> BN=64 ++ if D >= 256 and K <= 16: ++ BN = 128 ++ elif D >= 256 and K >= 32: ++ BN = 64 ++ elif K <= 4: ++ BN = 128 ++ elif K <= 8 and NB >= 30_000: ++ # Larger NB -> BN=128 saves enough HBM to overcome topk cost ++ BN = 128 ++ else: ++ BN = 64 ++ if K <= 4 and D >= 256: ++ BM = 64 # SMEM pressure at D=256 + BM=128 ++ elif K <= 4: ++ BM = 128 if (D <= 64 or BN == 128) else 64 ++ elif K <= 8 and BN == 128: ++ BM = 128 ++ elif K >= 64 and BN == 64 and D <= 128 and NB >= 30_000: ++ # klib carve-out: NB in [30K, ~200K] + BN=64 + K>=64 + ++ # D<=128 wins big at BM=128 splits=1 vs BM=64 splits=2. ++ # Empirically at (1, 48K, 48K, 64, 64) the upstream rule ++ # (BM=64, target_splits=2) runs 17.95 ms, while BM=128 ++ # splits=1 ns=2 runs 11.71 ms (-35%); at (1, 64K, 64K, 64, ++ # 64) 29.20 ms -> 18.91 ms (-35%). Bounded on K>=64 (K=32 ++ # autotune at this NB shows BM=64 splits=1 wins, ~1 ms over ++ # BM=128 splits=1) and on NB>=30K so the upstream rule ++ # still applies to smaller-NB shapes where its M-split ++ # target was directly autotune-derived. ++ BM = 128 ++ else: ++ BM = 64 ++ num_warps = 4 ++ ++ # ─────────────────────────────────────────────────────────────── ++ # Step 2: M_PER_SPLIT (target waves) ++ # ─────────────────────────────────────────────────────────────── ++ num_n_tiles = max(1, math.ceil(N / BN)) ++ ctas_no_split = num_n_tiles * B ++ ++ # Workaround for an insert-kernel correctness issue with the ++ # D-split path: when ``M_PER_SPLIT == BM`` (single M-loop ++ # iteration) and the d-loop has >=4 chunks, the result is wrong. ++ # Bumping the minimum mps to 2*BM keeps the M-loop at >=2 iters. ++ needs_min_2_iters = (D_INNER < D and (D + D_INNER - 1) // D_INNER >= 4) ++ min_mps = (2 * BM) if needs_min_2_iters else BM ++ ++ if is_large_n: ++ M_PER_SPLIT = ((M + BM - 1) // BM) * BM ++ elif ctas_no_split >= NUM_SMS * 8: ++ # Massively oversaturated (e.g. N=100K, BN=64, ctas=1563): one ++ # split fully utilises SMs. ++ M_PER_SPLIT = ((M + BM - 1) // BM) * BM ++ elif K >= 32 and BN == 64 and D <= 128 and NB >= 30_000 and not is_large_n: ++ # klib carve-out (pairs with the BM=128 K>=64 rule above ++ # plus K=32 BM=64 standalone): the K>=32 build regime at ++ # NB>=30K wants splits=1 -- 2 splits doubles the K-step ++ # argmin work per row and the Stage-2 reduce, neither of ++ # which the autotune at NB=10K K=8 accounted for. Empirically ++ # saves ~6 ms on (1, 48K, 48K, 64, 64) and ~10 ms on ++ # (1, 64K, 64K, 64, 64) vs target_splits=2, and ~1 ms on ++ # (1, 56K, 56K, 64, 32) vs target_splits=2. ++ M_PER_SPLIT = ((M + BM - 1) // BM) * BM ++ else: ++ if ctas_no_split >= NUM_SMS * 4: ++ # ctas >= 528: 2 splits for tail-fill at very large builds ++ # (autotune (1,100K,100K,64,4) BN=128 ctas=782 -> 2 splits) ++ target_splits = 2 ++ elif ctas_no_split >= NUM_SMS * 2: ++ # ctas in [264, 528): single split already 2-4 waves; adding ++ # splits hurts more than it helps. autotune ++ # (1,50K,50K,64,8) BN=128 ctas=391 -> 1 split. ++ target_splits = 1 ++ elif ctas_no_split >= NUM_SMS: ++ # ctas in [132, 264): 3 splits for better tail. autotune ++ # (1,10K,10K,64,16) BN=64 ctas=157 -> 3 splits. ++ target_splits = 3 ++ else: ++ # Under-saturated: target_waves by K (autotune-derived). ++ # Use FLOOR div: snaps to integer wave count, which matches ++ # autotune choices on shapes like NB=1024 K=32 ctas=16 ++ # (16 splits = 2 waves, 17 splits = 2.06 waves but bad ++ # tail). ++ if kernel_mode == "sortmerge": ++ target_waves = 2 ++ elif K >= 128 and NB >= 64: ++ # K=128 + medium NB: 2w wins (autotune ++ # (1,128,100K,128,128) picks 16 splits at ctas=16). ++ target_waves = 2 ++ elif K >= 128: ++ target_waves = 1 ++ elif K >= 64 and ctas_no_split <= 2: ++ # K=64 + tiny ctas: 1 wave wins (autotune ++ # (1,128,100K,128,64) picks 66 splits = 1w). ++ target_waves = 1 ++ elif K >= 64: ++ target_waves = 2 ++ elif K <= 4 and ctas_no_split == 1 and (NB <= 32 or M <= 1_000_000): ++ # K=4 single-CTA-stack: 4 waves at small NB or small M. ++ # autotune NB=1 K=4 -> 521 splits; (1,128,10M,64,4) ++ # however picks 264 splits (2w) -- too much M per CTA ++ # for more splits to help. ++ target_waves = 4 ++ elif K <= 8 and ctas_no_split >= 8 and ctas_no_split <= 32: ++ target_waves = 4 if M >= 500_000 else 2 ++ elif K <= 8 and ctas_no_split == 1 and M >= 5_000_000 and BN <= 32: ++ target_waves = 4 ++ elif B >= 2 and N <= 8 and K <= 8: ++ # Batched B>=2 + single-query + K<=8: 4w wins ++ target_waves = 4 ++ else: ++ target_waves = 2 ++ target_splits = max(1, NUM_SMS * target_waves // ctas_no_split) ++ target_splits = min(target_splits, max(1, math.ceil(M / BM))) ++ mps_raw = math.ceil(M / target_splits) ++ mps = math.ceil(mps_raw / BM) * BM ++ mps = max(min_mps, min(mps, ((M + BM - 1) // BM) * BM)) ++ mps = min(mps, MAX_MPS_TILES * BM) ++ M_PER_SPLIT = mps ++ ++ NUM_SPLITS = math.ceil(M / M_PER_SPLIT) ++ ++ # ─────────────────────────────────────────────────────────────── ++ # Step 3: NUM_STAGES_PIPE (M-loop pipeline depth) ++ # ─────────────────────────────────────────────────────────────── ++ # Derived from a 201-shape × 4-ns sweep on H200 (upstream ++ # ``bench/sweep_ns_by_k.py``), cross-checked against direct probes ++ # on N>=64 build shapes. Three competing effects: ++ # ++ # * Each extra pipeline stage adds another buffered C-tile, ++ # costing ``BM * D_INNER * 2`` bytes of SMEM/registers. Wide ++ # D-tiles amortise this well; narrow ones don't. ++ # * The K-step argmin inner loop holds live state proportional ++ # to ``BN * TOPK_PAD * 8`` bytes. For narrow rows (BN<=16) + ++ # large K, registers get tight and deeper pipelines spill. ++ # * Medium-N tiles (BN>=64) have abundant compute per tile to ++ # amortise prefetch, but get little benefit at very large K ++ # because the K loop dominates anyway. ++ # ++ # Empirical winners (mode_ns from the 201-shape sweep): ++ # ++ # D_INNER >= 256 (wide D) -> ns=2 (uniform across K) ++ # BN >= 128 (huge build N) -> ns=2 (registers tight) ++ # BN >= 64 + BM <= 128 + K <= 8 -> ns=3 (small tile + tiny K) ++ # BN >= 64 -> ns=2 (BM=256 SMEM-bound) ++ # BN ∈ [16,32] + BM=256 + K >= 64 -> ns=3 (big-tile large-K) ++ # otherwise (narrow tile + any K) -> ns=1 (avoid spills) ++ # ++ # Without this rule, K>=64 + D=256 shapes hit catastrophic ++ # regressions (e.g. (1,16,100K,256,128) 439us -> 2429us at ns=1, ++ # a 5.5x slowdown). Conversely K<=32 + huge-M shapes save 20-45% ++ # vs ns=2 (e.g. (1,1,10M,128,4) 1426us -> 785us at ns=1). ++ if D_INNER < D: ++ NUM_STAGES_PIPE = 1 # D-split path is hard-coded to ns=1 ++ elif (BN == 8 and BM == 64 and D_INNER == 256 ++ and K <= 4 and M >= 500_000): ++ # NB<=8 + K<=4 + D=256 + BM=64 + huge-M: ns=1 wins by 1.48-1.81x ++ # over the default ns=2. Direct verification at K ∈ {1, 2, 4} × ++ # M ∈ {1M, 5M, 10M, 30M}: (nw=4, ns=1) lifts %peak HBM from ++ # 39-48 % (ns=2) to 57-85 % (ns=1). The ns=2 default below was ++ # calibrated for wider-N tiles; at BN=8 the deeper pipeline ++ # spends too many registers on prefetch and the K=4 epilogue ++ # spills. Pairs with the nw=4 default the rule above falls into ++ # at D=256. ++ NUM_STAGES_PIPE = 1 ++ elif D_INNER >= 256: ++ NUM_STAGES_PIPE = 2 ++ elif BN >= 128: ++ NUM_STAGES_PIPE = 2 ++ elif BN >= 64 and BM <= 128 and K <= 8: ++ NUM_STAGES_PIPE = 3 ++ elif BN >= 64: ++ NUM_STAGES_PIPE = 2 ++ elif BN >= 16 and BM >= 256 and K >= 64: ++ NUM_STAGES_PIPE = 3 ++ else: ++ NUM_STAGES_PIPE = 1 ++ ++ cfg = { ++ "BN": BN, "BM": BM, "D_INNER": D_INNER, ++ "TOPK_PAD": ( ++ max(32, _next_pow2(K)) if kernel_mode == "sortmerge" ++ else _next_pow2(K) ++ ), ++ "kernel_mode": kernel_mode, "num_warps": num_warps, ++ "M_PER_SPLIT": M_PER_SPLIT, ++ "NUM_SPLITS": NUM_SPLITS, ++ "NUM_STAGES_PIPE": NUM_STAGES_PIPE, ++ } ++ ++ # Dtype-aware SMEM shrink. The shape-only heuristic above was tuned ++ # on bf16 and may pick (BN=128, BM=128, D_INNER=128) for shapes that ++ # cleanly fit ~210 KB at 2 bytes/elt but blow past the 227 KB H200 ++ # opt-in cap at 4 bytes/elt. The fitter shrinks BN/BM/NUM_STAGES_PIPE ++ # by powers of 2 until the kernel fits — leaving D_INNER and ++ # kernel_mode (which encode different perf regimes) alone. ++ if smem_limit is not None: ++ cfg = _fit_config_to_smem(cfg, D, K, dtype_bytes, smem_limit) ++ ++ return cfg ++ ++ ++# ── autotune (opt-in) ────────────────────────────────────────────────── ++ ++ ++_autotune_cache: dict = {} ++ ++ ++def _autotune(x, c, k, *, force_path=None): ++ """Brute-force autotune over the candidate config grid. ++ ++ Cached per ``(B, N, M, D, k, dtype, force_path)`` shape; first call ++ costs ~30-60 s of compile time, subsequent calls hit the cache in ++ sub-ms. ++ ++ The candidate grid is dtype-aware: ``_gen_configs`` filters out ++ configs whose SMEM exceeds the device opt-in limit at this dtype, ++ so fp32 inputs at large D never have an unrunnable config in the ++ search space. ++ """ ++ B, N, D = x.shape ++ M = c.shape[1] ++ dtype_bytes = x.element_size() ++ key = (B, N, M, D, k, x.dtype, force_path) ++ ++ if key in _autotune_cache: ++ return _autotune_cache[key] ++ ++ smem_limit = _smem_limit(x.device) ++ configs = _gen_configs(D, k, ++ dtype_bytes=dtype_bytes, ++ smem_limit=min(smem_limit, 220_000)) ++ if not configs: ++ raise RuntimeError( ++ f"No valid flash_knn config for D={D}, K={k}, " ++ f"dtype={x.dtype} (element_size={dtype_bytes})" ++ ) ++ ++ # Drop configs with BN > next_pow2(N) -- they would just pad N to BN ++ # and run the same kernel slower. ++ max_bn = max(16, _next_pow2(N)) ++ configs = [cfg for cfg in configs if cfg["BN"] <= max_bn] ++ ++ best_time = float('inf') ++ best_cfg = None ++ max_partial_bytes = 4 * 1024**3 ++ ++ for cfg in configs: ++ bn = cfg["BN"] ++ bm = cfg["BM"] ++ d_inner = cfg["D_INNER"] ++ topk_pad = cfg["TOPK_PAD"] ++ kernel_mode = cfg["kernel_mode"] ++ nw = cfg["num_warps"] ++ ns_pipe = cfg.get("NUM_STAGES_PIPE", 2) ++ ++ if force_path == "large_n": ++ mps_list = [((M + bm - 1) // bm) * bm] ++ else: ++ mps_list = _gen_m_splits(M, bm, B=B, N=N, BN=bn) ++ ++ for mps in mps_list: ++ num_splits = math.ceil(M / mps) ++ partial_bytes = B * num_splits * N * k * 8 ++ if partial_bytes > max_partial_bytes: ++ continue ++ ++ try: ++ pv = torch.empty((B, N, num_splits, k), ++ device=x.device, dtype=torch.float32) ++ pi = torch.empty((B, N, num_splits, k), ++ device=x.device, dtype=torch.int32) ++ grid = (num_splits, math.ceil(N / bn), B) ++ pv_s0, pv_s1, pv_s2, pv_s3 = pv.stride() ++ pi_s0, pi_s1, pi_s2, pi_s3 = pi.stride() ++ ++ if kernel_mode == "sortmerge": ++ def run(grid=grid, bn=bn, bm=bm, d_inner=d_inner, ++ topk_pad=topk_pad, mps=mps, nw=nw, ++ num_splits=num_splits, pv=pv, pi=pi, ++ ns_pipe=ns_pipe, ++ pv_s0=pv_s0, pv_s1=pv_s1, pv_s2=pv_s2, pv_s3=pv_s3, ++ pi_s0=pi_s0, pi_s1=pi_s1, pi_s2=pi_s2, pi_s3=pi_s3): ++ _flash_knn_sortmerge_kernel[grid]( ++ x, c, pv, pi, ++ x.stride(0), x.stride(1), x.stride(2), ++ c.stride(0), c.stride(1), c.stride(2), ++ pv_s0, pv_s2, pv_s1, pv_s3, ++ pi_s0, pi_s2, pi_s1, pi_s3, ++ N=N, M=M, D=D, K=k, M_PER_SPLIT=mps, ++ BN=bn, BM=bm, D_INNER=d_inner, ++ TOPK_PAD=topk_pad, ++ NUM_STAGES_PIPE=ns_pipe, ++ num_warps=nw, ++ ) ++ if num_splits > 1: ++ pv_flat = pv.view(B, N, -1) ++ pv_flat.topk(k, dim=-1, largest=False) ++ else: ++ max_steps = min(k, bm) ++ def run(grid=grid, bn=bn, bm=bm, d_inner=d_inner, ++ topk_pad=topk_pad, mps=mps, nw=nw, ++ num_splits=num_splits, pv=pv, pi=pi, ++ max_steps=max_steps, ns_pipe=ns_pipe, ++ pv_s0=pv_s0, pv_s1=pv_s1, pv_s2=pv_s2, pv_s3=pv_s3, ++ pi_s0=pi_s0, pi_s1=pi_s1, pi_s2=pi_s2, pi_s3=pi_s3): ++ _flash_knn_insert_kernel[grid]( ++ x, c, pv, pi, ++ x.stride(0), x.stride(1), x.stride(2), ++ c.stride(0), c.stride(1), c.stride(2), ++ pv_s0, pv_s2, pv_s1, pv_s3, ++ pi_s0, pi_s2, pi_s1, pi_s3, ++ N=N, M=M, D=D, K=k, M_PER_SPLIT=mps, ++ BN=bn, BM=bm, D_INNER=d_inner, ++ TOPK_PAD=topk_pad, MAX_STEPS=max_steps, ++ NUM_STAGES_PIPE=ns_pipe, ++ num_warps=nw, ++ ) ++ if num_splits > 1: ++ pv_flat = pv.view(B, N, -1) ++ pv_flat.topk(k, dim=-1, largest=False) ++ ++ t = _bench_quick(run) ++ if t < best_time: ++ best_time = t ++ best_cfg = {**cfg, "M_PER_SPLIT": mps, ++ "NUM_SPLITS": num_splits, ++ "NUM_STAGES_PIPE": ns_pipe} ++ ++ del pv, pi ++ except Exception: ++ continue ++ ++ if best_cfg is None: ++ raise RuntimeError( ++ f"All flash_knn configs failed for B={B}, N={N}, M={M}, D={D}, K={k}") ++ ++ _autotune_cache[key] = best_cfg ++ return best_cfg ++ ++ ++# ── runner ───────────────────────────────────────────────────────────── ++ ++ ++_OOR_FALLBACK_CACHE: dict = {} ++ ++ ++def _try_launch(x, c, cfg, B, N, M, D, k): ++ """Single launch attempt; returns idxs or raises ``triton.compiler.errors.OutOfResources``. ++ ++ Factored out of ``_run`` so the OOR-shrink retry loop can call it ++ repeatedly with a shrunk cfg. ++ """ ++ bn = cfg["BN"] ++ bm = cfg["BM"] ++ d_inner = cfg["D_INNER"] ++ topk_pad = cfg["TOPK_PAD"] ++ mps = cfg["M_PER_SPLIT"] ++ num_splits = cfg["NUM_SPLITS"] ++ kernel_mode = cfg["kernel_mode"] ++ nw = cfg["num_warps"] ++ num_stages_pipe = cfg.get("NUM_STAGES_PIPE", 2) ++ ++ partial_vals = torch.empty((B, N, num_splits, k), ++ device=x.device, dtype=torch.float32) ++ partial_idxs = torch.empty((B, N, num_splits, k), ++ device=x.device, dtype=torch.int32) ++ grid = (num_splits, math.ceil(N / bn), B) ++ ++ pv_s0, pv_s1, pv_s2, pv_s3 = partial_vals.stride() ++ pi_s0, pi_s1, pi_s2, pi_s3 = partial_idxs.stride() ++ ++ if kernel_mode == "sortmerge": ++ _flash_knn_sortmerge_kernel[grid]( ++ x, c, partial_vals, partial_idxs, ++ x.stride(0), x.stride(1), x.stride(2), ++ c.stride(0), c.stride(1), c.stride(2), ++ pv_s0, pv_s2, pv_s1, pv_s3, ++ pi_s0, pi_s2, pi_s1, pi_s3, ++ N=N, M=M, D=D, K=k, M_PER_SPLIT=mps, ++ BN=bn, BM=bm, D_INNER=d_inner, ++ TOPK_PAD=topk_pad, ++ NUM_STAGES_PIPE=num_stages_pipe, ++ num_warps=nw, ++ ) ++ else: ++ max_steps = min(k, bm) ++ _flash_knn_insert_kernel[grid]( ++ x, c, partial_vals, partial_idxs, ++ x.stride(0), x.stride(1), x.stride(2), ++ c.stride(0), c.stride(1), c.stride(2), ++ pv_s0, pv_s2, pv_s1, pv_s3, ++ pi_s0, pi_s2, pi_s1, pi_s3, ++ N=N, M=M, D=D, K=k, M_PER_SPLIT=mps, ++ BN=bn, BM=bm, D_INNER=d_inner, ++ TOPK_PAD=topk_pad, MAX_STEPS=max_steps, ++ NUM_STAGES_PIPE=num_stages_pipe, ++ num_warps=nw, ++ ) ++ ++ if num_splits == 1: ++ return partial_idxs[:, :, 0, :].contiguous() ++ ++ pv = partial_vals.view(B, N, -1) ++ pi = partial_idxs.view(B, N, -1) ++ _, sel = pv.topk(k, dim=-1, largest=False, sorted=True) ++ out_idxs = pi.gather(-1, sel.to(torch.int64)).to(torch.int32) ++ return out_idxs ++ ++ ++def _shrink_cfg_on_oor(cfg: dict) -> Optional[dict]: ++ """Halve the cheapest SMEM-bearing axis. ``None`` when nothing left. ++ ++ Order: ++ 1. NUM_STAGES_PIPE 3 → 2 → 1 (drops the c-tile pipeline buffer ++ — usually free or nearly-free at this depth). ++ 2. Tend toward a near-square tile by **halving the larger of ++ (BN, BM)** first. WGMMA shape efficiency drops rapidly once ++ BM<32 on Hopper (MMA atom is 64x16x16), so we prefer to ++ shrink BN when BN > BM. For BN=128 BM=64 (typical build ++ regime) this gives BN=64 BM=64 in one step, which empirically ++ matches the kmeans-style ``_fit_config_to_smem`` choice and ++ beats BN=128 BM=32 by ~20 % at D=1024 fp32. ++ 3. Floors: BN ≥ 8, BM ≥ 32 (both required by the WGMMA atom). ++ ++ sortmerge requires ``BM == TOPK_PAD``; skip BM shrinks and only ++ halve BN. ++ ++ Also recomputes ``M_PER_SPLIT`` so it remains a multiple of the ++ new BM (the kernel's inner ``tl.range`` walk requires this). ++ """ ++ ns = int(cfg.get("NUM_STAGES_PIPE", 2)) ++ bm = int(cfg["BM"]) ++ bn = int(cfg["BN"]) ++ sortmerge = cfg.get("kernel_mode") == "sortmerge" ++ ++ out = dict(cfg) ++ if ns > 1: ++ out["NUM_STAGES_PIPE"] = ns - 1 ++ return out ++ if sortmerge: ++ if bn > 8: ++ out["BN"] = bn // 2 ++ return out ++ return None ++ # Non-sortmerge: prefer shrinking the larger axis. ++ if bn > bm and bn > 8: ++ out["BN"] = bn // 2 ++ return out ++ if bm > 32: ++ out["BM"] = bm // 2 ++ new_bm = out["BM"] ++ out["M_PER_SPLIT"] = ((int(out["M_PER_SPLIT"]) + new_bm - 1) ++ // new_bm) * new_bm ++ return out ++ if bn > 8: ++ out["BN"] = bn // 2 ++ return out ++ return None ++ ++ ++def _run(x: torch.Tensor, c: torch.Tensor, k: int, *, ++ force_path: Optional[str] = None, ++ autotune: bool = False) -> torch.Tensor: ++ """Shared dispatcher -- launch stage 1, optional stage 2 reduce, return idxs. ++ ++ Returns indices only ``(B, N, k) int32``. The wrapper at ++ :func:`knnlib._kernels.primitives.knn.flash_knn` calls the gather kernel to ++ produce true squared distances per neighbour. ++ ++ On ``OutOfResources`` at launch, we shrink one SMEM-bearing axis ++ (NS → BM → BN) and retry, caching the surviving config per ++ ``(D, K, dtype, force_path)`` so subsequent calls with the same ++ shape skip the failed compiles. This is the source of truth for ++ "does this fit?"; ``_estimate_sram`` only provides a coarse hint ++ for the autotune candidate filter. ++ """ ++ assert x.is_cuda and c.is_cuda ++ assert x.dtype == c.dtype ++ B, N, D = x.shape ++ M = c.shape[1] ++ assert c.shape == (B, M, D) ++ assert 1 <= k <= M ++ ++ if autotune: ++ cfg = _autotune(x, c, k, force_path=force_path) ++ else: ++ cfg = _heuristic_config( ++ B, N, M, D, k, ++ force_path=force_path, ++ dtype_bytes=x.element_size(), ++ smem_limit=min(_smem_limit(x.device), 220_000), ++ ) ++ ++ # Surviving-config cache keyed on the parts the dispatcher can't ++ # vary (D, K, dtype, force_path). If a previous call already ++ # shrunk past an OOR for this shape, start with that smaller cfg. ++ fb_key = (D, k, x.dtype, force_path, ++ cfg["kernel_mode"], cfg.get("D_INNER")) ++ if fb_key in _OOR_FALLBACK_CACHE: ++ cached = _OOR_FALLBACK_CACHE[fb_key] ++ # Use the cached cfg only if it's smaller than the current one ++ # on at least one SMEM-bearing axis (otherwise the heuristic ++ # already picked something safe for this shape). ++ sm_cur = (cfg["BN"], cfg["BM"], cfg.get("NUM_STAGES_PIPE", 2)) ++ sm_cached = (cached["BN"], cached["BM"], ++ cached.get("NUM_STAGES_PIPE", 2)) ++ if sm_cached < sm_cur: ++ cfg = {**cfg, **{k_: cached[k_] for k_ in ++ ("BN", "BM", "NUM_STAGES_PIPE")}} ++ # Re-align mps to new BM ++ new_bm = cfg["BM"] ++ cfg["M_PER_SPLIT"] = ((cfg["M_PER_SPLIT"] + new_bm - 1) ++ // new_bm) * new_bm ++ cfg["NUM_SPLITS"] = math.ceil(M / cfg["M_PER_SPLIT"]) ++ ++ # Import lazily so we don't pay it on every dispatch. ++ try: ++ from triton.runtime.errors import OutOfResources ++ except ImportError: ++ try: ++ from triton.compiler.errors import OutOfResources ++ except ImportError: ++ OutOfResources = RuntimeError # safest fallback ++ ++ last_err = None ++ initial_cfg = cfg ++ shrunk = False ++ for _attempt in range(8): ++ try: ++ result = _try_launch(x, c, cfg, B, N, M, D, k) ++ # Cache the surviving cfg so subsequent calls with the ++ # same (D, K, dtype, ...) shape skip the failed compile. ++ if shrunk: ++ _OOR_FALLBACK_CACHE[fb_key] = { ++ "BN": cfg["BN"], "BM": cfg["BM"], ++ "NUM_STAGES_PIPE": cfg.get("NUM_STAGES_PIPE", 2), ++ } ++ return result ++ except OutOfResources as e: ++ last_err = e ++ new_cfg = _shrink_cfg_on_oor(cfg) ++ if new_cfg is None: ++ break ++ cfg = new_cfg ++ shrunk = True ++ continue ++ ++ raise last_err if last_err is not None else RuntimeError( ++ "flash_knn: exhausted OOR shrink attempts") ++ ++ ++# ── public API ───────────────────────────────────────────────────────── ++ ++ ++def flash_knn_triton_small_n(x: torch.Tensor, c: torch.Tensor, k: int, ++ *, autotune: bool = False) -> torch.Tensor: ++ """Force the M-split (search) path -- ``(B, N, k) int32`` indices.""" ++ return _run(x, c, k, force_path=None, autotune=autotune) ++ ++ ++def flash_knn_triton_large_n(x: torch.Tensor, c: torch.Tensor, k: int, ++ *, autotune: bool = False) -> torch.Tensor: ++ """Force the single-pass (build) path -- ``(B, N, k) int32`` indices.""" ++ return _run(x, c, k, force_path="large_n", autotune=autotune) ++ ++ ++def flash_knn_triton(x: torch.Tensor, c: torch.Tensor, k: int, ++ *, autotune: bool = False) -> torch.Tensor: ++ """Universal Triton dispatch -- the heuristic picks BN/BM/mode/ ++ mps/ns_pipe based on shape, including whether to single-pass or ++ M-split. The per-CTA-count check inside :func:`_heuristic_config` ++ handles both build and eval shapes without a host-side forced path. ++ ++ Neither path materialises an N×M cross or distance matrix to HBM, ++ and neither computes or loads ``x_sq`` -- the two hard contracts ++ klib's KNN imposes on every Triton entry point here. ++ ++ Args: ++ x: ``(B, N, D)`` bf16 / fp16 / fp32 query tensor. ++ c: ``(B, M, D)`` corpus, same dtype. ++ k: number of neighbours. ++ autotune: ``False`` (default) uses the shape-only heuristic -- ++ first call pays a single Triton compile (~0.5 s). ``True`` ++ runs the full brute-force sweep + caches per shape (~30-90 s ++ first call). ++ ++ Returns: ++ idxs: ``(B, N, k)`` int32 -- nearest-neighbour indices, sorted ++ by ascending true squared L2 distance (ties broken by index). ++ """ ++ return _run(x, c, k, force_path=None, autotune=autotune) +diff --git a/knnlib/_kernels/primitives/knn/triton/insert.py b/knnlib/_kernels/primitives/knn/triton/insert.py +new file mode 100644 +index 0000000..7c60023 +--- /dev/null ++++ b/knnlib/_kernels/primitives/knn/triton/insert.py +@@ -0,0 +1,204 @@ ++"""flash_knn — Triton iterative-insert top-K kernel (x²-free, signed score). ++ ++Same x²-free signed score as ``sortmerge.py`` — see that module's docstring ++for the rationale. The insert kernel differs in two ways: ++ ++ * ``BM`` is decoupled from ``TOPK_PAD`` (sortmerge requires ``BM == TOPK_PAD``). ++ The autotune sweeps ``BM ∈ {64, 128, 256}`` independently. ++ * Top-K is maintained as an unsorted ``(BN, TOPK_PAD)`` register array; ++ each chunk does at most ``MAX_STEPS = min(K, BM)`` argmin->insert ++ passes. fp32 ``<`` works natively on signed scores, so only the ++ trailing output sort uses the sortable-u32 transform. ++ ++Empirically the insert kernel wins on virtually every shape past the ++trivial ``K == 1`` case; sortmerge is kept for completeness (autotune ++still considers it, and it's the natural fit when ``K`` matches a ++hardware-sortable tile size). ++ ++Public name: ``_flash_knn_insert_kernel`` — called from ++``klib/primitives/knn/triton/dispatch.py``. ++""" ++import triton ++import triton.language as tl ++ ++from knnlib._kernels.primitives.knn.triton._common import ( ++ _fp32_to_sortable_u32, ++ _sortable_u32_to_fp32, ++) ++ ++ ++@triton.jit ++def _flash_knn_insert_kernel( ++ x_ptr, c_ptr, ++ partial_val_ptr, partial_idx_ptr, ++ stride_x_b, stride_x_n, stride_x_d, ++ stride_c_b, stride_c_m, stride_c_d, ++ stride_pv_b, stride_pv_s, stride_pv_n, stride_pv_k, ++ stride_pi_b, stride_pi_s, stride_pi_n, stride_pi_k, ++ N: tl.constexpr, M: tl.constexpr, ++ D: tl.constexpr, K: tl.constexpr, ++ M_PER_SPLIT: tl.constexpr, ++ BN: tl.constexpr, BM: tl.constexpr, ++ D_INNER: tl.constexpr, ++ TOPK_PAD: tl.constexpr, ++ MAX_STEPS: tl.constexpr, ++ NUM_STAGES_PIPE: tl.constexpr = 2, ++): ++ """Iterative-insert top-K on signed shifted distance. ++ ++ Grid: ``(num_m_splits, ceil(N/BN), B)``. ++ Two compile-time paths gated on ``D_INNER >= D`` (persistent x + ++ pipelined M-loop vs D-split with sequential M-loop). The M-loop ++ pipelining depth is exposed as ``NUM_STAGES_PIPE`` so the ++ dispatcher can tune it per shape (see :mod:`dispatch` for the ++ heuristic rules; defaults to 2 if not provided). ++ """ ++ pid_s = tl.program_id(0) ++ pid_n = tl.program_id(1) ++ pid_b = tl.program_id(2) ++ pid_b = pid_b.to(tl.int64) ++ ++ n_start = pid_n * BN ++ n_offs = (n_start + tl.arange(0, BN)).to(tl.int64) ++ n_mask = n_offs < N ++ ++ topk_vals = tl.full([BN, TOPK_PAD], float('inf'), dtype=tl.float32) ++ topk_idxs = tl.full([BN, TOPK_PAD], -1, dtype=tl.int32) ++ topk_max = tl.full([BN], float('inf'), dtype=tl.float32) ++ k_range = tl.arange(0, TOPK_PAD) ++ bm_range = tl.arange(0, BM) ++ m_base = pid_s.to(tl.int64) * M_PER_SPLIT ++ ++ if D_INNER >= D: ++ d_offs = tl.arange(0, D_INNER).to(tl.int64) ++ d_mask = d_offs < D ++ x_tile = tl.load( ++ x_ptr + pid_b * stride_x_b ++ + n_offs[:, None] * stride_x_n ++ + d_offs[None, :] * stride_x_d, ++ mask=n_mask[:, None] & d_mask[None, :], other=0.0) ++ ++ for m_local in tl.range(0, M_PER_SPLIT, BM, num_stages=NUM_STAGES_PIPE): ++ m_start = m_base + m_local ++ m_offs = m_start + bm_range.to(tl.int64) ++ m_mask = m_offs < M ++ ++ c_tile = tl.load( ++ c_ptr + pid_b * stride_c_b ++ + m_offs[:, None] * stride_c_m ++ + d_offs[None, :] * stride_c_d, ++ mask=m_mask[:, None] & d_mask[None, :], other=0.0) ++ ++ c_f = c_tile.to(tl.float32) ++ c_sq_tile = tl.sum(c_f * c_f, axis=1) ++ cross = tl.dot(x_tile, tl.trans(c_tile)).to(tl.float32) ++ score = c_sq_tile[None, :] - 2.0 * cross ++ score = tl.where(m_mask[None, :], score, float('inf')) ++ ++ chunk_best = tl.min(score) ++ threshold_worst = tl.max(topk_max) ++ if chunk_best < threshold_worst: ++ _active = tl.full([1], 1, dtype=tl.int32) ++ for _step in range(MAX_STEPS): ++ if tl.max(_active) > 0: ++ row_min = tl.min(score, axis=1) ++ row_argmin = tl.argmin(score, axis=1) ++ do_insert = row_min < topk_max ++ n_inserts = tl.sum(do_insert.to(tl.int32)) ++ if n_inserts > 0: ++ topk_argmax = tl.argmax(topk_vals, axis=1) ++ replace_mask = (k_range[None, :] == topk_argmax[:, None]) ++ insert_mask = do_insert[:, None] & replace_mask ++ topk_vals = tl.where(insert_mask, row_min[:, None], topk_vals) ++ topk_idxs = tl.where( ++ insert_mask, ++ (m_start + row_argmin.to(tl.int64))[:, None].to(tl.int32), ++ topk_idxs, ++ ) ++ topk_max = tl.max(topk_vals, axis=1) ++ used_mask = (bm_range[None, :] == row_argmin[:, None]) ++ score = tl.where(used_mask & do_insert[:, None], float('inf'), score) ++ _active = tl.where( ++ n_inserts > 0, ++ tl.full([1], 1, dtype=tl.int32), ++ tl.full([1], 0, dtype=tl.int32), ++ ) ++ ++ else: ++ for m_local in tl.range(0, M_PER_SPLIT, BM, num_stages=1): ++ m_start = m_base + m_local ++ m_offs = m_start + bm_range.to(tl.int64) ++ m_mask = m_offs < M ++ ++ cross = tl.zeros([BN, BM], dtype=tl.float32) ++ c_sq_tile = tl.zeros([BM], dtype=tl.float32) ++ for d_start in range(0, D, D_INNER): ++ d_offs = (d_start + tl.arange(0, D_INNER)).to(tl.int64) ++ d_mask = d_offs < D ++ ++ x_sub = tl.load( ++ x_ptr + pid_b * stride_x_b ++ + n_offs[:, None] * stride_x_n ++ + d_offs[None, :] * stride_x_d, ++ mask=n_mask[:, None] & d_mask[None, :], other=0.0) ++ ++ c_sub = tl.load( ++ c_ptr + pid_b * stride_c_b ++ + m_offs[:, None] * stride_c_m ++ + d_offs[None, :] * stride_c_d, ++ mask=m_mask[:, None] & d_mask[None, :], other=0.0) ++ ++ cross += tl.dot(x_sub, tl.trans(c_sub)).to(tl.float32) ++ c_f = c_sub.to(tl.float32) ++ c_sq_tile += tl.sum(c_f * c_f, axis=1) ++ ++ score = c_sq_tile[None, :] - 2.0 * cross ++ score = tl.where(m_mask[None, :], score, float('inf')) ++ ++ chunk_best = tl.min(score) ++ threshold_worst = tl.max(topk_max) ++ if chunk_best < threshold_worst: ++ _active = tl.full([1], 1, dtype=tl.int32) ++ for _step in range(MAX_STEPS): ++ if tl.max(_active) > 0: ++ row_min = tl.min(score, axis=1) ++ row_argmin = tl.argmin(score, axis=1) ++ do_insert = row_min < topk_max ++ n_inserts = tl.sum(do_insert.to(tl.int32)) ++ if n_inserts > 0: ++ topk_argmax = tl.argmax(topk_vals, axis=1) ++ replace_mask = (k_range[None, :] == topk_argmax[:, None]) ++ insert_mask = do_insert[:, None] & replace_mask ++ topk_vals = tl.where(insert_mask, row_min[:, None], topk_vals) ++ topk_idxs = tl.where( ++ insert_mask, ++ (m_start + row_argmin.to(tl.int64))[:, None].to(tl.int32), ++ topk_idxs, ++ ) ++ topk_max = tl.max(topk_vals, axis=1) ++ used_mask = (bm_range[None, :] == row_argmin[:, None]) ++ score = tl.where(used_mask & do_insert[:, None], float('inf'), score) ++ _active = tl.where( ++ n_inserts > 0, ++ tl.full([1], 1, dtype=tl.int32), ++ tl.full([1], 0, dtype=tl.int32), ++ ) ++ ++ val_sortable = _fp32_to_sortable_u32(topk_vals) ++ val_bits = val_sortable.to(tl.uint64) ++ idx_bits = topk_idxs.to(tl.uint32).to(tl.uint64) ++ packed = (val_bits << 32) | idx_bits ++ packed = tl.sort(packed, dim=1) ++ topk_score = _sortable_u32_to_fp32((packed >> 32).to(tl.uint32)) ++ topk_idx = (packed & 0xFFFFFFFF).to(tl.int32) ++ ++ k_offs = tl.arange(0, TOPK_PAD).to(tl.int64) ++ k_mask = k_offs < K ++ write_mask = n_mask[:, None] & k_mask[None, :] ++ pid_s_i64 = pid_s.to(tl.int64) ++ tl.store(partial_val_ptr + pid_b * stride_pv_b + pid_s_i64 * stride_pv_s ++ + n_offs[:, None] * stride_pv_n + k_offs[None, :] * stride_pv_k, ++ topk_score, mask=write_mask) ++ tl.store(partial_idx_ptr + pid_b * stride_pi_b + pid_s_i64 * stride_pi_s ++ + n_offs[:, None] * stride_pi_n + k_offs[None, :] * stride_pi_k, ++ topk_idx, mask=write_mask) +diff --git a/knnlib/_kernels/primitives/knn/triton/sortmerge.py b/knnlib/_kernels/primitives/knn/triton/sortmerge.py +new file mode 100644 +index 0000000..569caba +--- /dev/null ++++ b/knnlib/_kernels/primitives/knn/triton/sortmerge.py +@@ -0,0 +1,181 @@ ++"""flash_knn — Triton sort-merge top-K kernel (x²-free, signed score). ++ ++The kernel computes ``argmin-K`` over the *signed* shifted distance ++ ++ s(n, m) = ||c[m]||² − 2·⟨x[n], c[m]⟩ (== ||x[n] − c[m]||² − ||x[n]||²) ++ ++The ``||x||²`` term is constant per row and so does not affect the top-K ++selection over ``m``. Dropping it eliminates the ``x_sq`` HBM tensor, its ++precompute pass, its per-tile load, one fp32 ADD per accumulator element, ++and the ``dist >= 0`` underflow clamp. The trade-off is that ``s`` is no ++longer non-negative; the packed-uint64 sort-merge top-K therefore uses ++an IEEE-sortable u32 transform (positives flip just the sign bit, negatives ++flip every bit) so ascending uint32 order matches ascending fp32 order on ++signed inputs. ++ ++Public name: ``_flash_knn_sortmerge_kernel`` — called from ++``klib/primitives/knn/triton/dispatch.py``. Stage 1 writes signed ++fp32 scores + int32 idxs to a partial buffer; Stage 2 either returns ++``partial[:, 0]`` directly (single split) or runs ``torch.topk(..., ++largest=False)`` to reduce across splits. ++ ++Same kernel covers both ``small-N M-split flash-decode`` (multiple ++splits per query block) and ``large-N single-pass`` (M_PER_SPLIT == M) ++— the host-side dispatcher just chooses the grid + M_PER_SPLIT. ++""" ++import triton ++import triton.language as tl ++ ++from knnlib._kernels.primitives.knn.triton._common import ( ++ _INF_PACKED, ++ _fp32_to_sortable_u32, ++ _sortable_u32_to_fp32, ++) ++ ++ ++@triton.jit ++def _flash_knn_sortmerge_kernel( ++ x_ptr, c_ptr, ++ partial_val_ptr, partial_idx_ptr, ++ stride_x_b, stride_x_n, stride_x_d, ++ stride_c_b, stride_c_m, stride_c_d, ++ stride_pv_b, stride_pv_s, stride_pv_n, stride_pv_k, ++ stride_pi_b, stride_pi_s, stride_pi_n, stride_pi_k, ++ N: tl.constexpr, M: tl.constexpr, ++ D: tl.constexpr, K: tl.constexpr, ++ M_PER_SPLIT: tl.constexpr, ++ BN: tl.constexpr, BM: tl.constexpr, ++ D_INNER: tl.constexpr, ++ TOPK_PAD: tl.constexpr, ++ NUM_STAGES_PIPE: tl.constexpr = 2, ++): ++ """Packed-uint64 sort-merge top-K with IEEE-sortable u32 score bits. ++ ++ Grid: ``(num_m_splits, ceil(N/BN), B)``. ``BM == TOPK_PAD``. ++ ++ Two compile-time paths gated on ``D_INNER >= D``: ++ * persistent x tile + pipelined M-loop (single tile covers all of D) ++ * D-split with sequential M-loop (D-chunked GEMM for wide D) ++ ++ ``NUM_STAGES_PIPE`` controls the M-loop ``tl.range`` pipelining ++ depth (host-side dispatch tunes it per shape; defaults to 2 if not ++ provided -- matches the pre-port hard-coded value). ++ """ ++ pid_s = tl.program_id(0) ++ pid_n = tl.program_id(1) ++ pid_b = tl.program_id(2) ++ pid_b = pid_b.to(tl.int64) ++ ++ n_start = pid_n * BN ++ n_offs = (n_start + tl.arange(0, BN)).to(tl.int64) ++ n_mask = n_offs < N ++ ++ topk_packed = tl.full([BN, TOPK_PAD], _INF_PACKED, dtype=tl.uint64) ++ bm_range = tl.arange(0, BM) ++ m_base = pid_s.to(tl.int64) * M_PER_SPLIT ++ ++ if D_INNER >= D: ++ d_offs = tl.arange(0, D_INNER).to(tl.int64) ++ d_mask = d_offs < D ++ x_tile = tl.load( ++ x_ptr + pid_b * stride_x_b ++ + n_offs[:, None] * stride_x_n ++ + d_offs[None, :] * stride_x_d, ++ mask=n_mask[:, None] & d_mask[None, :], other=0.0) ++ ++ for m_local in tl.range(0, M_PER_SPLIT, BM, num_stages=NUM_STAGES_PIPE): ++ m_start = m_base + m_local ++ m_offs = m_start + bm_range.to(tl.int64) ++ m_mask = m_offs < M ++ ++ c_tile = tl.load( ++ c_ptr + pid_b * stride_c_b ++ + m_offs[:, None] * stride_c_m ++ + d_offs[None, :] * stride_c_d, ++ mask=m_mask[:, None] & d_mask[None, :], other=0.0) ++ ++ c_f = c_tile.to(tl.float32) ++ c_sq_tile = tl.sum(c_f * c_f, axis=1) ++ cross = tl.dot(x_tile, tl.trans(c_tile)).to(tl.float32) ++ score = c_sq_tile[None, :] - 2.0 * cross ++ score = tl.where(m_mask[None, :], score, float('inf')) ++ ++ chunk_best = tl.min(score) ++ topk_worst_packed = tl.max(topk_packed) ++ topk_worst_val = _sortable_u32_to_fp32( ++ (topk_worst_packed >> 32).to(tl.uint32)) ++ if chunk_best < topk_worst_val: ++ score_sortable = _fp32_to_sortable_u32(score) ++ idx_vals = (m_start + bm_range.to(tl.int64)).to(tl.int32) ++ idx_bits = idx_vals.to(tl.uint32).to(tl.uint64) ++ chunk_packed = (score_sortable.to(tl.uint64) << 32) | idx_bits[None, :] ++ ++ chunk_packed = chunk_packed ^ tl.full([1], 0xFFFFFFFFFFFFFFFF, dtype=tl.uint64) ++ chunk_packed = tl.sort(chunk_packed, dim=1) ++ chunk_packed = chunk_packed ^ tl.full([1], 0xFFFFFFFFFFFFFFFF, dtype=tl.uint64) ++ ++ merged = tl.minimum(topk_packed, chunk_packed) ++ topk_packed = tl.sort(merged, dim=1) ++ ++ else: ++ for m_local in tl.range(0, M_PER_SPLIT, BM, num_stages=1): ++ m_start = m_base + m_local ++ m_offs = m_start + bm_range.to(tl.int64) ++ m_mask = m_offs < M ++ ++ cross = tl.zeros([BN, BM], dtype=tl.float32) ++ c_sq_tile = tl.zeros([BM], dtype=tl.float32) ++ for d_start in range(0, D, D_INNER): ++ d_offs = (d_start + tl.arange(0, D_INNER)).to(tl.int64) ++ d_mask = d_offs < D ++ ++ x_sub = tl.load( ++ x_ptr + pid_b * stride_x_b ++ + n_offs[:, None] * stride_x_n ++ + d_offs[None, :] * stride_x_d, ++ mask=n_mask[:, None] & d_mask[None, :], other=0.0) ++ ++ c_sub = tl.load( ++ c_ptr + pid_b * stride_c_b ++ + m_offs[:, None] * stride_c_m ++ + d_offs[None, :] * stride_c_d, ++ mask=m_mask[:, None] & d_mask[None, :], other=0.0) ++ ++ cross += tl.dot(x_sub, tl.trans(c_sub)).to(tl.float32) ++ c_f = c_sub.to(tl.float32) ++ c_sq_tile += tl.sum(c_f * c_f, axis=1) ++ ++ score = c_sq_tile[None, :] - 2.0 * cross ++ score = tl.where(m_mask[None, :], score, float('inf')) ++ ++ chunk_best = tl.min(score) ++ topk_worst_packed = tl.max(topk_packed) ++ topk_worst_val = _sortable_u32_to_fp32( ++ (topk_worst_packed >> 32).to(tl.uint32)) ++ if chunk_best < topk_worst_val: ++ score_sortable = _fp32_to_sortable_u32(score) ++ idx_vals = (m_start + bm_range.to(tl.int64)).to(tl.int32) ++ idx_bits = idx_vals.to(tl.uint32).to(tl.uint64) ++ chunk_packed = (score_sortable.to(tl.uint64) << 32) | idx_bits[None, :] ++ ++ chunk_packed = chunk_packed ^ tl.full([1], 0xFFFFFFFFFFFFFFFF, dtype=tl.uint64) ++ chunk_packed = tl.sort(chunk_packed, dim=1) ++ chunk_packed = chunk_packed ^ tl.full([1], 0xFFFFFFFFFFFFFFFF, dtype=tl.uint64) ++ ++ merged = tl.minimum(topk_packed, chunk_packed) ++ topk_packed = tl.sort(merged, dim=1) ++ ++ topk_score_sortable = (topk_packed >> 32).to(tl.uint32) ++ topk_score = _sortable_u32_to_fp32(topk_score_sortable) ++ topk_idx = (topk_packed & 0xFFFFFFFF).to(tl.int32) ++ ++ k_offs = tl.arange(0, TOPK_PAD).to(tl.int64) ++ k_mask = k_offs < K ++ write_mask = n_mask[:, None] & k_mask[None, :] ++ pid_s_i64 = pid_s.to(tl.int64) ++ tl.store(partial_val_ptr + pid_b * stride_pv_b + pid_s_i64 * stride_pv_s ++ + n_offs[:, None] * stride_pv_n + k_offs[None, :] * stride_pv_k, ++ topk_score, mask=write_mask) ++ tl.store(partial_idx_ptr + pid_b * stride_pi_b + pid_s_i64 * stride_pi_s ++ + n_offs[:, None] * stride_pi_n + k_offs[None, :] * stride_pi_k, ++ topk_idx, mask=write_mask) +diff --git a/knnlib/knn.py b/knnlib/knn.py +index ad0c7b6..4668097 100644 +--- a/knnlib/knn.py ++++ b/knnlib/knn.py +@@ -1,13 +1,6 @@ +-"""Brute-force squared-L2 k-nearest-neighbours -- the reference you must optimise. ++"""Brute-force squared-L2 k-nearest-neighbours -- optimized Triton backend. + +-This implementation is intentionally naive. It computes the full ``(Q, M)`` +-pairwise distance matrix with :func:`torch.cdist` and materialises it in HBM, +-then runs :func:`torch.topk` over the whole matrix to pick the ``k`` nearest +-database points for every query. It is correct and deterministic, but it moves +-far more memory than necessary (the entire ``(Q, M)`` matrix) and cannot scale +-to large ``M`` without exhausting device memory. +- +-Contract (do NOT change): ++Public contract (unchanged): + + knn(queries, database, k) -> (distances, indices) + +@@ -19,21 +12,47 @@ Contract (do NOT change): + nearest database points, in ascending order (nearest first). + indices : (Q, k) int64 tensor of the corresponding database row indices. + +-You may add modules/kernels inside the ``knnlib`` package and rewrite the body +-of :func:`knn` freely (fused/streamed running top-k, tiled distance kernels, +-Triton), as long as the public contract above is preserved. In particular you +-should avoid ever materialising the full ``(Q, M)`` distance matrix. ++Instead of materialising the full ``(Q, M)`` pairwise distance matrix, this ++implementation dispatches to a fused Triton brute-force k-NN kernel that keeps ++a running top-k per query while streaming over the database in tiles (never ++writing an ``(Q, M)`` cross to HBM and never loading ``||x||^2``). The kernel ++ranks candidates with the shift-invariant score ``||c||^2 - 2 `` (same ++argmin-k as true squared L2); a second tiled gather pass then recovers the ++exact ``||x - c[idx]||^2`` in fp32 for the ``k`` selected neighbours. + """ + from __future__ import annotations + + import torch + ++from knnlib._kernels.primitives.knn.triton import flash_knn_triton ++from knnlib._kernels.kernels.distance.triton.knn_gather_l2sq import ( ++ triton_knn_gather_sqdist, ++) ++ + + def knn(queries, database, k): + """Return the (squared-L2) k nearest database points for each query. + +- Naive: materialise the full (Q, M) distance matrix, then top-k. ++ Fused: a Triton kNN pass returns the ``k`` nearest indices per query ++ (ranked by the x^2-free score), then a tiled gather kernel writes the ++ true squared-L2 distances for those indices. No ``(Q, M)`` matrix is ++ ever materialised. + """ +- d2 = torch.cdist(queries, database) ** 2 # (Q, M), materialized -- naive +- dist, idx = torch.topk(d2, k, dim=1, largest=False) +- return dist, idx.to(torch.long) ++ # flash_knn_triton / the gather kernel operate on batched (B, N, D) / ++ # (B, M, D) tensors. Add a batch axis of 1 and drop it on the way out. ++ qb = queries.unsqueeze(0) ++ cb = database.unsqueeze(0) ++ if not qb.is_contiguous(): ++ qb = qb.contiguous() ++ if not cb.is_contiguous(): ++ cb = cb.contiguous() ++ ++ # (1, Q, k) int32 indices, sorted by ascending true squared L2. ++ idxs = flash_knn_triton(qb, cb, k) ++ ++ # (1, Q, k) fp32 true ||x - c[idx]||^2, computed in fp32 throughout. ++ vals = triton_knn_gather_sqdist(qb, cb, idxs) ++ ++ dist = vals[0] ++ idx = idxs[0].to(torch.long) ++ return dist, idx