From 71756b902d060ed8a7a35fd5fedb94a0c2467bb7 Mon Sep 17 00:00:00 2001 From: EthanZero2Hero <1019419030@qq.com> Date: Sat, 25 Jul 2026 22:01:35 +0800 Subject: [PATCH] feat(kernels): add Triton batch-invariant attention --- rl_engine/kernels/gtest/operator_specs.py | 14 + .../kernels/ops/triton/attention/__init__.py | 14 + .../ops/triton/attention/standard_attn.py | 358 ++++++++++++++++++ .../test_triton_batch_invariant_attention.py | 219 +++++++++++ 4 files changed, 605 insertions(+) create mode 100644 rl_engine/kernels/ops/triton/attention/__init__.py create mode 100644 rl_engine/kernels/ops/triton/attention/standard_attn.py create mode 100644 tests/test_triton_batch_invariant_attention.py diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 55a4a203..e8c871be 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -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", diff --git a/rl_engine/kernels/ops/triton/attention/__init__.py b/rl_engine/kernels/ops/triton/attention/__init__.py new file mode 100644 index 00000000..220b6c95 --- /dev/null +++ b/rl_engine/kernels/ops/triton/attention/__init__.py @@ -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", +] diff --git a/rl_engine/kernels/ops/triton/attention/standard_attn.py b/rl_engine/kernels/ops/triton/attention/standard_attn.py new file mode 100644 index 00000000..938b66d0 --- /dev/null +++ b/rl_engine/kernels/ops/triton/attention/standard_attn.py @@ -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 + +_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, +): + 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 + + @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, + ) diff --git a/tests/test_triton_batch_invariant_attention.py b/tests/test_triton_batch_invariant_attention.py new file mode 100644 index 00000000..e0a2ddac --- /dev/null +++ b/tests/test_triton_batch_invariant_attention.py @@ -0,0 +1,219 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp +from rl_engine.kernels.ops.triton.attention.standard_attn import ( + TritonBatchInvariantAttentionOp, + triton_batch_invariant_attention_with_lse, +) + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def _qkv( + batch: int, + q_len: int, + kv_len: int, + *, + q_heads: int = 4, + kv_heads: int = 2, + head_dim: int = 32, + dtype: torch.dtype = torch.bfloat16, + seed: int = 0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + gen = torch.Generator(device="cuda").manual_seed(seed) + q = torch.randn((batch, q_heads, q_len, head_dim), generator=gen, device="cuda", dtype=dtype) + k = torch.randn((batch, kv_heads, kv_len, head_dim), generator=gen, device="cuda", dtype=dtype) + v = torch.randn((batch, kv_heads, kv_len, head_dim), generator=gen, device="cuda", dtype=dtype) + return q, k, v + + +def _run_backward(op, q, k, v, dy, *, causal=True, key_padding_mask=None): + q_req = q.detach().clone().requires_grad_(True) + k_req = k.detach().clone().requires_grad_(True) + v_req = v.detach().clone().requires_grad_(True) + out = op(q_req, k_req, v_req, causal=causal, key_padding_mask=key_padding_mask) + out.backward(dy) + return out.detach(), q_req.grad.detach(), k_req.grad.detach(), v_req.grad.detach() + + +@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) + + +@requires_cuda +def test_triton_attention_lse_matches_native_softmax_stats(): + q, k, v = _qkv(1, 9, 13, dtype=torch.float32, q_heads=2, kv_heads=1, head_dim=16) + mask = torch.ones((1, 13), device="cuda", dtype=torch.bool) + mask[:, 10:] = False + + out, lse = triton_batch_invariant_attention_with_lse( + q, k, v, causal=True, key_padding_mask=mask + ) + ref = NativeAttentionOp().forward_fp32(q, k, v, causal=True, key_padding_mask=mask) + + k_expanded = k.repeat_interleave(2, dim=1) + scores = torch.matmul(q.float(), k_expanded.float().transpose(-1, -2)) * (1.0 / 16**0.5) + causal_mask = torch.triu( + torch.ones(9, 13, dtype=torch.bool, device="cuda"), + diagonal=13 - 9 + 1, + ) + scores = scores.masked_fill(causal_mask, float("-inf")) + scores = scores.masked_fill(~mask[:, None, None, :], float("-inf")) + ref_lse = torch.logsumexp(scores, dim=-1) + + torch.testing.assert_close(out.float(), ref, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(lse, ref_lse, atol=1e-5, rtol=1e-5) + + +@requires_cuda +def test_triton_attention_batch_position_invariant(): + dtype = torch.bfloat16 + q_real, k_real, v_real = _qkv(1, 16, 16, dtype=dtype, seed=2) + op = TritonBatchInvariantAttentionOp() + + out_single = op(q_real, k_real, v_real, causal=True) + q_batch, k_batch, v_batch = _qkv(4, 16, 16, dtype=dtype, seed=3) + q_batch[2] = q_real[0] + k_batch[2] = k_real[0] + v_batch[2] = v_real[0] + out_batch = op(q_batch, k_batch, v_batch, causal=True) + + assert torch.equal(out_single[0], out_batch[2]) + + +@requires_cuda +def test_triton_attention_padding_layout_invariant(): + dtype = torch.bfloat16 + q, k_real, v_real = _qkv(1, 8, 8, dtype=dtype, seed=4) + op = TritonBatchInvariantAttentionOp() + + mask_a = torch.ones((1, 8), device="cuda", dtype=torch.bool) + out_a = op(q, k_real, v_real, causal=False, key_padding_mask=mask_a) + + _, k_pad, v_pad = _qkv(1, 8, 16, dtype=dtype, seed=5) + mask_b = torch.zeros((1, 16), device="cuda", dtype=torch.bool) + real_positions = [2 * i + 1 for i in range(8)] + for src, dst in enumerate(real_positions): + k_pad[:, :, dst] = k_real[:, :, src] + v_pad[:, :, dst] = v_real[:, :, src] + mask_b[:, dst] = True + out_b = op(q, k_pad, v_pad, causal=False, key_padding_mask=mask_b) + + torch.testing.assert_close(out_a.float(), out_b.float(), atol=5e-2, rtol=2e-2) + + +@requires_cuda +def test_triton_attention_lse_padding_layout_invariant(): + dtype = torch.bfloat16 + q, k_real, v_real = _qkv(1, 8, 8, dtype=dtype, seed=8) + + mask_a = torch.ones((1, 8), device="cuda", dtype=torch.bool) + out_a, lse_a = triton_batch_invariant_attention_with_lse( + q, + k_real, + v_real, + causal=False, + key_padding_mask=mask_a, + ) + + _, k_pad, v_pad = _qkv(1, 8, 17, dtype=dtype, seed=9) + mask_b = torch.zeros((1, 17), device="cuda", dtype=torch.bool) + real_positions = [2 * i for i in range(8)] + for src, dst in enumerate(real_positions): + k_pad[:, :, dst] = k_real[:, :, src] + v_pad[:, :, dst] = v_real[:, :, src] + mask_b[:, dst] = True + out_b, lse_b = triton_batch_invariant_attention_with_lse( + q, + k_pad, + v_pad, + causal=False, + key_padding_mask=mask_b, + ) + + torch.testing.assert_close(out_a.float(), out_b.float(), atol=5e-2, rtol=2e-2) + torch.testing.assert_close(lse_a, lse_b, atol=1e-5, rtol=1e-5) + + +@requires_cuda +def test_triton_attention_decode_matches_prefill_suffix_context(): + dtype = torch.bfloat16 + q_full, k_full, v_full = _qkv(1, 12, 12, dtype=dtype, seed=6) + op = TritonBatchInvariantAttentionOp() + + prefill = op(q_full, k_full, v_full, causal=True) + decode = op(q_full[:, :, -1:], k_full, v_full, causal=True) + + assert torch.equal(prefill[:, :, -1:], decode) + + +@requires_cuda +def test_triton_attention_backward_uses_reference_fallback(): + dtype = torch.bfloat16 + q, k, v = _qkv(1, 8, 8, dtype=dtype, seed=7) + dy = torch.randn_like(q) + op = TritonBatchInvariantAttentionOp() + native = NativeAttentionOp() + + out, dq, dk, dv = _run_backward(op, q, k, v, dy, causal=True) + ref_out, ref_dq, ref_dk, ref_dv = _run_backward(native, q, k, v, dy, causal=True) + + torch.testing.assert_close(out.float(), ref_out.float(), atol=5e-2, rtol=2e-2) + torch.testing.assert_close(dq.float(), ref_dq.float(), atol=5e-2, rtol=2e-2) + torch.testing.assert_close(dk.float(), ref_dk.float(), atol=5e-2, rtol=2e-2) + torch.testing.assert_close(dv.float(), ref_dv.float(), atol=5e-2, rtol=2e-2) + + +@requires_cuda +def test_triton_attention_backward_batch_position_invariant(): + dtype = torch.bfloat16 + q_real, k_real, v_real = _qkv(1, 8, 8, dtype=dtype, seed=10) + dy_real = torch.randn_like(q_real) + op = TritonBatchInvariantAttentionOp() + + _, dq_single, dk_single, dv_single = _run_backward( + op, + q_real, + k_real, + v_real, + dy_real, + causal=True, + ) + + q_batch, k_batch, v_batch = _qkv(4, 8, 8, dtype=dtype, seed=11) + dy_batch = torch.randn_like(q_batch) + q_batch[2] = q_real[0] + k_batch[2] = k_real[0] + v_batch[2] = v_real[0] + dy_batch[2] = dy_real[0] + _, dq_batch, dk_batch, dv_batch = _run_backward( + op, + q_batch, + k_batch, + v_batch, + dy_batch, + causal=True, + ) + + assert torch.equal(dq_single[0], dq_batch[2]) + assert torch.equal(dk_single[0], dk_batch[2]) + assert torch.equal(dv_single[0], dv_batch[2])