GitLab

using torch tools to change layout in bd pass

Optimized forward kernel for attention

https://github.com/NVIDIA/torch-harmonics/compare/main...azrael417:torch-harmonics:tkurth/mauro-rebase

* Added a new CSR array, psi_row_index, containing "ho" values sorted in descending order of CSR row length; this is used to process (ho, wo) points corresponding to longer rows before shorter ones, improving overlap and reducing the tail effect.

* Replaced PyTorch's slow permutation ops with custom kernels, significantly improving performance (especially on GB200).
* Split kernel into general and specialized versions for num_channel <= 16384, significantly reducing memory accesses.
* Enabled float4-based vectorized memory access when pointer alignment and channel size allow, improving throughput.
* Added runtime dispatch logic for kernel specialization.

Small fix in metric computation

Optimize bwd kernel: incremental qdot_max and alpha/integral/etc

adding lineinfo to optional debug flags

Leverage the same qdotk_max "trick" for the backward kernel. This avoids 1 loop
and saves about 20% of performance.

fixing attention perf test attempt 1

Optimized CUDA kernels for S2 Attention (forward and backward)

Merge branch 'mr/bwd-channel-permute-experiments' of https://github.com/rietmann-nv/torch-harmonics into mr/bwd-channel-permute-experiments

Also: Made fwd kernel use modified memory layout with standard shape

Also match the gradient output to the input, in terms of memory layout

Detect memory layout (B,C,H,W) (stride for C should be 1, if not, fix it)

This ensures that the backwards kernel is fast

putting qy in shared is a little faster

Changing internal memory layout means we can leave code in standard shape and
only change layout external to kernel