From 552678501a04dd3135937c12615c124f817b2a25 Mon Sep 17 00:00:00 2001 From: EthanZero2Hero <1019419030@qq.com> Date: Sat, 25 Jul 2026 22:01:35 +0800 Subject: [PATCH 1/2] feat(kernels): add Triton batch-invariant attention --- rl_engine/kernels/gtest/operator_specs.py | 14 + .../kernels/gtest/tolerance_contract.json | 5 + .../kernels/ops/triton/attention/__init__.py | 14 + .../ops/triton/attention/standard_attn.py | 359 ++++++++++++++++++ tests/test_tolerance_contract.py | 11 +- .../test_triton_batch_invariant_attention.py | 235 ++++++++++++ 6 files changed, 636 insertions(+), 2 deletions(-) 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..6f86df70 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="attention", + 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/gtest/tolerance_contract.json b/rl_engine/kernels/gtest/tolerance_contract.json index 7a2ce2b4..975ae450 100644 --- a/rl_engine/kernels/gtest/tolerance_contract.json +++ b/rl_engine/kernels/gtest/tolerance_contract.json @@ -16,6 +16,11 @@ "float32": {"atol": 1.0e-5, "rtol": 0.0}, "bfloat16": {"atol": 5.0e-2, "rtol": 0.0}, "float16": {"atol": 5.0e-3, "rtol": 0.0} + }, + "attention": { + "float32": {"atol": 1.0e-4, "rtol": 1.0e-4}, + "bfloat16": {"atol": 5.0e-2, "rtol": 2.0e-2}, + "float16": {"atol": 1.0e-3, "rtol": 1.0e-3} } }, "arch_overrides": { 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..e4d2d1b1 --- /dev/null +++ b/rl_engine/kernels/ops/triton/attention/standard_attn.py @@ -0,0 +1,359 @@ +# 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: + ctx.mark_non_differentiable(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_tolerance_contract.py b/tests/test_tolerance_contract.py index fb429d81..5eb75cbd 100644 --- a/tests/test_tolerance_contract.py +++ b/tests/test_tolerance_contract.py @@ -9,12 +9,12 @@ def test_load_contract_contains_expected_operator_classes(): contract = load_contract() accuracy = contract["accuracy"]["default"] - assert set(accuracy) == {"elementwise", "reduction", "logprob"} + assert set(accuracy) == {"elementwise", "reduction", "logprob", "attention"} def test_load_contract_contains_expected_dtypes(): contract = load_contract() - for op_class in ("elementwise", "reduction", "logprob"): + for op_class in ("elementwise", "reduction", "logprob", "attention"): assert set(contract["accuracy"]["default"][op_class]) == { "float32", "bfloat16", @@ -27,3 +27,10 @@ def test_logprob_bfloat16_tolerance_covers_observed_reference_drift(): tolerance = contract["accuracy"]["default"]["logprob"]["bfloat16"] assert tolerance["atol"] >= 5.0e-2 assert tolerance["rtol"] == 0.0 + + +def test_attention_bfloat16_tolerance_matches_contract(): + contract = load_contract() + tolerance = contract["accuracy"]["default"]["attention"]["bfloat16"] + assert tolerance["atol"] >= 5.0e-2 + assert tolerance["rtol"] >= 2.0e-2 diff --git a/tests/test_triton_batch_invariant_attention.py b/tests/test_triton_batch_invariant_attention.py new file mode 100644 index 00000000..45ecfa8b --- /dev/null +++ b/tests/test_triton_batch_invariant_attention.py @@ -0,0 +1,235 @@ +# 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_lse_is_non_differentiable(): + q, k, v = _qkv(1, 8, 8, dtype=torch.bfloat16, seed=12) + q.requires_grad_(True) + k.requires_grad_(True) + v.requires_grad_(True) + + out, lse = triton_batch_invariant_attention_with_lse(q, k, v, causal=True) + assert lse.requires_grad is False + (out.float().sum() + lse.float().sum()).backward() + + assert q.grad is not None + assert k.grad is not None + assert v.grad is not None + + +@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]) From 5d5c9648e27bf46a3ea30a52b69528fa624e0b5e Mon Sep 17 00:00:00 2001 From: EthanZero2Hero <1019419030@qq.com> Date: Wed, 29 Jul 2026 00:11:09 +0800 Subject: [PATCH 2/2] test(attention): cover Triton chunked prefill invariants --- .../test_triton_batch_invariant_attention.py | 125 ++++++++++++++++++ 1 file changed, 125 insertions(+) diff --git a/tests/test_triton_batch_invariant_attention.py b/tests/test_triton_batch_invariant_attention.py index 45ecfa8b..d4f195a4 100644 --- a/tests/test_triton_batch_invariant_attention.py +++ b/tests/test_triton_batch_invariant_attention.py @@ -42,6 +42,24 @@ def _run_backward(op, q, k, v, dy, *, causal=True, key_padding_mask=None): return out.detach(), q_req.grad.detach(), k_req.grad.detach(), v_req.grad.detach() +def _run_chunked_prefill(q, k, v, *, chunk_size: int, key_padding_mask=None): + outs = [] + lses = [] + for start in range(0, q.size(2), chunk_size): + end = min(q.size(2), start + chunk_size) + chunk_mask = None if key_padding_mask is None else key_padding_mask[:, :end] + out, lse = triton_batch_invariant_attention_with_lse( + q[:, :, start:end], + k[:, :, :end], + v[:, :, :end], + causal=True, + key_padding_mask=chunk_mask, + ) + outs.append(out) + lses.append(lse) + return torch.cat(outs, dim=2), torch.cat(lses, dim=2) + + @requires_cuda @pytest.mark.parametrize("causal", [False, True]) @pytest.mark.parametrize( @@ -170,6 +188,64 @@ def test_triton_attention_lse_padding_layout_invariant(): torch.testing.assert_close(lse_a, lse_b, atol=1e-5, rtol=1e-5) +@requires_cuda +@pytest.mark.parametrize( + "q_heads, kv_heads, head_dim", + [ + pytest.param(4, 4, 32, id="mha"), + pytest.param(4, 2, 32, id="gqa-small"), + pytest.param(8, 2, 64, id="gqa-wide"), + ], +) +@pytest.mark.parametrize("chunk_size", [1, 2, 3, 5, 13]) +def test_triton_attention_chunked_prefill_matches_full_prefill( + q_heads, kv_heads, head_dim, chunk_size +): + q, k, v = _qkv( + 2, + 13, + 13, + q_heads=q_heads, + kv_heads=kv_heads, + head_dim=head_dim, + dtype=torch.bfloat16, + seed=20, + ) + + full, full_lse = triton_batch_invariant_attention_with_lse(q, k, v, causal=True) + chunked, chunked_lse = _run_chunked_prefill(q, k, v, chunk_size=chunk_size) + + assert torch.equal(full, chunked) + assert torch.equal(full_lse, chunked_lse) + + +@requires_cuda +@pytest.mark.parametrize("chunk_size", [1, 4, 6]) +def test_triton_attention_chunked_prefill_with_key_padding_mask(chunk_size): + q, k, v = _qkv(2, 13, 13, dtype=torch.bfloat16, seed=21) + mask = torch.ones((2, 13), device="cuda", dtype=torch.bool) + mask[0, 9:] = False + mask[1, 11:] = False + + full, full_lse = triton_batch_invariant_attention_with_lse( + q, + k, + v, + causal=True, + key_padding_mask=mask, + ) + chunked, chunked_lse = _run_chunked_prefill( + q, + k, + v, + chunk_size=chunk_size, + key_padding_mask=mask, + ) + + assert torch.equal(full, chunked) + assert torch.equal(full_lse, chunked_lse) + + @requires_cuda def test_triton_attention_decode_matches_prefill_suffix_context(): dtype = torch.bfloat16 @@ -182,6 +258,55 @@ def test_triton_attention_decode_matches_prefill_suffix_context(): assert torch.equal(prefill[:, :, -1:], decode) +@requires_cuda +@pytest.mark.parametrize("s_new", [1, 3]) +def test_triton_attention_kv_cache_handoff_matches_prefill_suffix(s_new): + dtype = torch.bfloat16 + q_full, k_full, v_full = _qkv(2, 12, 12, dtype=dtype, seed=22) + op = TritonBatchInvariantAttentionOp() + + prefill = op(q_full, k_full, v_full, causal=True) + decode = op(q_full[:, :, -s_new:], k_full, v_full, causal=True) + + assert torch.equal(prefill[:, :, -s_new:], decode) + + +@requires_cuda +def test_triton_attention_all_false_key_padding_mask_row_matches_native(): + dtype = torch.bfloat16 + q, k, v = _qkv(2, 5, 5, dtype=dtype, seed=23) + mask = torch.ones((2, 5), device="cuda", dtype=torch.bool) + mask[1] = False + 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, lse = triton_batch_invariant_attention_with_lse( + q_req, + k_req, + v_req, + causal=False, + key_padding_mask=mask, + ) + ref = NativeAttentionOp().forward_fp32(q, k, v, causal=False, key_padding_mask=mask) + + assert torch.equal(out[1], torch.zeros_like(out[1])) + assert torch.isfinite(out).all() + assert torch.isneginf(lse[1]).all() + assert torch.equal(ref[1], torch.zeros_like(ref[1])) + torch.testing.assert_close(out.float(), ref, atol=5e-2, rtol=2e-2) + + dy = torch.randn_like(out) + out.backward(dy) + + assert torch.isfinite(q_req.grad).all() + assert torch.isfinite(k_req.grad).all() + assert torch.isfinite(v_req.grad).all() + assert torch.equal(q_req.grad[1], torch.zeros_like(q_req.grad[1])) + assert torch.equal(k_req.grad[1], torch.zeros_like(k_req.grad[1])) + assert torch.equal(v_req.grad[1], torch.zeros_like(v_req.grad[1])) + + @requires_cuda def test_triton_attention_backward_uses_reference_fallback(): dtype = torch.bfloat16