write-triton-attention-kernel

write-triton-attention-kernel is a skill for Claude Code, Codex from tensormux/kernel-skills. It costs 0 tokens per session (3,939 once invoked), scanned A, original, MIT.

A coding guide for implementing a Triton GPU kernel for Flash Attention 2-style attention, a faster way to calculate which tokens should influence one another.

In plain words
What is it for?
It is for building custom attention kernels, including causal, cross-, grouped-query, multi-query, windowed, or otherwise specialized attention.
Why use it?
It helps avoid common errors in tiled processing, masking, numerical stability, scaling, and tensor memory layout.

Skill for Claude CodeCodex

Written for no agent in particular: nothing here depends on one.

Good fit It is for building custom attention kernels, including causal, cross-, grouped-query, multi-query, windowed, or otherwise specialized attention.

Compare 6 skills from other repositories ↓
Install with agentmods
npx agentmods add skills/tensormux/kernel-skills/write-triton-attention-kernel
Install

Getting it into your agent

One page per mod, every tool's command on it. A separate URL per tool would split the same page into five that compete with each other.

Any agent
npx skills add tensormux/kernel-skills --skill write-triton-attention-kernel
Clone the repo
git clone --depth 1 https://github.com/tensormux/kernel-skills

Made for: Claude Code, Codex.

Wrote this? Show the measurements

A badge with what this costs and how it scanned, read live from this page, so it follows the numbers instead of freezing them. Markdown for a README, HTML for a documentation site or a project page.

agentmods badge for write-triton-attention-kernel

