diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 93c41a132..6e6d33232 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -64,6 +64,10 @@ jobs: run: | python -m pytest rl_engine/tests/test_dispatch.py -v + - name: Run Forward Invariance Tests + run: | + python -m pytest tests/test_forward_invariance.py -v + docs: runs-on: ubuntu-latest steps: diff --git a/docs/design/batch-invariant-elementwise-rope.md b/docs/design/batch-invariant-elementwise-rope.md new file mode 100644 index 000000000..02afcea14 --- /dev/null +++ b/docs/design/batch-invariant-elementwise-rope.md @@ -0,0 +1,103 @@ +# Batch-Invariant Elementwise and RoPE Audit + +Issue #149 audits the forward-path operations that should be pass-through with +respect to batch configuration. The goal is to prove that a row's output is +unchanged when unrelated rows are added, the row moves to a different batch +position, or padding/packed-sequence position ids are present. + +## Scope + +This audit covers Transformer pointwise operations and RoPE only. RMSNorm, +matmul, attention, selected logprob, and FP8 remain out of scope because they +have separate roadmap issues and reduction contracts. + +## Tested Contract + +The shared sweep helper lives in `rl_engine.testing.forward_invariance`: + +- `BatchInvariantConfig` +- `DEFAULT_BATCH_INVARIANT_SWEEP` +- `assert_batch_invariant_across_configs` + +The executable contract lives in `tests/test_forward_invariance.py` and calls +that helper for every audited elementwise operation and for fixed-position RoPE. + +Each audited operation is evaluated on a target row in isolation, then on the +same target row embedded into larger batches with unrelated noise rows. The +target output must be bitwise identical with `torch.equal` unless a tight +operator-specific tolerance is documented below. + +| Sweep case | Batch size | Target row | +| --- | ---: | ---: | +| `batch1` | 1 | 0 | +| `batch2-first` | 2 | 0 | +| `batch4-middle` | 4 | 2 | +| `batch9-last` | 9 | 8 | + +| Operation | Verdict | Coverage | +| --- | --- | --- | +| SiLU activation | Pass | Pointwise; no batch-dependent reduction. | +| GELU activation | Pass | Pointwise; uses `atol=1e-6, rtol=1e-6` for CPU libm/vector paths. | +| Residual add | Pass | Depends only on the matching residual row. | +| Scalar scaling | Pass | Scalar multiply has no accumulation. | +| Bias add | Pass | Broadcast bias is independent of batch position. | +| Mask fill | Pass | Uses only the value/mask pair for each element. | +| Explicit dtype cast | Pass | fp32/fp16/bf16 casts are batch-invariant. | + +CUDA runs cover fp32, fp16, and bf16 when the device supports bf16. CPU CI covers +fp32 for the arithmetic elementwise cases and fp32/fp16/bf16 for explicit dtype +casts. + +## RoPE Contract + +`rl_engine.testing.build_rope_cache` builds a deterministic table-lookup cache: + +1. Compute fp32 inverse frequencies from `base` and `head_dim`. +2. Build `[max_position, head_dim]` cos/sin tables. +3. `apply_rope_reference` gathers rows with explicit `position_ids`. +4. Gathered cos/sin values are cast to the Q/K dtype before the multiply/add. + +The reference path has no batch-dependent reduction, no launch-shape-dependent +accumulation, and no inline recomputation that can vary by batch shape. + +| RoPE case | Verdict | Coverage | +| --- | --- | --- | +| Fixed position | Pass | Same Q/K token and position id stays bitwise identical. | +| Batch position changes | Pass | Target row can move without drift. | +| Unrelated noise rows | Pass | Noise Q/K rows and position ids do not affect target output. | +| Padding | Pass | Valid tokens keep the same output when padding is inserted. | +| Packed-sequence reset | Pass | Local ids `[0, 1, 2]` match the standalone segment. | +| Position-id validation | Pass | Invalid position ids and cache shapes are rejected. | + +## Verification + +The CPU CI unit-test job runs `tests/test_forward_invariance.py` directly. The +standard `tests/` suite run by GPU CI also includes it when GPU CI is requested. + +Run the focused contract: + +```bash +python -m pytest tests/test_forward_invariance.py -q +``` + +Run the helper regression suite: + +```bash +python -m pytest tests/test_forward_invariance.py tests/test_reference_ops.py -q +``` + +Build the documentation after changing this page: + +```bash +mkdocs build --strict -f mkdocs.yaml +``` + +## Known Boundaries + +This audit adds the shared sweep helper used by the #149 tests, a PyTorch RoPE +reference contract, and documentation. It does not add a production RoPE kernel, +change runtime dispatch, or close issues assigned to RMSNorm, matmul, attention, +logprob, or FP8. + +If #108 later lands a broader model-level harness, these tests should bridge to +that harness without changing the pass/fail expectations above. diff --git a/rl_engine/testing/__init__.py b/rl_engine/testing/__init__.py index 42be8c1bb..9a602f399 100644 --- a/rl_engine/testing/__init__.py +++ b/rl_engine/testing/__init__.py @@ -3,6 +3,14 @@ """Testing helpers for RL-shaped kernel validation.""" +from .forward_invariance import ( + DEFAULT_BATCH_INVARIANT_SWEEP, + BatchInvariantConfig, + apply_rope_reference, + assert_batch_invariant_across_configs, + build_rope_cache, + rotate_half, +) from .reference_ops import ( active_token_count, compute_policy_ratio, @@ -15,13 +23,19 @@ from .rl_batch import SyntheticRLKernelBatch, make_synthetic_rl_kernel_batch __all__ = [ + "DEFAULT_BATCH_INVARIANT_SWEEP", + "BatchInvariantConfig", "SyntheticRLKernelBatch", "active_token_count", + "apply_rope_reference", + "assert_batch_invariant_across_configs", + "build_rope_cache", "compute_policy_ratio", "compute_reference_kl", "make_synthetic_rl_kernel_batch", "masked_mean", "masked_sum", + "rotate_half", "selected_logprobs_reference", "summarize_kernel_drift", ] diff --git a/rl_engine/testing/forward_invariance.py b/rl_engine/testing/forward_invariance.py new file mode 100644 index 000000000..515c66ca8 --- /dev/null +++ b/rl_engine/testing/forward_invariance.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import TypeAlias + +import torch + +BatchInvariantOutput: TypeAlias = torch.Tensor | tuple[torch.Tensor, ...] + + +@dataclass(frozen=True) +class BatchInvariantConfig: + """One target-row placement in a batch-invariance sweep.""" + + batch_size: int + target_index: int + seed: int + label: str + + def __post_init__(self) -> None: + if self.batch_size <= 0: + raise ValueError("batch_size must be greater than zero") + if not 0 <= self.target_index < self.batch_size: + raise ValueError("target_index must be within the batch") + + +DEFAULT_BATCH_INVARIANT_SWEEP: tuple[BatchInvariantConfig, ...] = ( + BatchInvariantConfig(batch_size=1, target_index=0, seed=14901, label="batch1"), + BatchInvariantConfig(batch_size=2, target_index=0, seed=14902, label="batch2-first"), + BatchInvariantConfig(batch_size=4, target_index=2, seed=14903, label="batch4-middle"), + BatchInvariantConfig(batch_size=9, target_index=8, seed=14904, label="batch9-last"), +) + +BatchInvariantOp = Callable[[Mapping[str, torch.Tensor]], BatchInvariantOutput] +BatchInputFactory = Callable[[BatchInvariantConfig, torch.Generator], Mapping[str, torch.Tensor]] + + +def assert_batch_invariant_across_configs( + reference_inputs: Mapping[str, torch.Tensor], + op: BatchInvariantOp, + make_batched_inputs: BatchInputFactory, + *, + configs: Sequence[BatchInvariantConfig] = DEFAULT_BATCH_INVARIANT_SWEEP, + reference_index: int = 0, + case_name: str = "batch-invariant op", + atol: float = 0.0, + rtol: float = 0.0, +) -> None: + """Assert that a target row is stable across batch placements.""" + + reference_output = _select_batch_output(op(reference_inputs), reference_index) + device = _first_tensor(reference_inputs).device + + for config in configs: + generator = torch.Generator(device=device) + generator.manual_seed(config.seed) + batched_inputs = make_batched_inputs(config, generator) + batched_output = op(batched_inputs) + actual = _select_batch_output(batched_output, config.target_index) + _assert_outputs_match( + actual, + reference_output, + case_name=f"{case_name}/{config.label}", + atol=atol, + rtol=rtol, + ) + + +def _select_batch_output( + output: BatchInvariantOutput, + index: int, +) -> BatchInvariantOutput: + if isinstance(output, torch.Tensor): + return output[index].clone() + return tuple(item[index].clone() for item in output) + + +def _assert_outputs_match( + actual: BatchInvariantOutput, + expected: BatchInvariantOutput, + *, + case_name: str, + atol: float, + rtol: float, +) -> None: + if isinstance(actual, torch.Tensor) and isinstance(expected, torch.Tensor): + if not _tensor_matches(actual, expected, atol=atol, rtol=rtol): + raise AssertionError(f"{case_name} drifted across batch configs") + return + + if not isinstance(actual, tuple) or not isinstance(expected, tuple): + raise TypeError("actual and expected outputs must have matching structures") + if len(actual) != len(expected): + raise AssertionError(f"{case_name} output arity changed across batch configs") + + for output_index, (actual_item, expected_item) in enumerate(zip(actual, expected, strict=True)): + if not _tensor_matches(actual_item, expected_item, atol=atol, rtol=rtol): + raise AssertionError(f"{case_name} output {output_index} drifted across batch configs") + + +def _tensor_matches( + actual: torch.Tensor, + expected: torch.Tensor, + *, + atol: float, + rtol: float, +) -> bool: + if actual.shape != expected.shape or actual.dtype != expected.dtype: + return False + if atol == 0.0 and rtol == 0.0: + return torch.equal(actual, expected) + return torch.allclose(actual, expected, atol=atol, rtol=rtol) + + +def _first_tensor(inputs: Mapping[str, torch.Tensor]) -> torch.Tensor: + for value in inputs.values(): + if isinstance(value, torch.Tensor): + return value + raise ValueError("reference_inputs must contain at least one tensor") + + +def rotate_half(x: torch.Tensor) -> torch.Tensor: + """Rotate the last dimension for LLaMA/Qwen-style RoPE.""" + + if x.size(-1) % 2 != 0: + raise ValueError("RoPE head_dim must be even") + x_first, x_second = x.chunk(2, dim=-1) + return torch.cat((-x_second, x_first), dim=-1) + + +def build_rope_cache( + max_position: int, + head_dim: int, + *, + base: float = 10000.0, + device: torch.device | str | None = None, + dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build a deterministic RoPE cos/sin lookup table.""" + + if max_position <= 0: + raise ValueError("max_position must be greater than zero") + if head_dim <= 0 or head_dim % 2 != 0: + raise ValueError("head_dim must be a positive even integer") + if base <= 0.0: + raise ValueError("base must be greater than zero") + if not torch.empty((), dtype=dtype).is_floating_point(): + raise ValueError("RoPE cache dtype must be a floating-point dtype") + + target_device = torch.device("cpu") if device is None else torch.device(device) + positions = torch.arange(max_position, device=target_device, dtype=torch.float32) + dims = torch.arange(0, head_dim, 2, device=target_device, dtype=torch.float32) + inv_freq = 1.0 / (base ** (dims / float(head_dim))) + freqs = torch.outer(positions, inv_freq) + angles = torch.cat((freqs, freqs), dim=-1) + + return angles.cos().to(dtype=dtype), angles.sin().to(dtype=dtype) + + +def apply_rope_reference( + query: torch.Tensor, + key: torch.Tensor, + position_ids: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply RoPE with explicit position-id lookup and no batch-dependent reduction.""" + + _validate_rope_inputs(query, key, position_ids, cos, sin) + + batch_size, seq_len, _, head_dim = query.shape + flat_position_ids = position_ids.reshape(-1).long() + cos_pos = cos.index_select(0, flat_position_ids).reshape(batch_size, seq_len, 1, head_dim) + sin_pos = sin.index_select(0, flat_position_ids).reshape(batch_size, seq_len, 1, head_dim) + cos_pos = cos_pos.to(dtype=query.dtype) + sin_pos = sin_pos.to(dtype=query.dtype) + + query_rotated = (query * cos_pos) + (rotate_half(query) * sin_pos) + key_rotated = (key * cos_pos) + (rotate_half(key) * sin_pos) + return query_rotated, key_rotated + + +def _validate_rope_inputs( + query: torch.Tensor, + key: torch.Tensor, + position_ids: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> None: + if query.shape != key.shape: + raise ValueError("query and key must have the same shape") + if query.dtype != key.dtype: + raise ValueError("query and key must have the same dtype") + if query.dim() != 4: + raise ValueError("query and key must have shape [batch, seq, heads, head_dim]") + + batch_size, seq_len, _, head_dim = query.shape + if head_dim % 2 != 0: + raise ValueError("RoPE head_dim must be even") + if tuple(position_ids.shape) != (batch_size, seq_len): + raise ValueError("position_ids must have shape [batch, seq]") + if position_ids.dtype == torch.bool or position_ids.is_floating_point(): + raise ValueError("position_ids must use an integer dtype") + + if cos.shape != sin.shape: + raise ValueError("cos and sin must have the same shape") + if cos.dtype != sin.dtype: + raise ValueError("cos and sin must have the same dtype") + if cos.dim() != 2 or cos.size(-1) != head_dim: + raise ValueError("cos and sin must have shape [max_position, head_dim]") + if cos.device != query.device or sin.device != query.device: + raise ValueError("cos and sin must be on the same device as query and key") + if position_ids.device != query.device: + raise ValueError("position_ids must be on the same device as query and key") + + if position_ids.numel() == 0: + return + + min_position = int(position_ids.min().item()) + max_position = int(position_ids.max().item()) + if min_position < 0: + raise ValueError("position_ids must be non-negative") + if max_position >= cos.size(0): + raise ValueError("position_ids exceed the RoPE cache length") diff --git a/tests/test_forward_invariance.py b/tests/test_forward_invariance.py new file mode 100644 index 000000000..175c3068e --- /dev/null +++ b/tests/test_forward_invariance.py @@ -0,0 +1,450 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.testing import ( + BatchInvariantConfig, + apply_rope_reference, + assert_batch_invariant_across_configs, + build_rope_cache, + rotate_half, +) + + +@dataclass(frozen=True) +class _ElementwiseCase: + name: str + apply: Callable[[dict[str, torch.Tensor]], torch.Tensor] + atol: float = 0.0 + rtol: float = 0.0 + + +_ELEMENTWISE_CASES = ( + _ElementwiseCase("silu activation", lambda tensors: F.silu(tensors["x"])), + _ElementwiseCase( + "gelu activation", + lambda tensors: F.gelu(tensors["x"], approximate="tanh"), + atol=1e-6, + rtol=1e-6, + ), + _ElementwiseCase("residual add", lambda tensors: tensors["x"] + tensors["residual"]), + _ElementwiseCase("scalar scaling", lambda tensors: tensors["x"] * tensors["scale"]), + _ElementwiseCase("bias add", lambda tensors: tensors["x"] + tensors["bias"]), + _ElementwiseCase( + "mask fill", + lambda tensors: tensors["x"].masked_fill(tensors["mask"], -0.75), + ), +) + + +def _available_devices() -> list[object]: + devices = [pytest.param(torch.device("cpu"), id="cpu")] + if torch.cuda.is_available(): + devices.append(pytest.param(torch.device("cuda"), id="cuda")) + return devices + + +def _skip_unsupported_dtype(device: torch.device, dtype: torch.dtype) -> None: + if device.type == "cpu" and dtype != torch.float32: + pytest.skip("low-precision elementwise CPU coverage is intentionally omitted") + if dtype == torch.bfloat16 and device.type == "cuda" and not torch.cuda.is_bf16_supported(): + pytest.skip("CUDA bfloat16 is not supported on this device") + + +def _generator(device: torch.device, seed: int) -> torch.Generator: + generator = torch.Generator(device=device) + generator.manual_seed(seed) + return generator + + +def _embed_item( + item: torch.Tensor, + *, + batch_size: int, + target_index: int, + generator: torch.Generator, +) -> torch.Tensor: + batch = torch.randn( + (batch_size, *item.shape[1:]), + device=item.device, + dtype=item.dtype, + generator=generator, + ) + batch[target_index].copy_(item[0]) + return batch + + +def _embed_bool_item( + item: torch.Tensor, + *, + batch_size: int, + target_index: int, + generator: torch.Generator, +) -> torch.Tensor: + batch = ( + torch.rand( + (batch_size, *item.shape[1:]), + device=item.device, + generator=generator, + ) + > 0.5 + ) + batch[target_index].copy_(item[0]) + return batch + + +def _embed_position_ids( + item: torch.Tensor, + *, + batch_size: int, + target_index: int, + max_position: int, + generator: torch.Generator, +) -> torch.Tensor: + batch = torch.randint( + low=0, + high=max_position, + size=(batch_size, *item.shape[1:]), + device=item.device, + dtype=item.dtype, + generator=generator, + ) + batch[target_index].copy_(item[0]) + return batch + + +@pytest.mark.parametrize("device", _available_devices()) +@pytest.mark.parametrize("dtype", (torch.float32, torch.float16, torch.bfloat16)) +@pytest.mark.parametrize("case", _ELEMENTWISE_CASES, ids=lambda case: case.name) +def test_forward_path_elementwise_ops_are_batch_invariant( + device: torch.device, + dtype: torch.dtype, + case: _ElementwiseCase, +): + _skip_unsupported_dtype(device, dtype) + generator = _generator(device, seed=149) + hidden_size = 16 + + target_x = torch.randn((1, 3, hidden_size), device=device, dtype=dtype, generator=generator) + target_residual = torch.randn( + target_x.shape, + device=device, + dtype=dtype, + generator=generator, + ) + target_mask = torch.tensor( + [[[False, True] * (hidden_size // 2)] * target_x.size(1)], + device=device, + ) + payload = { + "x": target_x, + "residual": target_residual, + "scale": torch.tensor(0.125, device=device, dtype=dtype), + "bias": torch.linspace(-0.5, 0.5, steps=hidden_size, device=device, dtype=dtype), + "mask": target_mask, + } + + def make_batched_inputs( + config: BatchInvariantConfig, + generator: torch.Generator, + ) -> dict[str, torch.Tensor]: + return { + "x": _embed_item( + target_x, + batch_size=config.batch_size, + target_index=config.target_index, + generator=generator, + ), + "residual": _embed_item( + target_residual, + batch_size=config.batch_size, + target_index=config.target_index, + generator=generator, + ), + "scale": payload["scale"], + "bias": payload["bias"], + "mask": _embed_bool_item( + target_mask, + batch_size=config.batch_size, + target_index=config.target_index, + generator=generator, + ), + } + + assert_batch_invariant_across_configs( + payload, + case.apply, + make_batched_inputs, + case_name=case.name, + atol=case.atol, + rtol=case.rtol, + ) + + +@pytest.mark.parametrize("device", _available_devices()) +@pytest.mark.parametrize("target_dtype", (torch.float32, torch.float16, torch.bfloat16)) +def test_dtype_casts_are_batch_invariant( + device: torch.device, + target_dtype: torch.dtype, +): + if ( + target_dtype == torch.bfloat16 + and device.type == "cuda" + and not torch.cuda.is_bf16_supported() + ): + pytest.skip("CUDA bfloat16 is not supported on this device") + + generator = _generator(device, seed=150) + target_x = torch.randn((1, 4, 9), device=device, dtype=torch.float32, generator=generator) + payload = {"x": target_x} + + def apply_dtype_cast(tensors: dict[str, torch.Tensor]) -> torch.Tensor: + return tensors["x"].to(dtype=target_dtype) + + def make_batched_inputs( + config: BatchInvariantConfig, + generator: torch.Generator, + ) -> dict[str, torch.Tensor]: + return { + "x": _embed_item( + target_x, + batch_size=config.batch_size, + target_index=config.target_index, + generator=generator, + ) + } + + assert_batch_invariant_across_configs( + payload, + apply_dtype_cast, + make_batched_inputs, + case_name=f"dtype cast to {target_dtype}", + ) + assert apply_dtype_cast(payload).dtype == target_dtype + + +@pytest.mark.parametrize("device", _available_devices()) +@pytest.mark.parametrize("dtype", (torch.float32, torch.float16, torch.bfloat16)) +def test_rope_fixed_positions_are_batch_invariant( + device: torch.device, + dtype: torch.dtype, +): + _skip_unsupported_dtype(device, dtype) + generator = _generator(device, seed=151) + seq_len = 4 + num_heads = 2 + head_dim = 8 + max_position = 16 + + target_q = torch.randn( + (1, seq_len, num_heads, head_dim), + device=device, + dtype=dtype, + generator=generator, + ) + target_k = torch.randn( + target_q.shape, + device=device, + dtype=dtype, + generator=generator, + ) + target_position_ids = torch.tensor([[0, 1, 4, 7]], device=device, dtype=torch.long) + cos, sin = build_rope_cache(max_position, head_dim, device=device, dtype=torch.float32) + payload = { + "query": target_q, + "key": target_k, + "position_ids": target_position_ids, + "cos": cos, + "sin": sin, + } + + def apply_rope(tensors: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + return apply_rope_reference( + tensors["query"], + tensors["key"], + tensors["position_ids"], + tensors["cos"], + tensors["sin"], + ) + + def make_batched_inputs( + config: BatchInvariantConfig, + generator: torch.Generator, + ) -> dict[str, torch.Tensor]: + return { + "query": _embed_item( + target_q, + batch_size=config.batch_size, + target_index=config.target_index, + generator=generator, + ), + "key": _embed_item( + target_k, + batch_size=config.batch_size, + target_index=config.target_index, + generator=generator, + ), + "position_ids": _embed_position_ids( + target_position_ids, + batch_size=config.batch_size, + target_index=config.target_index, + max_position=max_position, + generator=generator, + ), + "cos": cos, + "sin": sin, + } + + assert_batch_invariant_across_configs( + payload, + apply_rope, + make_batched_inputs, + case_name="RoPE fixed position", + ) + + +@pytest.mark.parametrize("device", _available_devices()) +def test_rope_padding_position_ids_do_not_shift_valid_tokens(device: torch.device): + generator = _generator(device, seed=152) + dtype = torch.float32 + num_heads = 2 + head_dim = 8 + valid_positions = torch.tensor([1, 3, 5], device=device) + compact_position_ids = torch.tensor([[0, 1, 2]], device=device, dtype=torch.long) + cos, sin = build_rope_cache(8, head_dim, device=device) + + compact_q = torch.randn( + (1, 3, num_heads, head_dim), + device=device, + dtype=dtype, + generator=generator, + ) + compact_k = torch.randn( + compact_q.shape, + device=device, + dtype=dtype, + generator=generator, + ) + expected_q, expected_k = apply_rope_reference( + compact_q, + compact_k, + compact_position_ids, + cos, + sin, + ) + + padded_q = torch.randn((1, 6, num_heads, head_dim), device=device, generator=generator) + padded_k = torch.randn(padded_q.shape, device=device, generator=generator) + padded_q[0, valid_positions] = compact_q[0] + padded_k[0, valid_positions] = compact_k[0] + padded_position_ids = torch.zeros((1, 6), device=device, dtype=torch.long) + padded_position_ids[0, valid_positions] = compact_position_ids[0] + + actual_q, actual_k = apply_rope_reference(padded_q, padded_k, padded_position_ids, cos, sin) + + assert torch.equal(actual_q[0, valid_positions], expected_q[0]) + assert torch.equal(actual_k[0, valid_positions], expected_k[0]) + + +@pytest.mark.parametrize("device", _available_devices()) +def test_rope_packed_sequence_position_reset_matches_standalone_segment(device: torch.device): + generator = _generator(device, seed=153) + dtype = torch.float32 + num_heads = 2 + head_dim = 8 + segment_len = 3 + cos, sin = build_rope_cache(8, head_dim, device=device) + local_position_ids = torch.tensor([[0, 1, 2]], device=device, dtype=torch.long) + + target_q = torch.randn( + (1, segment_len, num_heads, head_dim), + device=device, + dtype=dtype, + generator=generator, + ) + target_k = torch.randn( + target_q.shape, + device=device, + dtype=dtype, + generator=generator, + ) + expected_q, expected_k = apply_rope_reference(target_q, target_k, local_position_ids, cos, sin) + + packed_q = torch.randn( + (1, segment_len * 2, num_heads, head_dim), + device=device, + dtype=dtype, + generator=generator, + ) + packed_k = torch.randn( + packed_q.shape, + device=device, + dtype=dtype, + generator=generator, + ) + packed_q[:, segment_len:].copy_(target_q) + packed_k[:, segment_len:].copy_(target_k) + packed_position_ids = torch.tensor([[0, 1, 2, 0, 1, 2]], device=device, dtype=torch.long) + + actual_q, actual_k = apply_rope_reference(packed_q, packed_k, packed_position_ids, cos, sin) + + assert torch.equal(actual_q[:, segment_len:], expected_q) + assert torch.equal(actual_k[:, segment_len:], expected_k) + + +def test_rope_cache_rejects_invalid_shapes(): + with pytest.raises(ValueError, match="max_position"): + build_rope_cache(0, 8) + with pytest.raises(ValueError, match="head_dim"): + build_rope_cache(4, 7) + with pytest.raises(ValueError, match="base"): + build_rope_cache(4, 8, base=0.0) + + +@pytest.mark.parametrize("dtype", (torch.bool, torch.int64, torch.complex64)) +def test_rope_cache_rejects_non_floating_dtypes(dtype: torch.dtype): + with pytest.raises(ValueError, match="floating-point dtype"): + build_rope_cache(4, 8, dtype=dtype) + + +def test_rope_application_rejects_invalid_inputs(): + query = torch.randn(1, 2, 1, 8) + key = torch.randn_like(query) + cos, sin = build_rope_cache(4, 8) + position_ids = torch.tensor([[0, 1]]) + + with pytest.raises(ValueError, match="same shape"): + apply_rope_reference(query, key[:, :, :, :4], position_ids, cos, sin) + with pytest.raises(ValueError, match="same dtype"): + apply_rope_reference(query, key.to(dtype=torch.float64), position_ids, cos, sin) + with pytest.raises(ValueError, match="position_ids"): + apply_rope_reference(query, key, torch.tensor([[0.0, 1.0]]), cos, sin) + with pytest.raises(ValueError, match="cos and sin must have the same dtype"): + apply_rope_reference(query, key, position_ids, cos, sin.to(dtype=torch.float64)) + with pytest.raises(ValueError, match="non-negative"): + apply_rope_reference(query, key, torch.tensor([[0, -1]]), cos, sin) + with pytest.raises(ValueError, match="cache length"): + apply_rope_reference(query, key, torch.tensor([[0, 4]]), cos, sin) + with pytest.raises(ValueError, match="even"): + rotate_half(torch.randn(1, 7)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available") +def test_rope_application_rejects_device_mismatches(): + query = torch.randn(1, 2, 1, 8, device="cuda") + key = torch.randn_like(query) + cos, sin = build_rope_cache(4, 8, device="cuda") + position_ids = torch.tensor([[0, 1]], device="cuda") + + with pytest.raises(ValueError, match="position_ids must be on the same device"): + apply_rope_reference(query, key, position_ids.cpu(), cos, sin) + with pytest.raises(ValueError, match="cos and sin must be on the same device"): + apply_rope_reference(query, key, position_ids, cos.cpu(), sin.cpu())