Modular GDN forward / backward kernels (compatible with Kimi CP)

[icavan]

Description

Implement modular forward and backward CUDA kernels for the GDN (Gated Delta Network) algorithm, compatible with Kimi CP.

Context

GDN is a linear attention variant that improves expressiveness with gating and delta-style updates. Implementing modular forward/backward kernels compatible with Kimi CP would enable efficient distributed training via context parallelism.

Tasks

• Implement GDN forward kernels (modular, similar to KDA structure)
• Implement GDN backward kernels
• Ensure compatibility with Kimi CP's context parallelism protocol
• Add correctness tests against FLA reference implementation
• Add benchmarks

References

• GDN: https://arxiv.org/abs/2412.06464
• Kimi CP: https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/cp/KCP.md

[kingsleykimm]

I can take this issue - I was the one who commented on the Ant Ling Post with the h100 GDN kernels :). Right now I only have access to H100, so SM90 can be done first, and then I will try to move to Blackwell

[kingsleykimm]

@icavan - Should I duplicate all of the utility functions that kda/sm90 contains (like common.hpp files)? I could for now duplicate all the files, but it might be better in the future to refactor into a central utils/ or common/ directory if more linear attention models are added

[KevinZeng08]

@icavan - Should I duplicate all of the utility functions that kda/sm90 contains (like common.hpp files)? I could for now duplicate all the files, but it might be better in the future to refactor into a central utils/ or common/ directory if more linear attention models are added

Could you add gdn/sm90 directory under csrc for your implementation, making the code modular enough? Correspondingly, you can add cula/gdn directory for your python interface, maybe same with FLA interface.

[kingsleykimm]

Yes, I can do gdn/sm90 forward pass first, I'll try to keep as much of the KDA structure, and then I can push to a branch.

[kingsleykimm]

Sorry, I think I first implemented the fused Hopper forward pass for GDN, and also wrote tests for them, because that's what was in the kda/sm90 directory. I believe that the kda/sm100 structure contains the actual modular structure, is this correct? I've started working on those, but I also opened a pull request for the SM90 GDN fused forward pass, with GDN tests, #51. I'm going to start writing the SM90 kernel to match this FLA kernel.

[KevinZeng08]

Sorry, I think I first implemented the fused Hopper forward pass for GDN, and also wrote tests for them, because that's what was in the kda/sm90 directory. I believe that the kda/sm100 structure contains the actual modular structure, is this correct? I've started working on those, but I also opened a pull request for the SM90 GDN fused forward pass, with GDN tests, #51. I'm going to start writing the SM90 kernel to match this FLA kernel.

You're right, KDA in kda/sm100 is the modular forward impl.
For sm90 fused GDN forward, FlashInfer has an implementation, and the sm90 KDA is also modified from this code. We plan to port this code into cuLA and do optimizations. You can also try to compare with it and optimize something.
For sm90 modular GDN forward, it seems that there exists no CUDA optimized version currently, looking forward to your contribution!

[kingsleykimm]

Hi! Sorry for not updating for a bit - I've finished the chunk intra kernel and I am working on the chunk_fwd and fwd_o kernels right now - it should take around another week to finish and test these. Is it ok if I can first merge in unoptimized but correct kernel implementations, and then I can start optimizing the forward pass.