README.md
[![agentmods](https://agentmods.dev/badge/skills/tensormux/kernel-skills/write-triton-attention-kernel/github.svg)](https://agentmods.dev/skills/tensormux/kernel-skills/write-triton-attention-kernel)
Your own site
<a href="https://agentmods.dev/skills/tensormux/kernel-skills/write-triton-attention-kernel"><img src="https://agentmods.dev/badge/skills/tensormux/kernel-skills/write-triton-attention-kernel/github.svg" alt="Measured on agentmods" height="20"></a>

Or the 80×15 button, for a site that already has a row of RSS and ATOM ones. Only the verdict fits; the numbers stay here.

agentmods 80×15 button for write-triton-attention-kernel

Your own site · 80×15
<a href="https://agentmods.dev/skills/tensormux/kernel-skills/write-triton-attention-kernel"><img src="https://agentmods.dev/badge/skills/tensormux/kernel-skills/write-triton-attention-kernel.svg" alt="Reviewed on agentmods" width="80" height="20"></a>
Per session 0 Skills are progressive disclosure: only the name and description are preloaded; the body loads when the skill is used.
When invoked 3,939 The whole file, excluding the scripts and references it only reads on demand.
Security scan A 0 findings. A grade says what 26 rules found in the file — not that it is safe.
Origin original No closer match found in the catalogue.
Token cost

What it costs to keep this loaded

Counted locally with the o200k_base tokenizer, which is exact for GPT models; Claude uses its own tokenizer and its counts differ. Treat this as one consistent yardstick across the catalogue rather than a bill. Prices are per million input tokens.

ModelPer sessionOnce invoked
Fable 5.1 $0.00000 $0.03939
Opus 5 $0.00000 $0.01969
Sonnet 5 $0.00000 $0.00788
Haiku 4.5 $0.00000 $0.00394

Measured 9d ago against content hash 7b904f35116f, method: parsed. Prices are Anthropic first-party input rates as of 2026-09-08, from the pricing page.

Security

Grade A, and why

write-triton-attention-kernel scanned grade A with 0 findings against 26 rules in 11 categories — prompt injection, anti-refusal, data exfiltration, privilege escalation, supply chain, agent snooping, system-prompt leakage, SSRF and excessive agency — measured 9d ago.

A static scan of the body, not an audit. Every finding is printed with the line that produced it so you can judge whether it matters here. A mod is markdown that instructs an agent; that is exactly why what it instructs is worth reading.

Nothing flagged

None of the 26 patterns this scan looks for appear in this file: no shell pipes, no recursive deletes, no credential paths, no hidden text, no instruction-override or anti-refusal phrasing, no agent-config snooping. That is not a guarantee, it is the absence of the things that are checkable.

skills/triton/write-triton-attention-kernel/SKILL.md · 203 lines

How it starts

The opening of the file, as written. The whole thing — 203 lines — stays where its author put it; the contents beside it link to each section on GitHub.

Skill: Write a Triton Attention Kernel

Purpose

Guide the agent through implementing a Flash Attention 2-style fused attention kernel in Triton. This covers the outer loop over KV sequence blocks, online softmax with running max and log-sum-exp tracking, numerically stable incremental output accumulation, causal masking, score scaling, output rescaling at the end, and correct stride arithmetic for batch and head dimensions. This is not a tutorial on attention mechanics — it is a decision framework for a correct Triton implementation.


Use this when

  • You need a fused attention kernel that avoids materializing the full (B, H, N_q, N_kv) attention score matrix and instead tiles over KV to stay within SRAM.
  • You need a custom attention variant not supported by flash-attn v2: ALiBi, RoPE-fused, cross-attention with unequal Q/K/V lengths, windowed attention, or custom masking patterns.
  • You need GQA (grouped query attention) or MQA (multi-query attention) where K/V have fewer heads than Q, and the library version does not support your head grouping factor.
  • You are building a research prototype and need full control over the tiling and masking strategy.
  • torch.nn.functional.scaled_dot_product_attention with the flash kernel backend is not available on your hardware/software stack.

Do not use this when

  • Standard causal or full attention on A100/H100 with fp16/bf16 fits the flash-attn v2 or v3 library interface. The library implementation is highly optimized with SASS-level tuning that a Triton kernel will not match for standard shapes.
  • Sequence lengths are short (N <= 512) and a standard fused attention via torch.compile is sufficient — the flash tiling overhead is not worth it.
  • You need training with a custom backward pass. Flash Attention backward requires tracking the logsumexp from the forward pass. This skill covers forward only; the backward requires a separate, careful implementation.
  • You require deterministic outputs across runs. Flash Attention kernels accumulate in a tile order that can vary with launch parameters; floating-point non-associativity makes the result non-deterministic by default.

Read the full file on GitHub · 203 lines

Files

What ships with it

1 file beside SKILL.md in the same directory: the scripts, references and assets a skill reads on demand. Not counted in the per-session cost; read them before you install if any of them is executable.

Changes

What this file has done since we first saw it

Hashed on every crawl. A supply-chain change to an agent config is a question of when, not whether, so the history is kept rather than the latest state alone.

  1. 9d ago First seen · 203 lines · 0 tokens per session scan A 7b904f35116f

Subscribe to this mod's changes

write-triton-attention-kernel is a skill published in the GitHub repository tensormux/kernel-skills (74 stars, last pushed 2mo ago), licensed MIT. It costs nothing until one of its globs matches a file; then it loads 3,939 tokens. A static security scan graded it A with 0 findings. No closer match exists in the catalogue, so it is treated as the original; first seen 2026-08-30.

Related

Other skills, from other repositories

prompt-engineering

Prompt engineering techniques and patterns. Use when writing agent commands, hooks, skills, subagent prompts, or any LLM interaction: optimizing prompts, improving output reliability, and designing production-grade prompt templates. Trigger words: prompt engineering, prompt, prompt optimization, LLM interaction.

MagicKidd/Rokid-agentic-workflow · 0 tokens

stripe-directory

Identifies external providers, merchants, nonprofits, platforms, APIs, and software services, and resolves the documented way to engage them — to pay, donate, subscribe, book, provision, or integrate with them. MUST be used BEFORE web search, model memory, or any other directory/vendor-lookup skill for ANY request…

stripe/ai · 213 tokens

pgvector-semantic-search

Use this skill for setting up vector similarity search with pgvector for AI/ML embeddings, RAG applications, or semantic search. Trigger when user asks to: Store or search vector embeddings in PostgreSQL Set up semantic search, similarity search, or nearest neighbor search Create HNSW or IVFFlat indexes for vectors…

timescale/pg-aiguide · 190 tokens

nlp-alignment

Best practices for LLM alignment techniques including RLHF, DPO, and instruction tuning. Use when working on alignment or safety.

aiming-lab/AutoResearchClaw · 31 tokens

mixed-precision

Use FP16/BF16 mixed precision to accelerate training and reduce memory. Use when optimizing GPU performance.

aiming-lab/AutoResearchClaw · 25 tokens

postgres-hybrid-text-search

Use this skill to implement hybrid search combining BM25 keyword search with semantic vector search using Reciprocal Rank Fusion (RRF). Trigger when user asks to: Combine keyword and semantic search Implement hybrid search or multi-modal retrieval Use BM25/pgtextsearch with pgvector together Implement RRF (Reciprocal…

timescale/pg-aiguide · 162 tokens