[RFC] Fuse L2 norm into short conv?

[sustcsonglin]

Proposal

L2 norm on QK has been used in DeltaNet, Gated DeltaNet, KDA, and related models. However, this computation takes a non-negligible portion of the total running time, especially in backward where it recomputes the L2 norm to save activations. Since QK are often outputs of a short convolution, it would be interesting to fuse the L2 norm into the short convolution so that we avoid an additional kernel launch and activation recomputation or storage.

Rationale

No response

[zhiyuan1i]

We can also evaluate moving short_conv's computation inside the chunk kernel during training. This is because: although short_conv's backward pass currently recomputes its output to reduce peak memory usage, this output is simultaneously required by the chunk kernel and is saved for backward pass within it. As a result, during actual training, this approach not only fails to reduce memory usage but also increases recomputation time.

[sustcsonglin]

i guess we can follow mamba codebase to fuse both short conv and output gate into the kernel to have a "fast path" impl as well