From 523c9971c172f58b9887393a30741df3a2e2b0b1 Mon Sep 17 00:00:00 2001 From: Farhad Ramezanghorbani Date: Mon, 18 May 2026 16:36:32 -0600 Subject: [PATCH 01/12] feat(claude): add hyenand-retrofit Claude skill Adds a Claude Code / Cowork skill that automates the attention -> HyenaND swap. Targets two paths: - Native: in-repo configs that use the build_attention_net / build_hyena_net / build_hybrid_net builders. Swap is mechanical (change the builder call, flip compile_compatible_fftconv, pick a layer pattern for hybrids). - Foreign: external PyTorch models (timm ViT, HF transformers, hand- written) that use nn.MultiheadAttention or F.scaled_dot_product_ attention. Skill emits a sibling file with a full Hyena module constructed from paper-grounded vision/medical/genomics/PDE defaults. Paper-grounded defaults live in references/defaults.md (omega_0, register count, Gaussian mask, FFT padding, RoPE on/off by modality, plus the hybrid layer patterns from sec 5.2/5.3/5.5: HHHA x3 for ImageNet, HHAA for SwinUNETR, H_2 mixing for Evo2-style genomics LMs). Three eval cases under evals/ exercise the native pure-swap, native hybrid, and foreign-timm-vit paths. Iteration 1 scored 23/23 with the skill vs 20/22 without (baseline missed __main__ smoke-test blocks); foreign-path assertions tied -- a known gap to tighten next iteration. Co-Authored-By: Claude Opus 4.7 (1M context) --- .claude/skills/hyenand-retrofit/SKILL.md | 177 ++++++++++++++++++ .../skills/hyenand-retrofit/evals/evals.json | 57 ++++++ .../evals/inputs/tiny_vit_attention.py | 82 ++++++++ .../hyenand-retrofit/references/defaults.md | 136 ++++++++++++++ 4 files changed, 452 insertions(+) create mode 100644 .claude/skills/hyenand-retrofit/SKILL.md create mode 100644 .claude/skills/hyenand-retrofit/evals/evals.json create mode 100644 .claude/skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py create mode 100644 .claude/skills/hyenand-retrofit/references/defaults.md diff --git a/.claude/skills/hyenand-retrofit/SKILL.md b/.claude/skills/hyenand-retrofit/SKILL.md new file mode 100644 index 00000000..7f62f09d --- /dev/null +++ b/.claude/skills/hyenand-retrofit/SKILL.md @@ -0,0 +1,177 @@ +--- +name: hyenand-retrofit +description: Replace attention in a PyTorch model with HyenaND from the nvSubquadratic library, with paper-grounded default hyperparameters (SIREN ω₀, register count, per-axis Gaussian mask, FFT padding, hybrid layer pattern). Use this skill whenever the user wants to swap attention for a subquadratic operator, port a ViT / U-Net / diffusion / genomics LM to HyenaND, convert a `nn.MultiheadAttention` or `F.scaled_dot_product_attention` site to a Hyena mixer, build a striped Hyena LM, set up a HyenaND experiment config, or asks "how do I use nvSubquadratic with my model." Trigger even when the user does not explicitly name HyenaND — phrases like "make my ViT subquadratic," "subquadratic alternative to attention for my 3D U-Net," "Hyena layer for my transformer," "swap attention with FFT convolution," 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. + +## When to use + +- User has an attention-based model (ViT, U-Net, DiT, causal LM, hierarchical encoder) and wants a subquadratic alternative +- User explicitly mentions HyenaND, Hyena, nSubQ, nvSubquadratic, or "subquadratic attention" +- User is inside the `nvSubquadratic-private` repo and wants to add a new attention/Hyena variant +- User asks about scaling to long contexts (long genomes, high-resolution images, 3D volumes, PDE grids) + +If the user only wants conceptual explanation (no code), answer in chat. This skill is for producing a working file. + +## The two paths + +Decide which one applies *before* writing anything: + +1. **Native path** — the user's file already uses `nvsubquadratic` (`LazyConfig`, `build_attention_net`, `ViT5Attention`, `ViT5ClassificationNet`). The repo has matched `build_attention_net` / `build_hyena_net` / `build_hybrid_net` builders, so the swap is mechanical: replace the builder call, flip `compile_compatible_fftconv`, optionally pick a layer pattern. + +2. **Foreign path** — the user's file uses generic PyTorch (`nn.MultiheadAttention`, `F.scaled_dot_product_attention`, `timm`, `transformers`, `diffusers`). You must construct a full Hyena module from scratch and wire it in as a drop-in attention replacement. + +Look at the user's file. If you see `from nvsubquadratic` imports or `LazyConfig(ViT5Attention)`, take the native path. Otherwise foreign. + +## Native path workflow + +The repo factors attention vs Hyena into builder functions. The user's entry file is almost always just a thin wrapper around one of: + +- `build_attention_net(patch_size)` — pure attention +- `build_hyena_net(patch_size)` — pure Hyena +- `build_hybrid_net(layer_pattern="...", patch_size=...)` — mixed, where pattern is a string of `H`/`A` characters (one per block) + +### Step 1: ask the user which target + +Use AskUserQuestion if not already specified: + +- **Pure Hyena** — strongest when geometry + global structure dominate (PDE fields, high-resolution vision, long genomes) +- **Hybrid** — strongest when selectivity still matters (ImageNet classification, medical segmentation). Paper winners: + - 2D ImageNet ViT-Small: `(HA)×6` (best) or `(HHHA)×3` + - 3D medical SwinUNETR-style: `HHAA` (hierarchical — Hyena in early high-resolution stages, attention later) + - 1D genomics (1B striped LM, Evo2): one A per 4 blocks (`HHHA` repeat, H₂ mixing) was best in §5.2 + +### Step 2: write the sibling file + +Copy the entry-file structure verbatim, then change: + +- Import: replace `build_attention_net` with `build_hyena_net` (pure) or `build_hybrid_net` (hybrid). Drop the unused builder. +- Builder call inside `get_config()`: same swap. +- For hybrid: add a `LAYER_PATTERN = "..." * (NUM_BLOCKS // len_repeat)` line at module scope, pass it as `layer_pattern=LAYER_PATTERN`. +- Remove `config.compile_compatible_fftconv = False` if present (attention sets this False explicitly; the default is True, and HyenaND needs True). +- Update the module docstring and any inline comments to reflect the new grid math (Hyena adds registers — see `hyena_patch16.py` for the canonical comment style). +- Keep filename convention: `attention_patch16.py` → `hyena_patch16.py`, `full_attention.py` → `full_hyena.py` or `hybrid_hhha.py`. + +### Step 3: smoke-test stub + +Append a `__main__` block that constructs the config and prints `repr(config.net)`. The user can run this on CPU to verify imports resolve and the LazyConfig graph is well-formed before committing GPU time. + +```python +if __name__ == "__main__": + cfg = get_config() + print(cfg.net) +``` + +## Foreign path workflow + +The user has, e.g., a `timm` ViT, a HF transformers model, or a hand-written PyTorch transformer. You need to construct a Hyena module yourself. + +### Step 1: identify the attention sites + +Look for any of: + +- `nn.MultiheadAttention` +- `F.scaled_dot_product_attention` +- `flash_attn.*` calls +- Custom attention modules (look for `softmax(Q @ K.T / sqrt(d))` or `attn_drop`/`proj_drop` member names) +- timm `Attention` class, HF `*Attention` classes + +### Step 2: build the Hyena replacement + +For each attention site, emit a Hyena module config. The minimum spec: + +```python +from nvsubquadratic.lazy_config import LazyConfig +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.modules.rms_norm import RMSNorm +from nvsubquadratic.utils.qk_norm import L2Norm +import torch +``` + +The full Hyena config has many knobs — see `references/defaults.md` for the canonical table. The short version of the defaults by modality: + +| Modality | data_dim | mask | fft_padding | use_rope | gate_nonlinear | typical ω₀ | +|----------|----------|------|-------------|----------|----------------|------------| +| Vision (image classification, diffusion) | 2 | identity or per-axis Gaussian | circular or zero | False | SiLU | 10 | +| Medical 3D segmentation | 3 | per-axis Gaussian | zero | False | SiLU | 10 | +| Genomics / causal LM | 1 | exponential decay with causal zeroing | causal | True | SiLU | 10 | +| PDE fields | 2 or 3 | per-axis Gaussian | circular | False | SiLU | 10 | + +Read `references/defaults.md` for full parameter lists, init schemes, and the reasoning behind each choice. + +### Step 3: wire it in + +In the sibling file, subclass or monkey-patch the user's model to replace each attention module with the Hyena module. Prefer subclassing — it's easier to read and rollback. Example skeleton: + +```python +# my_model_hyenand.py — generated by hyenand-retrofit + +import torch +from nvsubquadratic.lazy_config import instantiate, LazyConfig +from nvsubquadratic.modules.hyena_nd import Hyena +# ... other imports per defaults.md ... + +from my_model import MyViT # user's original + +def build_hyena_mixer(hidden_dim: int, grid_h: int, grid_w: int) -> torch.nn.Module: + """Drop-in replacement for nn.MultiheadAttention(hidden_dim, num_heads).""" + cfg = LazyConfig(Hyena)( + global_conv_cfg=LazyConfig(CKConvND)( + data_dim=2, + hidden_dim=hidden_dim, + kernel_cfg=LazyConfig(SIRENKernelND)( + data_dim=2, out_dim=hidden_dim, mlp_hidden_dim=32, + num_layers=3, embedding_dim=32, omega_0=10.0, + L_cache=max(grid_h, grid_w), use_bias=True, hidden_omega_0=1.0, + ), + # ... see defaults.md ... + ), + # ... see defaults.md ... + ) + return instantiate(cfg) + +class MyViTHyenaND(MyViT): + def __init__(self, *args, grid_h: int, grid_w: int, **kwargs): + super().__init__(*args, **kwargs) + for block in self.blocks: + block.attn = build_hyena_mixer(self.embed_dim, grid_h, grid_w) +``` + +### Step 4: smoke-test stub + +Append a `__main__` block that constructs the model and runs one forward pass on synthetic input of the user's stated shape. This catches shape mismatches before the user invests in training. + +## Filename and location convention + +- Sibling file, same directory as the user's original +- Name: replace `attention` with `hyena` (or `hybrid_` for hybrid), keep all other tokens +- If the user's file has no `attention` token in the name, append `_hyenand` before the extension +- Do not edit the user's original — keep the diff trivial + +## Verification + +After writing the file: + +1. Read it back and confirm the changes you intended actually landed +2. Confirm the only difference from the user's file is the attention→Hyena swap plus any required toggles (`compile_compatible_fftconv`, layer pattern) +3. Confirm imports are syntactically correct (the user can `python -c "from import get_config"` to verify) +4. Confirm the smoke-test stub is present + +## What not to do + +- Do not modify the user's original file +- Do not invent new builders if the matched `build_*_net` pair already exists in `_base_config.py` +- Do not skip the smoke-test stub — silent shape mismatches at training start are the #1 retrofit failure +- Do not pick a hybrid pattern at random; either ask, or take the paper default for the modality (HHHA×3 for 2D vision, HHAA for 3D hierarchical, HHHA repeat for 1D genomics) +- Do not omit `use_rope=True` for 1D autoregressive — without RoPE, positional recall collapses +- Do not omit the per-axis Gaussian mask for ND≥2 unless the user explicitly opts out — it was the difference between Hyena and bidirectional Mamba in the §5.1 color_cond probes + +## References + +- `references/defaults.md` — full per-modality default tables, init schemes, and rationale diff --git a/.claude/skills/hyenand-retrofit/evals/evals.json b/.claude/skills/hyenand-retrofit/evals/evals.json new file mode 100644 index 00000000..2c038792 --- /dev/null +++ b/.claude/skills/hyenand-retrofit/evals/evals.json @@ -0,0 +1,57 @@ +{ + "skill_name": "hyenand-retrofit", + "evals": [ + { + "id": 1, + "name": "native-pure-swap", + "prompt": "I have an attention-based config at /sessions/gracious-trusting-pasteur/mnt/nvSubquadratic-private/examples/vit5_imagenet/v5_patch/attention_patch16.py. 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, does not set compile_compatible_fftconv=False, and contains a __main__ smoke-test block.", + "files": ["/sessions/gracious-trusting-pasteur/mnt/nvSubquadratic-private/examples/vit5_imagenet/v5_patch/attention_patch16.py", "/sessions/gracious-trusting-pasteur/mnt/nvSubquadratic-private/examples/vit5_imagenet/v5_patch/_base_config.py"], + "ground_truth": "/sessions/gracious-trusting-pasteur/mnt/nvSubquadratic-private/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 contains a __main__ smoke-test block", "check": "grep_output:if __name__"}, + {"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 /sessions/gracious-trusting-pasteur/mnt/nvSubquadratic-private/examples/vit5_imagenet/vit5_hybrid/full_attention.py. 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, and contains a __main__ smoke-test block.", + "files": ["/sessions/gracious-trusting-pasteur/mnt/nvSubquadratic-private/examples/vit5_imagenet/vit5_hybrid/full_attention.py", "/sessions/gracious-trusting-pasteur/mnt/nvSubquadratic-private/examples/vit5_imagenet/vit5_hybrid/_base_config.py"], + "ground_truth": "/sessions/gracious-trusting-pasteur/mnt/nvSubquadratic-private/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 contains a __main__ smoke-test block", "check": "grep_output:if __name__"}, + {"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 .claude/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), subclass or wrap TinyViT to swap each block's .attn for the Hyena mixer, and contain a __main__ smoke-test block that runs one forward pass.", + "files": [".claude/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 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"} + ] + } + ] +} diff --git a/.claude/skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py b/.claude/skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py new file mode 100644 index 00000000..51e32065 --- /dev/null +++ b/.claude/skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py @@ -0,0 +1,82 @@ +"""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): + def __init__(self, image_size: int = 64, patch_size: int = 8, in_channels: int = 3, hidden_dim: int = 128): + 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: + x = self.proj(x) # [B, C, H, W] + return x.flatten(2).transpose(1, 2) # [B, N, C] + + +class MLP(nn.Module): + def __init__(self, dim: int, mlp_ratio: float = 4.0): + 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: + return self.fc2(self.act(self.fc1(x))) + + +class TransformerBlock(nn.Module): + def __init__(self, dim: int = 128, num_heads: int = 4, mlp_ratio: float = 4.0): + 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: + 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): + 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: + 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] diff --git a/.claude/skills/hyenand-retrofit/references/defaults.md b/.claude/skills/hyenand-retrofit/references/defaults.md new file mode 100644 index 00000000..31d128bc --- /dev/null +++ b/.claude/skills/hyenand-retrofit/references/defaults.md @@ -0,0 +1,136 @@ +# HyenaND default hyperparameters by modality + +These defaults are grounded in the paper and the repo's reference configs (`examples/imagenet_classification/ccnn_7_512_hyena*.py`, `examples/vit5_imagenet/v5_patch/_base_config.py`, `examples/well/v2/*.py`, `examples/mnist_classification/ccnn_4_160_hyena_rope_qknorm.py`). Use them as the starting point for any retrofit; tune only if there's a specific reason. + +## Common (all modalities) + +| Knob | Value | Why | +|------|-------|-----| +| `omega_0` (SIREN first-layer ω) | 10.0 for vision/PDE, 100.0 for genomics | Frequency band; higher for 1D where positional recall matters more | +| `hidden_omega_0` | 1.0 | SIREN hidden-layer ω; preserves init scale | +| `mlp_hidden_dim` (SIREN width) | 32 | Small MLP; SIREN is shallow but wide enough to express smooth kernels | +| `num_layers` (SIREN depth) | 3 | Paper-standard depth | +| `embedding_dim` (SIREN coordinate embedding) | 32 | First-layer positional embedding | +| `use_bias` | True | | +| First-layer init | `U(-1/sqrt(N), 1/sqrt(N))` | Preserves unit-variance pre-activations (§3.2.3) | +| FiLM heads init | zero | Training starts from the unmodulated baseline (§3.2.3) | + +## Vision (2D, image classification or diffusion) + +```python +data_dim = 2 +use_rope = False +fft_padding = "circular" # or "zero" for CLS-row variants +mask_cfg = LazyConfig(GaussianModulationND)( + data_dim=2, + num_channels=hidden_dim, + min_attenuation_at_step=0.1, + max_attenuation_at_limit=0.95, + init_extent=1.0, + parametrization="direct", +) +gate_nonlinear_cfg = LazyConfig(torch.nn.SiLU)() +gate_nonlinear_2_cfg = LazyConfig(torch.nn.Sigmoid)() +qk_norm_cfg = LazyConfig(L2Norm)() +pixelhyena_norm_cfg = LazyConfig(torch.nn.GroupNorm)(num_groups=1, num_channels=hidden_dim) +short_conv = LazyConfig(torch.nn.Conv2d)( + in_channels=3 * hidden_dim, out_channels=3 * hidden_dim, + kernel_size=3, groups=3 * hidden_dim, padding=1, bias=False, +) +``` + +For ViT-5-style backbones with registers, use `ViT5HyenaAdapter` and `RegisterPooling` (see `examples/vit5_imagenet/v5_patch/_base_config.py` lines 270–353). + +**Hybrid patterns (12 blocks):** +- Best on ImageNet: `(HA)×6` (82.1 top-1) +- Second-best: `(HHHA)×3` (82.0 top-1) +- Pure: `H×12` (81.5 top-1, matches attention baseline) + +## Medical 3D segmentation (SwinUNETR-style hierarchical encoders) + +```python +data_dim = 3 +use_rope = False +fft_padding = "zero" +mask_cfg = LazyConfig(GaussianModulationND)(data_dim=3, num_channels=hidden_dim, ...) +short_conv = LazyConfig(torch.nn.Conv3d)( + in_channels=3 * hidden_dim, out_channels=3 * hidden_dim, + kernel_size=3, groups=3 * hidden_dim, padding=1, bias=False, +) +``` + +**Stage patterns (4-stage encoder, paper §5.5 PanTS, mean Dice):** + +| Pattern | Stage 1 | Stage 2 | Stage 3 | Stage 4 | Mean Dice | +|---------|---------|---------|---------|---------|-----------| +| `AAAA` (Swin baseline) | Attn | Attn | Attn | Attn | 0.7496 | +| `HHHH` (all-Hyena) | Hyena | Hyena | Hyena | Hyena | 0.7510 | +| `HAHA` (striped) | Hyena | Attn | Hyena | Attn | 0.7535 | +| `HHAA` (hierarchical, **best**) | Hyena | Hyena | Attn | Attn | **0.7559** | + +Recommend `HHAA` by default. Memory savings (~11%) are roughly invariant to placement. + +## Genomics / 1D causal LM (striped Hyena) + +```python +data_dim = 1 +use_rope = True # critical — without RoPE, positional recall collapses +rope_base = 10000.0 +fft_padding = "causal" # preserves autoregressive structure +mask_cfg = LazyConfig(GaussianModulationND)( # exponential decay with causal zeroing + data_dim=1, num_channels=hidden_dim, + min_attenuation_at_step=0.1, + max_attenuation_at_limit=0.95, + init_extent=1.0, + parametrization="exp_decay", +) +gate_nonlinear_cfg = LazyConfig(torch.nn.Identity)() # linear gating for AR +short_conv = LazyConfig(torch.nn.Conv1d)( + in_channels=3 * hidden_dim, out_channels=3 * hidden_dim, + kernel_size=3, groups=3 * hidden_dim, padding=1, bias=False, +) +``` + +**Mixing ratios (Evo2-1B, 8192-bp sequences, §5.2, perplexity lower is better):** + +| Config | Pattern (24 blocks) | Validation PPL | +|--------|---------------------|----------------| +| `T` (full transformer) | 24 attention | 2.9235 ± 0.0039 | +| `H₀` (full Hyena) | 24 Hyena | 2.8282 ± 0.0279 | +| `H₁` (1 MHA) | 23 H + 1 A | 2.8308 ± 0.0108 | +| `H₂` (2 MHA, **best**) | 22 H + 2 A | **2.7729 ± 0.0006** | +| `H₃` (3 MHA) | 21 H + 3 A | 2.8214 ± 0.0313 | +| `H₄` (4 MHA) | 20 H + 4 A | 2.8312 ± 0.0088 | + +Default for genomics retrofits: H₂-style pattern (sparse attention, ~1 A every 12 blocks for a 24-block model; for 12 blocks, place one A near the middle and one near the output). + +## PDE fields (The Well, 2D or 3D) + +Same as vision, with: + +- `fft_padding = "circular"` — physics is on a torus or with reflecting BCs +- Patch size matters: smaller patches (p=2, p=4) widen HyenaND's advantage; defaults from `examples/well/v2/` +- Mask: per-axis Gaussian, isotropic init + +## When to use registers + FiLM (input-dependent kernels) + +Always for vision and PDE (§3.2.2). Default `num_registers = 4` for ViT-5-Small; scale with hidden_dim. The `KernelFiLMGenerator` + `RegisterPooling` combo from `_base_config.py` lines 276–337 is the canonical recipe. + +For genomics, FiLM is optional — the Evo2 striped configs in the paper do not use it. + +## Things that look like knobs but aren't really + +- `grid_type`: use `"single"` for non-registered inputs, `"double"` for CLS-row + registers (see `_base_config.py` line 302) +- `L_cache`: set to the maximum sequence length you'll see + 1 (CLS row); the kernel caches at this length +- `parametrization` for Gaussian mask: `"direct"` for vision, `"exp_decay"` for 1D AR + +## Reference configs to copy from + +| Use case | File | +|----------|------| +| Smallest end-to-end Hyena example | `examples/mnist_classification/ccnn_4_160_hyena_rope_qknorm.py` | +| ImageNet ViT-5 with FiLM + registers | `examples/vit5_imagenet/v5_patch/_base_config.py` | +| ImageNet hybrid (pattern-driven) | `examples/vit5_imagenet/vit5_hybrid/_base_config.py` | +| Diffusion (HF diffusers retrofit) | `examples/imagenet_diffusion/ccnn_12_768_hyena_qknorm.py` | +| PDE (The Well) | `examples/well/v2/*.py` | +| CCNN-style (non-ViT) classification | `examples/imagenet_classification/ccnn_7_512_hyena_circular.py` | From 6af6a77ac0b8371b35f4e5f08f5540067662f70e Mon Sep 17 00:00:00 2001 From: Farhad Ramezanghorbani Date: Wed, 20 May 2026 10:16:36 -0700 Subject: [PATCH 02/12] update skill Signed-off-by: Farhad Ramezanghorbani --- .claude/skills/hyenand-retrofit/SKILL.md | 162 ++++++++++++++---- .../skills/hyenand-retrofit/evals/evals.json | 23 +-- .../hyenand-retrofit/references/defaults.md | 113 +++++++----- 3 files changed, 209 insertions(+), 89 deletions(-) diff --git a/.claude/skills/hyenand-retrofit/SKILL.md b/.claude/skills/hyenand-retrofit/SKILL.md index 7f62f09d..e513f48a 100644 --- a/.claude/skills/hyenand-retrofit/SKILL.md +++ b/.claude/skills/hyenand-retrofit/SKILL.md @@ -20,9 +20,14 @@ If the user only wants conceptual explanation (no code), answer in chat. This sk Decide which one applies *before* writing anything: -1. **Native path** — the user's file already uses `nvsubquadratic` (`LazyConfig`, `build_attention_net`, `ViT5Attention`, `ViT5ClassificationNet`). The repo has matched `build_attention_net` / `build_hyena_net` / `build_hybrid_net` builders, so the swap is mechanical: replace the builder call, flip `compile_compatible_fftconv`, optionally pick a layer pattern. +1. **Native path** — the user's file already uses `nvsubquadratic` (`LazyConfig`, `build_attention_net`, `ViT5Attention`, `ViT5ClassificationNet`). The repo has matched builders: -2. **Foreign path** — the user's file uses generic PyTorch (`nn.MultiheadAttention`, `F.scaled_dot_product_attention`, `timm`, `transformers`, `diffusers`). You must construct a full Hyena module from scratch and wire it in as a drop-in attention replacement. + - `build_attention_net` + `build_hyena_net` live in `examples/vit5_imagenet/v5_patch/_base_config.py` (pure variants, fixed pattern). + - `build_hybrid_net` lives in a separate file: `examples/vit5_imagenet/vit5_hybrid/_base_config.py` (pattern-driven). Switching from a pure-attention v5_patch entry to a hybrid means changing both the import path and the entry directory, not just the function name. + + With the right builder picked, the swap is mechanical: replace the builder call, flip `compile_compatible_fftconv`, optionally pick a layer pattern. + +1. **Foreign path** — the user's file uses generic PyTorch (`nn.MultiheadAttention`, `F.scaled_dot_product_attention`, `timm`, `transformers`, `diffusers`). You must construct a full Hyena module from scratch and wire it in as a drop-in attention replacement. Look at the user's file. If you see `from nvsubquadratic` imports or `LazyConfig(ViT5Attention)`, take the native path. Otherwise foreign. @@ -48,22 +53,14 @@ Use AskUserQuestion if not already specified: Copy the entry-file structure verbatim, then change: -- Import: replace `build_attention_net` with `build_hyena_net` (pure) or `build_hybrid_net` (hybrid). Drop the unused builder. +- Import: replace `build_attention_net` with `build_hyena_net` (pure). For hybrid, import `build_hybrid_net` from `vit5_hybrid._base_config` (different module). Drop the unused builder. - Builder call inside `get_config()`: same swap. - For hybrid: add a `LAYER_PATTERN = "..." * (NUM_BLOCKS // len_repeat)` line at module scope, pass it as `layer_pattern=LAYER_PATTERN`. - Remove `config.compile_compatible_fftconv = False` if present (attention sets this False explicitly; the default is True, and HyenaND needs True). - Update the module docstring and any inline comments to reflect the new grid math (Hyena adds registers — see `hyena_patch16.py` for the canonical comment style). - Keep filename convention: `attention_patch16.py` → `hyena_patch16.py`, `full_attention.py` → `full_hyena.py` or `hybrid_hhha.py`. -### Step 3: smoke-test stub - -Append a `__main__` block that constructs the config and prints `repr(config.net)`. The user can run this on CPU to verify imports resolve and the LazyConfig graph is well-formed before committing GPU time. - -```python -if __name__ == "__main__": - cfg = get_config() - print(cfg.net) -``` +Native sibling files in this repo do **not** carry an inline `__main__` smoke block — the canonical `hyena_patch16.py` and `hybrid_hhha.py` are tiny config shims. Don't add one. If the user wants to validate the LazyConfig graph, point them at the existing `examples/vit5_imagenet/v5_patch/_smoke_test.py` pattern (separate file, not inline). ## Foreign path workflow @@ -94,58 +91,150 @@ from nvsubquadratic.utils.qk_norm import L2Norm import torch ``` -The full Hyena config has many knobs — see `references/defaults.md` for the canonical table. The short version of the defaults by modality: +The full Hyena config has many knobs. Note that only `use_rope` and `gate_nonlinear` are direct `Hyena(...)` kwargs — the rest (`data_dim`, `mask_cfg`, `fft_padding`) belong to the `CKConvND` config that you pass in as `global_conv_cfg`. See `references/defaults.md` for the canonical table and parameter ownership. Short version: -| Modality | data_dim | mask | fft_padding | use_rope | gate_nonlinear | typical ω₀ | -|----------|----------|------|-------------|----------|----------------|------------| -| Vision (image classification, diffusion) | 2 | identity or per-axis Gaussian | circular or zero | False | SiLU | 10 | -| Medical 3D segmentation | 3 | per-axis Gaussian | zero | False | SiLU | 10 | -| Genomics / causal LM | 1 | exponential decay with causal zeroing | causal | True | SiLU | 10 | -| PDE fields | 2 or 3 | per-axis Gaussian | circular | False | SiLU | 10 | +| Modality | data_dim (CKConvND) | mask (CKConvND) | fft_padding (CKConvND) | use_rope (Hyena) | gate_nonlinear (Hyena) | ω₀ (SIREN kernel) | +| ---------------------------------------- | ------------------- | ------------------------------------- | ---------------------- | ---------------- | ---------------------- | ----------------- | +| Vision (image classification, diffusion) | 2 | identity or per-axis Gaussian | circular or zero | False | SiLU | 10 | +| Medical 3D segmentation | 3 | per-axis Gaussian | zero | False | SiLU | 10 | +| Genomics / causal LM | 1 | exponential decay with causal zeroing | causal | True | SiLU | 100 | +| PDE fields | 2 or 3 | per-axis Gaussian | circular | False | SiLU | 10 | Read `references/defaults.md` for full parameter lists, init schemes, and the reasoning behind each choice. -### Step 3: wire it in +### Step 3: wire it in via an adapter (do not assign Hyena directly) -In the sibling file, subclass or monkey-patch the user's model to replace each attention module with the Hyena module. Prefer subclassing — it's easier to read and rollback. Example skeleton: +`nn.MultiheadAttention` and `Hyena` have **three incompatible interfaces** that bite you on the first forward pass if you do `block.attn = Hyena(...)` naively: + +1. **Return shape.** `nn.MultiheadAttention(..., need_weights=False)` returns a `(out, weights)` tuple; callers commonly unpack `h, _ = self.attn(...)`. `Hyena.forward` returns a single tensor. +1. **Kwargs.** `nn.MultiheadAttention` accepts `need_weights=`, `key_padding_mask=`, `attn_mask=`. `Hyena.forward(query, key, value, cp_group=None, **mixer_kwargs)` does not — extra kwargs flow into the global conv and may crash. +1. **Input layout.** `nn.MultiheadAttention(batch_first=True)` takes `[B, N, C]` (flat token sequence). `Hyena.forward` requires `[B, *spatial, C]` channel-last and reshapes internally to `[B, C, H, W]` (or `[B, C, T]`, `[B, C, D, H, W]`). For ViT, this means undoing the patch flatten — and **handling the CLS token separately**, because `1 + grid_h*grid_w` is never a clean rectangular grid. + +You need an adapter wrapper. Skeleton (vision, 2D, with CLS): ```python # my_model_hyenand.py — generated by hyenand-retrofit import torch +import torch.nn as nn from nvsubquadratic.lazy_config import instantiate, LazyConfig from nvsubquadratic.modules.hyena_nd import Hyena -# ... other imports per defaults.md ... +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 from my_model import MyViT # user's original -def build_hyena_mixer(hidden_dim: int, grid_h: int, grid_w: int) -> torch.nn.Module: - """Drop-in replacement for nn.MultiheadAttention(hidden_dim, num_heads).""" + +def build_hyena_mixer(hidden_dim: int, grid_h: int, grid_w: int) -> nn.Module: + """Instantiate a 2D HyenaND mixer for a `grid_h × grid_w` patch grid. + + See references/defaults.md for the full per-modality knob list — this is + the minimum that passes a forward pass; tune mask/short_conv/qk_norm + against the canonical vision config in + examples/vit5_imagenet/v5_patch/_base_config.py. + """ cfg = LazyConfig(Hyena)( global_conv_cfg=LazyConfig(CKConvND)( data_dim=2, hidden_dim=hidden_dim, kernel_cfg=LazyConfig(SIRENKernelND)( - data_dim=2, out_dim=hidden_dim, mlp_hidden_dim=32, - num_layers=3, embedding_dim=32, omega_0=10.0, - L_cache=max(grid_h, grid_w), use_bias=True, hidden_omega_0=1.0, + data_dim=2, + out_dim=hidden_dim, + mlp_hidden_dim=32, + num_layers=3, + embedding_dim=32, + omega_0=10.0, + L_cache=max(grid_h, grid_w), + use_bias=True, + hidden_omega_0=1.0, ), - # ... see defaults.md ... + mask_cfg=LazyConfig(GaussianModulationND)( + data_dim=2, + num_channels=hidden_dim, + min_attenuation_at_step=0.1, + max_attenuation_at_limit=0.95, + init_extent=1.0, + parametrization="direct", + ), + fft_padding="circular", # see defaults.md for "zero" alternative + ), + short_conv_cfg=LazyConfig(nn.Conv2d)( + in_channels=3 * hidden_dim, + out_channels=3 * hidden_dim, + kernel_size=3, + groups=3 * hidden_dim, + padding=1, + bias=False, ), - # ... see defaults.md ... + gate_nonlinear_cfg=LazyConfig(nn.SiLU)(), + gate_nonlinear_2_cfg=LazyConfig(nn.Sigmoid)(), + pixelhyena_norm_cfg=LazyConfig(nn.GroupNorm)( + num_groups=1, num_channels=hidden_dim + ), + qk_norm_cfg=LazyConfig(L2Norm)(), + use_rope=False, ) return instantiate(cfg) + +class HyenaAttnAdapter(nn.Module): + """Drop-in replacement for ``nn.MultiheadAttention(dim, heads, batch_first=True)``. + + Bridges the three interface gaps: + - reshapes ``[B, N, C]`` token sequence to ``[B, H, W, C]`` channel-last + before calling Hyena, then flattens back + - peels CLS off the front (and any other prefix tokens) so the + remaining ``H * W`` tokens form a clean spatial grid + - returns ``(out, None)`` so existing ``h, _ = self.attn(...)`` callers + keep working + - swallows ``need_weights``/``attn_mask`` kwargs that Hyena doesn't take + """ + + def __init__( + self, + hyena_mixer: nn.Module, + grid_h: int, + grid_w: int, + num_prefix_tokens: int = 1, + ): + super().__init__() + self.mixer = hyena_mixer + self.grid_h = grid_h + self.grid_w = grid_w + self.num_prefix_tokens = num_prefix_tokens # e.g. 1 for CLS, 0 for none + + def forward(self, query, key=None, value=None, **_kwargs): + # Self-attention: query == key == value. We ignore key/value. + x = query + prefix, patches = x[:, : self.num_prefix_tokens], x[:, self.num_prefix_tokens :] + B, N, C = patches.shape + assert ( + N == self.grid_h * self.grid_w + ), f"HyenaAttnAdapter: expected {self.grid_h * self.grid_w} patch tokens, got {N}" + patches_2d = patches.view(B, self.grid_h, self.grid_w, C) # [B, H, W, C] + out_2d = self.mixer(patches_2d, patches_2d, patches_2d) # Q=K=V self-mix + out = out_2d.view(B, N, C) + out = torch.cat([prefix, out], dim=1) + return out, None # (attn_out, attn_weights) shape contract + + class MyViTHyenaND(MyViT): - def __init__(self, *args, grid_h: int, grid_w: int, **kwargs): + def __init__( + self, *args, grid_h: int, grid_w: int, num_prefix_tokens: int = 1, **kwargs + ): super().__init__(*args, **kwargs) for block in self.blocks: - block.attn = build_hyena_mixer(self.embed_dim, grid_h, grid_w) + mixer = build_hyena_mixer(self.embed_dim, grid_h, grid_w) + block.attn = HyenaAttnAdapter(mixer, grid_h, grid_w, num_prefix_tokens) ``` +If the host model's attention call site does *not* unpack a tuple (e.g. `x = self.attn(h, h, h)` with no `_`), drop the `(out, None)` tuple and just return `out` from the adapter. If it has no CLS/register prefix, pass `num_prefix_tokens=0`. + ### Step 4: smoke-test stub -Append a `__main__` block that constructs the model and runs one forward pass on synthetic input of the user's stated shape. This catches shape mismatches before the user invests in training. +Append a `__main__` block that constructs the model and runs one forward pass on synthetic input of the user's stated shape. This catches shape mismatches before the user invests in training. Keep this stub for the foreign path — the user has no pre-existing harness, unlike the native path where the LazyConfig graph is exercised by the experiment runner. ## Filename and location convention @@ -159,15 +248,16 @@ Append a `__main__` block that constructs the model and runs one forward pass on After writing the file: 1. Read it back and confirm the changes you intended actually landed -2. Confirm the only difference from the user's file is the attention→Hyena swap plus any required toggles (`compile_compatible_fftconv`, layer pattern) -3. Confirm imports are syntactically correct (the user can `python -c "from import get_config"` to verify) -4. Confirm the smoke-test stub is present +1. Confirm the only difference from the user's file is the attention→Hyena swap plus any required toggles (`compile_compatible_fftconv`, layer pattern) +1. Confirm imports are syntactically correct (the user can `python -c "from import get_config"` to verify) +1. **Foreign path only:** confirm the smoke-test stub is present *and* that the adapter wires CLS / prefix tokens and the tuple return correctly ## What not to do - Do not modify the user's original file - Do not invent new builders if the matched `build_*_net` pair already exists in `_base_config.py` -- Do not skip the smoke-test stub — silent shape mismatches at training start are the #1 retrofit failure +- Do not skip the foreign-path smoke-test stub — silent shape mismatches at training start are the #1 retrofit failure. (For the native path, omit the stub — the canonical configs don't carry one and the experiment runner exercises the graph.) +- Do not assign `Hyena(...)` directly to an `nn.MultiheadAttention` slot — always wrap with an adapter (see Foreign path Step 3) to bridge the tuple-return, kwargs, and shape mismatches - Do not pick a hybrid pattern at random; either ask, or take the paper default for the modality (HHHA×3 for 2D vision, HHAA for 3D hierarchical, HHHA repeat for 1D genomics) - Do not omit `use_rope=True` for 1D autoregressive — without RoPE, positional recall collapses - Do not omit the per-axis Gaussian mask for ND≥2 unless the user explicitly opts out — it was the difference between Hyena and bidirectional Mamba in the §5.1 color_cond probes diff --git a/.claude/skills/hyenand-retrofit/evals/evals.json b/.claude/skills/hyenand-retrofit/evals/evals.json index 2c038792..403b0b9d 100644 --- a/.claude/skills/hyenand-retrofit/evals/evals.json +++ b/.claude/skills/hyenand-retrofit/evals/evals.json @@ -4,16 +4,15 @@ { "id": 1, "name": "native-pure-swap", - "prompt": "I have an attention-based config at /sessions/gracious-trusting-pasteur/mnt/nvSubquadratic-private/examples/vit5_imagenet/v5_patch/attention_patch16.py. 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, does not set compile_compatible_fftconv=False, and contains a __main__ smoke-test block.", - "files": ["/sessions/gracious-trusting-pasteur/mnt/nvSubquadratic-private/examples/vit5_imagenet/v5_patch/attention_patch16.py", "/sessions/gracious-trusting-pasteur/mnt/nvSubquadratic-private/examples/vit5_imagenet/v5_patch/_base_config.py"], - "ground_truth": "/sessions/gracious-trusting-pasteur/mnt/nvSubquadratic-private/examples/vit5_imagenet/v5_patch/hyena_patch16.py", + "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 contains a __main__ smoke-test block", "check": "grep_output:if __name__"}, {"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"} ] @@ -21,16 +20,15 @@ { "id": 2, "name": "native-hybrid", - "prompt": "I have a pure-attention config at /sessions/gracious-trusting-pasteur/mnt/nvSubquadratic-private/examples/vit5_imagenet/vit5_hybrid/full_attention.py. 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, and contains a __main__ smoke-test block.", - "files": ["/sessions/gracious-trusting-pasteur/mnt/nvSubquadratic-private/examples/vit5_imagenet/vit5_hybrid/full_attention.py", "/sessions/gracious-trusting-pasteur/mnt/nvSubquadratic-private/examples/vit5_imagenet/vit5_hybrid/_base_config.py"], - "ground_truth": "/sessions/gracious-trusting-pasteur/mnt/nvSubquadratic-private/examples/vit5_imagenet/vit5_hybrid/hybrid_hhha.py", + "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 contains a __main__ smoke-test block", "check": "grep_output:if __name__"}, {"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"} ] @@ -39,7 +37,7 @@ "id": 3, "name": "foreign-timm-vit", "prompt": "I have a tiny standalone PyTorch ViT (not using nvsubquadratic conventions) at .claude/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), subclass or wrap TinyViT to swap each block's .attn for the Hyena mixer, and contain a __main__ smoke-test block that runs one forward pass.", + "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": [".claude/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"}, @@ -48,6 +46,9 @@ {"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"} diff --git a/.claude/skills/hyenand-retrofit/references/defaults.md b/.claude/skills/hyenand-retrofit/references/defaults.md index 31d128bc..bab70038 100644 --- a/.claude/skills/hyenand-retrofit/references/defaults.md +++ b/.claude/skills/hyenand-retrofit/references/defaults.md @@ -2,18 +2,31 @@ These defaults are grounded in the paper and the repo's reference configs (`examples/imagenet_classification/ccnn_7_512_hyena*.py`, `examples/vit5_imagenet/v5_patch/_base_config.py`, `examples/well/v2/*.py`, `examples/mnist_classification/ccnn_4_160_hyena_rope_qknorm.py`). Use them as the starting point for any retrofit; tune only if there's a specific reason. +## Where each knob actually lives + +A common foot-gun is passing the wrong kwarg to the wrong constructor. The class boundaries: + +| Knob | Class that owns it | +| ------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------- | +| `use_rope`, `rope_base`, `gate_nonlinear_cfg`, `gate_nonlinear_2_cfg`, `pixelhyena_norm_cfg`, `output_norm_cfg`, `qk_norm_cfg`, `short_conv_cfg`, `global_conv_cfg` | `Hyena` | +| `data_dim`, `hidden_dim`, `mask_cfg`, `fft_padding`, `kernel_cfg` | `CKConvND` (passed in as `Hyena.global_conv_cfg`) | +| `omega_0`, `hidden_omega_0`, `mlp_hidden_dim`, `num_layers`, `embedding_dim`, `L_cache`, `use_bias`, `out_dim` | `SIRENKernelND` (passed in as `CKConvND.kernel_cfg`) | +| `min_attenuation_at_step`, `max_attenuation_at_limit`, `init_extent`, `parametrization`, `num_channels` | `GaussianModulationND` (passed in as `CKConvND.mask_cfg`) | + +When the table below says "vision uses `use_rope=False`", that means the outer `Hyena(...)` constructor. When it says "vision uses `fft_padding=circular`", that means the inner `CKConvND(...)`. + ## Common (all modalities) -| Knob | Value | Why | -|------|-------|-----| -| `omega_0` (SIREN first-layer ω) | 10.0 for vision/PDE, 100.0 for genomics | Frequency band; higher for 1D where positional recall matters more | -| `hidden_omega_0` | 1.0 | SIREN hidden-layer ω; preserves init scale | -| `mlp_hidden_dim` (SIREN width) | 32 | Small MLP; SIREN is shallow but wide enough to express smooth kernels | -| `num_layers` (SIREN depth) | 3 | Paper-standard depth | -| `embedding_dim` (SIREN coordinate embedding) | 32 | First-layer positional embedding | -| `use_bias` | True | | -| First-layer init | `U(-1/sqrt(N), 1/sqrt(N))` | Preserves unit-variance pre-activations (§3.2.3) | -| FiLM heads init | zero | Training starts from the unmodulated baseline (§3.2.3) | +| Knob | Value | Why | +| -------------------------------------------- | --------------------------------------- | --------------------------------------------------------------------- | +| `omega_0` (SIREN first-layer ω) | 10.0 for vision/PDE, 100.0 for genomics | Frequency band; higher for 1D where positional recall matters more | +| `hidden_omega_0` | 1.0 | SIREN hidden-layer ω; preserves init scale | +| `mlp_hidden_dim` (SIREN width) | 32 | Small MLP; SIREN is shallow but wide enough to express smooth kernels | +| `num_layers` (SIREN depth) | 3 | Paper-standard depth | +| `embedding_dim` (SIREN coordinate embedding) | 32 | First-layer positional embedding | +| `use_bias` | True | | +| First-layer init | `U(-1/sqrt(N), 1/sqrt(N))` | Preserves unit-variance pre-activations (§3.2.3) | +| FiLM heads init | zero | Training starts from the unmodulated baseline (§3.2.3) | ## Vision (2D, image classification or diffusion) @@ -32,16 +45,23 @@ mask_cfg = LazyConfig(GaussianModulationND)( gate_nonlinear_cfg = LazyConfig(torch.nn.SiLU)() gate_nonlinear_2_cfg = LazyConfig(torch.nn.Sigmoid)() qk_norm_cfg = LazyConfig(L2Norm)() -pixelhyena_norm_cfg = LazyConfig(torch.nn.GroupNorm)(num_groups=1, num_channels=hidden_dim) +pixelhyena_norm_cfg = LazyConfig(torch.nn.GroupNorm)( + num_groups=1, num_channels=hidden_dim +) short_conv = LazyConfig(torch.nn.Conv2d)( - in_channels=3 * hidden_dim, out_channels=3 * hidden_dim, - kernel_size=3, groups=3 * hidden_dim, padding=1, bias=False, + in_channels=3 * hidden_dim, + out_channels=3 * hidden_dim, + kernel_size=3, + groups=3 * hidden_dim, + padding=1, + bias=False, ) ``` For ViT-5-style backbones with registers, use `ViT5HyenaAdapter` and `RegisterPooling` (see `examples/vit5_imagenet/v5_patch/_base_config.py` lines 270–353). **Hybrid patterns (12 blocks):** + - Best on ImageNet: `(HA)×6` (82.1 top-1) - Second-best: `(HHHA)×3` (82.0 top-1) - Pure: `H×12` (81.5 top-1, matches attention baseline) @@ -54,19 +74,23 @@ use_rope = False fft_padding = "zero" mask_cfg = LazyConfig(GaussianModulationND)(data_dim=3, num_channels=hidden_dim, ...) short_conv = LazyConfig(torch.nn.Conv3d)( - in_channels=3 * hidden_dim, out_channels=3 * hidden_dim, - kernel_size=3, groups=3 * hidden_dim, padding=1, bias=False, + in_channels=3 * hidden_dim, + out_channels=3 * hidden_dim, + kernel_size=3, + groups=3 * hidden_dim, + padding=1, + bias=False, ) ``` **Stage patterns (4-stage encoder, paper §5.5 PanTS, mean Dice):** -| Pattern | Stage 1 | Stage 2 | Stage 3 | Stage 4 | Mean Dice | -|---------|---------|---------|---------|---------|-----------| -| `AAAA` (Swin baseline) | Attn | Attn | Attn | Attn | 0.7496 | -| `HHHH` (all-Hyena) | Hyena | Hyena | Hyena | Hyena | 0.7510 | -| `HAHA` (striped) | Hyena | Attn | Hyena | Attn | 0.7535 | -| `HHAA` (hierarchical, **best**) | Hyena | Hyena | Attn | Attn | **0.7559** | +| Pattern | Stage 1 | Stage 2 | Stage 3 | Stage 4 | Mean Dice | +| ------------------------------- | ------- | ------- | ------- | ------- | ---------- | +| `AAAA` (Swin baseline) | Attn | Attn | Attn | Attn | 0.7496 | +| `HHHH` (all-Hyena) | Hyena | Hyena | Hyena | Hyena | 0.7510 | +| `HAHA` (striped) | Hyena | Attn | Hyena | Attn | 0.7535 | +| `HHAA` (hierarchical, **best**) | Hyena | Hyena | Attn | Attn | **0.7559** | Recommend `HHAA` by default. Memory savings (~11%) are roughly invariant to placement. @@ -74,11 +98,12 @@ Recommend `HHAA` by default. Memory savings (~11%) are roughly invariant to plac ```python data_dim = 1 -use_rope = True # critical — without RoPE, positional recall collapses +use_rope = True # critical — without RoPE, positional recall collapses rope_base = 10000.0 -fft_padding = "causal" # preserves autoregressive structure +fft_padding = "causal" # preserves autoregressive structure mask_cfg = LazyConfig(GaussianModulationND)( # exponential decay with causal zeroing - data_dim=1, num_channels=hidden_dim, + data_dim=1, + num_channels=hidden_dim, min_attenuation_at_step=0.1, max_attenuation_at_limit=0.95, init_extent=1.0, @@ -86,21 +111,25 @@ mask_cfg = LazyConfig(GaussianModulationND)( # exponential decay with causal ze ) gate_nonlinear_cfg = LazyConfig(torch.nn.Identity)() # linear gating for AR short_conv = LazyConfig(torch.nn.Conv1d)( - in_channels=3 * hidden_dim, out_channels=3 * hidden_dim, - kernel_size=3, groups=3 * hidden_dim, padding=1, bias=False, + in_channels=3 * hidden_dim, + out_channels=3 * hidden_dim, + kernel_size=3, + groups=3 * hidden_dim, + padding=1, + bias=False, ) ``` **Mixing ratios (Evo2-1B, 8192-bp sequences, §5.2, perplexity lower is better):** -| Config | Pattern (24 blocks) | Validation PPL | -|--------|---------------------|----------------| -| `T` (full transformer) | 24 attention | 2.9235 ± 0.0039 | -| `H₀` (full Hyena) | 24 Hyena | 2.8282 ± 0.0279 | -| `H₁` (1 MHA) | 23 H + 1 A | 2.8308 ± 0.0108 | -| `H₂` (2 MHA, **best**) | 22 H + 2 A | **2.7729 ± 0.0006** | -| `H₃` (3 MHA) | 21 H + 3 A | 2.8214 ± 0.0313 | -| `H₄` (4 MHA) | 20 H + 4 A | 2.8312 ± 0.0088 | +| Config | Pattern (24 blocks) | Validation PPL | +| ---------------------- | ------------------- | ------------------- | +| `T` (full transformer) | 24 attention | 2.9235 ± 0.0039 | +| `H₀` (full Hyena) | 24 Hyena | 2.8282 ± 0.0279 | +| `H₁` (1 MHA) | 23 H + 1 A | 2.8308 ± 0.0108 | +| `H₂` (2 MHA, **best**) | 22 H + 2 A | **2.7729 ± 0.0006** | +| `H₃` (3 MHA) | 21 H + 3 A | 2.8214 ± 0.0313 | +| `H₄` (4 MHA) | 20 H + 4 A | 2.8312 ± 0.0088 | Default for genomics retrofits: H₂-style pattern (sparse attention, ~1 A every 12 blocks for a 24-block model; for 12 blocks, place one A near the middle and one near the output). @@ -126,11 +155,11 @@ For genomics, FiLM is optional — the Evo2 striped configs in the paper do not ## Reference configs to copy from -| Use case | File | -|----------|------| -| Smallest end-to-end Hyena example | `examples/mnist_classification/ccnn_4_160_hyena_rope_qknorm.py` | -| ImageNet ViT-5 with FiLM + registers | `examples/vit5_imagenet/v5_patch/_base_config.py` | -| ImageNet hybrid (pattern-driven) | `examples/vit5_imagenet/vit5_hybrid/_base_config.py` | -| Diffusion (HF diffusers retrofit) | `examples/imagenet_diffusion/ccnn_12_768_hyena_qknorm.py` | -| PDE (The Well) | `examples/well/v2/*.py` | -| CCNN-style (non-ViT) classification | `examples/imagenet_classification/ccnn_7_512_hyena_circular.py` | +| Use case | File | +| ------------------------------------ | --------------------------------------------------------------- | +| Smallest end-to-end Hyena example | `examples/mnist_classification/ccnn_4_160_hyena_rope_qknorm.py` | +| ImageNet ViT-5 with FiLM + registers | `examples/vit5_imagenet/v5_patch/_base_config.py` | +| ImageNet hybrid (pattern-driven) | `examples/vit5_imagenet/vit5_hybrid/_base_config.py` | +| Diffusion (HF diffusers retrofit) | `examples/imagenet_diffusion/ccnn_12_768_hyena_qknorm.py` | +| PDE (The Well) | `examples/well/v2/*.py` | +| CCNN-style (non-ViT) classification | `examples/imagenet_classification/ccnn_7_512_hyena_circular.py` | From 17997533e5752c89c8f1903d6bcd37b263ac8908 Mon Sep 17 00:00:00 2001 From: Farhad Ramezanghorbani Date: Wed, 20 May 2026 11:12:56 -0700 Subject: [PATCH 03/12] generalize to all dims Signed-off-by: Farhad Ramezanghorbani --- .claude/skills/hyenand-retrofit/SKILL.md | 355 ++++++++---------- .../skills/hyenand-retrofit/evals/evals.json | 44 +++ .../hyenand-retrofit/references/defaults.md | 165 -------- 3 files changed, 193 insertions(+), 371 deletions(-) delete mode 100644 .claude/skills/hyenand-retrofit/references/defaults.md diff --git a/.claude/skills/hyenand-retrofit/SKILL.md b/.claude/skills/hyenand-retrofit/SKILL.md index e513f48a..3ed70bf2 100644 --- a/.claude/skills/hyenand-retrofit/SKILL.md +++ b/.claude/skills/hyenand-retrofit/SKILL.md @@ -1,267 +1,210 @@ --- name: hyenand-retrofit -description: Replace attention in a PyTorch model with HyenaND from the nvSubquadratic library, with paper-grounded default hyperparameters (SIREN ω₀, register count, per-axis Gaussian mask, FFT padding, hybrid layer pattern). Use this skill whenever the user wants to swap attention for a subquadratic operator, port a ViT / U-Net / diffusion / genomics LM to HyenaND, convert a `nn.MultiheadAttention` or `F.scaled_dot_product_attention` site to a Hyena mixer, build a striped Hyena LM, set up a HyenaND experiment config, or asks "how do I use nvSubquadratic with my model." Trigger even when the user does not explicitly name HyenaND — phrases like "make my ViT subquadratic," "subquadratic alternative to attention for my 3D U-Net," "Hyena layer for my transformer," "swap attention with FFT convolution," or "long-context model with O(L log L) scaling" should all activate this skill. +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. - -## When to use - -- User has an attention-based model (ViT, U-Net, DiT, causal LM, hierarchical encoder) and wants a subquadratic alternative -- User explicitly mentions HyenaND, Hyena, nSubQ, nvSubquadratic, or "subquadratic attention" -- User is inside the `nvSubquadratic-private` repo and wants to add a new attention/Hyena variant -- User asks about scaling to long contexts (long genomes, high-resolution images, 3D volumes, PDE grids) +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. -## The two paths - -Decide which one applies *before* writing anything: - -1. **Native path** — the user's file already uses `nvsubquadratic` (`LazyConfig`, `build_attention_net`, `ViT5Attention`, `ViT5ClassificationNet`). The repo has matched builders: - - - `build_attention_net` + `build_hyena_net` live in `examples/vit5_imagenet/v5_patch/_base_config.py` (pure variants, fixed pattern). - - `build_hybrid_net` lives in a separate file: `examples/vit5_imagenet/vit5_hybrid/_base_config.py` (pattern-driven). Switching from a pure-attention v5_patch entry to a hybrid means changing both the import path and the entry directory, not just the function name. - - With the right builder picked, the swap is mechanical: replace the builder call, flip `compile_compatible_fftconv`, optionally pick a layer pattern. +## Native path (user is already inside nvSubquadratic) -1. **Foreign path** — the user's file uses generic PyTorch (`nn.MultiheadAttention`, `F.scaled_dot_product_attention`, `timm`, `transformers`, `diffusers`). You must construct a full Hyena module from scratch and wire it in as a drop-in attention replacement. +If the user's file imports nvSubquadratic builders (`build_attention_net`, `LazyConfig(ViT5Attention)`, etc.), the swap is mechanical: -Look at the user's file. If you see `from nvsubquadratic` imports or `LazyConfig(ViT5Attention)`, take the native path. Otherwise foreign. +- 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 path workflow +Native sibling files are bare config shims — no `__main__` block; the experiment runner exercises the LazyConfig graph. -The repo factors attention vs Hyena into builder functions. The user's entry file is almost always just a thin wrapper around one of: +The rest of this skill covers the **foreign path** — generic PyTorch hosts using `nn.MultiheadAttention`, `F.scaled_dot_product_attention`, timm, HF, etc. -- `build_attention_net(patch_size)` — pure attention -- `build_hyena_net(patch_size)` — pure Hyena -- `build_hybrid_net(layer_pattern="...", patch_size=...)` — mixed, where pattern is a string of `H`/`A` characters (one per block) +## Decide four things up front -### Step 1: ask the user which target +These four axes are orthogonal. Fix them before writing. -Use AskUserQuestion if not already specified: +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`. -- **Pure Hyena** — strongest when geometry + global structure dominate (PDE fields, high-resolution vision, long genomes) -- **Hybrid** — strongest when selectivity still matters (ImageNet classification, medical segmentation). Paper winners: - - 2D ImageNet ViT-Small: `(HA)×6` (best) or `(HHHA)×3` - - 3D medical SwinUNETR-style: `HHAA` (hierarchical — Hyena in early high-resolution stages, attention later) - - 1D genomics (1B striped LM, Evo2): one A per 4 blocks (`HHHA` repeat, H₂ mixing) was best in §5.2 +1. **`causal ∈ {True, False}`** — autoregressive 1D LMs are causal; vision, segmentation, PDE are bidirectional. Causal sets `fft_padding="causal"`, `use_rope=True`, mask `parametrization="exp_decay"`, `omega_0=100`. Bidirectional sets `fft_padding ∈ {"zero", "circular"}`, `use_rope=False`, `parametrization="direct"`, `omega_0=10`. -### Step 2: write the sibling file +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]`. -Copy the entry-file structure verbatim, then change: +1. **Return contract** — `(out, None)` tuple if the call site is `h, _ = self.attn(...)` (matches `nn.MultiheadAttention`); bare tensor otherwise. -- Import: replace `build_attention_net` with `build_hyena_net` (pure). For hybrid, import `build_hybrid_net` from `vit5_hybrid._base_config` (different module). Drop the unused builder. -- Builder call inside `get_config()`: same swap. -- For hybrid: add a `LAYER_PATTERN = "..." * (NUM_BLOCKS // len_repeat)` line at module scope, pass it as `layer_pattern=LAYER_PATTERN`. -- Remove `config.compile_compatible_fftconv = False` if present (attention sets this False explicitly; the default is True, and HyenaND needs True). -- Update the module docstring and any inline comments to reflect the new grid math (Hyena adds registers — see `hyena_patch16.py` for the canonical comment style). -- Keep filename convention: `attention_patch16.py` → `hyena_patch16.py`, `full_attention.py` → `full_hyena.py` or `hybrid_hhha.py`. +## Pure or hybrid -Native sibling files in this repo do **not** carry an inline `__main__` smoke block — the canonical `hyena_patch16.py` and `hybrid_hhha.py` are tiny config shims. Don't add one. If the user wants to validate the LazyConfig graph, point them at the existing `examples/vit5_imagenet/v5_patch/_smoke_test.py` pattern (separate file, not inline). +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 via AskUserQuestion. -## Foreign path workflow +- **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 user has, e.g., a `timm` ViT, a HF transformers model, or a hand-written PyTorch transformer. You need to construct a Hyena module yourself. +## The Hyena module -### Step 1: identify the attention sites - -Look for any of: - -- `nn.MultiheadAttention` -- `F.scaled_dot_product_attention` -- `flash_attn.*` calls -- Custom attention modules (look for `softmax(Q @ K.T / sqrt(d))` or `attn_drop`/`proj_drop` member names) -- timm `Attention` class, HF `*Attention` classes - -### Step 2: build the Hyena replacement - -For each attention site, emit a Hyena module config. The minimum spec: +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 +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.modules.rms_norm import RMSNorm from nvsubquadratic.utils.qk_norm import L2Norm -import torch -``` - -The full Hyena config has many knobs. Note that only `use_rope` and `gate_nonlinear` are direct `Hyena(...)` kwargs — the rest (`data_dim`, `mask_cfg`, `fft_padding`) belong to the `CKConvND` config that you pass in as `global_conv_cfg`. See `references/defaults.md` for the canonical table and parameter ownership. Short version: - -| Modality | data_dim (CKConvND) | mask (CKConvND) | fft_padding (CKConvND) | use_rope (Hyena) | gate_nonlinear (Hyena) | ω₀ (SIREN kernel) | -| ---------------------------------------- | ------------------- | ------------------------------------- | ---------------------- | ---------------- | ---------------------- | ----------------- | -| Vision (image classification, diffusion) | 2 | identity or per-axis Gaussian | circular or zero | False | SiLU | 10 | -| Medical 3D segmentation | 3 | per-axis Gaussian | zero | False | SiLU | 10 | -| Genomics / causal LM | 1 | exponential decay with causal zeroing | causal | True | SiLU | 100 | -| PDE fields | 2 or 3 | per-axis Gaussian | circular | False | SiLU | 10 | - -Read `references/defaults.md` for full parameter lists, init schemes, and the reasoning behind each choice. - -### Step 3: wire it in via an adapter (do not assign Hyena directly) - -`nn.MultiheadAttention` and `Hyena` have **three incompatible interfaces** that bite you on the first forward pass if you do `block.attn = Hyena(...)` naively: - -1. **Return shape.** `nn.MultiheadAttention(..., need_weights=False)` returns a `(out, weights)` tuple; callers commonly unpack `h, _ = self.attn(...)`. `Hyena.forward` returns a single tensor. -1. **Kwargs.** `nn.MultiheadAttention` accepts `need_weights=`, `key_padding_mask=`, `attn_mask=`. `Hyena.forward(query, key, value, cp_group=None, **mixer_kwargs)` does not — extra kwargs flow into the global conv and may crash. -1. **Input layout.** `nn.MultiheadAttention(batch_first=True)` takes `[B, N, C]` (flat token sequence). `Hyena.forward` requires `[B, *spatial, C]` channel-last and reshapes internally to `[B, C, H, W]` (or `[B, C, T]`, `[B, C, D, H, W]`). For ViT, this means undoing the patch flatten — and **handling the CLS token separately**, because `1 + grid_h*grid_w` is never a clean rectangular grid. - -You need an adapter wrapper. Skeleton (vision, 2D, with CLS): - -```python -# my_model_hyenand.py — generated by hyenand-retrofit - import torch import torch.nn as nn -from nvsubquadratic.lazy_config import instantiate, LazyConfig -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 -from my_model import MyViT # user's original - - -def build_hyena_mixer(hidden_dim: int, grid_h: int, grid_w: int) -> nn.Module: - """Instantiate a 2D HyenaND mixer for a `grid_h × grid_w` patch grid. - - See references/defaults.md for the full per-modality knob list — this is - the minimum that passes a forward pass; tune mask/short_conv/qk_norm - against the canonical vision config in - examples/vit5_imagenet/v5_patch/_base_config.py. - """ - cfg = LazyConfig(Hyena)( - global_conv_cfg=LazyConfig(CKConvND)( - data_dim=2, - hidden_dim=hidden_dim, - kernel_cfg=LazyConfig(SIRENKernelND)( - data_dim=2, - out_dim=hidden_dim, - mlp_hidden_dim=32, - num_layers=3, - embedding_dim=32, - omega_0=10.0, - L_cache=max(grid_h, grid_w), - use_bias=True, - hidden_omega_0=1.0, - ), - mask_cfg=LazyConfig(GaussianModulationND)( - data_dim=2, - num_channels=hidden_dim, - min_attenuation_at_step=0.1, - max_attenuation_at_limit=0.95, - init_extent=1.0, - parametrization="direct", - ), - fft_padding="circular", # see defaults.md for "zero" alternative - ), - short_conv_cfg=LazyConfig(nn.Conv2d)( - in_channels=3 * hidden_dim, - out_channels=3 * hidden_dim, - kernel_size=3, - groups=3 * hidden_dim, - padding=1, - bias=False, - ), +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=( + "causal" if CAUSAL else "zero" + ), # "circular" is a valid bidirectional alt +) + +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)(), - gate_nonlinear_2_cfg=LazyConfig(nn.Sigmoid)(), pixelhyena_norm_cfg=LazyConfig(nn.GroupNorm)( - num_groups=1, num_channels=hidden_dim + num_groups=1, num_channels=HIDDEN_DIM ), qk_norm_cfg=LazyConfig(L2Norm)(), - use_rope=False, + use_rope=CAUSAL, + rope_base=10000.0, ) - return instantiate(cfg) +) +``` + +**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`, `kernel_cfg`. Also optional `grid_type` (`"single"`/`"double"`) and `use_chunked_fftconv` (memory optimization; requires `fft_padding="zero"`). +- `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`. + +## 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 replacement for ``nn.MultiheadAttention(dim, heads, batch_first=True)``. - - Bridges the three interface gaps: - - reshapes ``[B, N, C]`` token sequence to ``[B, H, W, C]`` channel-last - before calling Hyena, then flattens back - - peels CLS off the front (and any other prefix tokens) so the - remaining ``H * W`` tokens form a clean spatial grid - - returns ``(out, None)`` so existing ``h, _ = self.attn(...)`` callers - keep working - - swallows ``need_weights``/``attn_mask`` kwargs that Hyena doesn't take - """ + """Drop-in attention replacement. Parameters reflect the four-axis decision.""" def __init__( self, - hyena_mixer: nn.Module, - grid_h: int, - grid_w: int, - num_prefix_tokens: int = 1, + 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 = hyena_mixer - self.grid_h = grid_h - self.grid_w = grid_w - self.num_prefix_tokens = num_prefix_tokens # e.g. 1 for CLS, 0 for none - - def forward(self, query, key=None, value=None, **_kwargs): - # Self-attention: query == key == value. We ignore key/value. - x = query - prefix, patches = x[:, : self.num_prefix_tokens], x[:, self.num_prefix_tokens :] - B, N, C = patches.shape - assert ( - N == self.grid_h * self.grid_w - ), f"HyenaAttnAdapter: expected {self.grid_h * self.grid_w} patch tokens, got {N}" - patches_2d = patches.view(B, self.grid_h, self.grid_w, C) # [B, H, W, C] - out_2d = self.mixer(patches_2d, patches_2d, patches_2d) # Q=K=V self-mix - out = out_2d.view(B, N, C) - out = torch.cat([prefix, out], dim=1) - return out, None # (attn_out, attn_weights) shape contract - - -class MyViTHyenaND(MyViT): - def __init__( - self, *args, grid_h: int, grid_w: int, num_prefix_tokens: int = 1, **kwargs - ): - super().__init__(*args, **kwargs) - for block in self.blocks: - mixer = build_hyena_mixer(self.embed_dim, grid_h, grid_w) - block.attn = HyenaAttnAdapter(mixer, grid_h, grid_w, num_prefix_tokens) + 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 ``` -If the host model's attention call site does *not* unpack a tuple (e.g. `x = self.attn(h, h, h)` with no `_`), drop the `(out, None)` tuple and just return `out` from the adapter. If it has no CLS/register prefix, pass `num_prefix_tokens=0`. +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). + +## 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`. Fix: use `DepthwiseFFTConv3d` from nvSubquadratic for the short conv (eliminates `im2col`). 1D and 2D rarely hit this. -### Step 4: smoke-test stub +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. -Append a `__main__` block that constructs the model and runs one forward pass on synthetic input of the user's stated shape. This catches shape mismatches before the user invests in training. Keep this stub for the foreign path — the user has no pre-existing harness, unlike the native path where the LazyConfig graph is exercised by the experiment runner. +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. + +## 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 -- Name: replace `attention` with `hyena` (or `hybrid_` for hybrid), keep all other tokens -- If the user's file has no `attention` token in the name, append `_hyenand` before the extension -- Do not edit the user's original — keep the diff trivial +- 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 the file: - -1. Read it back and confirm the changes you intended actually landed -1. Confirm the only difference from the user's file is the attention→Hyena swap plus any required toggles (`compile_compatible_fftconv`, layer pattern) -1. Confirm imports are syntactically correct (the user can `python -c "from import get_config"` to verify) -1. **Foreign path only:** confirm the smoke-test stub is present *and* that the adapter wires CLS / prefix tokens and the tuple return correctly +After writing: -## What not to do +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. -- Do not modify the user's original file -- Do not invent new builders if the matched `build_*_net` pair already exists in `_base_config.py` -- Do not skip the foreign-path smoke-test stub — silent shape mismatches at training start are the #1 retrofit failure. (For the native path, omit the stub — the canonical configs don't carry one and the experiment runner exercises the graph.) -- Do not assign `Hyena(...)` directly to an `nn.MultiheadAttention` slot — always wrap with an adapter (see Foreign path Step 3) to bridge the tuple-return, kwargs, and shape mismatches -- Do not pick a hybrid pattern at random; either ask, or take the paper default for the modality (HHHA×3 for 2D vision, HHAA for 3D hierarchical, HHHA repeat for 1D genomics) -- Do not omit `use_rope=True` for 1D autoregressive — without RoPE, positional recall collapses -- Do not omit the per-axis Gaussian mask for ND≥2 unless the user explicitly opts out — it was the difference between Hyena and bidirectional Mamba in the §5.1 color_cond probes +## Reference configs in this repo -## References +Copy parameter values, not whole files: -- `references/defaults.md` — full per-modality default tables, init schemes, and rationale +| Use case | File | +| --------------------------------- | --------------------------------------------------------------- | +| Smallest end-to-end Hyena example | `examples/mnist_classification/ccnn_4_160_hyena_rope_qknorm.py` | +| ImageNet ViT-5 (FiLM + registers) | `examples/vit5_imagenet/v5_patch/_base_config.py` | +| ImageNet hybrid (pattern-driven) | `examples/vit5_imagenet/vit5_hybrid/_base_config.py` | +| Diffusion (HF diffusers retrofit) | `examples/imagenet_diffusion/ccnn_12_768_hyena_qknorm.py` | +| PDE fields (The Well) | `examples/well/v2/*.py` | diff --git a/.claude/skills/hyenand-retrofit/evals/evals.json b/.claude/skills/hyenand-retrofit/evals/evals.json index 403b0b9d..fb3dd476 100644 --- a/.claude/skills/hyenand-retrofit/evals/evals.json +++ b/.claude/skills/hyenand-retrofit/evals/evals.json @@ -53,6 +53,50 @@ {"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 .claude/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": [".claude/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 .claude/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: use_rope=True, fft_padding='causal', 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": [".claude/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 uses causal FFT padding", "check": "grep_output: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/.claude/skills/hyenand-retrofit/references/defaults.md b/.claude/skills/hyenand-retrofit/references/defaults.md deleted file mode 100644 index bab70038..00000000 --- a/.claude/skills/hyenand-retrofit/references/defaults.md +++ /dev/null @@ -1,165 +0,0 @@ -# HyenaND default hyperparameters by modality - -These defaults are grounded in the paper and the repo's reference configs (`examples/imagenet_classification/ccnn_7_512_hyena*.py`, `examples/vit5_imagenet/v5_patch/_base_config.py`, `examples/well/v2/*.py`, `examples/mnist_classification/ccnn_4_160_hyena_rope_qknorm.py`). Use them as the starting point for any retrofit; tune only if there's a specific reason. - -## Where each knob actually lives - -A common foot-gun is passing the wrong kwarg to the wrong constructor. The class boundaries: - -| Knob | Class that owns it | -| ------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------- | -| `use_rope`, `rope_base`, `gate_nonlinear_cfg`, `gate_nonlinear_2_cfg`, `pixelhyena_norm_cfg`, `output_norm_cfg`, `qk_norm_cfg`, `short_conv_cfg`, `global_conv_cfg` | `Hyena` | -| `data_dim`, `hidden_dim`, `mask_cfg`, `fft_padding`, `kernel_cfg` | `CKConvND` (passed in as `Hyena.global_conv_cfg`) | -| `omega_0`, `hidden_omega_0`, `mlp_hidden_dim`, `num_layers`, `embedding_dim`, `L_cache`, `use_bias`, `out_dim` | `SIRENKernelND` (passed in as `CKConvND.kernel_cfg`) | -| `min_attenuation_at_step`, `max_attenuation_at_limit`, `init_extent`, `parametrization`, `num_channels` | `GaussianModulationND` (passed in as `CKConvND.mask_cfg`) | - -When the table below says "vision uses `use_rope=False`", that means the outer `Hyena(...)` constructor. When it says "vision uses `fft_padding=circular`", that means the inner `CKConvND(...)`. - -## Common (all modalities) - -| Knob | Value | Why | -| -------------------------------------------- | --------------------------------------- | --------------------------------------------------------------------- | -| `omega_0` (SIREN first-layer ω) | 10.0 for vision/PDE, 100.0 for genomics | Frequency band; higher for 1D where positional recall matters more | -| `hidden_omega_0` | 1.0 | SIREN hidden-layer ω; preserves init scale | -| `mlp_hidden_dim` (SIREN width) | 32 | Small MLP; SIREN is shallow but wide enough to express smooth kernels | -| `num_layers` (SIREN depth) | 3 | Paper-standard depth | -| `embedding_dim` (SIREN coordinate embedding) | 32 | First-layer positional embedding | -| `use_bias` | True | | -| First-layer init | `U(-1/sqrt(N), 1/sqrt(N))` | Preserves unit-variance pre-activations (§3.2.3) | -| FiLM heads init | zero | Training starts from the unmodulated baseline (§3.2.3) | - -## Vision (2D, image classification or diffusion) - -```python -data_dim = 2 -use_rope = False -fft_padding = "circular" # or "zero" for CLS-row variants -mask_cfg = LazyConfig(GaussianModulationND)( - data_dim=2, - num_channels=hidden_dim, - min_attenuation_at_step=0.1, - max_attenuation_at_limit=0.95, - init_extent=1.0, - parametrization="direct", -) -gate_nonlinear_cfg = LazyConfig(torch.nn.SiLU)() -gate_nonlinear_2_cfg = LazyConfig(torch.nn.Sigmoid)() -qk_norm_cfg = LazyConfig(L2Norm)() -pixelhyena_norm_cfg = LazyConfig(torch.nn.GroupNorm)( - num_groups=1, num_channels=hidden_dim -) -short_conv = LazyConfig(torch.nn.Conv2d)( - in_channels=3 * hidden_dim, - out_channels=3 * hidden_dim, - kernel_size=3, - groups=3 * hidden_dim, - padding=1, - bias=False, -) -``` - -For ViT-5-style backbones with registers, use `ViT5HyenaAdapter` and `RegisterPooling` (see `examples/vit5_imagenet/v5_patch/_base_config.py` lines 270–353). - -**Hybrid patterns (12 blocks):** - -- Best on ImageNet: `(HA)×6` (82.1 top-1) -- Second-best: `(HHHA)×3` (82.0 top-1) -- Pure: `H×12` (81.5 top-1, matches attention baseline) - -## Medical 3D segmentation (SwinUNETR-style hierarchical encoders) - -```python -data_dim = 3 -use_rope = False -fft_padding = "zero" -mask_cfg = LazyConfig(GaussianModulationND)(data_dim=3, num_channels=hidden_dim, ...) -short_conv = LazyConfig(torch.nn.Conv3d)( - in_channels=3 * hidden_dim, - out_channels=3 * hidden_dim, - kernel_size=3, - groups=3 * hidden_dim, - padding=1, - bias=False, -) -``` - -**Stage patterns (4-stage encoder, paper §5.5 PanTS, mean Dice):** - -| Pattern | Stage 1 | Stage 2 | Stage 3 | Stage 4 | Mean Dice | -| ------------------------------- | ------- | ------- | ------- | ------- | ---------- | -| `AAAA` (Swin baseline) | Attn | Attn | Attn | Attn | 0.7496 | -| `HHHH` (all-Hyena) | Hyena | Hyena | Hyena | Hyena | 0.7510 | -| `HAHA` (striped) | Hyena | Attn | Hyena | Attn | 0.7535 | -| `HHAA` (hierarchical, **best**) | Hyena | Hyena | Attn | Attn | **0.7559** | - -Recommend `HHAA` by default. Memory savings (~11%) are roughly invariant to placement. - -## Genomics / 1D causal LM (striped Hyena) - -```python -data_dim = 1 -use_rope = True # critical — without RoPE, positional recall collapses -rope_base = 10000.0 -fft_padding = "causal" # preserves autoregressive structure -mask_cfg = LazyConfig(GaussianModulationND)( # exponential decay with causal zeroing - data_dim=1, - num_channels=hidden_dim, - min_attenuation_at_step=0.1, - max_attenuation_at_limit=0.95, - init_extent=1.0, - parametrization="exp_decay", -) -gate_nonlinear_cfg = LazyConfig(torch.nn.Identity)() # linear gating for AR -short_conv = LazyConfig(torch.nn.Conv1d)( - in_channels=3 * hidden_dim, - out_channels=3 * hidden_dim, - kernel_size=3, - groups=3 * hidden_dim, - padding=1, - bias=False, -) -``` - -**Mixing ratios (Evo2-1B, 8192-bp sequences, §5.2, perplexity lower is better):** - -| Config | Pattern (24 blocks) | Validation PPL | -| ---------------------- | ------------------- | ------------------- | -| `T` (full transformer) | 24 attention | 2.9235 ± 0.0039 | -| `H₀` (full Hyena) | 24 Hyena | 2.8282 ± 0.0279 | -| `H₁` (1 MHA) | 23 H + 1 A | 2.8308 ± 0.0108 | -| `H₂` (2 MHA, **best**) | 22 H + 2 A | **2.7729 ± 0.0006** | -| `H₃` (3 MHA) | 21 H + 3 A | 2.8214 ± 0.0313 | -| `H₄` (4 MHA) | 20 H + 4 A | 2.8312 ± 0.0088 | - -Default for genomics retrofits: H₂-style pattern (sparse attention, ~1 A every 12 blocks for a 24-block model; for 12 blocks, place one A near the middle and one near the output). - -## PDE fields (The Well, 2D or 3D) - -Same as vision, with: - -- `fft_padding = "circular"` — physics is on a torus or with reflecting BCs -- Patch size matters: smaller patches (p=2, p=4) widen HyenaND's advantage; defaults from `examples/well/v2/` -- Mask: per-axis Gaussian, isotropic init - -## When to use registers + FiLM (input-dependent kernels) - -Always for vision and PDE (§3.2.2). Default `num_registers = 4` for ViT-5-Small; scale with hidden_dim. The `KernelFiLMGenerator` + `RegisterPooling` combo from `_base_config.py` lines 276–337 is the canonical recipe. - -For genomics, FiLM is optional — the Evo2 striped configs in the paper do not use it. - -## Things that look like knobs but aren't really - -- `grid_type`: use `"single"` for non-registered inputs, `"double"` for CLS-row + registers (see `_base_config.py` line 302) -- `L_cache`: set to the maximum sequence length you'll see + 1 (CLS row); the kernel caches at this length -- `parametrization` for Gaussian mask: `"direct"` for vision, `"exp_decay"` for 1D AR - -## Reference configs to copy from - -| Use case | File | -| ------------------------------------ | --------------------------------------------------------------- | -| Smallest end-to-end Hyena example | `examples/mnist_classification/ccnn_4_160_hyena_rope_qknorm.py` | -| ImageNet ViT-5 with FiLM + registers | `examples/vit5_imagenet/v5_patch/_base_config.py` | -| ImageNet hybrid (pattern-driven) | `examples/vit5_imagenet/vit5_hybrid/_base_config.py` | -| Diffusion (HF diffusers retrofit) | `examples/imagenet_diffusion/ccnn_12_768_hyena_qknorm.py` | -| PDE (The Well) | `examples/well/v2/*.py` | -| CCNN-style (non-ViT) classification | `examples/imagenet_classification/ccnn_7_512_hyena_circular.py` | From 8ee115f5c30caed0daa6d04eb7565387d3d568cd Mon Sep 17 00:00:00 2001 From: Farhad Ramezanghorbani Date: Wed, 20 May 2026 11:29:18 -0700 Subject: [PATCH 04/12] future harness plans Signed-off-by: Farhad Ramezanghorbani --- .../skills/hyenand-retrofit/evals/README.md | 203 ++++++++++++++++++ .../evals/inputs/tiny_charlm_attention.py | 74 +++++++ .../evals/inputs/tiny_unet3d_attention.py | 97 +++++++++ 3 files changed, 374 insertions(+) create mode 100644 .claude/skills/hyenand-retrofit/evals/README.md create mode 100644 .claude/skills/hyenand-retrofit/evals/inputs/tiny_charlm_attention.py create mode 100644 .claude/skills/hyenand-retrofit/evals/inputs/tiny_unet3d_attention.py diff --git a/.claude/skills/hyenand-retrofit/evals/README.md b/.claude/skills/hyenand-retrofit/evals/README.md new file mode 100644 index 00000000..b53daded --- /dev/null +++ b/.claude/skills/hyenand-retrofit/evals/README.md @@ -0,0 +1,203 @@ +# 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: ['.claude/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 +# .claude/skills/hyenand-retrofit/evals/run_evals.py +import json, re, subprocess, pathlib, sys + +REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] +EVALS = json.load(open(pathlib.Path(__file__).parent / "evals.json"))["evals"] + + +def run_one(eval_case): + # Spawn claude -p with the skill auto-loaded via .claude/ discovery + 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/.claude/skills/hyenand-retrofit/evals/inputs/tiny_charlm_attention.py b/.claude/skills/hyenand-retrofit/evals/inputs/tiny_charlm_attention.py new file mode 100644 index 00000000..42f3af1b --- /dev/null +++ b/.claude/skills/hyenand-retrofit/evals/inputs/tiny_charlm_attention.py @@ -0,0 +1,74 @@ +"""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/.claude/skills/hyenand-retrofit/evals/inputs/tiny_unet3d_attention.py b/.claude/skills/hyenand-retrofit/evals/inputs/tiny_unet3d_attention.py new file mode 100644 index 00000000..e07c3d83 --- /dev/null +++ b/.claude/skills/hyenand-retrofit/evals/inputs/tiny_unet3d_attention.py @@ -0,0 +1,97 @@ +"""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] From f0717e56cc4eaa2221dc6c77f4b5590dd2ccbb61 Mon Sep 17 00:00:00 2001 From: Farhad Ramezanghorbani Date: Wed, 20 May 2026 12:16:17 -0700 Subject: [PATCH 05/12] subq ops 2D support foot gun Signed-off-by: Farhad Ramezanghorbani --- .claude/skills/hyenand-retrofit/SKILL.md | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/.claude/skills/hyenand-retrofit/SKILL.md b/.claude/skills/hyenand-retrofit/SKILL.md index 3ed70bf2..22045be2 100644 --- a/.claude/skills/hyenand-retrofit/SKILL.md +++ b/.claude/skills/hyenand-retrofit/SKILL.md @@ -82,6 +82,9 @@ global_conv_cfg = LazyConfig(CKConvND)( fft_padding=( "causal" if CAUSAL else "zero" ), # "circular" is a valid bidirectional alt + 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] @@ -112,10 +115,19 @@ mixer = instantiate( **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`, `kernel_cfg`. Also optional `grid_type` (`"single"`/`"double"`) and `use_chunked_fftconv` (memory optimization; requires `fft_padding="zero"`). +- `CKConvND(...)` (passed as `global_conv_cfg`): `data_dim`, `hidden_dim`, `mask_cfg`, `fft_padding`, `kernel_cfg`, `fft_backend`. Also optional `grid_type` (`"single"`/`"double"`), `use_chunked_fftconv` (memory optimization; `zero`/`causal` padding only — circular has no chunked variant by design), and `use_fp16_fft` (memory; `circular` requires power-of-2 spatial dims; not allowed with `subq_ops`). - `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 selection + +`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"`, non-causal, `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", "causal"}` (causal is 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. + ## 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: @@ -179,6 +191,8 @@ These break the first forward pass or the first large-input forward pass. 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"`, `is_causal=True`, `use_fp16_fft=True`. The canonical 2D vision case is fine; extending to 3D, circular padding, 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. From 87e2a95e4281544b05499b530f2189837fb35ec2 Mon Sep 17 00:00:00 2001 From: Farhad Ramezanghorbani Date: Wed, 20 May 2026 12:22:33 -0700 Subject: [PATCH 06/12] fix ruff Signed-off-by: Farhad Ramezanghorbani --- .../evals/inputs/tiny_vit_attention.py | 29 +++++++++++++++---- 1 file changed, 24 insertions(+), 5 deletions(-) diff --git a/.claude/skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py b/.claude/skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py index 51e32065..387215ac 100644 --- a/.claude/skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py +++ b/.claude/skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py @@ -10,7 +10,10 @@ 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 @@ -18,12 +21,16 @@ def __init__(self, image_size: int = 64, patch_size: int = 8, in_channels: int = 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) @@ -31,11 +38,15 @@ def __init__(self, dim: int, mlp_ratio: float = 4.0): 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) @@ -43,6 +54,7 @@ def __init__(self, dim: int = 128, num_heads: int = 4, mlp_ratio: float = 4.0): 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 @@ -53,19 +65,26 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: 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): + 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.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 From 36cf1bedf41432d33839b9b98dafd9b824da30ba Mon Sep 17 00:00:00 2001 From: Farhad Ramezanghorbani Date: Wed, 27 May 2026 12:17:46 -0700 Subject: [PATCH 07/12] docs: escape bare $ in skill evals README for mdformat-myst Co-Authored-By: Claude Opus 4.7 (1M context) --- .claude/skills/hyenand-retrofit/evals/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.claude/skills/hyenand-retrofit/evals/README.md b/.claude/skills/hyenand-retrofit/evals/README.md index b53daded..c2c179f6 100644 --- a/.claude/skills/hyenand-retrofit/evals/README.md +++ b/.claude/skills/hyenand-retrofit/evals/README.md @@ -69,7 +69,7 @@ PRs don't pay for skill evals. | 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) | +| **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 @@ -108,7 +108,7 @@ deliberately tolerant of surface-level variation (`data_dim\s*=\s*3`, not `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 +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? From ce51f39975a88d8e6d3e7534991aace5817d426e Mon Sep 17 00:00:00 2001 From: Farhad Ramezanghorbani Date: Wed, 27 May 2026 13:59:33 -0600 Subject: [PATCH 08/12] update skill using docs Signed-off-by: Farhad Ramezanghorbani --- .claude/skills/hyenand-retrofit/SKILL.md | 42 +++++++++++-------- .../skills/hyenand-retrofit/evals/evals.json | 5 ++- 2 files changed, 27 insertions(+), 20 deletions(-) diff --git a/.claude/skills/hyenand-retrofit/SKILL.md b/.claude/skills/hyenand-retrofit/SKILL.md index 22045be2..8a6daec3 100644 --- a/.claude/skills/hyenand-retrofit/SKILL.md +++ b/.claude/skills/hyenand-retrofit/SKILL.md @@ -26,7 +26,7 @@ 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 `fft_padding="causal"`, `use_rope=True`, mask `parametrization="exp_decay"`, `omega_0=100`. Bidirectional sets `fft_padding ∈ {"zero", "circular"}`, `use_rope=False`, `parametrization="direct"`, `omega_0=10`. +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]`. @@ -79,9 +79,8 @@ global_conv_cfg = LazyConfig(CKConvND)( hidden_dim=HIDDEN_DIM, kernel_cfg=kernel_cfg, mask_cfg=mask_cfg, - fft_padding=( - "causal" if CAUSAL else "zero" - ), # "circular" is a valid bidirectional alt + 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 @@ -115,19 +114,21 @@ mixer = instantiate( **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`, `kernel_cfg`, `fft_backend`. Also optional `grid_type` (`"single"`/`"double"`), `use_chunked_fftconv` (memory optimization; `zero`/`causal` padding only — circular has no chunked variant by design), and `use_fp16_fft` (memory; `circular` requires power-of-2 spatial dims; not allowed with `subq_ops`). +- `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 selection +## 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"`, non-causal, `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", "causal"}` (causal is 1D-only), optional fp16 and chunked variants. +- **`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: @@ -177,11 +178,13 @@ class HyenaAttnAdapter(nn.Module): 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`. Fix: use `DepthwiseFFTConv3d` from nvSubquadratic for the short conv (eliminates `im2col`). 1D and 2D rarely hit this. +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. @@ -191,7 +194,7 @@ These break the first forward pass or the first large-input forward pass. 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"`, `is_causal=True`, `use_fp16_fft=True`. The canonical 2D vision case is fine; extending to 3D, circular padding, or causal LMs requires flipping to `fft_backend="torch_fft"`. Easy to miss when copy-pasting a 2D config as a 3D starting point. +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 @@ -213,12 +216,15 @@ After writing: ## Reference configs in this repo -Copy parameter values, not whole files: +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` | -| Use case | File | -| --------------------------------- | --------------------------------------------------------------- | -| Smallest end-to-end Hyena example | `examples/mnist_classification/ccnn_4_160_hyena_rope_qknorm.py` | -| ImageNet ViT-5 (FiLM + registers) | `examples/vit5_imagenet/v5_patch/_base_config.py` | -| ImageNet hybrid (pattern-driven) | `examples/vit5_imagenet/vit5_hybrid/_base_config.py` | -| Diffusion (HF diffusers retrofit) | `examples/imagenet_diffusion/ccnn_12_768_hyena_qknorm.py` | -| PDE fields (The Well) | `examples/well/v2/*.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/.claude/skills/hyenand-retrofit/evals/evals.json b/.claude/skills/hyenand-retrofit/evals/evals.json index fb3dd476..43ff7a53 100644 --- a/.claude/skills/hyenand-retrofit/evals/evals.json +++ b/.claude/skills/hyenand-retrofit/evals/evals.json @@ -78,7 +78,7 @@ "id": 5, "name": "foreign-1d-causal-lm", "prompt": "I have a tiny causal char-level LM at .claude/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: use_rope=True, fft_padding='causal', 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.", + "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": [".claude/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"}, @@ -87,7 +87,8 @@ {"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 uses causal FFT padding", "check": "grep_output:fft_padding\\s*=\\s*['\\\"]causal['\\\"]"}, + {"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"}, From 9b78d6f4bc291cf62b0177024d519be33c8b0998 Mon Sep 17 00:00:00 2001 From: Farhad Ramezanghorbani Date: Wed, 27 May 2026 13:02:43 -0700 Subject: [PATCH 09/12] docs: re-align table columns in SKILL.md per mdformat Co-Authored-By: Claude Opus 4.7 (1M context) --- .claude/skills/hyenand-retrofit/SKILL.md | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/.claude/skills/hyenand-retrofit/SKILL.md b/.claude/skills/hyenand-retrofit/SKILL.md index 8a6daec3..e079d2a5 100644 --- a/.claude/skills/hyenand-retrofit/SKILL.md +++ b/.claude/skills/hyenand-retrofit/SKILL.md @@ -218,13 +218,13 @@ After writing: 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") | +| 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` | +| 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. From 3a36280c63d47cd2d4778aa177f07341e02a23ca Mon Sep 17 00:00:00 2001 From: Farhad Ramezanghorbani Date: Wed, 3 Jun 2026 12:07:41 -0700 Subject: [PATCH 10/12] Move hyenand-retrofit skill to root skills/ and generalize Signed-off-by: Farhad Ramezanghorbani --- {.claude/skills => skills}/hyenand-retrofit/SKILL.md | 2 +- .../hyenand-retrofit/evals/README.md | 12 +++++++----- .../hyenand-retrofit/evals/evals.json | 12 ++++++------ .../evals/inputs/tiny_charlm_attention.py | 0 .../evals/inputs/tiny_unet3d_attention.py | 0 .../evals/inputs/tiny_vit_attention.py | 0 6 files changed, 14 insertions(+), 12 deletions(-) rename {.claude/skills => skills}/hyenand-retrofit/SKILL.md (99%) rename {.claude/skills => skills}/hyenand-retrofit/evals/README.md (94%) rename {.claude/skills => skills}/hyenand-retrofit/evals/evals.json (85%) rename {.claude/skills => skills}/hyenand-retrofit/evals/inputs/tiny_charlm_attention.py (100%) rename {.claude/skills => skills}/hyenand-retrofit/evals/inputs/tiny_unet3d_attention.py (100%) rename {.claude/skills => skills}/hyenand-retrofit/evals/inputs/tiny_vit_attention.py (100%) diff --git a/.claude/skills/hyenand-retrofit/SKILL.md b/skills/hyenand-retrofit/SKILL.md similarity index 99% rename from .claude/skills/hyenand-retrofit/SKILL.md rename to skills/hyenand-retrofit/SKILL.md index e079d2a5..2290a7fe 100644 --- a/.claude/skills/hyenand-retrofit/SKILL.md +++ b/skills/hyenand-retrofit/SKILL.md @@ -34,7 +34,7 @@ These four axes are orthogonal. Fix them before writing. ## 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 via AskUserQuestion. +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. diff --git a/.claude/skills/hyenand-retrofit/evals/README.md b/skills/hyenand-retrofit/evals/README.md similarity index 94% rename from .claude/skills/hyenand-retrofit/evals/README.md rename to skills/hyenand-retrofit/evals/README.md index c2c179f6..65022ac1 100644 --- a/.claude/skills/hyenand-retrofit/evals/README.md +++ b/skills/hyenand-retrofit/evals/README.md @@ -56,7 +56,7 @@ Separate workflow, not bolted onto the main GPU CI in ```yaml on: pull_request: - paths: ['.claude/skills/hyenand-retrofit/**'] + 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 @@ -123,15 +123,17 @@ deliberately tolerant of surface-level variation (`data_dim\s*=\s*3`, not The full eval harness is ~50 lines of Python. Core loop: ```python -# .claude/skills/hyenand-retrofit/evals/run_evals.py +# skills/hyenand-retrofit/evals/run_evals.py import json, re, subprocess, pathlib, sys -REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] +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 with the skill auto-loaded via .claude/ discovery + # 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"], @@ -184,7 +186,7 @@ Caveats this sketch glosses over: - 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) +[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. diff --git a/.claude/skills/hyenand-retrofit/evals/evals.json b/skills/hyenand-retrofit/evals/evals.json similarity index 85% rename from .claude/skills/hyenand-retrofit/evals/evals.json rename to skills/hyenand-retrofit/evals/evals.json index 43ff7a53..e32e83c0 100644 --- a/.claude/skills/hyenand-retrofit/evals/evals.json +++ b/skills/hyenand-retrofit/evals/evals.json @@ -36,9 +36,9 @@ { "id": 3, "name": "foreign-timm-vit", - "prompt": "I have a tiny standalone PyTorch ViT (not using nvsubquadratic conventions) at .claude/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.", + "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": [".claude/skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py"], + "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"}, @@ -57,9 +57,9 @@ { "id": 4, "name": "foreign-3d-feature-map", - "prompt": "I have a small standalone 3D U-Net at .claude/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.", + "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": [".claude/skills/hyenand-retrofit/evals/inputs/tiny_unet3d_attention.py"], + "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"}, @@ -77,9 +77,9 @@ { "id": 5, "name": "foreign-1d-causal-lm", - "prompt": "I have a tiny causal char-level LM at .claude/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.", + "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": [".claude/skills/hyenand-retrofit/evals/inputs/tiny_charlm_attention.py"], + "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"}, diff --git a/.claude/skills/hyenand-retrofit/evals/inputs/tiny_charlm_attention.py b/skills/hyenand-retrofit/evals/inputs/tiny_charlm_attention.py similarity index 100% rename from .claude/skills/hyenand-retrofit/evals/inputs/tiny_charlm_attention.py rename to skills/hyenand-retrofit/evals/inputs/tiny_charlm_attention.py diff --git a/.claude/skills/hyenand-retrofit/evals/inputs/tiny_unet3d_attention.py b/skills/hyenand-retrofit/evals/inputs/tiny_unet3d_attention.py similarity index 100% rename from .claude/skills/hyenand-retrofit/evals/inputs/tiny_unet3d_attention.py rename to skills/hyenand-retrofit/evals/inputs/tiny_unet3d_attention.py diff --git a/.claude/skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py b/skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py similarity index 100% rename from .claude/skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py rename to skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py From fd05b9f57987b04f252b45a524381baf09e8daa5 Mon Sep 17 00:00:00 2001 From: Farhad Ramezanghorbani Date: Wed, 3 Jun 2026 12:12:37 -0700 Subject: [PATCH 11/12] update Signed-off-by: Farhad Ramezanghorbani --- .../evals/inputs/tiny_charlm_attention.py | 15 +++++++++++++++ .../evals/inputs/tiny_unet3d_attention.py | 15 +++++++++++++++ .../evals/inputs/tiny_vit_attention.py | 15 +++++++++++++++ 3 files changed, 45 insertions(+) diff --git a/skills/hyenand-retrofit/evals/inputs/tiny_charlm_attention.py b/skills/hyenand-retrofit/evals/inputs/tiny_charlm_attention.py index 42f3af1b..d6a6e079 100644 --- a/skills/hyenand-retrofit/evals/inputs/tiny_charlm_attention.py +++ b/skills/hyenand-retrofit/evals/inputs/tiny_charlm_attention.py @@ -1,3 +1,18 @@ +# 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, diff --git a/skills/hyenand-retrofit/evals/inputs/tiny_unet3d_attention.py b/skills/hyenand-retrofit/evals/inputs/tiny_unet3d_attention.py index e07c3d83..371f9b09 100644 --- a/skills/hyenand-retrofit/evals/inputs/tiny_unet3d_attention.py +++ b/skills/hyenand-retrofit/evals/inputs/tiny_unet3d_attention.py @@ -1,3 +1,18 @@ +# 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 diff --git a/skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py b/skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py index 387215ac..10716901 100644 --- a/skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py +++ b/skills/hyenand-retrofit/evals/inputs/tiny_vit_attention.py @@ -1,3 +1,18 @@ +# 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, From b747e36004446901b3afa1353b092078f603e31f Mon Sep 17 00:00:00 2001 From: Farhad Ramezanghorbani Date: Wed, 3 Jun 2026 12:38:18 -0700 Subject: [PATCH 12/12] ci: always run gpu-tests (drop path filters) Signed-off-by: Farhad Ramezanghorbani --- .github/workflows/gpu-tests.yml | 33 +++++---------------------------- 1 file changed, 5 insertions(+), 28 deletions(-) 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