valoomba/rdna3-qwen-attention

RDNA3 Qwen Attention

ROCm cached-decode GQA attention kernel for Qwen3.6-style full-attention layers on RDNA3 (gfx1100) GPUs.

The package exposes:

• attention_decode(query, key, value, num_key_value_groups=None, scaling=None)
• rdna3_gqa_attention_forward(...), a Transformers AttentionInterface-style wrapper that uses the native kernel for supported cached-decode calls and falls back to PyTorch SDPA otherwise.

How to use

from kernels import get_kernel

kernel = get_kernel("valoomba/rdna3-qwen-attention", version=1)
out = kernel.attention_decode(query, key, value)

For Transformers integration, register kernel.rdna3_gqa_attention_forward with the model's attention registry in the calling project.

Benchmarks

Benchmarking script is available for this kernel:

kernels benchmark valoomba/rdna3-qwen-attention --version 1