feat(kernels): add Triton batch-invariant attention - #246
Conversation
📝 WalkthroughWalkthroughAdds a Triton batch-invariant standard-softmax attention operator with causal and padding masks, LSE output, native-autograd backward fallback, public exports, operator registration, tolerance configuration, 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 reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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.
323c840 to
71756b9
Compare
Flink-ddd
left a comment
Flink-ddd
left a comment
There was a problem hiding this comment.
Left a few comments, otherwise LGTM!
71756b9 to
7ae3d07
Compare
7ae3d07 to
5526785
Compare
maxiaosong1124
left a comment
maxiaosong1124
left a comment
There was a problem hiding this comment.
Left some comments,please resolve these problems, thanks!
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/test_triton_batch_invariant_attention.py (1)
310-361: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winMissing backward coverage: causal + partial
key_padding_maskcombined.All backward tests here use
causal=Truewith no mask, orcausal=Falsewith an all-False mask (Line 274-308). None exercisecausal=Truetogether with a partial (some valid/some padded)key_padding_mask— the combination most representative of real training batches with left-padding. A bug in how the backward's native-reference fallback threads a partial mask through causal masking would go undetected.Consider adding a parametrized backward test mirroring
test_triton_attention_matches_native_forward(Lines 68-77) but calling.backward()and comparingdq/dk/dvagainst the native reference.🤖 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 310 - 361, Add a parametrized backward test alongside the existing Triton attention backward tests that uses causal=True with a partially populated key_padding_mask containing both valid and padded positions. Run _run_backward with TritonBatchInvariantAttentionOp and NativeAttentionOp using identical q, k, v, dy, and mask inputs, then compare dq, dk, and dv with appropriate tolerances, covering the native-reference fallback’s combined causal and mask behavior.
🤖 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 `@tests/test_triton_batch_invariant_attention.py`:
- Around line 63-77: The float16 tolerance in
test_triton_attention_matches_native_forward must match the attention accuracy
contract. Update the float16 parameterized case from 3e-3 for both atol and rtol
to 1e-3, while preserving the existing bfloat16 tolerances and test behavior.
---
Nitpick comments:
In `@tests/test_triton_batch_invariant_attention.py`:
- Around line 310-361: Add a parametrized backward test alongside the existing
Triton attention backward tests that uses causal=True with a partially populated
key_padding_mask containing both valid and padded positions. Run _run_backward
with TritonBatchInvariantAttentionOp and NativeAttentionOp using identical q, k,
v, dy, and mask inputs, then compare dq, dk, and dv with appropriate tolerances,
covering the native-reference fallback’s combined causal and mask behavior.
🪄 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: 22b84beb-a809-4799-8000-8d387603a1fd
📒 Files selected for processing (6)
rl_engine/kernels/gtest/operator_specs.pyrl_engine/kernels/gtest/tolerance_contract.jsonrl_engine/kernels/ops/triton/attention/__init__.pyrl_engine/kernels/ops/triton/attention/standard_attn.pytests/test_tolerance_contract.pytests/test_triton_batch_invariant_attention.py
🚧 Files skipped from review as they are similar to previous changes (4)
- rl_engine/kernels/gtest/tolerance_contract.json
- rl_engine/kernels/ops/triton/attention/init.py
- rl_engine/kernels/gtest/operator_specs.py
- rl_engine/kernels/ops/triton/attention/standard_attn.py
| @requires_cuda | ||
| @pytest.mark.parametrize("causal", [False, True]) | ||
| @pytest.mark.parametrize( | ||
| "dtype, atol, rtol", [(torch.float16, 3e-3, 3e-3), (torch.bfloat16, 5e-2, 2e-2)] | ||
| ) | ||
| def test_triton_attention_matches_native_forward(causal, dtype, atol, rtol): | ||
| q, k, v = _qkv(2, 17, 17, dtype=dtype) | ||
| mask = torch.ones((2, 17), device="cuda", dtype=torch.bool) | ||
| mask[1, 13:] = False | ||
|
|
||
| out = TritonBatchInvariantAttentionOp()(q, k, v, causal=causal, key_padding_mask=mask) | ||
| ref = NativeAttentionOp().forward_fp32(q, k, v, causal=causal, key_padding_mask=mask) | ||
|
|
||
| assert out.dtype == dtype | ||
| torch.testing.assert_close(out.float(), ref, atol=atol, rtol=rtol) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm the current declared float16 tolerance for attention in the contract.
fd tolerance_contract.json --exec cat {}Repository: RL-Align/RL-Kernel
Length of output: 1047
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Locate relevant files and occurrences"
fd -i 'test_triton_batch_invariant_attention.py|accuracy.*attention|tolerance_contract.json' .
echo
rg -n "test_triton_attention_matches_native_forward|float16.*atol=3e-|atol=3e-3|attention/triton|attention/triton|attention.*triton|key_padding_mask|NativeAttentionOp|TritonBatchInvariantAttentionOp" .
echo
echo "Relevant test file section"
file="$(fd test_triton_batch_invariant_attention.py . | head -n1)"
if [ -n "${file:-}" ]; then
wc -l "$file"
cat -n "$file" | sed -n '1,120p'
fiRepository: RL-Align/RL-Kernel
Length of output: 22585
Align the float16 forward test tolerance with the attention contract.
tolerance_contract.json uses accuracy.default.attention.float16: {atol: 1.0e-3, rtol: 1.0e-3}, but this test uses atol=3e-3, rtol=3e-3. Tighten the float16 case to 1e-3 so the baseline still passes, or update the contract if the observed drift genuinely requires a looser threshold.
🤖 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 63 - 77, The
float16 tolerance in test_triton_attention_matches_native_forward must match the
attention accuracy contract. Update the float16 parameterized case from 3e-3 for
both atol and rtol to 1e-3, while preserving the existing bfloat16 tolerances
and test behavior.
Flink-ddd
left a comment
Flink-ddd
left a comment
There was a problem hiding this comment.
LGTM now, Thank you for update.
maxiaosong1124
left a comment
maxiaosong1124
left a comment
There was a problem hiding this comment.
LGTM now! thanks for your update!
3 participants
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
Chores