LICENSE
README.md
pyproject.toml
hiera_optim/__init__.py
hiera_optim/checkpoint.py
hiera_optim/fp8.py
hiera_optim/lightning.py
hiera_optim/patch.py
hiera_optim.egg-info/PKG-INFO
hiera_optim.egg-info/SOURCES.txt
hiera_optim.egg-info/dependency_links.txt
hiera_optim.egg-info/requires.txt
hiera_optim.egg-info/top_level.txt
hiera_optim/adapters/__init__.py
hiera_optim/adapters/hiera.py
hiera_optim/attention/__init__.py
hiera_optim/attention/mask_unit.py
hiera_optim/attention/mask_unit_fused.py
hiera_optim/kernels/__init__.py
hiera_optim/kernels/flash_qpool.py
hiera_optim/kernels/fused_mlp_block.py
hiera_optim/kernels/fused_proj_epilogue.py
hiera_optim/kernels/fused_proj_epilogue_bwd.py
hiera_optim/kernels/fused_qkv_preamble.py
hiera_optim/kernels/fused_qkv_preamble_bwd.py
hiera_optim/kernels/gpu_config.py
hiera_optim/kernels/layernorm.py
hiera_optim/kernels/mask_gather.py
hiera_optim/kernels/pack_qkv_grad.py
hiera_optim/kernels/pack_z.py
hiera_optim/kernels/residual_layernorm.py
hiera_optim/ops/__init__.py
hiera_optim/ops/mask_gather.py
tests/test_bwd_kernels.py
tests/test_compile_reduce_overhead.py
tests/test_e2e_equivalence.py
tests/test_flash_qpool.py
tests/test_fused_mlp_block.py
tests/test_graph_safe.py
tests/test_graph_safe_matrix.py
tests/test_lightning_callback.py
tests/test_mask_gather.py
tests/test_mask_unit_attention.py
tests/test_matrix.py
tests/test_pack_kernels.py
tests/test_sdpa_backend.py
tests/test_selective_checkpoint.py