Skip to content

feat(kernels): add Triton batch-invariant attention #246

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
EthanZero2Hero wants to merge 1 commit into
base: main
Choose a base branch
from
+605 −0
Open

feat(kernels): add Triton batch-invariant attention #246

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions rl_engine/kernels/gtest/operator_specs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,20 @@ def _load_object(path: str) -> Any:


OP_SPECS = {
"attention": OperatorSpec(
name="attention",
op_class="reduction",
gold_path="rl_engine.kernels.ops.pytorch.attention.standard_attn.NativeAttentionOp",
gold_method="forward_fp32",
candidate_paths={
"pytorch": "rl_engine.kernels.ops.pytorch.attention.standard_attn.NativeAttentionOp",
"triton": (
"rl_engine.kernels.ops.triton.attention.standard_attn."
"TritonBatchInvariantAttentionOp"
),
},
grad_input_names=("q", "k", "v"),
),
"logp": OperatorSpec(
name="logp",
op_class="logprob",
Expand Down
14 changes: 14 additions & 0 deletions rl_engine/kernels/ops/triton/attention/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
# SPDX-License-Identifier: Apache-2.0
# Copyright (c) 2026 RL-Kernel Contributors

from rl_engine.kernels.ops.triton.attention.standard_attn import (
TritonBatchInvariantAttentionOp,
triton_batch_invariant_attention,
triton_batch_invariant_attention_with_lse,
)

__all__ = [
"TritonBatchInvariantAttentionOp",
"triton_batch_invariant_attention",
"triton_batch_invariant_attention_with_lse",
]
358 changes: 358 additions & 0 deletions rl_engine/kernels/ops/triton/attention/standard_attn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,358 @@
# SPDX-License-Identifier: Apache-2.0
# Copyright (c) 2026 RL-Kernel Contributors

from __future__ import annotations

import math
from typing import Optional

import torch
import triton
import triton.language as tl

from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp
Comment thread
coderabbitai[bot] marked this conversation as resolved.

_BLOCK_N = 64


def _next_power_of_2(value: int) -> int:
return 1 << (value - 1).bit_length()


@triton.jit
def _standard_attn_fwd_kernel(
q_ptr,
k_ptr,
v_ptr,
mask_ptr,
out_ptr,
lse_ptr,
B: tl.constexpr,
H_Q: tl.constexpr,
H_KV: tl.constexpr,
S_Q: tl.constexpr,
S_KV: tl.constexpr,
D: tl.constexpr,
stride_qb: tl.constexpr,
stride_qh: tl.constexpr,
stride_qs: tl.constexpr,
stride_qd: tl.constexpr,
stride_kb: tl.constexpr,
stride_kh: tl.constexpr,
stride_ks: tl.constexpr,
stride_kd: tl.constexpr,
stride_vb: tl.constexpr,
stride_vh: tl.constexpr,
stride_vs: tl.constexpr,
stride_vd: tl.constexpr,
stride_ob: tl.constexpr,
stride_oh: tl.constexpr,
stride_os: tl.constexpr,
stride_od: tl.constexpr,
sm_scale: tl.constexpr,
BLOCK_N: tl.constexpr,
BLOCK_D: tl.constexpr,
CAUSAL: tl.constexpr,
HAS_KEY_PADDING_MASK: tl.constexpr,
Comment on lines +30 to +56

@coderabbitai coderabbitai Bot Jul 25, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy lift

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

file="rl_engine/kernels/ops/triton/attention/standard_attn.py"

if [ -f "$file" ]; then
  echo "== file size =="
  wc -l "$file"
  echo
  echo "== outline =="
  ast-grep outline "$file" 2>/dev/null || true
  echo
  echo "== relevant source =="
  sed -n '1,240p' "$file" | cat -n
else
  echo "File not found: $file"
  echo "Matching files:"
  fd -i 'standard_attn.py' .
fi

echo
echo "== usages / launches of standard_attn or fwd function name =="
rg -n "def .*standard|standard_attn|triton\.jit|stride_qb|stride_qh|stride_qs|stride_qd|HAS_KEY_PADDING_MASK|CAUSAL" .

echo
echo "== triton dependency specification =="
rg -n "triton|vllm-ascend" pyproject.toml requirements*.txt setup.py setup.cfg 2>/dev/null || true

Repository: RL-Align/RL-Kernel

Length of output: 14722

🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
import ast, re
from pathlib import Path

path = Path("rl_engine/kernels/ops/triton/attention/standard_attn.py")
src = path.read_text()
tree = ast.parse(src)

fn = tree.body[1]
assert isinstance(fn, ast.FunctionDef) and fn.name == "_standard_attn_fwd_kernel"
params = [(a.arg, isinstance(a.annotation, ast.Attribute) and
          isinstance(a.annotation.value, ast.Name) and
          a.annotation.value.id == "tl" and
          a.annotation.attr == "constexpr")
          for a in fn.args.args]

name_count = {}
for node in ast.walk(fn):
    if isinstance(node, ast.Name):
        name_count[node.id] = name_count.get(node.id, 0) + 1

print("parameter_constexpr_dict=", dict(params))
print("unused_params=", [name for name, annotated in params if name not in name_count])
print("B_use_count=", name_count.get("B", 0))

launch = None
for node in tree.body:
    if isinstance(node, ast.ClassDef) and node.name == "_TritonBatchInvariantAttention":
        for item in node.body:
            if isinstance(item, ast.FunctionDef) and item.name == "forward":
                for subnode in ast.walk(item):
                    if isinstance(subnode, ast.Call):
                        try:
                            name = ast.unparse(subnode.func)
                        except Exception:
                            name = ""
                        if name == "_standard_attn_fwd_kernel[grid]":
                            launch = subnode
print("launch_found=", launch is not None)
if launch:
    positional_names = []
    for node in launch.args.args:
        try:
            positional_names.append(ast.unparse(node))
        except Exception:
            positional_names.append("<expr>")
    print("launch_positional_keywords=", [ast.unparse(node) for node in launch.keywords])
    print("launch_positional_args=", positional_names)
PY

Repository: RL-Align/RL-Kernel

Length of output: 244

🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
import ast, re
from pathlib import Path

path = Path("rl_engine/kernels/ops/triton/attention/standard_attn.py")
src = path.read_text()
tree = ast.parse(src)

for item in tree.body:
    if isinstance(item, ast.FunctionDef):
        print("FUNCTION", item.lineno, item.name, item.args.args[0].lineno if item.args.args else None)
    if isinstance(item, ast.ClassDef):
        print("CLASS", item.lineno, item.name)

# Locate target function by source text.
match = re.search(r'def\r?\n[^()]*\(_standard_attn_fwd_kernel\s*\(', src, re.S)
print("text_target_found=", bool(match))

fn = tree.body[1] if len(tree.body) > 1 else None
params = [(a.arg, isinstance(a.annotation, ast.Attribute) and
          isinstance(a.annotation.value, ast.Name) and
          a.annotation.value.id == "tl" and
          a.annotation.attr == "constexpr")
          for a in fn.args.args] if fn else []

name_count = {}
if fn:
    for node in ast.walk(fn):
        if isinstance(node, ast.Name):
            name_count[node.id] = name_count.get(node.id, 0) + 1

print("parameter_constexpr_dict=", dict(params))
print("unused_params=", [name for name, annotated in params if annotated and name not in name_count])
print("B_use_count=", name_count.get("B", 0))
PY

echo "== launch call structure =="
python3 - <<'PY'
import ast, re
from pathlib import Path
path = Path("rl_engine/kernels/ops/triton/attention/standard_attn.py")
src = path.read_text()
tree = ast.parse(src)

# Print the exact function signature as declared in source.
lines = src.splitlines()
for start in [idx for idx,l in enumerate(lines,1) if 'def _standard_attn_fwd_kernel(' in l][0:1]:
    print("\n".join(f"{i:4}: {lines[i-1]}" for i in range(start, start+45)))

# Scan source text for index call with no keyword matching B.
for i,line in enumerate(lines):
    if '_standard_attn_fwd_kernel[grid](' in line:
        print("INDEX_CALL_LINE", i+1, line.strip())
        brace=lines[i].index('_standard_attn_fwd_kernel[grid](')
        depth=0
        text=''; idx=-1
        end=len(lines)
        while idx < len(lines)-1:
            idx += 1
            chunk=lines[idx]
            depth += chunk.count('(') - chunk.count(')')
            text += '\n'+chunk
            if depth <= 0:
                break
        print("INDEX_CALL_BODY")
        print(text)
PY

Repository: RL-Align/RL-Kernel

Length of output: 608

🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
import ast
from pathlib import Path

path = Path("rl_engine/kernels/ops/triton/attention/standard_attn.py")
src = path.read_text()
tree = ast.parse(src)

# Find target function by name.
fn = next(
    item for item in tree.body if isinstance(item, ast.FunctionDef) and item.name == "_standard_attn_fwd_kernel"
)

params = dict(
    (a.arg, isinstance(a.annotation, ast.Attribute) and
      isinstance(a.annotation.value, ast.Name) and
      a.annotation.value.id == "tl" and a.annotation.attr == "constexpr")
    for a in fn.args.args
)

name_count = {}
for node in ast.walk(fn):
    if isinstance(node, ast.Name):
        name_count[node.id] = name_count.get(node.id, 0) + 1

print("parameter_constexpr_dict=")
for k, v in params.items():
    print(f"  {k}: {v}")
print("unused_params=", [name for name, annotated in params.items() if annotated and name not in name_count])
print("B_use_count=", name_count.get("B", 0))

# Find target launch call by exact function name.
launch = next(
    (node for node in ast.walk(tree)
     if isinstance(node, ast.Call)
     and all(getattr(getattr(node.func, "value", None), "id", None) == "_standard_attn_fwd_kernel"
             and getattr(node.func, "attr", None) == "__getitem__"
             and getattr(getattr(node.func, "slice", None), "id", None) == "grid")),
    None
)
print("launch_found=", launch is not None)
if launch:
    positional_names = []
    for node in launch.args.args:
        try:
            positional_names.append(ast.unparse(node))
        except Exception:
            positional_names.append("<expr>")
    print("launch_positional_args=", positional_names)
    print("launch_keyword_names=", [kw.arg for kw in launch.keywords])
    print("launch_keyword_names_without_B=", [kw.arg for kw in launch.keywords if kw.arg != "B"])
PY

Repository: RL-Align/RL-Kernel

Length of output: 921

🌐 Web query:

Triton jit constexpr parameters compile-time constants each distinct value creates separate LLVM IR specialization

💡 Result:

Yes, in Triton, parameters annotated with tl.constexpr are treated as compile-time constants, and each distinct value passed to these parameters triggers a separate kernel recompilation, resulting in a unique LLVM IR specialization [1][2]. Because tl.constexpr values are "baked into" the compiled program, they allow the Triton compiler to perform aggressive optimizations based on these fixed values [3][1]. Consequently, whenever a kernel is called with a new or different value for a constexpr parameter, the Triton JIT runtime detects this as a change to the function's signature and compiles a new version of the kernel [4][1]. This behavior is distinct from standard kernel parameters, which are designed to be dynamic; using tl.constexpr for values that change frequently can lead to a "compilation explosion," where excessive recompilation significantly increases the total execution time [5][1]. If a parameter needs to change frequently without triggering recompilation, developers should avoid the tl.constexpr annotation and instead use standard parameters or rely on the do_not_specialize flag, which prevents the compiler from baking the value into the IR [1][6].

Citations:

De-constexpr runtime shapes, strides, and sm_scale.

B, S_Q, S_KV, the 16 tensor strides, and sm_scale are runtime values that vary with rollout batch/sequence/layout/scale. Marking them tl.constexpr creates a separate specialized Triton binary for each distinct value, while B is not referenced in the kernel at all. Keep tile sizes and boolean compilation flags as constexpr; pass the rest as runtime arguments.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@rl_engine/kernels/ops/triton/attention/standard_attn.py` around lines 31 -
57, Update the kernel parameter declarations to remove tl.constexpr from B, S_Q,
S_KV, all 16 tensor stride parameters, and sm_scale; leave BLOCK_N, BLOCK_D, and
boolean flags such as CAUSAL and HAS_KEY_PADDING_MASK constexpr. Since B is
unused in the kernel, remove it from the signature and update its launch
arguments accordingly, while passing the remaining shapes, strides, and scale as
runtime values.

):
row = tl.program_id(0)
q_head = tl.program_id(1)
batch = tl.program_id(2)
kv_head = q_head // (H_Q // H_KV)

offs_d = tl.arange(0, BLOCK_D)
d_mask = offs_d < D
q = tl.load(
q_ptr + batch * stride_qb + q_head * stride_qh + row * stride_qs + offs_d * stride_qd,
mask=d_mask,
other=0.0,
).to(tl.float32)

max_score = -float("inf")
for start_n in range(0, S_KV, BLOCK_N):
cols = start_n + tl.arange(0, BLOCK_N)
col_mask = cols < S_KV
k = tl.load(
k_ptr
+ batch * stride_kb
+ kv_head * stride_kh
+ cols[:, None] * stride_ks
+ offs_d[None, :] * stride_kd,
mask=col_mask[:, None] & d_mask[None, :],
other=0.0,
).to(tl.float32)
scores = tl.sum(k * q[None, :], axis=1) * sm_scale
scores = tl.where(col_mask, scores, -float("inf"))

if CAUSAL:
causal_keep = cols <= (row + S_KV - S_Q)
scores = tl.where(causal_keep, scores, -float("inf"))

if HAS_KEY_PADDING_MASK:
keep = tl.load(mask_ptr + batch * S_KV + cols, mask=col_mask, other=0)
scores = tl.where(keep != 0, scores, -float("inf"))

max_score = tl.maximum(max_score, tl.max(scores, axis=0))

denom = 0.0
acc = tl.zeros((BLOCK_D,), dtype=tl.float32)
for start_n in range(0, S_KV, BLOCK_N):
cols = start_n + tl.arange(0, BLOCK_N)
col_mask = cols < S_KV
k = tl.load(
k_ptr
+ batch * stride_kb
+ kv_head * stride_kh
+ cols[:, None] * stride_ks
+ offs_d[None, :] * stride_kd,
mask=col_mask[:, None] & d_mask[None, :],
other=0.0,
).to(tl.float32)
scores = tl.sum(k * q[None, :], axis=1) * sm_scale
scores = tl.where(col_mask, scores, -float("inf"))

if CAUSAL:
causal_keep = cols <= (row + S_KV - S_Q)
scores = tl.where(causal_keep, scores, -float("inf"))

if HAS_KEY_PADDING_MASK:
keep = tl.load(mask_ptr + batch * S_KV + cols, mask=col_mask, other=0)
scores = tl.where(keep != 0, scores, -float("inf"))

probs = tl.exp(scores - max_score)
probs = tl.where(max_score == -float("inf"), 0.0, probs)
denom += tl.sum(probs, axis=0)

v = tl.load(
v_ptr
+ batch * stride_vb
+ kv_head * stride_vh
+ cols[:, None] * stride_vs
+ offs_d[None, :] * stride_vd,
mask=col_mask[:, None] & d_mask[None, :],
other=0.0,
).to(tl.float32)
acc += tl.sum(probs[:, None] * v, axis=0)

out = tl.where(denom > 0.0, acc / denom, 0.0)
tl.store(
out_ptr + batch * stride_ob + q_head * stride_oh + row * stride_os + offs_d * stride_od,
out,
mask=d_mask,
)
tl.store(lse_ptr + (batch * H_Q + q_head) * S_Q + row, max_score + tl.log(denom))


class _TritonBatchInvariantAttention(torch.autograd.Function):
@staticmethod
def forward(
ctx,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
key_padding_mask: Optional[torch.Tensor],
causal: bool,
scale: float,
return_lse: bool,
):
q = q.contiguous()
k = k.contiguous()
v = v.contiguous()
if key_padding_mask is not None:
key_padding_mask = key_padding_mask.contiguous()

batch, q_heads, q_len, head_dim = q.shape
kv_batch, kv_heads, kv_len, k_dim = k.shape
if v.shape != k.shape:
raise ValueError("k and v must have the same shape")
if kv_batch != batch:
raise ValueError("q, k, and v must have the same batch size")
if k_dim != head_dim:
raise ValueError("q, k, and v must have the same head dimension")
if q_heads % kv_heads != 0:
raise ValueError("q heads must be divisible by k/v heads for GQA/MQA")
if head_dim > 256:
raise ValueError("Triton batch-invariant attention supports head_dim <= 256")
if key_padding_mask is not None and key_padding_mask.shape != (batch, kv_len):
raise ValueError("key_padding_mask must have shape [batch, key_seq_len]")

out = torch.empty_like(q)
lse = torch.empty((batch, q_heads, q_len), device=q.device, dtype=torch.float32)
block_d = _next_power_of_2(head_dim)
grid = (q_len, q_heads, batch)
dummy_mask = (
key_padding_mask
if key_padding_mask is not None
else q.new_empty((1,), dtype=torch.bool)
)

_standard_attn_fwd_kernel[grid](
q,
k,
v,
dummy_mask,
out,
lse,
batch,
q_heads,
kv_heads,
q_len,
kv_len,
head_dim,
q.stride(0),
q.stride(1),
q.stride(2),
q.stride(3),
k.stride(0),
k.stride(1),
k.stride(2),
k.stride(3),
v.stride(0),
v.stride(1),
v.stride(2),
v.stride(3),
out.stride(0),
out.stride(1),
out.stride(2),
out.stride(3),
scale,
BLOCK_N=_BLOCK_N,
BLOCK_D=block_d,
CAUSAL=causal,
HAS_KEY_PADDING_MASK=key_padding_mask is not None,
num_warps=8,
)

ctx.save_for_backward(q, k, v, key_padding_mask)
ctx.causal = causal
ctx.scale = scale
ctx.has_key_padding_mask = key_padding_mask is not None
if return_lse:
return out, lse
return out
Comment on lines +230 to +232

@coderabbitai coderabbitai Bot Jul 25, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Mark lse non-differentiable; grads flowing into it are silently dropped.

lse is returned from an autograd.Function without ctx.mark_non_differentiable, so it carries requires_grad=True when q/k/v do. backward only consumes grad_outputs[0], so if a caller uses lse in a loss the contribution is discarded silently instead of erroring.

🐛 Proposed fix 
         if return_lse:
+            ctx.mark_non_differentiable(lse)
             return out, lse
         return out
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
if return_lse:
return out, lse
return out
if return_lse:
ctx.mark_non_differentiable(lse)
return out, lse
return out
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@rl_engine/kernels/ops/triton/attention/standard_attn.py` around lines 235 -
237, Update the autograd Function forward path around the return_lse branch to
call ctx.mark_non_differentiable on lse before returning it with out. Keep out
differentiable and preserve the existing return behavior when return_lse is
false.


@staticmethod
def backward(ctx, *grad_outputs):
q, k, v, key_padding_mask = ctx.saved_tensors
grad_out = grad_outputs[0]
with torch.enable_grad():
q_ref = q.detach().requires_grad_(True)
k_ref = k.detach().requires_grad_(True)
v_ref = v.detach().requires_grad_(True)
out = NativeAttentionOp().forward(
q_ref,
k_ref,
v_ref,
causal=ctx.causal,
scale=ctx.scale,
key_padding_mask=key_padding_mask if ctx.has_key_padding_mask else None,
)
dq, dk, dv = torch.autograd.grad(out, (q_ref, k_ref, v_ref), grad_out)
return dq, dk, dv, None, None, None, None


def _resolve_scale(
head_dim: int,
*,
scale: Optional[float],
softmax_scale: Optional[float],
) -> float:
if scale is not None and softmax_scale is not None:
raise ValueError("set only one of scale or softmax_scale")
value = scale if scale is not None else softmax_scale
return float(value) if value is not None else 1.0 / math.sqrt(head_dim)


def triton_batch_invariant_attention(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
*,
causal: bool = True,
scale: Optional[float] = None,
softmax_scale: Optional[float] = None,
key_padding_mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
scale_value = _resolve_scale(q.shape[-1], scale=scale, softmax_scale=softmax_scale)
return _TritonBatchInvariantAttention.apply(
q,
k,
v,
key_padding_mask,
causal,
scale_value,
False,
)


def triton_batch_invariant_attention_with_lse(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
*,
causal: bool = True,
scale: Optional[float] = None,
softmax_scale: Optional[float] = None,
key_padding_mask: Optional[torch.Tensor] = None,
) -> tuple[torch.Tensor, torch.Tensor]:
scale_value = _resolve_scale(q.shape[-1], scale=scale, softmax_scale=softmax_scale)
return _TritonBatchInvariantAttention.apply(
q,
k,
v,
key_padding_mask,
causal,
scale_value,
True,
)


class TritonBatchInvariantAttentionOp:
"""Triton standard-softmax attention with fixed key-order reductions."""

def __call__(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
*,
causal: bool = True,
scale: Optional[float] = None,
softmax_scale: Optional[float] = None,
key_padding_mask: Optional[torch.Tensor] = None,
dropout_p: float = 0.0,
) -> torch.Tensor:
return self.forward(
q,
k,
v,
causal=causal,
scale=scale,
softmax_scale=softmax_scale,
key_padding_mask=key_padding_mask,
dropout_p=dropout_p,
)

def forward(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
*,
causal: bool = True,
scale: Optional[float] = None,
softmax_scale: Optional[float] = None,
key_padding_mask: Optional[torch.Tensor] = None,
dropout_p: float = 0.0,
) -> torch.Tensor:
if dropout_p != 0.0:
raise ValueError("batch-invariant attention does not support dropout")
return triton_batch_invariant_attention(
q,
k,
v,
causal=causal,
scale=scale,
softmax_scale=softmax_scale,
key_padding_mask=key_padding_mask,
)
Loading
Loading