Overview

FlashKDA is a cutting-edge library delivering high-performance Kimi Delta Attention (KDA) kernels. Built upon the robust CUTLASS framework, FlashKDA is engineered for speed and efficiency, targeting modern GPU architectures (SM90 and above). It seamlessly integrates with PyTorch, enabling accelerated KDA operations within your machine learning workflows.

Key Features

• High Performance: Leverages highly optimized CUDA kernels to achieve significant speedups in KDA computations.
• CUTLASS Integration: Built on the battle-tested CUTLASS library for advanced CUDA programming.
• PyTorch Compatibility: Designed for easy integration with PyTorch (2.4 and above).
• Flexible State Handling: Supports both stateless and stateful recurrent operations, with options for bf16 and fp32 precision for recurrent states.
• Variable Length Batching: Efficiently handles sequences of varying lengths using cumulative sequence lengths (cu_seqlens).
• Configurable Kernel Options: Fine-grained control over kernel behavior, including gating, L2 normalization, and sigmoid activation.

Typical Use Cases

FlashKDA is ideal for applications requiring efficient and accelerated processing of Kimi Delta Attention. This includes:

• Large Language Models (LLMs): Speeding up attention mechanisms in transformer architectures for LLM training and inference.
• Recurrent Neural Networks (RNNs): Enhancing the performance of stateful recurrent computations.
• Sequence Modeling: Accelerating various sequence-to-sequence tasks and time-series analysis.
• Any PyTorch application that can benefit from optimized KDA kernels, particularly when using the flash-linear-attention library as a backend.