feat(kernels): add Triton batch-invariant attention#246
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (4)
📝 WalkthroughWalkthroughAdds a Triton batch-invariant standard-softmax attention operator with causal and padding masks, LSE output, native-autograd backward fallback, public exports, operator registration, and CUDA tests. ChangesAttention operator
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant TritonBatchInvariantAttentionOp
participant TritonAttentionKernel
participant NativeAttentionOp
Caller->>TritonBatchInvariantAttentionOp: submit q, k, v and masks
TritonBatchInvariantAttentionOp->>TritonAttentionKernel: compute attention output and LSE
TritonAttentionKernel-->>Caller: return output
Caller->>TritonBatchInvariantAttentionOp: request backward gradients
TritonBatchInvariantAttentionOp->>NativeAttentionOp: recompute forward with autograd
NativeAttentionOp-->>Caller: return q, k, and v gradients
Possibly related issues
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (3)
tests/test_triton_batch_invariant_attention.py (1)
46-58: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd coverage for a fully-masked query row.
The kernel has dedicated branches for that path —
tl.where(max_score == -inf, 0.0, probs)(Line 128) andtl.where(denom > 0.0, acc / denom, 0.0)(Line 142) — but no test drives them.mask[1, 13:] = Falsestill leaves valid keys in every row. A row whose every key is padded is exactly where NaN would leak, and it's also where the kernel'slse = -infconvention diverges from anything the native op returns.Suggest a case with an all-
Falsemask row asserting finite zero output (matchingNativeAttentionOp's fully-masked-row convention) and-infLSE. Want me to draft it?Also applies to: 102-153
🤖 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 `@tests/test_triton_batch_invariant_attention.py` around lines 46 - 58, Add coverage in test_triton_attention_matches_native_forward for a query row whose key_padding_mask is entirely False, using the existing Triton and NativeAttentionOp paths. Assert that the fully masked row produces finite zero output matching the native convention, and verify its LSE value is -inf while preserving the existing dtype and tolerance checks.rl_engine/kernels/ops/triton/attention/standard_attn.py (2)
239-256: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winDocument the fallback's memory profile and honor
ctx.needs_input_grad.
NativeAttentionOp.forwardmaterializes the full[B, Hq, Sq, Skv]score and probability tensors, so backward peak memory scales quadratically in sequence length even though the forward kernel does not — it will OOM at production prefill shapes where the forward succeeds. Also, grads are always computed for all three inputs regardless ofneeds_input_grad, andcreate_graph/double-backward isn't supported.Suggest a class docstring stating this is a validation-only backward, plus:
♻️ Honor `needs_input_grad`
- dq, dk, dv = torch.autograd.grad(out, (q_ref, k_ref, v_ref), grad_out) - return dq, dk, dv, None, None, None, None + wanted = [ + t for t, need in zip((q_ref, k_ref, v_ref), ctx.needs_input_grad[:3]) if need + ] + grads = iter(torch.autograd.grad(out, wanted, grad_out) if wanted else ()) + dq, dk, dv = ( + next(grads) if need else None for need in ctx.needs_input_grad[:3] + ) + return dq, dk, dv, None, None, None, None🤖 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 239 - 256, Update StandardAttention’s backward fallback around NativeAttentionOp.forward to document that it is validation-only, uses quadratic [B, Hq, Sq, Skv] memory, and does not support create_graph or double backward. Honor ctx.needs_input_grad by enabling gradients and requesting torch.autograd.grad only for required q, k, and v inputs, while returning None for inputs that do not need gradients and preserving the existing non-tensor argument returns.
191-191: 🚀 Performance & Scalability | 🔵 Trivial | ⚖️ Poor tradeoffConsider tiling queries instead of one program per query row.
grid = (q_len, q_heads, batch)withnum_warps=8gives each program a single query vector and a(BLOCK_N, BLOCK_D)tile, so all K/V traffic is re-read per row and the warps are likely underutilized on the D-reduction. ABLOCK_Mquery tile would amortize K/V loads across rows while keeping the key-order reduction fixed (batch invariance holds as long asBLOCK_M/BLOCK_Nare shape-independent constants). Worth deferring if this baseline is only for validation, not throughput.Also applies to: 228-228
🤖 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` at line 191, Update the attention kernel launch and corresponding kernel indexing around the grid definitions to process a shape-independent BLOCK_M tile of query rows rather than one program per query row. Adjust the query, output, and masking logic to handle the tiled M dimension while preserving the fixed key-order reduction and batch/head behavior; apply the same change to the second grid definition referenced in the review.
🤖 Prompt for all review comments with 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.
Inline comments:
In `@rl_engine/kernels/ops/triton/attention/standard_attn.py`:
- Around line 6-13: Apply the isort-generated ordering to the import block in
rl_engine/kernels/ops/triton/attention/standard_attn.py (lines 6-13), then apply
the corresponding isort ordering to the import block in
tests/test_triton_batch_invariant_attention.py (lines 6-13). Run pre-commit
across all files to verify both import blocks pass linting.
- Around line 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.
- Around line 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.
---
Nitpick comments:
In `@rl_engine/kernels/ops/triton/attention/standard_attn.py`:
- Around line 239-256: Update StandardAttention’s backward fallback around
NativeAttentionOp.forward to document that it is validation-only, uses quadratic
[B, Hq, Sq, Skv] memory, and does not support create_graph or double backward.
Honor ctx.needs_input_grad by enabling gradients and requesting
torch.autograd.grad only for required q, k, and v inputs, while returning None
for inputs that do not need gradients and preserving the existing non-tensor
argument returns.
- Line 191: Update the attention kernel launch and corresponding kernel indexing
around the grid definitions to process a shape-independent BLOCK_M tile of query
rows rather than one program per query row. Adjust the query, output, and
masking logic to handle the tiled M dimension while preserving the fixed
key-order reduction and batch/head behavior; apply the same change to the second
grid definition referenced in the review.
In `@tests/test_triton_batch_invariant_attention.py`:
- Around line 46-58: Add coverage in
test_triton_attention_matches_native_forward for a query row whose
key_padding_mask is entirely False, using the existing Triton and
NativeAttentionOp paths. Assert that the fully masked row produces finite zero
output matching the native convention, and verify its LSE value is -inf while
preserving the existing dtype and tolerance checks.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 4f5c745f-b212-4a99-8fa7-45457bf53c02
📒 Files selected for processing (4)
rl_engine/kernels/gtest/operator_specs.pyrl_engine/kernels/ops/triton/attention/__init__.pyrl_engine/kernels/ops/triton/attention/standard_attn.pytests/test_triton_batch_invariant_attention.py
| 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, |
There was a problem hiding this comment.
🚀 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 || trueRepository: 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)
PYRepository: 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)
PYRepository: 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"])
PYRepository: 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:
- 1: different between do_not_specialize and tl.constexpr triton-lang/triton#5000
- 2: Update kernel name to include the constexprs names and vals triton-lang/triton#6164
- 3: [BACKEND] Dedup elementwise in LLVM IR based on constancy triton-lang/triton#2512
- 4: https://github.com/triton-lang/triton/blob/main/python/triton/runtime/jit.py
- 5: [RFC] Support specialization on
tl.constexprparameters triton-lang/triton#5274 - 6: do_not_specialize not preventing new compilations triton-lang/triton#8273
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.
| if return_lse: | ||
| return out, lse | ||
| return out |
There was a problem hiding this comment.
🎯 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.
| 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.
323c840 to
71756b9
Compare
1 participant
Closes #147.
This PR adds a Triton standard-softmax attention baseline for WS1 batch-invariant forward validation. The kernel uses a fixed key-order reduction for max, sum-exp/LSE, and P @ V, supports causal and key-padding masks, and is wired into the #108 operator harness as the attention/triton candidate.
Summary by CodeRabbit
New Features
Tests