GitLab

Use 64-bit for pointer offsets

Updated pointer offset calculations to use 64-bit integers to prevent overflow with large batch or image sizes.

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