Triton-Ascend Fused Attention

What it does

When your agent uses it

• Writing a forward scaled dot-product attention kernel for Ascend NPU
• Supporting causal and non-causal masking in a Triton-Ascend kernel
• Handling large head dimensions without overflowing on-chip memory
• Checking a custom attention kernel against torch_npu.npu_fusion_attention

Example prompts

• “Write a causal fused attention forward kernel in Triton-Ascend for BNSD tensors with head dimension 128.”
• “My attention kernel runs out of on-chip memory at head dimension 256, so apply the scratch accumulator approach.”
• “Add a test that compares our kernel with torch_npu.npu_fusion_attention in float16 and bfloat16.”
• “Choose BLOCK_M and BLOCK_N values that divide the sequence length evenly.”

Requirements

• Triton-Ascend on an Ascend NPU
• `torch_npu` for the reference comparison

Workflow steps

6 steps, taken from the first numbered list in SKILL.md.

1. 1Validate matching Q, K, and V shapes, compatible strides, supported dtype, and HEAD_DIM in {16, 32, 64, 128, 256}.
2. 2Choose BLOCK_M for query rows and BLOCK_N for key/value rows. Require N_CTX to be divisible by both in the compact implementation.
3. 3Map each task to a batch, head, and query-block tuple and create block pointers for Q, K, V, and output.
4. 4Load one Q block and stream K/V blocks through an online-softmax inner loop.
5. 5For causal attention, process preceding blocks and the diagonal block as separate stages; apply the triangular mask only to the diagonal…
6. 6Normalize the accumulated output by the running denominator and store it in the output dtype.

What it can do on your machine

Read from SKILL.md and the folder at commit 4ab5ee7. It shows what the files ask for, not the result of running them.

• Tool permissions

  Pre-approves nothing: there is no allowed-tools line, so your agent's usual permission prompts apply.

  From allowed-tools in the SKILL.md frontmatter.

• Runs code

  No scripts in the folder and no shell commands in SKILL.md.

  From the folder's file list and the shell code blocks in SKILL.md.

• Network

  No URLs in SKILL.md.

  From URLs in SKILL.md, links to its own repository left out.

• Credentials

  Names no API keys, tokens, secrets or passwords.

  From names ending in _API_KEY, _TOKEN, _SECRET, _KEY or _PASSWORD in SKILL.md.

Context cost

Triton-Ascend Fused Attention loads about 753 tokens when it runs. Until then it costs about 106 tokens; SKILL.md has 328 words of instructions outside code blocks.

Always · name and description, kept in context so the agent knows when to use it

~106

When it runs · the whole SKILL.md, loaded when a task matches

~753

Estimates: characters ÷ 4, the usual rule of thumb; real counts depend on the model's tokenizer. Scripts and assets cost tokens only if the agent reads them.

Safety

Automated static check — not a guarantee. Review scripts before installing. It scans the text of SKILL.md for risky patterns (piping downloads into a shell, reading credential files, hidden Unicode, destructive commands); files beside SKILL.md are not scanned.

SKILL.md

The full file from Krusty84/triton-ascend-agent-dev-kit at commit 4ab5ee7, republished under its Apache-2.0 licence (© Krusty84). 328 words, ~753 tokens.

name

build-triton-ascend-fused-attention

description

Build and adapt a FlashAttention-v2-style fused forward attention kernel for Triton-Ascend using tiled Q/K/V block pointers, online softmax, causal staging, and float32 accumulation. Use when an agent needs forward scaled dot-product attention in BNSD layout on Ascend NPU, must support causal or non-causal masking, or must handle large head dimensions without overflowing on-chip memory.

scores = dot(q, transpose(k)) * scale
scores = scores + causal_mask_if_diagonal
m_new = max(m_i, row_max(scores))
p = exp(scores - m_new)
alpha = exp(m_i - m_new)
l_i = l_i * alpha + row_sum(p)
acc = acc * alpha + dot(cast(p, k.dtype), v)
m_i = m_new
output = acc / l_i

© Krusty84, Apache-2.0. Rendered from Markdown: HTML in the file is shown as text, images as links, and headings moved down two levels. Raw file

Files

Just SKILL.md in skills/build-triton-ascend-fused-attention of Krusty84/triton-ascend-agent-dev-kit.

Open the folder on GitHubat commit 4ab5ee7