Repository navigation
feat: unify hierarchical patch merging and experiment recipes - #141
Merged
farhadrgh merged 80 commits intoSep 16, 2026
Merged
Conversation
added 30 commits
May 25, 2026 16:04
…es are Hyena short-conv configs The module docstring claimed CausalConv1D was used in mamba_nd.py, but it is not imported there at all. The actual call sites are the Hyena short-conv configuration helpers: examples/spatial_recall_v2/mixer_defaults.py examples/spatial_recall_1d/mixer_defaults.py Update the 'Use in …' section header and body to reflect reality. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…s div-by-zero Copilot flagged that the docstring implied drop_prob=1.0 would zero every sample cleanly, but worried the keep_prob=0 division would produce inf/NaN first. The implementation already has an explicit 'if keep_prob > 0.0' guard (line 67) that skips the rescaling division, so Bernoulli(0) produces an all-zero mask and x*0=0 with no numerical issue. Update the docstring to document this guarantee. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…ored CONVENTIONS.md documented that the pre-commit ruff hook catches D100, but pyproject.toml had D100 in the global ignore list, so both the pre-commit hook and the CI diff-check silently skipped missing module docstrings (ruff applies config-file ignores on top of CLI --select). Fix: - Remove D100 from pyproject.toml ignore so the hook matches the docs - Add missing module docstring to nvsubquadratic/parallel/utils.py (the only file in scope that was missing one) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…forced Four files were missing module-level docstrings that ruff D100 now catches (after removing D100 from the global ignore list): - docs/conf.py — Sphinx configuration file - examples/imagenet_diffusion/ccnn_jit_baseline.py — CCNN-Hyena JiT-B-matched diffusion baseline - examples/imagenet_diffusion/hf_uvit_baseline.py — HuggingFace UViT diffusion baseline - examples/imagenet_diffusion/jit_baseline.py — JiT-B flow-matching diffusion baseline Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Implements 2x2 spatial patch merging for hierarchical ViT-5/Hyena networks. Key features: - Pure-spatial layout: [B, H*W, C] -> [B, (H/2)*(W/2), out_dim] - Register-row layout: passes the leading grid_w register tokens through a dedicated reg_proj Linear so FiLM conditioning survives the channel-dim change, then repacks them to width grid_w//2 - Post-concat norm (configurable via LazyConfig) + bias-free reduction Linear, both trunc_normal_ initialised (std=0.02) - flop_count() helper for FLOP bookkeeping - Full test suite covering shape, zero-pad, independent paths, error guards, and FLOP count Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Swin-style 4-stage hierarchical ViT-5 classifier with PatchMerging between stages. Two readout layouts: - pure: flat patch grid [B, H*W, C] + GAP - register_row: FiLM register tokens prepended as first grid row at every stage; GAP excludes them Key properties: - Per-stage dims/depths fully configurable via LazyConfig - flop_count() aggregates patch-embed + blocks + merges + head - Backward tested: patch_embed and reg_proj grads both non-zero ImageNet configs for 4-stage Swin-T-like Hyena hierarchy (p=4, dims [96,192,384,768], depths [2,2,6,2]) in both 'pure' and 'register_row' FiLM variants, plus CIFAR-10 patch and capacity ablation configs. Mark patch_merging.py and vit5_hierarchical_classification.py as [x]. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…r; align ImageNet configs with vit5_hybrid reference
## CIFAR-10 subfolder
Move all CIFAR-10 experiment files into examples/vit5_imagenet/v6_hierarchical/cifar10/:
_cifar10_patch_ablation_base.py → cifar10/_base.py
cifar10_{flat,hier}_p{4,8,16}.py → cifar10/{flat,hier}_p{4,8,16}.py
cifar10_hyena_{hier,flat}.py → cifar10/hyena_{hier,flat}.py
Leaf configs import updated to reference cifar10._base.
## _base_config.py — align with vit5_hybrid/_film.py + _learnable_omega.py
- Mask: swap torch.nn.Identity for BlockAlignedGaussianModulationND
(data_dim=2, extent=1.0, direct parametrization) matching
apply_learnable_omega_blockdiag_overrides in the reference config.
- FiLM constants: add FILM_INIT_TYPE='identity', FILM_WEIGHT_DECAY=5e-3,
FILM_AFTER_POS_EMBED=True matching vit5_hybrid/_film.py best-run defaults.
- KernelFiLMGenerator: add init_type + no_weight_decay; bump num_film_layers
from KERNEL_NUM_LAYERS-1 to KERNEL_NUM_LAYERS (3) since film_after_pos_embed
adds one extra pair for the positional-embedding sine.
- _siren_kernel_cfg: set cfg.film_after_pos_embed=True when film_cfg provided.
## hyena_hier_p4_{pure,film}.py — match full_hyena_learnable_omega_blockdiag.py
- Append MaskMonitorCallback + OmegaScaleMonitorCallback (log every 50 steps).
- Set distinct wandb.job_group: 'v6_hier_pure' / 'v6_hier_film'.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…far10 configs The v6_hierarchical CIFAR-10 example configs added in this PR import `experiments.datamodules.cifar10.CIFAR10DataModule` but the file was local-only. Add it to the PR with docstrings on all public methods so ruff D102/D107 pass under the new docstring CI. Also apply mdformat to docs-tracker.md to satisfy the pre-commit hook (column-width normalization).
Preserve PR NVIDIA-BioNeMo#122 feature history while resolving squash-history conflicts to main. Recover the spatial-grid implementation and recipes from ed0bcbc and 74dd91c; retain the two architectures and CIFAR data recipes under explicit APIs. Share 2x2 merging, fix autocast padding, update stale configs and validate equivalence, gradients and both real Hyena layouts. 69 focused CPU tests and all pre-commit hooks pass. Signed-off-by: dafidofff <davidwessels15@gmail.com>
There was a problem hiding this comment.
🟡 Changes recommended
One or more issues must be addressed before approval.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
This PR synchronizes the patch-merging branch with main and unifies token/register-row and spatial-grid implementations, along with corresponding hierarchical models, recipes, tests, and documentation.
Changes:
- Added shared 2×2 patch-merging behavior and dtype-safe register padding.
- Recovered hierarchical model implementations and ImageNet/CIFAR recipes.
- Added equivalence, gradient, dtype, and recipe validation tests.
File summaries
| File | Description |
|---|---|
| VERSION | Updated as part of this pull request. |
| tests/test_well_staging_ddp.py | Updated as part of this pull request. |
| tests/test_well_dataloader_parity.py | Updated as part of this pull request. |
| tests/test_safe_eval_resolver.py | Updated as part of this pull request. |
| tests/test_param_groups.py | Updated as part of this pull request. |
| tests/test_nightly_well_validation.py | Updated as part of this pull request. |
| tests/test_nightly_validation.py | Updated as part of this pull request. |
| tests/test_iteration_speed_callback.py | Updated as part of this pull request. |
| tests/test_image_grid_callback.py | Updated as part of this pull request. |
| tests/test_h5_caching.py | Updated as part of this pull request. |
| tests/test_dali_rand_augment.py | Updated as part of this pull request. |
| tests/test_dali_epoch_transition.py | Updated as part of this pull request. |
| tests/test_compile_lcache.py | Updated as part of this pull request. |
| tests/test_classification_loss.py | Updated as part of this pull request. |
| tests/test_checkpoint_resume.py | Updated as part of this pull request. |
| tests/test_basic.py | Updated as part of this pull request. |
| tests/test_autoregressive_wrapper.py | Updated as part of this pull request. |
| tests/README.md | Updated as part of this pull request. |
| tests/parallel/test_a2a_comms.py | Updated as part of this pull request. |
| tests/parallel/init.py | Updated as part of this pull request. |
| tests/ops/test_subq_ops_fft_conv2d.py | Updated as part of this pull request. |
| tests/ops/test_mixed_fftconv.py | Updated as part of this pull request. |
| tests/ops/test_fftconv.py | Updated as part of this pull request. |
| tests/ops/test_fftconv_custom.py | Updated as part of this pull request. |
| tests/ops/test_fftconv_custom_1d.py | Updated as part of this pull request. |
| tests/ops/test_fftconv_chunked.py | Updated as part of this pull request. |
| tests/ops/test_circular_vs_fftconv_perf.py | Updated as part of this pull request. |
| tests/ops/test_circular_fftconv.py | Updated as part of this pull request. |
| tests/ops/test_causal_conv1d_custom.py | Updated as part of this pull request. |
| tests/ops/init.py | Updated as part of this pull request. |
| tests/networks/baselines/init.py | Updated as part of this pull request. |
| tests/networks/init.py | Updated as part of this pull request. |
| tests/modules/torchrun_sequence_mixer_cp_test.py | Updated as part of this pull request. |
| tests/modules/test_vit5_components.py | Updated as part of this pull request. |
| tests/modules/test_subq_ops_causal_conv1d.py | Updated as part of this pull request. |
| tests/modules/test_self_attention.py | Updated as part of this pull request. |
| tests/modules/test_patchify.py | Updated as part of this pull request. |
| tests/modules/test_multi_omega_siren.py | Updated as part of this pull request. |
| tests/modules/test_mlp.py | Updated as part of this pull request. |
| tests/modules/test_masks_nd.py | Updated as part of this pull request. |
| tests/modules/test_kernels_nd.py | Updated as part of this pull request. |
| tests/modules/test_hyena_film.py | Updated as part of this pull request. |
| tests/modules/test_grn.py | Updated as part of this pull request. |
| tests/modules/test_distributed_depthwise_conv_nd.py | Updated as part of this pull request. |
| tests/modules/test_causality_hyena.py | Updated as part of this pull request. |
| tests/modules/test_causality_attn_mamba.py | Updated as part of this pull request. |
| tests/modules/init.py | Updated as part of this pull request. |
| tests/conftest.py | Updated as part of this pull request. |
| scripts/visualization/visualize_repeated_aug.py | Updated as part of this pull request. |
| scripts/verify_v2_samples_per_epoch.py | Updated as part of this pull request. |
| scripts/run_cifar10_comparison.sh | Updated as part of this pull request. |
| scripts/data/extract_kernel_data.py | Updated as part of this pull request. |
| scripts/data/extract_imagenet_to_folder.py | Updated as part of this pull request. |
| scripts/data/compute_imagenet_stats.py | Updated as part of this pull request. |
| scripts/check_gpu_availability.py | Updated as part of this pull request. |
| scripts/_wandb_check_ablation.py | Updated as part of this pull request. |
| reports/vit5_imagenet_dataloader_profiling/profile_step_breakdown.py | Updated as part of this pull request. |
| reports/siren_omega0_dimensional_scaling/omega0_scaling.py | Updated as part of this pull request. |
| reports/ckconv_block_diagonal_kernel/production_vs_prototype.py | Updated as part of this pull request. |
| reports/ckconv_block_diagonal_kernel/multiomega_classes.py | Updated as part of this pull request. |
| reports/ckconv_block_diagonal_kernel/multiomega_block_diag.py | Updated as part of this pull request. |
| reports/ckconv_block_diagonal_kernel/block_diag_default_choice.py | Updated as part of this pull request. |
| reports/ckconv_block_diagonal_kernel/analyze_kernel_spectrum.py | Updated as part of this pull request. |
| reports/ckconv_block_diagonal_kernel/.gitignore | Updated as part of this pull request. |
| nvsubquadratic/utils/rope.py | Updated as part of this pull request. |
| nvsubquadratic/utils/quack_utils.py | Updated as part of this pull request. |
| nvsubquadratic/utils/qk_norm.py | Updated as part of this pull request. |
| nvsubquadratic/utils/init.py | Updated as part of this pull request. |
| nvsubquadratic/testing/utils.py | Updated as part of this pull request. |
| nvsubquadratic/testing/init.py | Updated as part of this pull request. |
| nvsubquadratic/parallel/a2a_comms.py | Updated as part of this pull request. |
| nvsubquadratic/parallel/init.py | Updated as part of this pull request. |
| nvsubquadratic/ops/fftconv_chunked.py | Updated as part of this pull request. |
| nvsubquadratic/ops/circular_fftconv.py | Updated as part of this pull request. |
| nvsubquadratic/ops/init.py | Updated as part of this pull request. |
| nvsubquadratic/networks/vit5_classification.py | Updated as part of this pull request. |
| nvsubquadratic/modules/vit5_attention.py | Updated as part of this pull request. |
| nvsubquadratic/modules/subq_ops_causal_conv1d.py | Updated as part of this pull request. |
| nvsubquadratic/modules/sequence_mixer.py | Updated as part of this pull request. |
| nvsubquadratic/modules/schedulers.py | Updated as part of this pull request. |
| nvsubquadratic/modules/rms_norm.py | Updated as part of this pull request. |
| nvsubquadratic/modules/position_encoding.py | Updated as part of this pull request. |
| nvsubquadratic/modules/patchify.py | Updated as part of this pull request. |
| nvsubquadratic/modules/mlp.py | Updated as part of this pull request. |
| nvsubquadratic/modules/masks_nd.py | Updated as part of this pull request. |
| nvsubquadratic/modules/mamba_nd.py | Updated as part of this pull request. |
| nvsubquadratic/modules/hyena_nd.py | Updated as part of this pull request. |
| nvsubquadratic/modules/grn.py | Updated as part of this pull request. |
| nvsubquadratic/modules/drop_path.py | Updated as part of this pull request. |
| nvsubquadratic/modules/distributed_depthwise_conv_nd.py | Updated as part of this pull request. |
| nvsubquadratic/modules/condition_mixer.py | Updated as part of this pull request. |
| nvsubquadratic/modules/causal_conv1d.py | Updated as part of this pull request. |
| nvsubquadratic/modules/attention.py | Updated as part of this pull request. |
| nvsubquadratic/modules/_channels_first_utils.py | Updated as part of this pull request. |
| nvsubquadratic/modules/init.py | Updated as part of this pull request. |
| experiments/utils/cli.py | Updated as part of this pull request. |
| experiments/utils/checkpointing.py | Updated as part of this pull request. |
| experiments/lightning_wrappers/well_lightning_wrapper.py | Updated as part of this pull request. |
| experiments/datamodules/utils/dali_rand_augment.py | Updated as part of this pull request. |
| experiments/datamodules/utils/init.py | Updated as part of this pull request. |
| experiments/datamodules/tinyimagenet.py | Updated as part of this pull request. |
| experiments/datamodules/spatial_recall_dataset.py | Updated as part of this pull request. |
| experiments/datamodules/pde/well.py | Updated as part of this pull request. |
| experiments/datamodules/pde/init.py | Updated as part of this pull request. |
| experiments/datamodules/mnist.py | Updated as part of this pull request. |
| experiments/datamodules/emnist.py | Updated as part of this pull request. |
| experiments/datamodules/dali_imagenet_fused.py | Updated as part of this pull request. |
| experiments/datamodules/cifar10.py | Updated as part of this pull request. |
| experiments/datamodules/_deprecated/ref_imagenet.py | Updated as part of this pull request. |
| experiments/datamodules/_deprecated/dali_imagenet.py | Updated as part of this pull request. |
| experiments/datamodules/_deprecated/dali_imagenet_optimized.py | Updated as part of this pull request. |
| experiments/datamodules/_deprecated/init.py | Updated as part of this pull request. |
| experiments/datamodules/init.py | Updated as part of this pull request. |
| experiments/callbacks/wandb_cache_cleanup.py | Updated as part of this pull request. |
| experiments/callbacks/walltime_checkpointer.py | Updated as part of this pull request. |
| experiments/callbacks/sequence_visualization_1d.py | Updated as part of this pull request. |
| experiments/callbacks/omega_scale_monitor.py | Updated as part of this pull request. |
| experiments/callbacks/model_ema.py | Updated as part of this pull request. |
| experiments/callbacks/mask_monitor.py | Updated as part of this pull request. |
| experiments/callbacks/iteration_speed.py | Updated as part of this pull request. |
| experiments/callbacks/image_grid_val_visualization.py | Updated as part of this pull request. |
| experiments/callbacks/film_monitor.py | Updated as part of this pull request. |
| experiments/init.py | Updated as part of this pull request. |
| examples/well/v2/viscoelastic_instability/unet_convnext.py | Updated as part of this pull request. |
| examples/well/v2/viscoelastic_instability/_base.py | Updated as part of this pull request. |
| examples/well/v2/turbulence_gravity_cooling/_base.py | Updated as part of this pull request. |
| examples/well/v2/TRACKER.md | Updated as part of this pull request. |
| examples/well/v2/supernova_explosion_64/attention.py | Updated as part of this pull request. |
| examples/well/v2/supernova_explosion_64/_base.py | Updated as part of this pull request. |
| examples/well/v2/shear_flow/unet_convnext.py | Updated as part of this pull request. |
| examples/well/v2/shear_flow/cfg_unet_convnext.py | Updated as part of this pull request. |
| examples/well/v2/shear_flow/_base.py | Updated as part of this pull request. |
| examples/well/v2/rayleigh_taylor_instability/_base.py | Updated as part of this pull request. |
| examples/well/v2/rayleigh_benard/unet_convnext.py | Updated as part of this pull request. |
| examples/well/v2/rayleigh_benard/_base.py | Updated as part of this pull request. |
| examples/well/v2/MHD_64/unet_convnext.py | Updated as part of this pull request. |
| examples/well/v2/MHD_64/hyena_gaussian_mask.py | Updated as part of this pull request. |
| examples/well/v2/MHD_64/attention.py | Updated as part of this pull request. |
| examples/well/v2/MHD_64/_base.py | Updated as part of this pull request. |
| examples/well/v2/helmholtz_staircase/_base.py | Updated as part of this pull request. |
| examples/well/v2/gray_scott_reaction_diffusion/attention.py | Updated as part of this pull request. |
| examples/well/v2/gray_scott_reaction_diffusion/_base.py | Updated as part of this pull request. |
| examples/well/v2/euler_multi_quadrants_periodicBC/_base.py | Updated as part of this pull request. |
| examples/well/v2/active_matter/unet_convnext.py | Updated as part of this pull request. |
| examples/well/v2/active_matter/cfg_unet_convnext.py | Updated as part of this pull request. |
| examples/well/v2/active_matter/attention.py | Updated as part of this pull request. |
| examples/well/v2/acoustic_scattering_maze/unet_convnext.py | Updated as part of this pull request. |
| examples/well/v2/acoustic_scattering_maze/attention.py | Updated as part of this pull request. |
| examples/well/v2/acoustic_scattering_maze/_base.py | Updated as part of this pull request. |
| examples/well/v1/TRACKER.md | Updated as part of this pull request. |
| examples/well/v1/MHD_64/cfg_attention.py | Updated as part of this pull request. |
| examples/well/v1/gray_scott_reaction_diffusion/_base.py | Updated as part of this pull request. |
| examples/well/v1/gray_scott_reaction_diffusion/init.py | Updated as part of this pull request. |
| examples/well/v1/euler_multi_quadrants_periodicBC/_base.py | Updated as part of this pull request. |
| examples/well/v1/euler_multi_quadrants_periodicBC/init.py | Updated as part of this pull request. |
| examples/well/v1/active_matter/cfg_hyena.py | Updated as part of this pull request. |
| examples/well/v1/active_matter/cfg_ckconv.py | Updated as part of this pull request. |
| examples/well/v1/active_matter/cfg_attention.py | Updated as part of this pull request. |
| examples/well/v1/init.py | Updated as part of this pull request. |
| examples/well/README.md | Updated as part of this pull request. |
| examples/well/init.py | Updated as part of this pull request. |
| examples/vit5_imagenet/vit5_hybrid/hybrid_ha.py | Updated as part of this pull request. |
| examples/vit5_imagenet/vit5_hybrid/full_mamba.py | Updated as part of this pull request. |
| examples/vit5_imagenet/vit5_hybrid/full_hyena.py | Updated as part of this pull request. |
| examples/vit5_imagenet/vit5_hybrid/full_attention.py | Updated as part of this pull request. |
| examples/vit5_imagenet/vit5_hybrid/_film.py | Updated as part of this pull request. |
| examples/vit5_imagenet/vit5_hybrid/_blockdiag.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v6_hierarchical/cifar10/init.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v6_hierarchical/_smoke_test.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v6_hierarchical/init.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v5/attention_pretrain.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v5/_base.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v5_patchmerge/init.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v5_patch/README.md | Updated as part of this pull request. |
| examples/vit5_imagenet/v5_patch/attention_patch8.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v5_patch/attention_patch4.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v5_patch/attention_patch2.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v5_patch/attention_patch16.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v5_patch/attention_patch1.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v5_patch/_smoke_test.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v5_patch/_base_config.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v5_patch/init.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v4/_base.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v4/init.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v3/gap_film_regs/init.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v3/_pretrain_base.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v2/wsd_ft_ablation/_base.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v2/wsd_ft_ablation/init.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v2/init.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v1/vit5_small_pretrain.py | Updated as part of this pull request. |
| examples/vit5_imagenet/v1/init.py | Updated as part of this pull request. |
| examples/vit5_imagenet/local_comparison/init.py | Updated as part of this pull request. |
| examples/vit5_imagenet/init.py | Updated as part of this pull request. |
| examples/vit_b_benchmark/hyena.py | Updated as part of this pull request. |
| examples/vit_b_benchmark/hyena_patchify.py | Updated as part of this pull request. |
| examples/vit_b_benchmark/attention.py | Updated as part of this pull request. |
| examples/vit_b_benchmark/attention_patchify.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/simple_copy_3d/mamba.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/simple_copy_3d/mamba_patch.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/simple_copy_3d/hyena.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/simple_copy_3d/hyena_patch.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/simple_copy_3d/attention.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/simple_copy_3d/attention_patch.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/simple_copy_3d/_base.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/simple_copy_2d/mamba.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/simple_copy_2d/mamba_patch.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/simple_copy_2d/hyena.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/simple_copy_2d/hyena_patch.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/simple_copy_2d/attention.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/simple_copy_2d/attention_patch.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/simple_copy_2d/_base.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/simple_copy_1d/mamba.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/simple_copy_1d/mamba_bidir.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/simple_copy_1d/hyena.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/simple_copy_1d/hyena_patch.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/simple_copy_1d/attention.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/simple_copy_1d/attention_patch.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/simple_copy_1d/_base.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/mixer_defaults.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/color_conditioning_3d/mamba.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/color_conditioning_3d/mamba_patch.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/color_conditioning_3d/hyena.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/color_conditioning_3d/hyena_patch.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/color_conditioning_3d/attention.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/color_conditioning_3d/attention_patch.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/color_conditioning_3d/_base.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/color_conditioning_3d/init.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/color_conditioning_2d/mamba.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/color_conditioning_2d/mamba_patch.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/color_conditioning_2d/hyena.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/color_conditioning_2d/hyena_patch.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/color_conditioning_2d/attention.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/color_conditioning_2d/attention_patch.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/color_conditioning_2d/_base.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/color_conditioning_1d/mamba.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/color_conditioning_1d/hyena.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/color_conditioning_1d/hyena_patch.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/color_conditioning_1d/attention.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/color_conditioning_1d/attention_patch.py | Updated as part of this pull request. |
| examples/spatial_recall_v2/color_conditioning_1d/_base.py | Updated as part of this pull request. |
| examples/spatial_recall_3d/mixer_defaults.py | Updated as part of this pull request. |
| examples/spatial_recall_3d/base_config.py | Updated as part of this pull request. |
| examples/spatial_recall_2d/mixer_defaults.py | Updated as part of this pull request. |
| examples/spatial_recall_2d/base_config.py | Updated as part of this pull request. |
| examples/spatial_recall_1d/mixer_defaults.py | Updated as part of this pull request. |
| examples/spatial_recall_1d/base_config.py | Updated as part of this pull request. |
| examples/patch_merging/cifar10/init.py | Updated as part of this pull request. |
| examples/patch_merging/init.py | Updated as part of this pull request. |
| examples/imagenet_classification/ccnn_7_512_hyena.py | Updated as part of this pull request. |
| examples/imagenet_classification/ccnn_7_512_attention.py | Updated as part of this pull request. |
| examples/init.py | Updated as part of this pull request. |
| docs/ops/MIXED_BC_PLAN.md | Updated as part of this pull request. |
| docs/api_reference/parallel.rst | Updated as part of this pull request. |
| docs/api_reference/index.rst | Updated as part of this pull request. |
| benchmarks/well/profile_training_loop.py | Updated as part of this pull request. |
| benchmarks/well/bench_training_step.py | Updated as part of this pull request. |
| benchmarks/well/bench_dataloader.py | Updated as part of this pull request. |
| benchmarks/well/bench_ab_comparison.py | Updated as part of this pull request. |
| benchmarks/vit5_imagenet/verify_dali_fused.py | Updated as part of this pull request. |
| benchmarks/vit5_imagenet/validate_checkpoint.py | Updated as part of this pull request. |
| benchmarks/vit5_imagenet/bench_profile.sh | Updated as part of this pull request. |
| benchmarks/vit5_imagenet/bench_optimized.sh | Updated as part of this pull request. |
| benchmarks/vit5_imagenet/bench_compile.sh | Updated as part of this pull request. |
| benchmarks/README.md | Updated as part of this pull request. |
| benchmarks/ops/bench_fftconv2d.py | Updated as part of this pull request. |
| benchmarks/compare_flops.py | Updated as part of this pull request. |
| benchmarks/benchmark_patch_size_2d.py | Updated as part of this pull request. |
| .github/pull_request_template.md | Updated as part of this pull request. |
| .devcontainer/postCreateCommand.sh | Updated as part of this pull request. |
Review details
Suppressed comments (2)
docs/api_reference/modules.rst:93
- The new API reference's patching section omits the public
PatchMerging,PatchMerging2D, andPatchEmbedHierarchicalclasses that this repair documents and tests. Add them here so users can discover the central patch-merging APIs from the generated reference.
docs/api_reference/networks.rst:21 - The classification API reference does not list either hierarchical network introduced/recovered by this PR (
ViT5HierarchicalNetandViT5HierarchicalClassificationNet). Add both entries to the curated network reference so the documented architectures are actually reachable from the API index.
- Files reviewed: 289/913 changed files
- Comments generated: 1
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Dafidofff
changed the title
farhadrgh
reviewed
Sep 14, 2026
farhadrgh
left a comment
farhadrgh
left a comment
Collaborator
There was a problem hiding this comment.
Review pass on the unified patch-merging recipes. The structural work looks solid — 69 new tests pass, and every leaf config builds.
My main concern is that the hierarchical-vs-flat/isotropic comparison this PR is built to enable has four independent confounds as currently configured, flagged inline below. Each one alone would make a hier-vs-flat result hard to interpret; together I don't think the runs would be comparable.
Separately (not inline, happy to expand): stage_grid_sides is floor-halved and never cross-checked against the supplied patch_merge_cfgs, so a mismatched merger silently corrupts the register_row GAP boundary instead of raising — and initial_grid=6, num_stages=3 computes [6, 3, 1] for a 3x3 stage that can't be 2x2-merged at all. Also: a single-block net gets the maximum drop-path rate instead of 0 under the linear schedule; the stage-count assert vanishes under python -O and then IndexErrors in forward via the hardcoded i < 3; and out_proj being a read-only property alias means DropKeysFromCheckpoint(prefixes=('network.out_proj',)) matches zero keys, since the state_dict keys are head.*.
farhadrgh
approved these changes
Sep 15, 2026
farhadrgh
added a commit
that referenced
this pull request
Sep 16, 2026
The 0.2.0 section was written when #139 was the only thing in the release and dated 2026-09-08. Since then #140 and #141 merged to main and came into this branch via e505327, so both ship inside this tag and belong in the notes. - Dates the section 2026-09-16. - Adds #141: PatchMerging / PatchMerging2D / PatchEmbedHierarchical, the ViT5HierarchicalNet + StageSpec / ViT5HierarchicalClassificationNet pair built on them, the CIFAR-10 and ImageNet recipes, and the two new CIFAR-10 datamodules (torchvision and HuggingFace). - Adds #140: SpatialRecall3DMotionDataset / SpatialRecall3DMotionDataModule and the motion_3d example configs. Every class named above was checked to exist at this commit rather than transcribed from the PR descriptions. pre-commit clean, mdformat idempotent over two passes. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
3 of 4 tasks
farhadrgh
added a commit
that referenced
this pull request
Sep 16, 2026
… 13 migration (#139) * feat(ops): fused FFT-conv2d backend, torch.compile lowering, and CUDA 13 migration Adds fft_backend="subq_ops_fused" on CKConvND, backed by subquadratic_ops_torch.fused_fft_conv2d: one cuFFTDx launch for the whole rfft2 → multiply → irfft2 pipeline, running natively in fp32/fp16/bf16 instead of upcasting. Ships with a torch.compile pre-grad pass (nvsubquadratic.ops.fftconv_lowering) that rewrites the existing torch_fft chain onto the same kernel, so models pick it up without a config change. Also migrates the CUDA runtime 12.9 → 13.0: cu130 torch wheels, torch floor raised to >=2.12.0, [cuda] extra retargeted to subquadratic-ops-torch-cu13, [dali] to nvidia-dali-cuda130, plus the Dockerfile/enroot build updates and a clean-install CI matrix. The [cuda] extra pins >=0.3.0 rather than 0.2.2: 0.2.2 was published only to the internal GitLab registry, so pinning it would have broken `pip install nvsubquadratic[cuda]` for public users. 0.3.0 reached public PyPI on 2026-08-28 with the same fused_fft_conv2d entry point, which lets the Dockerfile and enroot scripts drop their GITLAB_TOKEN/BuildKit-secret plumbing and resolve everything from PyPI. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * ci: fail the build when torch/CUDA pins drift from the pyproject range pyproject.toml declares the supported range, but the Docker image, the conda bootstrap, the SLURM builder and the install docs each repeat a concrete version and wheel index. Nothing links them — build_sqsh_slurm.sh replays the Dockerfile rather than parsing it, and the docs are prose — so they drift when only one is edited, which is exactly how setup_conda_env.sh ended up pointing at the cu130 index while still pinning torch 2.10.0. That drift is not cosmetic: installing a torch outside the pyproject range means the later `pip install -e .` re-resolves torch from PyPI and replaces the CUDA-matched wheel after apex/mamba/causal-conv1d were already compiled against it, leaving extensions built for a torch that is no longer installed. scripts/check_version_pins.py parses the range from pyproject.toml and asserts across all 14 pin sites that every exact pin satisfies it, that the pins agree with each other, that every PyTorch wheel index names the same cuXYZ build, that the [cuda] and [dali] extras target that same CUDA major, and that files mirroring the specifier quote it verbatim. It also fails loudly rather than silently passing when a pin is reworded out from under its pattern. Wired in as a pre-commit local hook, so CI's `pre-commit run --all-files` picks it up with no workflow change. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix(deps): raise torch floor to 2.14 so [cuda] resolves; guard and document SM90+ `pip install nvsubquadratic[cuda]` could not resolve at all on the previous 2.12 floor. torch pins nvidia-cudnn-cu13 exactly, and subquadratic-ops-torch-cu13>=0.3.0 requires nvidia-cudnn-cu13>=9.24.0.43: torch 2.12/2.13 -> nvidia-cudnn-cu13==9.20.0.48 subq-ops 0.3.0 -> nvidia-cudnn-cu13>=9.24.0.43 torch 2.14.0 pins ==9.24.0.43 and satisfies both, so 0.3.0 was in fact built against 2.14. In the Docker build the conflict surfaced misleadingly: pip backtracked through ~26 versions of datasets/jinja2/MarkupSafe hunting for an escape and died on MarkupSafe 1.0's setup.py under --no-build-isolation. Falling back is not an option — 0.2.1 is the newest cu13 release without the cudnn floor and its wheel carries no fused_fft_conv2d module. Verified with `pip install --dry-run` on py3.12 against the cu130 index that the whole [all] closure now resolves: torch 2.14.0+cu130, torchvision 0.29.0+cu130, nvidia-cudnn-cu13 9.24.0.43, subquadratic-ops-torch-cu13 0.3.0, DALI 1.53.0, megatron-core, timm, quack-kernels. torch 2.14.0+cu130 is published for both x86_64 and aarch64, and is still CUDA 13.0, so the 13.0.3-devel base still satisfies apex/mamba's exact-CUDA-match requirement. Also drops the torchaudio pin from setup_env.sh: nothing in the repo imports it, and the cu130 index has no build past 2.11. SM90+ gating: The 128 FFT tile needs more shared memory than SM80/SM86 provide, but that gate only existed in the torch.compile lowering pass. The eager fft_backend="subq_ops_fused" path had no check, while every user-facing doc described the limit as purely spatial — so a 64x64 input on an A100/A6000 (which resolves to the 128 tile) failed inside the kernel with no explanation. Adds fused_fftconv2d_arch_supported() plus a raise at the single call site, with the constant now shared with the lowering pass rather than duplicated, and states the requirement in README, docs/ops/README.md, ops.rst and the CHANGELOG. Docs: - docs/index.rst named subquadratic-ops-torch-cu12 and CUDA Toolkit 12.0+, which this branch made wrong; also syncs the backend paragraph with README. - docs-tracker.md gains the fftconv_lowering.py row required by CONVENTIONS.md and the PR template. - scripts/slurm/enroot/README.md documents build_sqsh_slurm.sh, which the branch added without mentioning it anywhere. - FusedFFTConv2dLowering gets an API page, so the three docstrings pointing at it are no longer dead text; adds the two missing chunked variants to ops.rst. `sphinx -W --keep-going` builds clean; 276 tests pass. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * test: gate the 64x64 fused cases on SM90+ The first gpu-tests run that got far enough to execute the suite reported 1126 passed / 5 failed, and all five were the `[64]` parameter: test_ckconv_nd_subq_fused.py::test_spatial_sizes_within_cap[64] test_fftconv_lowering.py::test_rewrites_across_supported_spatial_sizes[64] test_fused_fftconv2d.py::TestForwardMatchesReference::test_double_grid_kernel[64-dtype0/1/2] The CI runner reports compute capability 8.6, and resolve_fused_fft_size escalates any extent above 32 per axis to the 128 FFT tile, which needs more shared memory than SM80/SM86 provide. So these tests encoded an assumption — that the documented 64-per-axis cap holds on every GPU — that is only true on Hopper/Blackwell. On Ampere the effective cap is 32. Four of the five surfaced as the new arch guard raising; the fifth (test_rewrites_across_supported_spatial_sizes[64]) is the lowering pass correctly declining to rewrite via its own pre-existing SM90 check, so that one would have failed on this hardware regardless of the guard. Neither is a regression — this is the first run in which the fused suite executed at all. Adds a shared `requires_sm90` marker in tests/conftest.py and applies it to the individual 64 parameters rather than whole files, so the sub-64 coverage keeps running everywhere and the 64 cases still execute on an H100/B200 runner. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix(ops): scope the lowering pass's erase; attempt every FFT-conv chain Replaces fftconv_lowering.py and its test file with the revision Alireza posted on PR #139, which fixes the two correctness bugs raised in review plus a matcher gap neither of us had caught. Whole-graph `graph.eliminate_dead_code()` ran after every successful rewrite with no is_impure_node override. torch.fx's Node.is_impure() only inspects call_function schemas, so an in-place call_method elsewhere in the same compiled graph (`state.add_(1)`) is treated as pure and silently deleted when its result is only reachable through the mutated variable — wrong numerics, no error. _erase_chain now erases a known consumers-first list of just the nodes the rewrite orphaned. `any(_try_rewrite(...) for node in candidates)` short-circuited on the first success, so a graph with two independent FFT-conv chains fused only the first and silently left the rest on the eager cuFFT path (and under-counted lowering_stats()). Every candidate is now attempted. The matcher additionally declines when a spectrum or the multiply result has users outside the chain — previously that extra math was dropped, or the cuFFT chain was left running beside the fused kernel while the pass reported success. That gate is what makes the erase unconditionally safe. Also declined rather than mis-rewritten: an in-place mul_ ordered after the irfft2, non-default norm, fp64/cross-device casts, kernel dtype/device mismatches, an output whose dtype/device does not match x, and symbolic slice bounds under dynamic shapes. uuid() hashes this file's contents instead of a hand-bumped version constant. L2_TOL/L2_TOL_GRAD/l2_rel/assert_l2_close move to tests/conftest.py so the three fused test files share one definition rather than three drifting copies; test_ckconv_nd_subq_fused keeps its own looser gradient table, since gradients through the full CKConvND stack accumulate more error than the bare op. 702 passed, 177 skipped, 77 xfailed, 3 xpassed across tests/ops and tests/modules; pre-commit clean. The fused path itself still does not execute here — this env has subquadratic-ops-torch-cu12 0.1.1 (below the 0.2.2 gate) and SM86 hardware — so the rewrite still needs an H100/B200 run. Co-Authored-By: Alireza Moradzadeh <amoradzadeh@nvidia.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * release: v0.2.0 Bumps VERSION 0.1.1 -> 0.2.0 and closes the changelog section. release.yml fails fast when the tag does not match the VERSION file, so this has to land before the tag is created. 0.2.0 rather than 0.1.2: this release is dependency-breaking for every 0.1.1 user. The torch floor moves >=2.10,<2.11 -> >=2.14,<2.15, [cuda] retargets from subquadratic-ops-torch-cu12 to -cu13 (the cu12 line tops out at 0.2.1 and has no fused_fft_conv2d), and [dali] moves to nvidia-dali-cuda130. Under PEP 440 both `~=0.1.1` and `>=0.1,<0.2` reject 0.2.0, which is what protects downstream pins from resolving into an environment that cannot build. Also records three things the changelog was missing: - SubqOpsCausalConv1d could not run under torch.autocast in 0.1.0/0.1.1. The kernel selects its specialisation from the input dtype and then demands an exact match, so bf16/fp16 activations against fp32 parameters raised ValueError. This is a genuine Fixed entry — unlike the fftconv_lowering correctness fixes, which never shipped and so are folded into that feature's own description rather than listed as regressions. - The lowering pass's decline conditions, which are user-visible through lowering_stats() skip reasons, and the scoped erase that replaced whole-graph dead-code elimination. - That CI validates the fused path only up to 32 extents per axis (64 FFT tile) on SM86; the 64-per-axis 128-tile path is guarded but has never run on Hopper/Blackwell. Better stated in the release notes than discovered by a user with an H100. scripts/check_version_pins.py: 14 pins consistent. pre-commit clean. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * docs(release): finish the cu12->cu13 sweep and document the CUDA 13 host requirements Three things this release got wrong that become immutable once 0.2.0 is on PyPI. The [quack] extra was missed by the cu12->cu13 sweep. pyproject.toml's `cuda` and `dali` extras were retargeted at CUDA 13, but `quack = ["quack-kernels"]` still resolved nvidia-cutlass-dsl's CUDA 12 runtime libraries — libs-cu12 is an unconditional dependency, so only the [cu13] extra adds libs-cu13. Now `quack-kernels[cu13]`, with a note to install via pip rather than uv (quack's README documents a uv ordering race on that path, NVIDIA/cutlass#3259). nvsubquadratic/modules/mlp.py imports quack.linear for fused GEMM+activation, so this path is more load-bearing than the RMSNorm-only framing suggested. Upgrading a 0.1.1 GPU environment silently corrupts it. cu12 and cu13 both install a top-level subquadratic_ops_torch/ package, and both DALI builds install nvidia/, so `pip install -U` overwrites files without removing the old distribution and a later uninstall of either breaks the other. Adds a Notes block with the uninstall-first sequence. Also records that a plain `pip install nvsubquadratic[cuda]` resolves from public PyPI unaided — torch 2.14.0 there is already a CUDA 13.0 build — so the cu130 index is only needed when compiling Apex/mamba against a specific torch. The NVIDIA driver >= 580 floor was documented only in a Dockerfile comment. A host on a 12.x-era driver fails at runtime with no documented cause. Added to the requirements lists in README.md, docs/getting_started.md and docs/index.rst. While in docs/index.rst, dropped the claim that [cuda] needs nvcc to build: subquadratic-ops-torch-cu13 0.3.0 ships prebuilt manylinux wheels for x86_64 and aarch64, so nothing is compiled at install time. Not touched: the "3-4x faster" claims in README/CHANGELOG/docs/ckconv_nd, which have no reproducing benchmark in the tree. Tracked separately. pre-commit clean; 14 version pins consistent; sphinx -W --keep-going green. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * docs: cite the measurement behind the fused-kernel speedup claims "Roughly 3-4x faster than torch_fft" appeared in four public places — README, CHANGELOG, docs/ops/README.md and the CKConvND docstring — with nothing in the tree to back it. benchmarks/ has no fused arm: bench_fftconv2d.py covers torch_fft and subq_ops only, and on farhadr/2d_bench the 2D sweep records one backend per resolution (subq_ops_fused at R=16/32/64, subq_ops at R>=128), so no same-resolution head-to-head exists there either. The numbers are not invented — they trace to ee5146e, which measured 3.6-3.9x over torch_fft and 1.2-2.4x over subq_ops on an H200 at B=8, hidden=768, forward+backward, bf16. That configuration existed only in a commit message no user will ever read, and "3-4x" alone reads as a property of the kernel rather than of one shape on one GPU. Replaces the rounded figure with the measured range plus its conditions at all four sites, and notes that speedups are shape- and hardware-dependent. Nothing here changes what ships; it makes a public performance claim checkable. An H100/B200 re-measurement is planned. If it disagrees, these four sites are where to correct it. pre-commit clean (mdformat idempotent over two passes on docs/ops/README.md); sphinx -W --keep-going green. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix(pyproject): correct the [quack] comment — the MLP backend is unreachable The comment added in 5a3a985 described [quack] as providing "fused RMSNorm + GEMM/activation kernels". The GEMM/activation half is not reachable: _validate_quack_backend (modules/mlp.py:285) raises NotImplementedError unconditionally, so MLP(backend="quack") always fails at __init__ and the quack.linear imports below it are dead code. Nothing in the tree passes backend="quack". The live consumers are RMSNorm and RMSNormChannelFirst, both of which gate on _cuda_supports_quack (SM90+) and fall back to pure PyTorch otherwise. This does not change the dependency — quack-kernels[cu13] is still the right target for a CUDA 13 release — only the comment, which ships in the sdist. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * bench: add the fused-vs-subq_ops-vs-torch_fft 2D comparison The README, CHANGELOG, docs/ops/README.md and the CKConvND docstring all cite 3.6-3.9x over torch_fft and 1.2-2.4x over subq_ops, but nothing in benchmarks/ produced those numbers: bench_fftconv2d.py compares torch_fft against subq_ops and never touches the fused kernel, and bench_subquadratic_fftconv.py is 1D causal. The figures traced to a commit message (ee5146e), not to runnable code. This script reproduces them. It defaults to the configuration the claim was measured at — B=8, hidden=768, bf16, forward+backward, median of 100 iters — and sweeps spatial 16/32/64 with the CKConvND double-grid kernel K=2N-1, which resolves to FFT tiles 32/64/128 respectively. The 128 tile means the spatial-64 row needs SM90+. Two traps it avoids, both worth knowing: - fftconv2d_fp32_bhl exists in BOTH ops/fftconv.py and ops/fftconv_chunked.py. Importing the wrong one silently benchmarks the chunked implementation. - subq_ops takes the kernel without the leading batch-1 dim ([H, Kx, Ky]), while torch_fft and the fused wrapper take [1|B, H, Kx, Ky]. Passing the same tensor to all three either throws or measures the wrong shape. Not in the wheel: [tool.setuptools.packages.find] includes only nvsubquadratic*, so this cannot affect the published artifact. Has not been run on SM90+ yet — that is the point of adding it. If the measured numbers disagree with the citations, the four sites to correct are README.md:70, CHANGELOG.md:17, docs/ops/README.md:127 and modules/ckconv_nd.py:621. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix(bench): time subq_ops through the wrapper, not the fp32-only raw op The script called subquadratic_ops_torch.fft_conv2d directly, which raises NotImplementedError: fft_conv2d supports float32 only on bf16/fp16 input. Since the benchmark's whole subject is native-dtype execution, the "vs subq_ops" column could never be produced — it failed at every spatial size on the first H100 run and printed n/a throughout. Uses nvsubquadratic.ops.fftconv_custom.fftconv2d_bhl instead, the repo's documented drop-in for fftconv2d_fp32_bhl. It performs the fp32 upcast/downcast that the raw op requires, which is also the honest comparison: subq_ops really does pay an upcast that the fused kernel does not, and that cost belongs in the measurement rather than being defined away. It also takes the same [1|B, H, Kx, Ky] kernel as torch_fft and the fused wrapper, so the special-cased squeeze for subq_ops is gone along with the docstring note warning about it. Found by the first H100 run of this script. The torch_fft column was unaffected and is reported separately. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * docs: replace the fused-kernel speedup claims with measured H100 numbers The cited 3.6-3.9x over torch_fft does not hold. Measured on an H100 (B=8, hidden=768, bf16, fwd+bwd, median of 100, via the now-fixed benchmarks/ops/bench_fused_fftconv2d.py): spatial tile torch_fft subq_ops fused vs torch vs subq 16 32 0.620ms 0.466ms 0.329ms 1.88x 1.42x 32 64 2.124ms 1.214ms 0.482ms 4.41x 2.52x 64 128 8.266ms 2.248ms 1.689ms 4.89x 1.33x Reproduced across two runs within 1-2%, so the spread is real rather than measurement noise. vs torch_fft is 1.9-4.9x, not 3.6-3.9x. The direction that matters is spatial=16: 1.88x is less than half the previously advertised floor, so the old text overclaimed for small extents — the case where the fused kernel is weakest, since its advantage grows with the FFT tile. The 3.6-3.9x figure was likely a single shape or a narrow sweep, which is why it read as a tight range. vs subq_ops is 1.3-2.5x, close enough to the previous 1.2-2.4x, but it does NOT vary monotonically with extent (1.42x, 2.52x, 1.33x), so it cannot be summarised as "larger is better" the way torch_fft can. All four sites now quote the range with its conditions and point at the script, so the claim is checkable rather than inherited. The H200 attribution is dropped: it is not reproducible from this tree, whereas these numbers are. The benchmark records the reference run in its docstring so a future regression is visible as a diff against a committed baseline. Measured on H100 80GB, torch 2.14.0+cu130, subquadratic-ops-torch-cu13 0.3.0. pre-commit clean (mdformat stable over two passes on the new CHANGELOG table); sphinx -W --keep-going green. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * docs(release): date 0.2.0 and cover #140 / #141 in the release notes The 0.2.0 section was written when #139 was the only thing in the release and dated 2026-09-08. Since then #140 and #141 merged to main and came into this branch via e505327, so both ship inside this tag and belong in the notes. - Dates the section 2026-09-16. - Adds #141: PatchMerging / PatchMerging2D / PatchEmbedHierarchical, the ViT5HierarchicalNet + StageSpec / ViT5HierarchicalClassificationNet pair built on them, the CIFAR-10 and ImageNet recipes, and the two new CIFAR-10 datamodules (torchvision and HuggingFace). - Adds #140: SpatialRecall3DMotionDataset / SpatialRecall3DMotionDataModule and the motion_3d example configs. Every class named above was checked to exist at this commit rather than transcribed from the PR descriptions. pre-commit clean, mdformat idempotent over two passes. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Alireza Moradzadeh <amoradzadeh@nvidia.com>
farhadrgh
added a commit
that referenced
this pull request
Sep 16, 2026
NVIDIA runs CodeRabbit self-hosted (coderabbit.nvidia.com), approved for EC 3e991, with code disposed after each review. The app is installed through https://github-onboarding.nvidia.com/install-github-app rather than the public GitHub Marketplace; this file only configures it once installed. Both top-level settings match the in-org precedent, NVIDIA-BioNeMo/bionemo-recipes, whose entire config is the same four lines: profile: chill - balanced rather than assertive, which reads as nitpicky and is the fastest way to get an AI reviewer muted. auto_review: false - on request only, via `@coderabbitai review` or `@coderabbitai full review`. Keeps the tool opt-in while the team forms a view of its signal-to-noise. Neither costs CI time: CodeRabbit is a hosted app reacting to webhooks, not a GitHub Action, so it never triggers the ~36-minute self-hosted gpu-tests job. The path_instructions are the part worth reviewing. Each encodes a defect that actually occurred in this repo rather than generic advice: - the worker_init_fn DDP-seeding trap, which silently collapses augmentation diversity to 1/world_size (PR #140, still open) - config.net needing to be a LazyConfig, and iterations-per-epoch needing the effective batch rather than the micro-batch (PR #141) - torch.fx passes calling eliminate_dead_code without an is_impure_node override, which can delete unrelated in-place ops (PR #139) - the two colliding fftconv2d_fp32_bhl definitions in ops/fftconv.py and ops/fftconv_chunked.py - the raw subquadratic_ops_torch.fft_conv2d op being fp32-only - SM90+ gating for fused spatial extents above 32 per axis - the requires_sm90 marker convention and the shared L2 tolerance helpers - display math nested under a bullet in docs/ops/, which makes mdformat non-idempotent and loops pre-commit forever - new public classes needing both a docs-tracker.md row and a Sphinx entry Ordering note: the pin-consistency instruction names scripts/check_version_pins.py, which lands with #139 and is not on main yet. The invariant it describes holds regardless; only the enforcing script is pending. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>
farhadrgh
added a commit
that referenced
this pull request
Sep 16, 2026
main now contains the curated extraction of this branch's library work, shipped as v0.2.0 via PR #139, plus #140 and #141. Every file both sides touched was resolved to main's version, because main is authoritative for all of it: - fftconv_lowering.py: this branch carried the original 533-line pass. main has the 631-line revision with _erase_chain, which fixes two correctness bugs — an unguarded whole-graph eliminate_dead_code that could delete unrelated in-place ops, and an any() short-circuit that fused only the first chain per graph. Keeping this branch's copy would reintroduce both. - pyproject.toml: torch >=2.12 here vs >=2.14 on main. The 2.14 floor is not a preference — torch pins nvidia-cudnn-cu13 exactly, and subquadratic-ops-torch-cu13 >=0.3.0 needs >=9.24.0.43, which only 2.14 satisfies. 2.12 makes `pip install nvsubquadratic[cuda]` unresolvable. - ckconv_nd.py, fftconv_custom.py, tests/conftest.py, the Dockerfile and enroot scripts, and the docs: same work, later revision on main. 21 files resolved this way. What this branch uniquely owns is untouched — the ND benchmark harness, the sweep JSONLs and results digest, the visualization and submit scripts, docs/mamba2_limits.md, and the FA4/FlexAttention additions to attention.py and sequence_mixer.py. Merged rather than rebased: 37 commits against 21 overlapping files would replay the same conflicts repeatedly, and Alireza has commits here, so rewriting the branch's history has a cost a merge does not. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>
farhadrgh
added a commit
that referenced
this pull request
Sep 16, 2026
…ess (#142) * ci: add CodeRabbit configuration NVIDIA runs CodeRabbit self-hosted (coderabbit.nvidia.com), approved for EC 3e991, with code disposed after each review. The app is installed through https://github-onboarding.nvidia.com/install-github-app rather than the public GitHub Marketplace; this file only configures it once installed. Both top-level settings match the in-org precedent, NVIDIA-BioNeMo/bionemo-recipes, whose entire config is the same four lines: profile: chill - balanced rather than assertive, which reads as nitpicky and is the fastest way to get an AI reviewer muted. auto_review: false - on request only, via `@coderabbitai review` or `@coderabbitai full review`. Keeps the tool opt-in while the team forms a view of its signal-to-noise. Neither costs CI time: CodeRabbit is a hosted app reacting to webhooks, not a GitHub Action, so it never triggers the ~36-minute self-hosted gpu-tests job. The path_instructions are the part worth reviewing. Each encodes a defect that actually occurred in this repo rather than generic advice: - the worker_init_fn DDP-seeding trap, which silently collapses augmentation diversity to 1/world_size (PR #140, still open) - config.net needing to be a LazyConfig, and iterations-per-epoch needing the effective batch rather than the micro-batch (PR #141) - torch.fx passes calling eliminate_dead_code without an is_impure_node override, which can delete unrelated in-place ops (PR #139) - the two colliding fftconv2d_fp32_bhl definitions in ops/fftconv.py and ops/fftconv_chunked.py - the raw subquadratic_ops_torch.fft_conv2d op being fp32-only - SM90+ gating for fused spatial extents above 32 per axis - the requires_sm90 marker convention and the shared L2 tolerance helpers - display math nested under a bullet in docs/ops/, which makes mdformat non-idempotent and loops pre-commit forever - new public classes needing both a docs-tracker.md row and a Sphinx entry Ordering note: the pin-consistency instruction names scripts/check_version_pins.py, which lands with #139 and is not on main yet. The invariant it describes holds regardless; only the enforcing script is pending. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com> * docs: document the CodeRabbit review request process in the PR template auto_review is disabled, so CodeRabbit does nothing unless asked. Without this the config lands and nobody uses it — the trigger comments are invisible to anyone who has not read .coderabbit.yaml. Adds a Review section with two checkboxes: request the AI first pass, then request a codeowner. That order is the point — a codeowner's time is better spent on design than on what a bot would have caught. Records two ruleset facts that are easy to learn the hard way: - A CodeRabbit review does not count toward the required approving review. main requires a codeowner approval and bypass_actors is empty, so there is no override. - required_review_thread_resolution is true, so unresolved CodeRabbit threads block merge. They have to be resolved or replied to, not ignored. Also lists the other trigger commands, and points at .coderabbit.yaml as the place to correct systematic bad advice rather than repeating the correction on every PR. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com> * ci: rewrite the CodeRabbit path instructions against main The first version was written against the PR #139 branch and committed to a branch off main, so a third of it described code that does not exist here. Verified: resolve_fused_fft_size, requires_sm90, assert_l2_close, L2_TOL, torch.fx/eliminate_dead_code and scripts/check_version_pins.py each return zero hits in the tree — their only occurrence was .coderabbit.yaml citing itself. It also instructed on *.rs files, of which the repo has none. Two of those were worse than stale. tests/conftest.py on main defines only _subq_ops_version, requires_subq_ops_v2, device and dtype_fixture, so telling a reviewer to require requires_sm90 and assert_l2_close imports was advice that breaks test collection if followed. Those belong in #139, which introduces them. Also cut the SPDX rule as pure CI redundancy: scripts/license_check.py inserts the header itself and lint.yml runs pre-commit --all-files, so a missing header can never reach a human reviewer. What replaces them are invariants a general reviewer cannot infer from a diff: - COMPILE_COMPATIBLE and the chunking flags are flipped at runtime after import, so `from ... import` snapshots them and the flip silently no-ops. - fftconv_chunked.py shadows fftconv.py's public names deliberately; the defect is a new op with no matching shim, not the shadowing. - Transform length and crop offset are one contract — change one and the output is merely shifted, which a symmetric probe kernel will not catch. - OmegaConf returns ListConfig/DictConfig, neither of which subclasses list/dict, so isinstance checks on config-supplied values silently take the wrong branch. - Buffers derived from config use persistent=False so old checkpoints survive a resolution change. - QKVSequenceMixer passes cp_group positionally; a mixer reading it from **kwargs gets None and runs without cross-rank communication. - Only torch/einops/omegaconf/numpy may be imported at module scope; extras need a lazy getter and an autodoc_mock_imports entry. Every symbol and path cited above was checked to exist at this commit. Reviewed by four independent critics (over-fitting, CI redundancy, staleness, coverage gaps) and an adversarial verifier that re-derived each claim. 8 entries, 10.9k chars, well under the 20k per-entry limit. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com> --------- Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Add hierarchical patch merging to main, combining the implementation from #122 with the complementary local spatial-grid implementation. This PR supersedes #122 and targets main directly.
Instead of keeping one spatial resolution throughout a classifier, the new networks can repeatedly merge neighbouring 2×2 patches: for example, a 224×224 image becomes a 56×56 patch grid and then progresses through 28×28, 14×14 and 7×7 stages, typically doubling channel width at each transition. This supports hierarchical Hyena experiments and comparisons with flat architectures.
What changes
Shared merging operation with two interfaces
nvsubquadratic/modules/patch_merging.pyadds:PatchMerging: flat token input, optional prepended register row, configurable output width and normalization. Requires even spatial grids. Register tokens take a separate learned projection, preserving their conditioning role as patch resolution changes.PatchMerging2D: channels-last spatial grids, LayerNorm and C → 2C projection, with zero padding for odd heights/widths.PatchEmbedHierarchical: stride-four convolution and LayerNorm for the spatial hierarchy's input stem.Both merger interfaces share the same 2×2 concatenation order. Tests verify identical outputs and gradients with compatible shared weights.
Two hierarchical classifiers
ViT5HierarchicalClassificationNet, from feat(modules): Swin-style PatchMerging with register-row support #122: configurable stage widths/depths, absolute position embeddings, and either pure patch tokens or register-row FiLM conditioning. Global average pooling excludes the register row.ViT5HierarchicalNet, recovered locally: four spatial-grid stages, normalized convolutional stem, no position/register tokens, configurable stochastic-depth schedule and odd-grid padding.These are distinct experiment architectures; their checkpoints are not interchangeable. Existing parameter names within each implementation are preserved.
Experiment recipes and data adapters
Includes the original
v6_hierarchicalImageNet/CIFAR recipes, recoveredv5_patchmergespatial hierarchies and block-diagonal kernel variant, local ImageNet comparisons, and CIFAR hierarchical-versus-flat experiments.CIFAR data recipes remain explicit:
cifar10.pyuses torchvision, whilecifar10_hf.pyretains the local Hugging Face source and augmentation settings. This avoids silently changing the meaning of either experiment family.Correctness and compatibility fixes
Hyena(use_rope=False)arguments from recovered configs.Review corrections
network.headwhen replacing the spatial classifier, preserving existing checkpoint keys. Restore the networks section in the test README.Scope and review order
The main-relative diff contains 43 files, primarily recovered source, experiment configs and tests. Review the merger module first, then the two networks, then the recipes/data adapters and their tests.
The branch includes current main and preserves #122's feature history. There is no dependency on merging #122 first. Local benchmark scripts already exist on main with identical Python ASTs and are not duplicated. Unrelated spatial-recall Gaussian changes, global FFT/mask edits, generated figures and historical result trackers are outside this PR.
Validation
python -O, and classifier replacement with ordinary/compiled checkpoint keys.Local validation used Python 3.12 / Torch 2.9.0+cu128 on CPU. Real Hyena CPU bf16 encounters an existing FFT shortcut-dtype assertion, so real-mixer CPU tests use float32; bf16 merging and hierarchy wiring are tested independently. ImageNet lazy-recipe tests substitute the unrelated Apex/DALI training-base dependency. No new training runs or accuracy claims are included.