FlashKDA is a CUTLASS-based implementation of Kimi Delta Attention (KDA) kernels developed by Moonshot AI, designed to accelerate the prefill phase of large language model inference through optimized GPU computation. The system provides significant performance improvements over existing attention implementations while maintaining compatibility with standard model backends.

FlashKDA implements the Kimi Delta Attention mechanism using NVIDIA's CUTLASS (Custom Architecture for Specialized Tensor Operations) framework, enabling efficient computation on modern GPU architectures. The implementation focuses on optimizing the prefill stage of transformer inference, where input tokens are processed to generate initial key-value cache states before autoregressive token generation begins ¹⁾.

Delta Attention represents an optimization approach to the standard multi-head attention mechanism, reducing computational redundancy during the prefill phase through specialized kernel implementations. The CUTLASS-based approach allows fine-grained control over memory access patterns and computation scheduling on GPUs, enabling significant speedup compared to naive or baseline attention implementations.

FlashKDA demonstrates substantial performance improvements across multiple metrics. On NVIDIA H20 GPUs, the implementation achieves 1.72x to 2.22x prefill speedup compared to flash-linear-attention baselines, a significant acceleration for the computationally intensive initial inference phase ²⁾.

At scale, FlashKDA enables 508 tokens per second (tok/s) throughput on systems configured with 8x AMD MI300X GPUs ³⁾. This throughput metric reflects end-to-end inference performance across both prefill and decode phases on multi-GPU systems, representing practical production-level performance for serving large language models.

A key design principle of FlashKDA is drop-in backend compatibility, enabling integration into existing model serving infrastructure without requiring modifications to higher-level code. This compatibility approach reduces deployment friction and allows organizations to adopt the performance improvements with minimal engineering effort. The implementation maintains standard attention interface contracts while replacing the underlying computational kernels with optimized CUTLASS implementations.

FlashKDA addresses a critical bottleneck in LLM inference pipelines—the prefill phase processing latency. For applications involving long context windows, batch processing of multiple prompts, or serving scenarios with varied input lengths, prefill optimization directly impacts end-to-end latency and throughput metrics. The 1.72x-2.22x speedup improvement translates to measurable reductions in time-to-first-token for end users and increased batch processing capacity for server deployments.

The demonstrated 508 tok/s throughput on 8x MI300X systems indicates FlashKDA's effectiveness in large-scale deployment scenarios, where multiple GPUs are typically coordinated for tensor parallelism or data parallelism across batches.

FlashKDA builds upon the lineage of flash attention optimizations, including foundational work on IO-aware attention kernel design. The progression from standard attention implementations to flash variants to specialized delta attention approaches reflects ongoing GPU kernel optimization efforts in the ML infrastructure space. CUTLASS-based implementations benefit from compiler infrastructure that understands modern GPU memory hierarchies and execution models, enabling automatic tuning and optimization across different GPU families.

• Moonshot FlashKDA
• MLA (Multi-head Latent Attention) Kernels
• Attention Kernel Optimization
• DSA Kernels
• Kimi Linear