Skip to content

Add hyenaND-retrofit Claude skill#115

Merged
farhadrgh merged 16 commits into
mainfrom
farhadr/retrofit_skill
Jun 3, 2026
Merged

Add hyenaND-retrofit Claude skill#115
farhadrgh merged 16 commits into
mainfrom
farhadr/retrofit_skill

Conversation

@farhadrgh

@farhadrgh farhadrgh commented May 18, 2026

Copy link
Copy Markdown
Collaborator

Adds a Claude Code skill that automates the attention -> HyenaND swap. Two paths:

  • Native: in-repo configs that already use the build_attention_net / build_hyena_net / build_hybrid_net builders. Swap is mechanical — change the builder call, drop compile_compatible_fftconv = False, pick a layer pattern for hybrids. Sibling files are bare config shims (no __main__ block, matching the canonical hyena_patch16.py / hybrid_hhha.py).
  • Foreign: external PyTorch models (timm ViT, HF transformers, hand-written) using nn.MultiheadAttention or F.scaled_dot_product_attention. Skill emits a sibling file with a full Hyena module plus a drop-in HyenaAttnAdapter that bridges the three nn.MHA ↔ Hyena interface gaps: tuple return, kwarg surface (need_weights/attn_mask/key_padding_mask), and channel layout / CLS handling.

The foreign path is organized around four orthogonal axes the user fixes up front:

  1. data_dim ∈ {1, 2, 3} — picks ConvNd, sets CKConvND / SIRENKernelND / GaussianModulationND.
  2. causal ∈ {True, False} — sets fft_padding, use_rope, mask parametrization, omega_0.
  3. Host layout — tokens [B, N, C] vs feature_map [B, C, *S].
  4. Return contract — (out, None) tuple vs bare tensor.

One parameterized adapter covers all sixteen combinations. FFT backend is selected automatically from the axes: fft_backend="subq_ops" (optimized CUDA kernel from subquadratic_ops_torch) when data_dim == 2 and not causal, else fft_backend="torch_fft" — the latter is required for 1D, 3D, circular padding, fp16, and causal paths because subq_ops is 2D-only in the current nvSubquadratic wiring (the upstream package ships fft_causal_conv1d and causal_conv1d but they are not exposed).

A foot-guns section calls out the failures that bite in practice: INT32 im2col overflow in large 3D short conv (≥160³ trips it; fix is DepthwiseFFTConv3d), RoPE divisibility (% 2 / % 4 / % 6 for 1D / 2D / 3D, must hold at every hierarchical stage), L_cache memory blowup ((2L-1)^D × D × 4 bytes — D=3 / L=512 ≈ 12.8 GB), state_dict-prefix mismatch on pretrained-attention loading, and the subq_ops constraint set (asserts loudly on data_dim ≠ 2, fft_padding ≠ "zero", causal, or fp16).

Hierarchical encoders with a use_hyena=True flag are an explicit exception to the sibling-file rule — the natural API is an in-place edit.

Five eval cases under evals/ cover the matrix:

  • native-pure-swap, native-hybrid — v5_patch and vit5_hybrid configs.
  • foreign-timm-vit — 2D, tokens, with CLS.
  • foreign-3d-feature-map — 3D, feature_map, no CLS, F.scaled_dot_product_attention on [B, C, D, H, W].
  • foreign-1d-causal-lm — 1D, causal tokens, attn_mask call site.

Together they exercise every (data_dim, causal, host_layout, prefix-tokens) combination the skill is meant to handle.

farhadrgh and others added 6 commits May 18, 2026 16:36
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) <noreply@anthropic.com>
Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>
Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>
Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>
Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>
Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>
farhadrgh and others added 5 commits May 27, 2026 11:44
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Farhad Ramezanghorbani <farhad.ghorbani@gmail.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@farhadrgh farhadrgh requested a review from tshimko-nv May 28, 2026 21:36

@tshimko-nv tshimko-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good - one small comment about Claude vs. universal skill implementation, but otherwise, should be good.

Comment thread .claude/skills/hyenand-retrofit/SKILL.md Outdated
farhadrgh and others added 4 commits June 3, 2026 12:07
Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>
Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>
Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>
@farhadrgh farhadrgh merged commit d74ec18 into main Jun 3, 2026
3 checks passed
@farhadrgh farhadrgh deleted the farhadr/retrofit_skill branch June 3, 2026 20:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants