diff --git a/.github/workflows/gpu-tests.yml b/.github/workflows/gpu-tests.yml index 9775b1c6..4a6c45ca 100644 --- a/.github/workflows/gpu-tests.yml +++ b/.github/workflows/gpu-tests.yml @@ -1,37 +1,14 @@ name: GPU tests +# `gpu-tests` is a *required* status check on `main` (repo ruleset). It runs on +# every PR (no path filter) so the required status is always reported — a +# path-filtered required check never reports on non-matching PRs and blocks +# merge forever ("Expected — waiting for status to be reported"). + on: pull_request: - paths: - - "nvsubquadratic/**" - - "tests/**" - - "benchmarks/**" - - "examples/**" - - "experiments/**" - - "scripts/**" - - "pyproject.toml" - - "Pipfile" - - "Pipfile.lock" - - "requirements-dev.txt" - - "Dockerfile" - - "Dockerfile.*" - - ".github/workflows/gpu-tests.yml" push: branches: [main] - paths: - - "nvsubquadratic/**" - - "tests/**" - - "benchmarks/**" - - "examples/**" - - "experiments/**" - - "scripts/**" - - "pyproject.toml" - - "Pipfile" - - "Pipfile.lock" - - "requirements-dev.txt" - - "Dockerfile" - - "Dockerfile.*" - - ".github/workflows/gpu-tests.yml" workflow_dispatch: # Cancel in-flight runs for the same PR/branch when a new commit is pushed diff --git a/skills/hyenand-retrofit/SKILL.md b/skills/hyenand-retrofit/SKILL.md new file mode 100644 index 00000000..2290a7fe --- /dev/null +++ b/skills/hyenand-retrofit/SKILL.md @@ -0,0 +1,230 @@ +--- +name: hyenand-retrofit +description: Replace attention in a PyTorch model with HyenaND from the nvSubquadratic library. Covers 1D / 2D / 3D hosts (ViT, U-Net, diffusion, causal LM, hierarchical encoders). Trigger when the user wants a subquadratic alternative to attention, ports a model to HyenaND, swaps `nn.MultiheadAttention` or `F.scaled_dot_product_attention` for a Hyena mixer, builds a striped Hyena LM, or asks "how do I use nvSubquadratic with my model." Phrases like "make my ViT subquadratic," "Hyena layer for my U-Net," "swap attention with FFT convolution," "subquadratic alternative for my 3D segmentation network," or "long-context model with O(L log L) scaling" should all activate this skill. +--- + +# hyenand-retrofit + +Replace attention in a user's model with HyenaND from the nvSubquadratic library. The output is a runnable sibling file alongside the user's original — the original is not modified, with one exception: hierarchical hosts whose natural API is a conditional swap (`use_hyena=True` flag inside the host class) edit in place. + +If the user only wants conceptual explanation (no code), answer in chat. This skill is for producing a working file. + +## Native path (user is already inside nvSubquadratic) + +If the user's file imports nvSubquadratic builders (`build_attention_net`, `LazyConfig(ViT5Attention)`, etc.), the swap is mechanical: + +- Pure Hyena: replace `build_attention_net` with `build_hyena_net`, drop `compile_compatible_fftconv = False` if present (Hyena needs the default `True`). +- Hybrid: import `build_hybrid_net` from the matching `_base_config.py`, pass `layer_pattern=...`. + +Native sibling files are bare config shims — no `__main__` block; the experiment runner exercises the LazyConfig graph. + +The rest of this skill covers the **foreign path** — generic PyTorch hosts using `nn.MultiheadAttention`, `F.scaled_dot_product_attention`, timm, HF, etc. + +## Decide four things up front + +These four axes are orthogonal. Fix them before writing. + +1. **`data_dim ∈ {1, 2, 3}`** — number of spatial axes the mixer sees. Picks `Conv1d/2d/3d` for the short conv and sets `data_dim` on `CKConvND`, `SIRENKernelND`, and `GaussianModulationND`. + +1. **`causal ∈ {True, False}`** — autoregressive 1D LMs are causal; vision, segmentation, PDE are bidirectional. Causal sets `is_causal=True` on `CKConvND` (1D only) plus `use_rope=True`, mask `parametrization="exp_decay"`, `omega_0=100`. Bidirectional sets `is_causal=False`, `use_rope=False`, `parametrization="direct"`, `omega_0=10`. Boundary condition (`fft_padding`) is a separate axis-4-style choice — see the FFT-backend section. + +1. **Host layout** — `tokens [B, N, C]` (most ViTs, causal LMs) or `feature_map [B, C, *spatial]` (CNNs, U-Nets, hierarchical encoders). Spatial dim count does *not* change this — a 1D, 2D, or 3D feature-map host all use `[B, C, *S] -> [B, *S, C]`. + +1. **Return contract** — `(out, None)` tuple if the call site is `h, _ = self.attn(...)` (matches `nn.MultiheadAttention`); bare tensor otherwise. + +## Pure or hybrid + +Separate decision: replace every attention site (pure) or leave some as attention (hybrid). Hybrids are common in vision and genomics LMs because attention's selectivity complements Hyena's global mixing. If unsure, ask the user. + +- **Pure** — swap every site. Smallest change, cleanest comparison. +- **Hybrid** — *which* sites stay attention is itself an ablation, not a settled choice. Pick any reasonable starting point (e.g., alternate, or hold attention in the deepest stages) and treat the pattern as a knob to sweep. For hierarchical encoders, prefer a per-stage `bool` list (`[True, True, False, False]`) over a `"HHAA"` string — it maps cleanly onto the encoder's stage construction loop. + +## The Hyena module + +Knobs below that don't depend on the four axes are the dim-agnostic default. Substitute `DATA_DIM`, `CAUSAL`, and `HIDDEN_DIM` from the four-axis decision. + +```python +from nvsubquadratic.lazy_config import LazyConfig, instantiate +from nvsubquadratic.modules.hyena_nd import Hyena +from nvsubquadratic.modules.ckconv_nd import CKConvND +from nvsubquadratic.modules.kernels_nd import SIRENKernelND +from nvsubquadratic.modules.masks_nd import GaussianModulationND +from nvsubquadratic.utils.qk_norm import L2Norm +import torch +import torch.nn as nn + +kernel_cfg = LazyConfig(SIRENKernelND)( + data_dim=DATA_DIM, + out_dim=HIDDEN_DIM, + mlp_hidden_dim=32, + num_layers=3, + embedding_dim=32, + omega_0=100.0 if CAUSAL else 10.0, + hidden_omega_0=1.0, + L_cache=MAX_SPATIAL, # see foot-gun #3 + use_bias=True, +) + +mask_cfg = LazyConfig(GaussianModulationND)( + data_dim=DATA_DIM, + num_channels=HIDDEN_DIM, + min_attenuation_at_step=0.1, + max_attenuation_at_limit=0.95, + init_extent=1.0, + parametrization="exp_decay" if CAUSAL else "direct", +) + +global_conv_cfg = LazyConfig(CKConvND)( + data_dim=DATA_DIM, + hidden_dim=HIDDEN_DIM, + kernel_cfg=kernel_cfg, + mask_cfg=mask_cfg, + fft_padding="zero", # "circular" for periodic BCs; per-axis list for mixed + is_causal=CAUSAL, # 1D-only; ignored when CAUSAL is False + fft_backend=( + "subq_ops" if (DATA_DIM == 2 and not CAUSAL) else "torch_fft" + ), # see "FFT backend selection" below +) + +ConvND = {1: nn.Conv1d, 2: nn.Conv2d, 3: nn.Conv3d}[DATA_DIM] +short_conv_cfg = LazyConfig(ConvND)( + in_channels=3 * HIDDEN_DIM, + out_channels=3 * HIDDEN_DIM, + kernel_size=3, + groups=3 * HIDDEN_DIM, + padding=1, + bias=False, +) + +mixer = instantiate( + LazyConfig(Hyena)( + global_conv_cfg=global_conv_cfg, + short_conv_cfg=short_conv_cfg, + gate_nonlinear_cfg=LazyConfig(nn.SiLU)(), + pixelhyena_norm_cfg=LazyConfig(nn.GroupNorm)( + num_groups=1, num_channels=HIDDEN_DIM + ), + qk_norm_cfg=LazyConfig(L2Norm)(), + use_rope=CAUSAL, + rope_base=10000.0, + ) +) +``` + +**Knob ownership** (common foot-gun): + +- `Hyena(...)`: `use_rope`, `rope_base`, `gate_nonlinear_cfg`, `pixelhyena_norm_cfg`, `qk_norm_cfg`, `short_conv_cfg`, `global_conv_cfg`. +- `CKConvND(...)` (passed as `global_conv_cfg`): `data_dim`, `hidden_dim`, `mask_cfg`, `fft_padding`, `is_causal`, `kernel_cfg`, `fft_backend`. Also optional `grid_type` (`"single"`/`"double"`), `use_chunked_fftconv` (memory optimization; works with zero/causal padding only — circular has no chunked variant by design), and `use_fp16_fft` (memory; requires power-of-2 spatial dims for circular padding; not allowed with `subq_ops`). `fft_padding` accepts a single mode string (`"zero"`/`"circular"`) or a per-axis list (`["circular", "zero", ...]`) — see the FFT-backend section. +- `SIRENKernelND(...)` (passed as `kernel_cfg`): `omega_0`, `hidden_omega_0`, `mlp_hidden_dim`, `num_layers`, `embedding_dim`, `L_cache`, `use_bias`, `out_dim`. +- `GaussianModulationND(...)` (passed as `mask_cfg`): `data_dim`, `num_channels`, `min_attenuation_at_step`, `max_attenuation_at_limit`, `init_extent`, `parametrization`. + +## FFT backend and boundary conditions + +`CKConvND` has two backends. The skeleton above picks one from the four axes; the rule: + +- **`fft_backend="subq_ops"`** — optimized CUDA kernel from the optional `subquadratic_ops_torch` package (`pip install subquadratic_ops_torch`). Faster on H100/A100 for the 2D vision case. Constraints (all asserted at construction): `data_dim == 2`, `fft_padding == "zero"`, `is_causal == False`, `use_fp16_fft == False`. The canonical 2D ImageNet configs (`vit5_hybrid`, `v5/hyena_gap_pretrain.py`) use this path. +- **`fft_backend="torch_fft"`** (default) — pure `torch.fft`, supports the full matrix: `data_dim ∈ {1, 2, 3}`, `fft_padding ∈ {"zero", "circular"}` or per-axis list, `is_causal=True` (1D only), optional fp16 and chunked variants. + +The skeleton's `fft_backend="subq_ops" if (DATA_DIM == 2 and not CAUSAL) else "torch_fft"` picks the fast path when eligible and falls back otherwise. If the user has not installed `subquadratic_ops_torch`, the import fails at first forward — either swap to `"torch_fft"` or tell them to install it. + +**Per-axis mixed boundary conditions.** `fft_padding` accepts a list of modes — one per spatial axis — when the user's domain has different BCs on different axes. Example: a PDE on a periodic channel with hard walls in the cross-stream direction is `fft_padding=["circular", "zero"]` for 2D. This routes through `nvsubquadratic.ops.mixed_fftconv` and forces `fft_backend="torch_fft"` (the subq_ops kernel is single-mode). Default to a single mode if the user hasn't specified mixed; ask if the domain has clearly heterogeneous boundaries. + +## Wire it in + +`nn.MultiheadAttention` and `Hyena` have three incompatible interfaces — tuple return, kwargs, and input layout. Assigning `Hyena` directly into an attention slot fails. One adapter, parameterized by the four axes: + +```python +class HyenaAttnAdapter(nn.Module): + """Drop-in attention replacement. Parameters reflect the four-axis decision.""" + + def __init__( + self, + mixer: nn.Module, + spatial_shape: tuple[int, ...], # (T,), (H, W), or (D, H, W) + host_layout: str, # "tokens" or "feature_map" + num_prefix_tokens: int = 0, # CLS / registers; only meaningful for "tokens" + return_tuple: bool = True, # (out, None) for nn.MHA contract + ): + super().__init__() + self.mixer = mixer + self.spatial_shape = spatial_shape + self.host_layout = host_layout + self.num_prefix_tokens = num_prefix_tokens + self.return_tuple = return_tuple + + def forward(self, query, key=None, value=None, **_ignored_kwargs): + x = query # self-attention: query == key == value + if self.host_layout == "tokens": + # [B, N, C] -> peel prefix -> [B, *spatial, C] -> mix -> flatten back -> re-attach prefix + prefix, patches = ( + x[:, : self.num_prefix_tokens], + x[:, self.num_prefix_tokens :], + ) + B, _, C = patches.shape + assert ( + patches.shape[1] == torch.prod(torch.tensor(self.spatial_shape)).item() + ), f"expected {self.spatial_shape} flattened, got {patches.shape[1]} tokens" + patches_nd = patches.view(B, *self.spatial_shape, C) + out_nd = self.mixer(patches_nd, patches_nd, patches_nd) + out = out_nd.view(B, -1, C) + out = torch.cat([prefix, out], dim=1) + else: # "feature_map": [B, C, *S] -> [B, *S, C] -> mix -> back + x_cl = x.moveaxis(1, -1).contiguous() + out_cl = self.mixer(x_cl, x_cl, x_cl) + out = out_cl.moveaxis(-1, 1).contiguous() + + return (out, None) if self.return_tuple else out +``` + +For hierarchical hosts where the natural API is a `use_hyena=True` flag inside the host class, edit in place — the sibling-file rule doesn't apply. Use an outer `nn.Linear(C, 3*C)` for QKV projection and an `nn.Linear(C, C)` for output projection if you want q/k/v streams to diverge (rather than self-mixing on a single tensor). + +**Layout convention.** The library uses `BHL` (channels-first, `[B, H, *spatial]`) as the fast path internally and exposes `_w_reshape` wrappers for channels-last (`BLH`) callers. The adapter above does the explicit `moveaxis`, which works and is the safest pattern for retrofitting. If the user is constructing directly from `nvsubquadratic.ops.*` (rather than going through the full `Hyena` module), prefer the `*_w_reshape` variants instead — see `docs/architecture.md` for the full BHL/BLH convention. + +## Foot-guns + +These break the first forward pass or the first large-input forward pass. + +1. **INT32 unfold overflow in large 3D short conv.** `F.conv3d` uses `im2col` with INT32 indexing. For typical channel counts, spatial extent ≥ 160³ overflows. Symptom: `RuntimeError: Input tensor is too large`. There is no drop-in fix in this repo today (real 3D Hyena configs in `examples/well/v2/*/hyena.py` run at 64³ with plain `nn.Conv3d` and don't hit this); options at larger extent are (a) patch-merge to a smaller working grid, (b) replace the short conv with `nn.Identity()` (the global Hyena conv still mixes spatial information; only fine-grained QKV smoothing is lost), or (c) chunk the short conv along the batch dim. 1D and 2D rarely hit this. + +1. **RoPE divisibility.** With `use_rope=True`, the per-block hidden dim must be divisible by `2 * data_dim`: `% 2` for 1D, `% 4` for 2D, `% 6` for 3D. In hierarchical encoders, this must hold at every stage that uses Hyena (`embed_dim * 2^stage`). Validate at construction; fail loudly with a helpful message. + +1. **`L_cache` memory in higher dim.** The SIREN coordinate cache allocates an `L^D` volumetric buffer; memory scales roughly as `(2L − 1)^D × D × 4` bytes. For D=3: L=32 → ~3 MB, L=256 → ~1.6 GB, L=512 → ~12.8 GB. Set `L_cache` to the minimum spatial extent you need; the grid expands automatically beyond it. + +1. **Shape contract.** `nn.MultiheadAttention(batch_first=True)` returns `(out, attn_weights)` and accepts `need_weights=`/`attn_mask=`/`key_padding_mask=`. `Hyena.forward` returns a single tensor and rejects extra kwargs. The adapter above handles both; never assign `Hyena(...)` directly to an `nn.MultiheadAttention` slot. + +1. **Pretrained-weight loading.** Hyena blocks have an entirely different `state_dict` prefix than attention blocks. Loading an attention checkpoint into a Hyena variant produces a wall of missing/unexpected keys. Either filter to shared submodules (patch_embed, downsample, MLP, norms), or skip pretrained loading for Hyena variants. + +1. **`subq_ops` backend constraint set.** `fft_backend="subq_ops"` asserts loudly on any one of: `data_dim ≠ 2`, `fft_padding ≠ "zero"` (so circular and per-axis-mixed both disqualify), `is_causal=True`, `use_fp16_fft=True`. The canonical 2D vision case is fine; extending to 3D, circular padding, mixed BCs, or causal LMs requires flipping to `fft_backend="torch_fft"`. Easy to miss when copy-pasting a 2D config as a 3D starting point. + +## Smoke-test stub + +Append a `__main__` block that constructs the model and runs one forward pass at the user's stated input shape. This catches axis-1 (`data_dim`) and axis-3 (host layout) mismatches before training. + +## Filename and location convention + +- Sibling file, same directory as the user's original. +- Replace `attention` with `hyena` (or `hybrid` if mixed); keep all other tokens. If the host has no `attention` token in the filename, append `_hyenand` before the extension. +- Exception: conditional-swap hosts (`use_hyena=True` flag inside the host) edit in place. + +## Verification + +After writing: + +1. Re-read the file and confirm the four-axis choices are reflected (`data_dim`, `causal`-derived knobs, layout, return type). +1. Confirm imports are syntactically correct. +1. Run the smoke-test stub on synthetic input of the stated shape. + +## Reference configs in this repo + +Copy parameter values, not whole files. The repo's curated index lives at `docs/examples/index.md` and `docs/repository_overview.md`; the highest-leverage pointers per axis combination: + +| Axes (data_dim, causal, layout) | File | +| ------------------------------- | ---------------------------------------------------------------------------------------- | +| 2D, non-causal, tokens (ViT) | `examples/vit5_imagenet/v5_patch/_base_config.py` | +| 2D, non-causal, hybrid pattern | `examples/vit5_imagenet/vit5_hybrid/_base_config.py` | +| 2D, non-causal, feature_map | `examples/well/v2/active_matter/hyena.py` (FFT_PADDING="circular") | +| 3D, non-causal, feature_map | `examples/well/v2/supernova_explosion_64/hyena.py` (zero) / `MHD_64/hyena.py` (circular) | +| 1D, non-causal, smallest | `examples/mnist_classification/ccnn_4_160_hyena_rope_qknorm.py` | +| Diffusion (HF diffusers) | `examples/imagenet_diffusion/ccnn_12_768_hyena_qknorm.py` | + +For full per-module API and the math primer on the FFT ops, see `docs/api_reference/modules.rst` and `docs/architecture.md` — both are auto-built from the rich inline docstrings. diff --git a/skills/hyenand-retrofit/evals/README.md b/skills/hyenand-retrofit/evals/README.md new file mode 100644 index 00000000..65022ac1 --- /dev/null +++ b/skills/hyenand-retrofit/evals/README.md @@ -0,0 +1,205 @@ +# Skill evals — future CI integration + +This directory holds eval cases for the `hyenand-retrofit` skill. As of writing +they function as a **spec** for the skill's intended behavior, not a runnable +test suite — there is no harness in this repo that executes them. This README +is the design for wiring that up later. + +## What's here + +| File | Purpose | +| ------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `evals.json` | Five eval cases (native pure swap, native hybrid, foreign 2D ViT, foreign 3D feature-map U-Net, foreign 1D causal LM). Each case has a prompt, the input files the agent should see, an `expected_output` description, and a list of grep-based assertions against the file the agent writes. | +| `inputs/*.py` | Standalone test fixtures the agent retrofits. Each runs end-to-end (`python ` produces a forward-pass shape) so the harness can sanity-check the inputs themselves before evaluating the agent's output. | + +The evals together cover the four-axis grid from `SKILL.md`: + +| Eval | data_dim | causal | host layout | prefix tokens | +| ------------------------ | -------- | ------ | ----------- | ------------- | +| `foreign-timm-vit` | 2 | False | tokens | 1 (CLS) | +| `foreign-3d-feature-map` | 3 | False | feature_map | 0 | +| `foreign-1d-causal-lm` | 1 | True | tokens | 0 | + +`native-pure-swap` and `native-hybrid` cover the native path where the host +already uses nvSubquadratic builders. + +## Why CI for these + +Three reasons, in order of value: + +1. **Regression guard.** When SKILL.md is edited, the description, the four-axis + decision tree, or the adapter skeleton may drift in a way that breaks a + retrofit pattern that used to work. Running the evals on every skill change + catches that early. +1. **Model-drift catch.** As Claude models update (4.7 → 4.8 → ...), the same + skill text may be interpreted slightly differently. A weekly cron picks this + up before the next user does. +1. **Documented behavior.** The asserts encode "what counts as a correct + retrofit" in a machine-checkable form. Reviewers don't have to read SKILL.md + to know what the skill claims to do — they can read evals.json. + +## Recommended shape + +Separate workflow, not bolted onto the main GPU CI in +[../../../.github/workflows/ci.yml](../../../.github/workflows/ci.yml). Reasons: + +- The skill evals don't need a GPU — they only produce code, then grep it. + `ubuntu-latest` is enough. No reason to bloat the Colossus-runner queue. +- The skill evals are LLM-driven and inherently flakier than the existing + pytest suite. Keeping them in a separate workflow means a flake here doesn't + block a code PR. +- The triggers are different: the GPU pipeline runs on every code change; the + skill evals only need to run when the skill itself changes. + +### Triggers + +```yaml +on: + pull_request: + paths: ['skills/hyenand-retrofit/**'] + workflow_dispatch: # manual button for ad-hoc runs + schedule: + - cron: '0 6 * * 1' # weekly Monday 06:00 UTC — catches model drift +``` + +The `paths:` filter keeps cost down: skill changes are infrequent, code-only +PRs don't pay for skill evals. + +### Two-layer cheap-first design + +| Layer | Cost | When to run | Tool | +| -------------------------------------------------------------------------------------------------- | ------------------------------- | ----------------------------------------------------------- | ----------------------------------------------------- | +| **Trigger eval** — does the skill description cause Claude to load the skill for the eval prompts? | ~5 sec, ~\$0.001 per query | Every PR touching the skill | skill-creator's `run_eval.py` (already exists; reuse) | +| **Full eval** — does the agent produce a file that passes the asserts? | ~1-5 min, ~$0.05-$0.20 per eval | `pull_request` to main and `workflow_dispatch`; weekly cron | Custom harness, see sketch below | + +Trigger eval is cheap enough to run on every skill PR. Full eval is reserved +for main-PR and manual triggers because a single run is ~$0.50-$1 total and +takes 5-25 minutes serial (or ~1-5 min with a GHA matrix). + +### Parallelization + +Run the five evals as a GHA matrix: + +```yaml +strategy: + fail-fast: false # report all failures, not just the first + matrix: + eval: [native-pure-swap, native-hybrid, foreign-timm-vit, foreign-3d-feature-map, foreign-1d-causal-lm] +``` + +Each matrix job runs one eval. Wall-clock is bounded by the slowest single eval +(~5 min) at the cost of 5× concurrent ubuntu-latest runners. + +### Flake handling + +LLM outputs are not deterministic. The grep patterns in `evals.json` are +deliberately tolerant of surface-level variation (`data_dim\s*=\s*3`, not +`data_dim=3` exactly), but flakes will still happen. Two mitigations: + +- **Retry once on failure.** A retry that still fails is a real failure; + a retry that passes is a flake (note it for tracking, don't block the PR). +- **Treat the workflow as advisory at first.** Mark it non-blocking + (`continue-on-error: true` or a separate non-required check) until you have + a few weeks of data on its flake rate. Promote to required once stable. + +## Open questions to resolve before building + +1. **API key as a GHA secret.** Does the NVIDIA org policy allow + `ANTHROPIC_API_KEY` in GitHub Actions secrets? If not, the full-eval layer + is dead in the water and only trigger evals (which can run via the local + `claude` CLI configured with a personal token) are viable. +1. **Cost budget.** At ~\$1/full-run × N PRs/week + cron, is the budget + acceptable? If not, drop the per-PR trigger and run only on + `workflow_dispatch` + weekly cron. +1. **Blocking vs advisory.** Should a failing eval block a PR merge? + Recommend: advisory for the first few months, then promote to required if + the flake rate is low. +1. **Skill maintenance frequency.** If the skill is touched ~once a quarter, + the CI overhead may not be worth it vs. running `python run_evals.py` + manually as part of any skill edit. + +## Minimal harness sketch + +The full eval harness is ~50 lines of Python. Core loop: + +```python +# skills/hyenand-retrofit/evals/run_evals.py +import json, re, subprocess, pathlib, sys + +REPO_ROOT = pathlib.Path(__file__).resolve().parents[3] +EVALS = json.load(open(pathlib.Path(__file__).parent / "evals.json"))["evals"] + + +def run_one(eval_case): + # Spawn claude -p. NOTE: the skill now lives in skills/ (not .claude/skills/), + # so it is NOT auto-discovered — reference SKILL.md explicitly in the prompt, + # or symlink/copy it under .claude/skills/ for the duration of the eval run. + prompt = eval_case["prompt"] + result = subprocess.run( + ["claude", "-p", prompt, "--output-format", "json"], + cwd=REPO_ROOT, + capture_output=True, + text=True, + timeout=600, + ) + # Find the file the agent wrote (delta vs git HEAD). + written = subprocess.run( + ["git", "status", "--porcelain"], cwd=REPO_ROOT, capture_output=True, text=True + ).stdout + new_files = [ + line.split()[-1] for line in written.splitlines() if line.startswith("?? ") + ] + output_path = REPO_ROOT / new_files[0] + text = output_path.read_text() + + fails = [] + for a in eval_case["assertions"]: + kind, pattern = ( + a["check"].split(":", 1) if ":" in a["check"] else (a["check"], "") + ) + if kind == "grep_output" and not re.search(pattern, text): + fails.append(a["text"]) + elif kind == "grep_output_negative" and re.search(pattern, text): + fails.append(f"(negative) {a['text']}") + return fails, output_path + + +if __name__ == "__main__": + failed = 0 + for e in EVALS: + fails, path = run_one(e) + status = "PASS" if not fails else "FAIL" + print(f"[{status}] #{e['id']} {e['name']} ({path.name})") + for f in fails: + print(f" - {f}") + failed += bool(fails) + sys.exit(1 if failed else 0) +``` + +Caveats this sketch glosses over: + +- Cleaning up the agent's written file between evals (so the next eval's + "files written since HEAD" doesn't include the previous output) +- Handling the `original_unmodified` and `sibling_location` assertion kinds + (which need filesystem inspection, not grep) +- Per-eval working-directory isolation if evals are matrix-parallelized +- Passing the eval's `files` input list to the agent so it knows what to read + +The skill-creator plugin's +[scripts/run_eval.py](../../../.claude/plugins/marketplaces/claude-plugins-official/plugins/skill-creator/skills/skill-creator/scripts/run_eval.py) +already handles the `claude -p` spawn correctly (uuid-named command file for +skill discovery, partial-message stream parsing for fast trigger detection). +Cannibalize it rather than re-implementing. + +## Does this interfere with the skill itself? + +No. The skill loader (Claude Code's skill discovery) reads `SKILL.md` matched +by its frontmatter. README files, eval definitions, and test fixtures in +`evals/` are inert from the skill's perspective. You can freely add docs, +helper scripts, and CI config under this directory without affecting how the +skill triggers or executes. + +The only files in this directory that the skill *might* surface to a running +agent are those the agent itself navigates to (e.g., reading +`inputs/tiny_vit_attention.py` because the eval prompt named it). That's +intended — the inputs are part of the eval's surface area. diff --git a/skills/hyenand-retrofit/evals/evals.json b/skills/hyenand-retrofit/evals/evals.json new file mode 100644 index 00000000..e32e83c0 --- /dev/null +++ b/skills/hyenand-retrofit/evals/evals.json @@ -0,0 +1,103 @@ +{ + "skill_name": "hyenand-retrofit", + "evals": [ + { + "id": 1, + "name": "native-pure-swap", + "prompt": "I have an attention-based config at examples/vit5_imagenet/v5_patch/attention_patch16.py (relative to the repo root). Make a sibling file that swaps attention for pure HyenaND. Use the repo's existing builder abstractions — don't reinvent the Hyena module. Write the output sibling file into the same directory as the original.", + "expected_output": "A new file hyena_patch16.py (or similar) in the v5_patch/ directory that calls build_hyena_net instead of build_attention_net and does not set compile_compatible_fftconv=False. Should NOT include an inline __main__ stub — the canonical sibling files in this repo are bare config shims.", + "files": ["examples/vit5_imagenet/v5_patch/attention_patch16.py", "examples/vit5_imagenet/v5_patch/_base_config.py"], + "ground_truth": "examples/vit5_imagenet/v5_patch/hyena_patch16.py", + "assertions": [ + {"text": "Output file imports build_hyena_net from v5_patch._base_config", "check": "grep_output:from examples.vit5_imagenet.v5_patch._base_config import.*build_hyena_net"}, + {"text": "Output file does NOT import build_attention_net", "check": "grep_output_negative:from examples.vit5_imagenet.v5_patch._base_config import.*build_attention_net"}, + {"text": "Output file calls build_hyena_net(PATCH_SIZE) inside get_config", "check": "grep_output:build_hyena_net\\s*\\(\\s*PATCH_SIZE"}, + {"text": "Output file does NOT set compile_compatible_fftconv=False", "check": "grep_output_negative:compile_compatible_fftconv\\s*=\\s*False"}, + {"text": "Output file does not modify the original attention_patch16.py", "check": "original_unmodified"}, + {"text": "Output file is a sibling in the same directory as the original", "check": "sibling_location"} + ] + }, + { + "id": 2, + "name": "native-hybrid", + "prompt": "I have a pure-attention config at examples/vit5_imagenet/vit5_hybrid/full_attention.py (relative to the repo root). I want to produce a hybrid HyenaND+attention version using the paper-recommended HHHA pattern (3 Hyena blocks followed by 1 attention, repeated for the 12 blocks). Write a sibling file. Use the repo's existing build_hybrid_net builder.", + "expected_output": "A new file (e.g., hybrid_hhha.py) in the vit5_hybrid/ directory that imports build_hybrid_net, defines LAYER_PATTERN = 'HHHA' * 3 (or equivalent multiplication of NUM_BLOCKS // 4), calls build_hybrid_net with layer_pattern=LAYER_PATTERN. Should NOT include an inline __main__ stub.", + "files": ["examples/vit5_imagenet/vit5_hybrid/full_attention.py", "examples/vit5_imagenet/vit5_hybrid/_base_config.py"], + "ground_truth": "examples/vit5_imagenet/vit5_hybrid/hybrid_hhha.py", + "assertions": [ + {"text": "Output file imports build_hybrid_net", "check": "grep_output:from examples.vit5_imagenet.vit5_hybrid._base_config import.*build_hybrid_net"}, + {"text": "Output file defines a LAYER_PATTERN variable", "check": "grep_output:LAYER_PATTERN\\s*="}, + {"text": "LAYER_PATTERN contains the substring 'HHHA'", "check": "grep_output:HHHA"}, + {"text": "Output file passes layer_pattern to build_hybrid_net", "check": "grep_output:build_hybrid_net\\s*\\([^)]*layer_pattern"}, + {"text": "Output file does not modify the original full_attention.py", "check": "original_unmodified"}, + {"text": "Output file is a sibling in the same directory as the original", "check": "sibling_location"} + ] + }, + { + "id": 3, + "name": "foreign-timm-vit", + "prompt": "I have a tiny standalone PyTorch ViT (not using nvsubquadratic conventions) at skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py (path is relative to the repo root). It uses nn.MultiheadAttention inside its TransformerBlock. Produce a sibling file that swaps out the attention for HyenaND from the nvsubquadratic library, keeping the rest of the model intact. The model handles 64x64 RGB images with 8x8 patches (8x8 grid = 64 patches + 1 CLS = 65 tokens). Apply reasonable vision-modality defaults from the library.", + "expected_output": "A new file (e.g., tiny_vit_hyenand.py or tiny_vit_attention_hyenand.py) in the same inputs/ directory. It must import Hyena from nvsubquadratic, construct a Hyena mixer with vision defaults (data_dim=2, use_rope=False, per-axis Gaussian mask or identity, SiLU gate), wrap each block's .attn with an ADAPTER module that (a) handles the [B, 65, C] -> [B, 8, 8, C] reshape with CLS peeled off, (b) returns (out, None) so `h, _ = self.attn(...)` keeps working, (c) swallows need_weights/attn_mask kwargs. Must contain a __main__ smoke-test block that runs one forward pass.", + "files": ["skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py"], + "assertions": [ + {"text": "Output file imports Hyena from nvsubquadratic.modules.hyena_nd", "check": "grep_output:from nvsubquadratic.modules.hyena_nd import.*Hyena"}, + {"text": "Output file imports CKConvND", "check": "grep_output:from nvsubquadratic.modules.ckconv_nd import.*CKConvND"}, + {"text": "Output file imports SIRENKernelND", "check": "grep_output:from nvsubquadratic.modules.kernels_nd import.*SIRENKernelND"}, + {"text": "Output file references the original TinyViT class (subclass or wrap)", "check": "grep_output:TinyViT"}, + {"text": "Output file does not enable RoPE (use_rope=False for 2D vision)", "check": "grep_output_negative:use_rope\\s*=\\s*True"}, + {"text": "Output file uses data_dim=2", "check": "grep_output:data_dim\\s*=\\s*2"}, + {"text": "Output file defines an adapter class (look for Adapter or Wrapper in a class name)", "check": "grep_output:class\\s+\\w*(Adapter|Wrapper|Mixer)\\w*\\("}, + {"text": "Adapter returns a 2-tuple to match nn.MultiheadAttention contract", "check": "grep_output:return\\s+\\w+\\s*,\\s*None"}, + {"text": "Adapter handles CLS / prefix tokens (slice before / after index 1)", "check": "grep_output:\\[:,\\s*:\\s*1\\s*\\]|\\[:,\\s*1\\s*:\\s*\\]|num_prefix_tokens"}, + {"text": "Output file contains a __main__ smoke-test block", "check": "grep_output:if __name__"}, + {"text": "Output file does not modify the original tiny_vit_attention.py", "check": "original_unmodified"}, + {"text": "Output file is a sibling in the same directory as the original", "check": "sibling_location"} + ] + }, + { + "id": 4, + "name": "foreign-3d-feature-map", + "prompt": "I have a small standalone 3D U-Net at skills/hyenand-retrofit/evals/inputs/tiny_unet3d_attention.py (path relative to repo root). It has a SpatialAttention3D bottleneck module that operates on [B, C, D, H, W] channel-first feature maps and internally uses F.scaled_dot_product_attention — there is no nn.MultiheadAttention slot to swap. Produce a sibling file that replaces the bottleneck with HyenaND from the nvSubquadratic library, keeping the rest of the model intact. The bottleneck operates on 16x16x16 feature maps with 48 channels. Apply reasonable 3D bidirectional defaults.", + "expected_output": "A new file (e.g., tiny_unet3d_hyenand.py) in the same inputs/ directory that imports Hyena from nvsubquadratic with data_dim=3 and uses Conv3d for the short conv. The replacement must handle the channel-first [B, C, D, H, W] -> channel-last [B, D, H, W, C] layout (feature_map host, not tokens). Must NOT slice CLS / prefix tokens. Must contain a __main__ smoke-test block.", + "files": ["skills/hyenand-retrofit/evals/inputs/tiny_unet3d_attention.py"], + "assertions": [ + {"text": "Output file imports Hyena from nvsubquadratic.modules.hyena_nd", "check": "grep_output:from nvsubquadratic.modules.hyena_nd import.*Hyena"}, + {"text": "Output file imports CKConvND", "check": "grep_output:from nvsubquadratic.modules.ckconv_nd import.*CKConvND"}, + {"text": "Output file imports SIRENKernelND", "check": "grep_output:from nvsubquadratic.modules.kernels_nd import.*SIRENKernelND"}, + {"text": "Output file uses data_dim=3", "check": "grep_output:data_dim\\s*=\\s*3"}, + {"text": "Output file uses Conv3d for the short conv", "check": "grep_output:Conv3d"}, + {"text": "Output file references original TinyUNet3D class", "check": "grep_output:TinyUNet3D"}, + {"text": "Output file does NOT introduce a non-zero prefix-token count (no CLS in feature-map host)", "check": "grep_output_negative:num_prefix_tokens\\s*=\\s*[1-9]"}, + {"text": "Output file handles channel-first <-> channel-last layout", "check": "grep_output:moveaxis|permute|rearrange"}, + {"text": "Output file contains a __main__ smoke-test block", "check": "grep_output:if __name__"}, + {"text": "Output file does not modify the original tiny_unet3d_attention.py", "check": "original_unmodified"}, + {"text": "Output file is a sibling in the same directory as the original", "check": "sibling_location"} + ] + }, + { + "id": 5, + "name": "foreign-1d-causal-lm", + "prompt": "I have a tiny causal char-level LM at skills/hyenand-retrofit/evals/inputs/tiny_charlm_attention.py (path relative to repo root). It uses causal nn.MultiheadAttention inside CausalBlock with an explicit attn_mask. Produce a sibling file that swaps the attention for HyenaND from the nvSubquadratic library, keeping the rest of the model intact. The model processes sequences of length 256 with hidden_dim=128. Apply reasonable 1D causal defaults.", + "expected_output": "A new file (e.g., tiny_charlm_hyenand.py) in the same inputs/ directory that imports Hyena from nvsubquadratic with data_dim=1 and causal defaults: is_causal=True on CKConvND (NOT fft_padding='causal' — that string is rejected by CKConvND validation), use_rope=True, fft_padding='zero' (or another valid mode string), mask parametrization='exp_decay', omega_0=100. Uses Conv1d for the short conv. Wraps each block's .attn with an adapter that handles the tokens [B, T, C] layout without prefix tokens, returns (out, None) for the nn.MultiheadAttention contract, and swallows attn_mask/need_weights kwargs. Must contain a __main__ smoke-test block.", + "files": ["skills/hyenand-retrofit/evals/inputs/tiny_charlm_attention.py"], + "assertions": [ + {"text": "Output file imports Hyena from nvsubquadratic.modules.hyena_nd", "check": "grep_output:from nvsubquadratic.modules.hyena_nd import.*Hyena"}, + {"text": "Output file imports CKConvND", "check": "grep_output:from nvsubquadratic.modules.ckconv_nd import.*CKConvND"}, + {"text": "Output file imports SIRENKernelND", "check": "grep_output:from nvsubquadratic.modules.kernels_nd import.*SIRENKernelND"}, + {"text": "Output file uses data_dim=1", "check": "grep_output:data_dim\\s*=\\s*1"}, + {"text": "Output file uses Conv1d for the short conv", "check": "grep_output:Conv1d"}, + {"text": "Output file enables RoPE for causal LM", "check": "grep_output:use_rope\\s*=\\s*True"}, + {"text": "Output file sets is_causal=True on CKConvND (causal LM)", "check": "grep_output:is_causal\\s*=\\s*True"}, + {"text": "Output file does NOT use the obsolete fft_padding='causal' value (it's rejected by CKConvND validation)", "check": "grep_output_negative:fft_padding\\s*=\\s*['\\\"]causal['\\\"]"}, + {"text": "Output file uses exp_decay mask parametrization", "check": "grep_output:parametrization\\s*=\\s*['\\\"]exp_decay['\\\"]"}, + {"text": "Output file uses omega_0=100 (causal default)", "check": "grep_output:omega_0\\s*=\\s*100"}, + {"text": "Output file references original TinyCharLM class", "check": "grep_output:TinyCharLM"}, + {"text": "Adapter returns 2-tuple to match nn.MultiheadAttention contract", "check": "grep_output:return\\s+\\w+\\s*,\\s*None"}, + {"text": "Adapter does NOT slice CLS / prefix tokens (causal LM has no CLS)", "check": "grep_output_negative:num_prefix_tokens\\s*=\\s*[1-9]"}, + {"text": "Output file contains a __main__ smoke-test block", "check": "grep_output:if __name__"}, + {"text": "Output file does not modify the original tiny_charlm_attention.py", "check": "original_unmodified"}, + {"text": "Output file is a sibling in the same directory as the original", "check": "sibling_location"} + ] + } + ] +} diff --git a/skills/hyenand-retrofit/evals/inputs/tiny_charlm_attention.py b/skills/hyenand-retrofit/evals/inputs/tiny_charlm_attention.py new file mode 100644 index 00000000..d6a6e079 --- /dev/null +++ b/skills/hyenand-retrofit/evals/inputs/tiny_charlm_attention.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tiny char-level causal LM with multi-head self-attention. + +4 transformer blocks, hidden_dim=128, num_heads=4, sequence length 256, +vocab_size=256. Each block applies a causal attention mask. Used as a +small test target for attention -> HyenaND retrofits on 1D causal hosts. + +No CLS / prefix tokens. Retrofit must use causal-LM defaults: +``use_rope=True``, ``fft_padding="causal"``, mask +``parametrization="exp_decay"``, ``omega_0=100``. +""" + +import torch +import torch.nn as nn + + +class CausalBlock(nn.Module): + """Pre-norm transformer block with causal multi-head self-attention.""" + + def __init__(self, dim: int = 128, num_heads: int = 4, mlp_ratio: float = 4.0): + """Configure attention and MLP sublayers for hidden size ``dim``.""" + super().__init__() + self.norm1 = nn.LayerNorm(dim) + self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True) + self.norm2 = nn.LayerNorm(dim) + hidden = int(dim * mlp_ratio) + self.mlp = nn.Sequential( + nn.Linear(dim, hidden), + nn.GELU(), + nn.Linear(hidden, dim), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Apply causal attention and MLP; input shape ``[B, T, C]``.""" + T = x.size(1) + causal_mask = torch.triu(torch.full((T, T), float("-inf"), device=x.device), diagonal=1) + h = self.norm1(x) + h, _ = self.attn(h, h, h, attn_mask=causal_mask, need_weights=False) + x = x + h + x = x + self.mlp(self.norm2(x)) + return x + + +class TinyCharLM(nn.Module): + """Small causal character LM for HyenaND retrofit evals.""" + + def __init__( + self, + vocab_size: int = 256, + seq_len: int = 256, + dim: int = 128, + num_blocks: int = 4, + num_heads: int = 4, + ): + """Build embeddings, ``num_blocks`` causal blocks, and output head.""" + super().__init__() + self.tok_embed = nn.Embedding(vocab_size, dim) + self.pos_embed = nn.Parameter(torch.zeros(1, seq_len, dim)) + self.blocks = nn.ModuleList([CausalBlock(dim, num_heads) for _ in range(num_blocks)]) + self.norm = nn.LayerNorm(dim) + self.head = nn.Linear(dim, vocab_size) + + def forward(self, idx: torch.Tensor) -> torch.Tensor: + """Return logits for token indices ``[B, T]``.""" + x = self.tok_embed(idx) + self.pos_embed[:, : idx.size(1)] + for blk in self.blocks: + x = blk(x) + return self.head(self.norm(x)) + + +if __name__ == "__main__": + model = TinyCharLM() + idx = torch.randint(0, 256, (2, 256)) + y = model(idx) + print(y.shape) # [2, 256, 256] diff --git a/skills/hyenand-retrofit/evals/inputs/tiny_unet3d_attention.py b/skills/hyenand-retrofit/evals/inputs/tiny_unet3d_attention.py new file mode 100644 index 00000000..371f9b09 --- /dev/null +++ b/skills/hyenand-retrofit/evals/inputs/tiny_unet3d_attention.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tiny 3D U-Net with bottleneck self-attention over spatial positions. + +Used as a small test target for attention -> HyenaND retrofits on +feature-map hosts. The bottleneck attention operates directly on +[B, C, D, H, W] channel-first feature maps (no token sequence in the +host API). Internally it uses ``F.scaled_dot_product_attention`` rather +than ``nn.MultiheadAttention``, so there is no MHA slot to swap — the +retrofit must replace the whole ``SpatialAttention3D`` module. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class ConvBlock3D(nn.Module): + """Two 3x3x3 convolutions with GroupNorm and SiLU.""" + + def __init__(self, in_ch: int, out_ch: int): + """Stack conv layers from ``in_ch`` to ``out_ch`` channels.""" + super().__init__() + self.net = nn.Sequential( + nn.Conv3d(in_ch, out_ch, kernel_size=3, padding=1), + nn.GroupNorm(num_groups=1, num_channels=out_ch), + nn.SiLU(), + nn.Conv3d(out_ch, out_ch, kernel_size=3, padding=1), + nn.GroupNorm(num_groups=1, num_channels=out_ch), + nn.SiLU(), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Run the conv block on ``[B, C, D, H, W]``.""" + return self.net(x) + + +class SpatialAttention3D(nn.Module): + """Self-attention over spatial positions of a 3D feature map. + + Input/output: [B, C, D, H, W]. Retrofit target: replace this whole + module with a Hyena mixer that operates on channel-last 3D feature + maps. No CLS / prefix tokens. + """ + + def __init__(self, dim: int, num_heads: int = 4): + """Build QKV projection and output proj for ``dim`` channels.""" + super().__init__() + assert dim % num_heads == 0 + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.norm = nn.LayerNorm(dim) + self.qkv = nn.Linear(dim, 3 * dim, bias=False) + self.proj = nn.Linear(dim, dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Apply full-spatial self-attention; preserves ``[B, C, D, H, W]``.""" + B, C, D, H, W = x.shape + N = D * H * W + h = x.flatten(2).transpose(1, 2) # [B, N, C] + h = self.norm(h) + qkv = self.qkv(h).view(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) + q, k, v = qkv.unbind(0) # [B, heads, N, head_dim] + out = F.scaled_dot_product_attention(q, k, v) + out = out.transpose(1, 2).reshape(B, N, C) + out = self.proj(out) + return out.transpose(1, 2).view(B, C, D, H, W) + + +class TinyUNet3D(nn.Module): + """Tiny 3D U-Net for 32x32x32 volumes, 2-class segmentation.""" + + def __init__(self, in_channels: int = 1, base_ch: int = 24, out_channels: int = 2): + """Wire encoder, bottleneck attention, decoder, and segmentation head.""" + super().__init__() + self.enc1 = ConvBlock3D(in_channels, base_ch) + self.pool = nn.MaxPool3d(2) + self.enc2 = ConvBlock3D(base_ch, base_ch * 2) + self.bottleneck = SpatialAttention3D(base_ch * 2, num_heads=4) + self.up = nn.ConvTranspose3d(base_ch * 2, base_ch, kernel_size=2, stride=2) + self.dec1 = ConvBlock3D(base_ch * 2, base_ch) + self.head = nn.Conv3d(base_ch, out_channels, kernel_size=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Encode, attend at bottleneck, decode; return per-voxel logits.""" + e1 = self.enc1(x) # [B, base_ch, 32, 32, 32] + e2 = self.enc2(self.pool(e1)) # [B, 2*base_ch, 16, 16, 16] + b = e2 + self.bottleneck(e2) # residual attention + u = self.up(b) # [B, base_ch, 32, 32, 32] + d = self.dec1(torch.cat([u, e1], dim=1)) + return self.head(d) + + +if __name__ == "__main__": + model = TinyUNet3D() + x = torch.randn(2, 1, 32, 32, 32) + y = model(x) + print(y.shape) # [2, 2, 32, 32, 32] diff --git a/skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py b/skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py new file mode 100644 index 00000000..10716901 --- /dev/null +++ b/skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tiny ViT with standard multi-head self-attention. + +Hand-written transformer for 64x64 RGB images, 4 transformer blocks, +hidden_dim=128, num_heads=4. Patches are 8x8 (so 8x8 = 64 patches). +Used as a small test target for attention -> HyenaND retrofits. +""" + +import torch +import torch.nn as nn + + +class PatchEmbed(nn.Module): + """Conv-based patch embedding that flattens 2D patches into a token sequence.""" + + def __init__(self, image_size: int = 64, patch_size: int = 8, in_channels: int = 3, hidden_dim: int = 128): + """Configure the patch-projection conv for the given image / patch size.""" + super().__init__() + self.grid_h = image_size // patch_size + self.grid_w = image_size // patch_size + self.num_patches = self.grid_h * self.grid_w + self.proj = nn.Conv2d(in_channels, hidden_dim, kernel_size=patch_size, stride=patch_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Project image ``[B, C, H, W]`` to patch tokens ``[B, N, C]``.""" + x = self.proj(x) # [B, C, H, W] + return x.flatten(2).transpose(1, 2) # [B, N, C] + + +class MLP(nn.Module): + """Two-layer GELU MLP used as the transformer feed-forward sublayer.""" + + def __init__(self, dim: int, mlp_ratio: float = 4.0): + """Build a ``dim -> dim * mlp_ratio -> dim`` MLP with GELU activation.""" + super().__init__() + hidden = int(dim * mlp_ratio) + self.fc1 = nn.Linear(dim, hidden) + self.act = nn.GELU() + self.fc2 = nn.Linear(hidden, dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Apply ``fc1 -> GELU -> fc2`` to ``[B, N, C]``.""" + return self.fc2(self.act(self.fc1(x))) + + +class TransformerBlock(nn.Module): + """Pre-norm transformer block with multi-head self-attention and an MLP.""" + + def __init__(self, dim: int = 128, num_heads: int = 4, mlp_ratio: float = 4.0): + """Configure attention and MLP sublayers for hidden size ``dim``.""" + super().__init__() + self.norm1 = nn.LayerNorm(dim) + self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True) + self.norm2 = nn.LayerNorm(dim) + self.mlp = MLP(dim, mlp_ratio) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Apply self-attention and MLP residual sublayers to ``[B, N, C]``.""" + h = self.norm1(x) + h, _ = self.attn(h, h, h, need_weights=False) + x = x + h + x = x + self.mlp(self.norm2(x)) + return x + + +class TinyViT(nn.Module): + """4-block ViT for 64x64 images, 10-class classification.""" + + def __init__( + self, + image_size: int = 64, + patch_size: int = 8, + hidden_dim: int = 128, + num_blocks: int = 4, + num_heads: int = 4, + num_classes: int = 10, + ): + """Build patch embedding, CLS token, ``num_blocks`` transformer blocks, and classifier head.""" + super().__init__() + self.embed = PatchEmbed(image_size, patch_size, in_channels=3, hidden_dim=hidden_dim) + self.cls_token = nn.Parameter(torch.zeros(1, 1, hidden_dim)) + self.pos_embed = nn.Parameter(torch.zeros(1, self.embed.num_patches + 1, hidden_dim)) + self.blocks = nn.ModuleList([TransformerBlock(hidden_dim, num_heads) for _ in range(num_blocks)]) + self.norm = nn.LayerNorm(hidden_dim) + self.head = nn.Linear(hidden_dim, num_classes) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Return class logits for input images ``[B, 3, H, W]``.""" + x = self.embed(x) + cls = self.cls_token.expand(x.size(0), -1, -1) + x = torch.cat([cls, x], dim=1) + self.pos_embed + for blk in self.blocks: + x = blk(x) + x = self.norm(x[:, 0]) + return self.head(x) + + +if __name__ == "__main__": + model = TinyViT() + x = torch.randn(2, 3, 64, 64) + y = model(x) + print(y.shape) # [2, 10]