MoonshotAI/FlashKDA

FlashKDA: high-performance Kimi Delta Attention kernels

何を解決するか

FlashKDA は、Kimi Delta Attention (KDA) のための高性能 CUDA カーネルを提供します。これは線形アテンション機構であり、この特定のアテンション変種を使用するモデルの速度とメモリ効率を向上させるために、効率的でハードウェア加速された計算の必要性に対応しています。

動作方法

このプロジェクトは、CUTLASS ライブラリに基づいた KDA カーネルを実装しており、特に NVIDIA SM90 (H100) 以降のアーキテクチャを対象としています。flash-linear-attention (FLA) ライブラリのバックエンドとして統合され、ユーザーが標準の Triton ベース実装の代わりに、より高いパフォーマンスを発揮する最適化された C++ カーネルを使用できるようにします。

対象ユーザー

線形アテンションモデルを扱っている AI 研究者やエンジニア、および flash-linear-attention フレームワークを利用して、最新の NVIDIA GPU 上でモデルの推論およびトレーニングパフォーマンスを最適化したい方々を対象としています。

特徴

  • SM90+ GPU 向けに CUTLASS に基づいた高性能カーネル。
  • flash-linear-attention ライブラリとの自動ディスパッチによるシームレスな統合。
  • 累積シーケンス長 (cu_seqlens) を使用した可変長バッチ処理のサポート。
  • 状態なしおよび状態あり（初期/最終状態）の再帰的計算の両方をサポート。

関連

  • プロジェクト
    inclusionAI/cuLA

    KDAやLightning Attentionなどの線形アテンションのバリエーション向けに、NVIDIA BlackwellおよびHopper GPU向けに特別に最適化された、高性能CUDAカーネルのライブラリ。

  • Dispatch
    Kimi Linearは、メモリ使用量を75%削減し、6倍高速なデコーディングを実現しつつ、フルアテンションを凌駕する

    Kimi Linearのハイブリッド・リニア・アテンション・アーキテクチャは、KV-cacheのメモリ使用量を最大75%削減し、1Mトークンのコンテキストに対して最大6倍高速なデコーディングを実現しながら、さまざまなタスクにおいてフルアテンションを凌駕します。

  • プロジェクト
    NVIDIA/cutlass

    A collection of CUDA C++ and Python DSL abstractions for implementing high-performance matrix-matrix multiplication (GEMM) and the related linear algebra operations on NVIDIA GPUs.

  • Dispatch
    Kimi Delta Attention: 導出、実装、そしてコミュニティの反応

    Kimi Delta Attention (KDA) は、スカラー保持をチャネルごとの対角行列に置き換えることで DeltaNet を改良し、キー・チャネルを独立に忘却できるようにしつつ delta‑rule 書き込みを保持します。融合再帰カーネルまたはチャンク単位 Triton カーネルで効率的に実行可能です。

  • プロジェクト
    Tencent/hpc-ops

    LLM推論の生産環境向けに最適化されたCUDAカーネルを提供する、高性能オペレータライブラリ。アテンション、MoE、サンプリングの各処理を高速化し、サーバーのレイテンシを低減します。