From aradotso-trending-skills-37
Implements high-performance CUDA kernels for Kimi Delta Attention (KDA) using FlashKDA and CUTLASS on SM90+ GPUs. Integrates as backend for flash-linear-attention chunk_kda in PyTorch for recurrent state inference.
npx claudepluginhub joshuarweaver/cascade-ai-ml-agents-misc-1 --plugin aradotso-trending-skills-37This skill uses the workspace's default tool permissions.
> Skill by [ara.so](https://ara.so) — Daily 2026 Skills collection.
Guides Next.js Cache Components and Partial Prerendering (PPR) with cacheComponents enabled. Implements 'use cache', cacheLife(), cacheTag(), revalidateTag(), static/dynamic optimization, and cache debugging.
Guides building MCP servers enabling LLMs to interact with external services via tools. Covers best practices, TypeScript/Node (MCP SDK), Python (FastMCP).
Generates original PNG/PDF visual art via design philosophy manifestos for posters, graphics, and static designs on user request.
Skill by ara.so — Daily 2026 Skills collection.
FlashKDA provides high-performance CUDA kernels for Kimi Delta Attention (KDA) built on CUTLASS. It targets SM90+ GPUs (H100/H20 class) and integrates as a drop-in backend for flash-linear-attention's chunk_kda operation.
git clone https://github.com/MoonshotAI/FlashKDA.git flash-kda
cd flash-kda
git submodule update --init --recursive
pip install -v .
Install the FLA integration (optional but recommended):
pip install -U flash-linear-attention # >= 0.5.0
flash_kda.fwdThe primary low-level kernel call:
import torch
import flash_kda
flash_kda.fwd(
q, k, v, g, beta, scale, out,
A_log, dt_bias, lower_bound,
initial_state=None,
final_state=None,
cu_seqlens=None
)
Tensor shapes and dtypes:
| Parameter | Dtype | Shape | Notes |
|---|---|---|---|
q | bf16 | [B, T, H, K] | Query; K must be 128 |
k | bf16 | [B, T, H, K] | Key; K must be 128 |
v | bf16 | [B, T, H, V] | Value; V must be 128 |
g | bf16 | [B, T, H, K] | Gate logits (sigmoid/activation applied internally) |
beta | bf16 | [B, T, H] | Beta logits (sigmoid applied internally) |
scale | float | scalar | Attention scale factor |
out | bf16 | [B, T, H, V] | Pre-allocated output tensor |
A_log | fp32 | [H] | Per-head log-gate parameter |
dt_bias | fp32 | [H, K] | Per-head gate bias |
lower_bound | float | scalar | Gate lower bound, range [-5.0, 0] |
initial_state | bf16/fp32/None | [B, H, V, K] or [N, H, V, K] | Optional initial recurrent state |
final_state | bf16/fp32/None | [B, H, V, K] or [N, H, V, K] | Optional output final state |
cu_seqlens | int64 | [N+1] | Optional cumulative seq lengths for varlen |
Constraints:
K == V == 128 requiredcu_seqlens is provided, B must be 1 and T is total tokens across all sequencesinitial_state and final_state dtypes must match when both providedFlashKDA auto-dispatches from FLA's chunk_kda when installed:
import torch
import logging
from fla.ops.kda import chunk_kda
# Optional: see dispatch decisions
logging.basicConfig(level=logging.INFO)
B, T, H, K, V = 2, 2048, 16, 128, 128
q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')
k = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')
v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda')
g = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')
beta = torch.randn(B, T, H, dtype=torch.bfloat16, device='cuda')
A_log = torch.randn(H, dtype=torch.float32, device='cuda')
dt_bias = torch.zeros(H, K, dtype=torch.float32, device='cuda')
h0 = torch.zeros(B, H, V, K, dtype=torch.float32, device='cuda')
scale = K ** -0.5
lower_bound = -5.0
with torch.inference_mode():
out, final_state = chunk_kda(
q=q, k=k, v=v, g=g, beta=beta,
scale=scale,
initial_state=h0,
output_final_state=True,
use_gate_in_kernel=True,
use_qk_l2norm_in_kernel=True,
use_beta_sigmoid_in_kernel=True,
safe_gate=True,
A_log=A_log,
dt_bias=dt_bias,
lower_bound=lower_bound,
transpose_state_layout=True,
)
# out: [B, T, H, V], final_state: [B, H, V, K]
import torch
import flash_kda
def run_flash_kda(
q, k, v, g, beta,
A_log, dt_bias,
lower_bound=-5.0,
initial_state=None,
):
B, T, H, K = q.shape
V = v.shape[-1]
scale = K ** -0.5
out = torch.empty(B, T, H, V, dtype=torch.bfloat16, device=q.device)
final_state = torch.zeros(B, H, V, K, dtype=torch.float32, device=q.device)
flash_kda.fwd(
q, k, v, g, beta,
scale, out,
A_log, dt_bias, lower_bound,
initial_state=initial_state,
final_state=final_state,
cu_seqlens=None,
)
return out, final_state
B, T, H, K = 1, 4096, 8, 128
device = 'cuda'
dtype = torch.bfloat16
q = torch.randn(B, T, H, K, device=device, dtype=dtype)
k = torch.randn(B, T, H, K, device=device, dtype=dtype)
v = torch.randn(B, T, H, K, device=device, dtype=dtype) # V==K==128
g = torch.randn(B, T, H, K, device=device, dtype=dtype)
beta = torch.randn(B, T, H, device=device, dtype=dtype)
A_log = torch.full((H,), -0.1, device=device, dtype=torch.float32)
dt_bias = torch.zeros(H, K, device=device, dtype=torch.float32)
with torch.inference_mode():
out, state = run_flash_kda(q, k, v, g, beta, A_log, dt_bias)
print(out.shape) # [1, 4096, 8, 128]
print(state.shape) # [1, 8, 128, 128]
Use cu_seqlens for variable-length sequences packed into a single batch dimension:
import torch
import flash_kda
# Two sequences of lengths 512 and 768, packed together
seq_lens = [512, 768]
T_total = sum(seq_lens)
N = len(seq_lens)
H, K, V = 16, 128, 128
cu_seqlens = torch.tensor([0, 512, 1280], dtype=torch.int64, device='cuda')
# B must be 1 for varlen mode
q = torch.randn(1, T_total, H, K, dtype=torch.bfloat16, device='cuda')
k = torch.randn(1, T_total, H, K, dtype=torch.bfloat16, device='cuda')
v = torch.randn(1, T_total, H, V, dtype=torch.bfloat16, device='cuda')
g = torch.randn(1, T_total, H, K, dtype=torch.bfloat16, device='cuda')
beta = torch.randn(1, T_total, H, dtype=torch.bfloat16, device='cuda')
A_log = torch.zeros(H, dtype=torch.float32, device='cuda')
dt_bias = torch.zeros(H, K, dtype=torch.float32, device='cuda')
out = torch.empty(1, T_total, H, V, dtype=torch.bfloat16, device='cuda')
# State shape is [N, H, V, K] in varlen mode
final_state = torch.zeros(N, H, V, K, dtype=torch.float32, device='cuda')
scale = K ** -0.5
with torch.inference_mode():
flash_kda.fwd(
q, k, v, g, beta,
scale, out,
A_log, dt_bias, lower_bound=-5.0,
initial_state=None,
final_state=final_state,
cu_seqlens=cu_seqlens,
)
print(out.shape) # [1, 1280, 16, 128]
print(final_state.shape) # [2, 16, 128, 128]
Pass initial_state from a previous call to maintain recurrent state across chunks:
import torch
import flash_kda
H, K, V = 16, 128, 128
B = 2
scale = K ** -0.5
def inference_step(q, k, v, g, beta, A_log, dt_bias, state=None):
T = q.shape[1]
out = torch.empty(B, T, H, V, dtype=torch.bfloat16, device='cuda')
new_state = torch.zeros(B, H, V, K, dtype=torch.float32, device='cuda')
flash_kda.fwd(
q, k, v, g, beta, scale, out,
A_log, dt_bias, lower_bound=-5.0,
initial_state=state,
final_state=new_state,
cu_seqlens=None,
)
return out, new_state
A_log = torch.zeros(H, dtype=torch.float32, device='cuda')
dt_bias = torch.zeros(H, K, dtype=torch.float32, device='cuda')
state = None
for chunk_idx in range(4):
q = torch.randn(B, 256, H, K, dtype=torch.bfloat16, device='cuda')
k = torch.randn(B, 256, H, K, dtype=torch.bfloat16, device='cuda')
v = torch.randn(B, 256, H, V, dtype=torch.bfloat16, device='cuda')
g = torch.randn(B, 256, H, K, dtype=torch.bfloat16, device='cuda')
beta = torch.randn(B, 256, H, dtype=torch.bfloat16, device='cuda')
with torch.inference_mode():
out, state = inference_step(q, k, v, g, beta, A_log, dt_bias, state)
print(f"Chunk {chunk_idx}: out={out.shape}, state={state.shape}")
| Variable | Values | Effect |
|---|---|---|
FLA_FLASH_KDA | 0 / 1 | Set to 0 to force Triton fallback in FLA |
# Disable FlashKDA, use Triton path
FLA_FLASH_KDA=0 python your_script.py
bash tests/test.sh
tests/test_fwd.py — correctness tests against PyTorch reference and flash-linear-attentionimport logging
logging.basicConfig(level=logging.INFO)
# Successful: [FLA Backend] kda.chunk_kda -> flashkda
# Rejected: [FLA Backend] kda.chunk_kda rejected: <reason>
import torch
cap = torch.cuda.get_device_capability()
assert cap >= (9, 0), f"FlashKDA requires SM90+, got SM{cap[0]}{cap[1]}"
# WRONG — will error
q = torch.randn(1, 512, 8, 64, ...) # K=64 not supported
# CORRECT
q = torch.randn(1, 512, 8, 128, ...) # K=128 required
torch.inference_mode() not torch.no_grad()# FlashKDA requires inference_mode for FLA dispatch
with torch.inference_mode():
out, state = chunk_kda(...)
# initial_state and final_state must have matching dtypes
initial = torch.zeros(B, H, V, K, dtype=torch.float32, device='cuda')
final = torch.zeros(B, H, V, K, dtype=torch.float32, device='cuda') # must match
# bf16 initial + fp32 final → error
lower_bound valid rangelower_bound = -5.0 # valid: range is [-5.0, 0]
lower_bound = -2.5 # valid
lower_bound = 0.0 # valid boundary
lower_bound = -10.0 # out of spec — use -5.0 as safe minimum
bash setup_clangd.sh
# Generates .clangd with correct include paths for CUDA/CUTLASS sources