diff --git a/docs/design/ws2-attention-decode-replay.md b/docs/design/ws2-attention-decode-replay.md new file mode 100644 index 00000000..60583d5b --- /dev/null +++ b/docs/design/ws2-attention-decode-replay.md @@ -0,0 +1,83 @@ +# WS2 attention decode replay + +PR6 of issue #235 extends the single-GPU attention attribution harness with a +correctness-first decode-stage KV-cache replay. The reference compares each +decode query with the same query evaluated over a fully materialized logical KV +sequence. Both paths export attention output and attention-domain LSE. + +## Cache identity + +`DecodeKVCacheMetadata` separates logical sequence identity from physical cache +layout: + +- `cache_position` and `query_position_ids` identify each decode query; +- `kv_seq_lens` records the active cached sequence length; +- `block_table` maps logical blocks to physical pages; +- `global_token_positions` identifies every populated physical cache slot; +- `key_position_ids` binds cached K to its RoPE positions; +- `q_rope_state` and `k_cache_rope_state` declare whether tensors are before or + after RoPE; +- `prefix_cache_key`, `prefix_length`, and `prefix_cache_fingerprint` identify a + reused logical prefix and bind it to the cached K/V content; +- `cp_block_owners` records logical CP ownership without changing merge order. + +Metadata is validated before attention runs. Missing pages, duplicated active +pages, non-canonical `-1` page/owner tails, out-of-range positions, mismatched +RoPE positions, or inconsistent prefix-cache identity fail with an explicit +error. Logical positions must be non-negative and strictly increasing, but may +start at a nonzero global offset (for example after a sliding-window eviction). +Prefix identity is verified by recomputing a physical-layout-invariant +SHA-256 fingerprint over logical prefix positions, cached K/V content, storage +dtypes, and the cache-side RoPE materialization configuration (`theta`, rotary +dimension, cast boundary, and output dtype). The prefix key and enabled state are +part of that fingerprint. Every replay report also includes a physical-layout- +invariant full cache execution fingerprint over logical positions, page size, CP +ownership, prefix/RoPE identity, dtypes, and active K/V content. The current reference requires +`rope_cast_at="after_rope"` and `rope_rotary_dim == head_dim`; partial rotary is +not supported yet. Q and cached K retain separate RoPE output dtype contracts so +mixed rollout query/KV storage dtypes are represented faithfully. + +## Deterministic replay + +The replay restores logical KV order from the block table, computes one FP32 +partial `(out, lse)` state per logical block, and merges partial states in +`global_block_index` order. CP ownership and physical page order never determine +the numerical reduction order. Both the full logical-KV reference and paged +replay accumulate and export LSE in FP32; output downcast occurs only at final +write. + +For each decode query at logical position `t`, only cached positions less than +or equal to `t` participate. This supports `Sq=1` and few-query replay. + +The reference path removes physical page boundaries after reconstructing the +same logical KV sequence. Therefore the primary check is: + +```text +decode_paged_kv(query_t) == full_logical_kv(query_t) +``` + +The report reuses the PR2 drift format and includes maximum, mean, p95, and p99 +absolute drift for output and LSE, plus decode/cache/RoPE provenance. + +## Transformer Engine + +Transformer Engine remains optional. When requested, the replay passes the same +sorted partial states to the capability-probed TE context-parallel correction +helpers already used by the PR2 harness. TE validates merge arithmetic only; it +does not interpret cache metadata, choose logical order, or replace RL-Kernel's +reference semantics. The comparison always reports the RL-Kernel result; when +TE is requested but unavailable or incompatible, it records the reason in the +report instead of failing the core decode harness. + +## Current boundary + +The harness models CP block ownership on one device so cache construction, +logical ordering, RoPE identity, and deterministic merging can be attributed +without communication. It does not implement P2P NCCL transport, production +custom CUDA all-gather/reduce-scatter, FlashInfer runtime execution, or training +backward. Reports explicitly record `single_device_logical_reference`, +`runtime_verified=false`, `supports_backward=false`, and `communication=none`; +they are not valid Megatron/vLLM runtime readbacks. A distributed caller can gather the same partial states using the PR3 +transport layer; numerical merging must still happen in the fixed logical order +described above. This PR targets the shared `test` integration branch and has a +logical dependency on PR2/#253. diff --git a/docs/design/ws2-attention-single-gpu-harness.md b/docs/design/ws2-attention-single-gpu-harness.md new file mode 100644 index 00000000..acca4101 --- /dev/null +++ b/docs/design/ws2-attention-single-gpu-harness.md @@ -0,0 +1,129 @@ +# WS2 Attention Single-GPU Comparison Harness + +Status: PR2 harness for [#235](https://github.com/RL-Align/RL-Kernel/issues/235) + +## Scope + +This harness compares attention materializations on one device before CP +communication is introduced. It is diagnostic infrastructure: it does not launch +collectives and does not replace the deterministic CP reference planned in PR3. + +Implemented paths: + +- `full_prefill`: training-style full-sequence softmax attention; +- `chunked_prefill`: rollout-style query chunk replay over full KV; +- `rl_kernel_paged_kv`: rollout-style KV page replay with fp32 attention-domain + LSE merge by logical KV block order; +- `transformer_engine_paged_kv`: optional oracle that reuses NVIDIA Transformer + Engine's context-parallel PyTorch correction helpers when TE is installed. + +RoPE scope: + +- `unfused_rope_attention`: canonical `RoPE -> Attention` path; +- `fused_like_rope_attention`: semantic `RoPE+Attention` path that applies the + same canonical RoPE rules before attention, then records the fused boundary in + provenance. + +The RoPE path is still single-GPU attribution. It proves that both sides agree +on post-RoPE Q/K, `out`, attention-domain `lse`, and optional active-token +`dlogp` before CP communication or production fused kernels are introduced. + +## Report + +`rl_engine.testing.attention_comparison.compare_single_gpu_attention` emits a +structured report with: + +- `out` max / mean / p95 / p99 absolute drift; +- attention-domain `lse` max / mean / p95 / p99 absolute drift; +- optional active-token-only `dlogp` drift when `lm_head_weight`, `target_ids`, + and an active token mask are provided; +- per-path provenance including chunk/page sizes, KV page bounds, merge backend, + merge order, and LSE domain; +- optional-backend unavailability reasons. + +`compare_single_gpu_rope_attention` emits the same drift schema and additionally +reports post-RoPE Q/K drift. Its provenance records: + +- Q/K state as `post_rope`; +- `position_ids` shape and range; +- `rope_theta`, `rotary_dim`, `rope_cast_at`, and `rope_output_dtype`; +- `fusion_boundary` as either `unfused_rope_attention` or + `fused_rope_attention`. + +The selected-logprob convention follows #207: + +```text +dlogp = candidate selected logp - full_prefill selected logp +``` + +## Transformer Engine Reuse + +The harness does not make Transformer Engine a runtime dependency. When +available, it lazily imports: + +```text +transformer_engine.pytorch.attention.dot_product_attention.context_parallel +``` + +and calls: + +```text +flash_attn_fwd_softmax_lse_correction +flash_attn_fwd_out_correction_init +flash_attn_fwd_out_correction +``` + +Those helpers provide an industrial implementation oracle for the same fp32 +`(out, lse)` online-softmax merge policy that later CP/fused paths must match. +When TE is not installed, the TE path is reported as unavailable and the local +RL-Kernel paths still run. + +## CLI Registration + +The existing generic operator harness now registers `attention`, so a local +candidate smoke can run with: + +```bash +python scripts/check_operator.py --op attention --candidate pytorch --dtype fp32 +``` + +The attention-specific WS2 comparison entry point is Python-first for now: + +```python +from rl_engine.testing.attention_comparison import ( + AttentionComparisonInputs, + compare_single_gpu_rope_attention, + compare_single_gpu_attention, +) + +report = compare_single_gpu_attention( + AttentionComparisonInputs(q=q, k=k, v=v, target_ids=target_ids, lm_head_weight=w), + query_chunk_size=512, + kv_page_size=512, + include_transformer_engine=True, +) +print(report.to_dict()) + +rope_report = compare_single_gpu_rope_attention( + AttentionComparisonInputs( + q=q, + k=k, + v=v, + rope_positions=torch.arange(q.size(2), device=q.device), + target_ids=target_ids, + lm_head_weight=w, + ) +) +print(rope_report.to_dict()) +``` + +## Validation + +```bash +python -m pytest tests/test_attention_comparison.py -q +``` + +The tests cover full vs chunked/paged equivalence, active-token `dlogp` drift, +optional TE correction-helper reuse through a fake TE module, JSON-compatible +reports, RoPE+Attention post-RoPE Q/K attribution, and `attention` registration +in the generic operator comparison specs. diff --git a/docs/design/ws2-attention-transformer-engine-reuse-plan.md b/docs/design/ws2-attention-transformer-engine-reuse-plan.md new file mode 100644 index 00000000..76c5bb40 --- /dev/null +++ b/docs/design/ws2-attention-transformer-engine-reuse-plan.md @@ -0,0 +1,106 @@ +# WS2 Attention Transformer Engine 复用方案 + +Status: #235 设计补充 + +## 设计结论 + +Transformer Engine（TE）在 #235 中只能是显式 opt-in 的 validation oracle +或 backend candidate，不是 RL-Kernel attention 语义的可信源。可信源仍然是 +RL-Kernel 自己的 `AttentionContract`、RoPE/cache metadata、 +attention-domain `lse`、固定 `global_block_index` merge 顺序、 +deterministic reference 和 drift report。 + +TE 复用分为三层： + +| 层级 | TE 角色 | 允许范围 | +| --- | --- | --- | +| Merge oracle | 复用 TE context-parallel correction helpers 校验 `(out, lse)` online-softmax merge | PR2、PR3、PR5、PR6 | +| Fused forward candidate | 评估 `DotProductAttention` 作为 opt-in 生产后端候选 | 仅 PR7 | +| Backward oracle | 仅在 TE 暴露兼容 saved forward state 时，通过 autograd/backward 对比 `dq/dk/dv` | 仅 PR8 | + +## Merge Oracle Contract + +对任意 Q row，RL-Kernel 先按逻辑 KV block 生成 partial states： + +```text +state_i = (out_i, lse_i, global_block_index_i) +``` + +其中 `out_i` 是本地 KV block 内已经归一化的 attention output，`lse_i` +是 attention-domain LSE，shape 为 `[B, Hq, Sq]`。所有 state 必须按 +`global_block_index` 排序后再合并： + +```text +lse_new = logaddexp(lse_prev, lse_i) +out_new = exp(lse_prev - lse_new) * out_prev + + exp(lse_i - lse_new) * out_i +``` + +TE helper 可以负责 correction arithmetic，但语义输入必须由 RL-Kernel 提供： + +```text +TE_merge(sorted(RL-Kernel partial states)) == RL-Kernel_merge(sorted(partial states)) +``` + +调用 TE 前，RL-Kernel 必须保证： + +- merge accumulation 使用 FP32，只有 `final_write` 才 downcast； +- merge 顺序来自逻辑 `global_block_index`，不是通信 arrival order； +- all-masked / empty-KV row 保持 `lse = -inf`、`out = 0`，不能产生 NaN； +- TE adapter 启用前必须完成 capability probe：module/symbol 存在、helper + signature 兼容、tiny numeric merge smoke 通过； +- RoPE state、causal/padding mask、packed/varlen boundary、cache position 已经对齐。 + +## PR-level TE Plan + +| PR | RL-Kernel 核心功能 | TE 复用方式 | 精确 TE API | RL-Kernel 必须准备 | Gate / fallback | +| --- | --- | --- | --- | --- | --- | +| PR1 / #236 | 定义 attention contract、sharding/reduction metadata、RoPE/cache 字段 | 不调用 TE；只预留 `transformer_engine` 作为未来显式 backend 名称 | 无 | backend、reduction、`lse_domain`、`merge_order`、RoPE/cache identity 字段 | 不依赖 TE；metadata 缺失仍由 RL-Kernel contract fail | +| PR2 / #253 | 单 GPU full/chunked/paged-KV attention comparison harness | optional paged-KV merge oracle | `transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py`；`transformer_engine.pytorch.attention.dot_product_attention.context_parallel`；`flash_attn_fwd_softmax_lse_correction`；`flash_attn_fwd_out_correction_init`；`flash_attn_fwd_out_correction` | 相同 Q/K/V、相同 causal/padding metadata、相同 KV page order、RL-Kernel partial states `(out_i, lse_i)` | 对比 `TE_merge(partials)` 和 `RL-Kernel_merge(partials)` 的 `out/lse`；TE 不可用时 report `unavailable` | +| PR3 / #238 | post-RoPE Q/K 上的 deterministic CP attention reference | optional CP merge oracle test | 同 PR2 的 `context_parallel.py` module/functions | post-RoPE Q/K boundary、CP partial states、不重叠 global KV block ranges、固定 merge order | TE 不可用时 skip；TE 不定义 reference path | +| PR4 | Qwen3-8B TP=2 CP=2 BF16 cross-config 集成和 backend provenance | policy/provenance only | 不新增 TE 调用 | runtime descriptor 可 request `transformer_engine`，但默认执行仍是 deterministic reference | 记录 requested backend、actual backend、fallback reason、TE availability；禁止 silent fallback | +| PR5 | 分布式 prefill/chunked-prefill drift benchmark 和 report artifacts | benchmark merge oracle | 通过 `TEContextParallelMergeAdapter` 调用同 PR2 的 `context_parallel.py` module/functions | 与 RL-Kernel merge 完全相同的 gathered CP partial states、per-rank block metadata hash、FP32 merge dtype | 报告 `merge_drift = drift(TE_merge(partials), RL-Kernel_merge(partials))`；benchmark 可 provenance fallback | +| PR6 | decode-stage KV-cache CP attention replay | decode / paged-KV merge oracle only | 通过 decode TE merge adapter 调用同 PR2 的 `context_parallel.py` module/functions | `cache_position`、`kv_seq_lens`、page table、prefix-cache identity、global token positions、RoPE cache state、sorted logical page/block order | TE 只验证 `(out, lse)` merge；cache/page identity 不一致时，在调用 TE 前 fail | +| PR7 | deterministic reference 稳定后的 fused prefill/decode backend alignment | full fused forward backend candidate | `transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py`；`transformer_engine.pytorch.DotProductAttention`；actual backend 可观测时记录为 `FlashAttention` / `FusedAttention` / `UnfusedDotProductAttention` | 精确 layout / `qkv_format`、mask mode、RoPE fusion boundary、dtype、scale placement、dropout=0 correctness mode、deterministic controls、LSE export capability、actual-backend 观测方式 | 只有 TE output、attention-domain LSE、actual backend provenance 都能对齐/记录时才可作为 production candidate；如果不能导出 LSE 或不能观测 actual backend，只能算 exploratory，并记录原因 | +| PR8 | training backward CP attention reference 和 gradient drift validation | optional backward oracle | `DotProductAttention` autograd/backward path，仅当 compatible saved forward state 暴露时使用 | 与 RL-Kernel reference 相同的 forward inputs/metadata：`out`、attention-domain `lse`、masks、RoPE state、sequence/cache metadata、CP block ownership | 对比 `dq/dk/dv`；没有兼容 TE backward state 时明确写 `not used`，不能宣称复用 TE backward | + +## Capability / Provenance Checklist + +任何 PR 只要提到 TE，都必须写清： + +```text +te_available, te_version, te_module, te_symbols +te_capability_probe, te_signature_checked, te_numeric_selftest +requested_backend, actual_backend, actual_backend_source +fallback, fallback_reason +attention_mode, dtype, layout/qkv_format, mask_alignment +lse_domain, lse_exported, merge_order, accum_dtype, downcast_at +split_kv_policy, paged_kv_policy, cp_block_metadata_hash +scale_placement, deterministic_controls, dropout_policy, te_env_controls +``` + +fallback 策略： + +| 场景 | TE 不可用 / capability 不匹配时 | +| --- | --- | +| optional oracle test | skip / report unavailable | +| benchmark exploration | provenance fallback 到 deterministic reference | +| correctness gate | fail closed | +| production backend | fail closed 或显式 provenance fallback；禁止 silent fallback | + +## 不宣称的事 + +- 不把 TE 设为 #235 的硬依赖。 +- 不用 TE API 反向定义 RL-Kernel contract。 +- 不在 metadata 不完整时 silent fallback 到 TE。 +- 不用 NCCL / TE arrival order 决定 attention merge 数值顺序。 +- 不在 PR7 前把 TE fused path 宣称为默认生产路径。 +- PR7 如果拿不到 attention-domain LSE，不宣称完整 correctness closure。 +- PR8 如果拿不到兼容 backward state，不宣称复用 TE backward。 + +## 最终判断标准 + +TE 可以帮助验证和加速，但 #235 的正确性仍由 RL-Kernel 自己的 contract、 +metadata、deterministic reference 和 drift report 保证。当前最值得复用的是 +TE context-parallel correction helper；完整 `DotProductAttention` 路径只有在 +显式声明 capability 并满足 RL-Kernel 语义契约后，才允许作为生产候选后端。 diff --git a/rl_engine/testing/__init__.py b/rl_engine/testing/__init__.py index 8a66f5de..602c875b 100644 --- a/rl_engine/testing/__init__.py +++ b/rl_engine/testing/__init__.py @@ -3,6 +3,31 @@ """Testing helpers for RL-shaped kernel validation.""" +from .attention_comparison import ( + AttentionComparisonInputs, + AttentionComparisonReport, + AttentionPathDrift, + AttentionPathResult, + DecodeAttentionInputs, + DecodeKVCacheMetadata, + DriftStats, + TransformerEngineUnavailable, + compare_decode_kv_replay, + compare_single_gpu_attention, + compare_single_gpu_rope_attention, + decode_kv_cache_fingerprint, + decode_rope_identity_fingerprint, + decode_prefix_cache_fingerprint, + run_chunked_query_attention, + run_decode_full_prefill_reference, + run_decode_kv_replay, + run_decode_strict_shared_core, + run_full_attention, + run_fused_like_rope_attention, + run_paged_kv_attention, + run_unfused_rope_attention, + transformer_engine_context_parallel_available, +) from .reference_ops import ( active_token_count, compute_policy_ratio, @@ -30,6 +55,22 @@ ) __all__ = [ + "AttentionComparisonInputs", + "AttentionComparisonReport", + "AttentionPathDrift", + "AttentionPathResult", + "DecodeAttentionInputs", + "DecodeKVCacheMetadata", + "DriftStats", + "SyntheticRLKernelBatch", + "TransformerEngineUnavailable", + "active_token_count", + "compare_single_gpu_rope_attention", + "compare_single_gpu_attention", + "compare_decode_kv_replay", + "decode_kv_cache_fingerprint", + "decode_rope_identity_fingerprint", + "decode_prefix_cache_fingerprint", "LogicalBatch", "PhysicalLayout", "SyntheticRLKernelBatch", @@ -47,9 +88,18 @@ "make_synthetic_rl_kernel_batch", "masked_mean", "masked_sum", + "run_chunked_query_attention", + "run_decode_full_prefill_reference", + "run_decode_kv_replay", + "run_decode_strict_shared_core", + "run_fused_like_rope_attention", + "run_full_attention", + "run_paged_kv_attention", + "run_unfused_rope_attention", "reference_payload", "restore_logical_order", "restore_logical_order_from_padded", "selected_logprobs_reference", "summarize_kernel_drift", + "transformer_engine_context_parallel_available", ] diff --git a/rl_engine/testing/attention_comparison.py b/rl_engine/testing/attention_comparison.py new file mode 100644 index 00000000..def7a449 --- /dev/null +++ b/rl_engine/testing/attention_comparison.py @@ -0,0 +1,1745 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Single-GPU WS2 attention cross-implementation comparison harness. + +This module compares logically equivalent attention materializations before CP +communication is introduced. The full path is the training-style reference. +The chunked-query and paged-KV paths emulate rollout-style prefill layouts on a +single device while preserving global causal positions and attention-domain LSE. +""" + +from __future__ import annotations + +import hashlib +import importlib +import importlib.metadata as importlib_metadata +import inspect +import math +from dataclasses import dataclass +from numbers import Integral +from typing import Any, Literal + +import torch + +from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp +from rl_engine.testing.reference_ops import selected_logprobs_reference + +MergeBackend = Literal["rl_kernel", "transformer_engine"] +RoPEState = Literal["pre_rope", "post_rope"] + +_TE_CONTEXT_PARALLEL_MODULE = ( + "transformer_engine.pytorch.attention.dot_product_attention.context_parallel" +) +_TE_CONTEXT_PARALLEL_HELPERS = { + "flash_attn_fwd_softmax_lse_correction": ("softmax_lse", "softmax_lse_per_step"), + "flash_attn_fwd_out_correction_init": ( + "out_init_step", + "softmax_lse", + "softmax_lse_init_step", + "seq_dim", + ), + "flash_attn_fwd_out_correction": ( + "out", + "out_per_step", + "softmax_lse", + "softmax_lse_per_step", + "seq_dim", + ), +} + + +class TransformerEngineUnavailable(RuntimeError): + """Raised when the optional Transformer Engine oracle cannot be imported.""" + + +@dataclass(frozen=True) +class AttentionComparisonInputs: + """Inputs shared by every single-GPU attention comparison path.""" + + q: torch.Tensor + k: torch.Tensor + v: torch.Tensor + causal: bool = True + scale: float | None = None + key_padding_mask: torch.Tensor | None = None + lm_head_weight: torch.Tensor | None = None + target_ids: torch.Tensor | None = None + active_token_mask: torch.Tensor | None = None + output_dtype: torch.dtype = torch.float32 + rope_positions: torch.Tensor | None = None + rope_theta: float = 1_000_000.0 + rope_rotary_dim: int | None = None + rope_cast_at: str = "after_rope" + rope_output_dtype: torch.dtype | None = None + + +@dataclass(frozen=True) +class DecodeKVCacheMetadata: + """Logical identity and physical layout for decode-stage cached KV. + + ``block_table`` maps logical KV blocks to physical cache pages. Positions + are stored per physical cache slot; unused slots must contain ``-1``. + Keeping both mappings explicit lets the harness distinguish layout changes + from changes to the logical token sequence. + """ + + cache_position: torch.Tensor + kv_seq_lens: torch.Tensor + block_table: torch.Tensor + global_token_positions: torch.Tensor + query_position_ids: torch.Tensor + key_position_ids: torch.Tensor + page_size: int + prefix_cache_key: str | None = None + prefix_cache_enabled: bool = False + prefix_length: int = 0 + prefix_cache_fingerprint: str | None = None + q_rope_state: RoPEState = "post_rope" + k_cache_rope_state: RoPEState = "post_rope" + cp_block_owners: torch.Tensor | None = None + + +@dataclass(frozen=True) +class DecodeAttentionInputs: + """Decode queries and physically paged KV cache used by the PR6 harness.""" + + q: torch.Tensor + k_cache: torch.Tensor + v_cache: torch.Tensor + metadata: DecodeKVCacheMetadata + scale: float | None = None + output_dtype: torch.dtype = torch.float32 + rope_theta: float = 1_000_000.0 + rope_rotary_dim: int | None = None + rope_cast_at: str = "after_rope" + q_rope_output_dtype: torch.dtype | None = None + k_cache_rope_output_dtype: torch.dtype | None = None + lm_head_weight: torch.Tensor | None = None + target_ids: torch.Tensor | None = None + active_token_mask: torch.Tensor | None = None + + +@dataclass(frozen=True) +class AttentionPathResult: + """One materialized attention path result.""" + + name: str + out: torch.Tensor + lse: torch.Tensor + provenance: dict[str, Any] + post_rope_q: torch.Tensor | None = None + post_rope_k: torch.Tensor | None = None + + +@dataclass(frozen=True) +class DriftStats: + """Shape-aware absolute drift summary.""" + + max_abs: float + mean_abs: float + p95_abs: float + p99_abs: float + active_count: int + + def to_dict(self) -> dict[str, Any]: + return { + "max_abs": self.max_abs, + "mean_abs": self.mean_abs, + "p95_abs": self.p95_abs, + "p99_abs": self.p99_abs, + "active_count": self.active_count, + } + + +@dataclass(frozen=True) +class AttentionPathDrift: + """Candidate-vs-reference drift for one attention path.""" + + candidate_name: str + out: DriftStats + lse: DriftStats + dlogp: DriftStats | None + provenance: dict[str, Any] + post_rope_q: DriftStats | None = None + post_rope_k: DriftStats | None = None + + def to_dict(self) -> dict[str, Any]: + return { + "candidate_name": self.candidate_name, + "out": self.out.to_dict(), + "lse": self.lse.to_dict(), + "dlogp": None if self.dlogp is None else self.dlogp.to_dict(), + "post_rope_q": (None if self.post_rope_q is None else self.post_rope_q.to_dict()), + "post_rope_k": (None if self.post_rope_k is None else self.post_rope_k.to_dict()), + "provenance": self.provenance, + } + + +@dataclass(frozen=True) +class AttentionComparisonReport: + """Structured report for PR2 single-GPU attention attribution.""" + + reference_name: str + drifts: tuple[AttentionPathDrift, ...] + unavailable: tuple[str, ...] = () + + def to_dict(self) -> dict[str, Any]: + return { + "reference_name": self.reference_name, + "drifts": [drift.to_dict() for drift in self.drifts], + "unavailable": list(self.unavailable), + } + + +@dataclass(frozen=True) +class _PartialAttentionState: + out: torch.Tensor + lse: torch.Tensor + block_start: int + block_end: int + + +def compare_single_gpu_attention( + inputs: AttentionComparisonInputs, + *, + query_chunk_size: int | None = None, + kv_page_size: int | None = None, + include_transformer_engine: bool = False, +) -> AttentionComparisonReport: + """Compare full attention with chunked/paged single-GPU materializations. + + If ``lm_head_weight`` and ``target_ids`` are provided, the report also + includes active-token selected-logprob drift using the #207 convention: + candidate logp minus reference logp. + """ + + _validate_comparison_inputs(inputs) + reference = run_full_attention(inputs) + candidates = [ + run_chunked_query_attention(inputs, query_chunk_size=query_chunk_size), + run_paged_kv_attention(inputs, kv_page_size=kv_page_size, merge_backend="rl_kernel"), + ] + unavailable: list[str] = [] + if include_transformer_engine: + try: + candidates.append( + run_paged_kv_attention( + inputs, + kv_page_size=kv_page_size, + merge_backend="transformer_engine", + ) + ) + except TransformerEngineUnavailable as exc: + unavailable.append(f"transformer_engine_paged_kv: {exc}") + + drifts = tuple(_compare_path(candidate, reference, inputs) for candidate in candidates) + return AttentionComparisonReport( + reference_name=reference.name, + drifts=drifts, + unavailable=tuple(unavailable), + ) + + +def compare_single_gpu_rope_attention( + inputs: AttentionComparisonInputs, +) -> AttentionComparisonReport: + """Compare canonical unfused RoPE+Attention with fused-like materialization. + + This attribution path keeps the computation on one device and checks the + boundary that matters before CP communication: post-RoPE Q/K identity and + the resulting attention ``out`` / attention-domain ``lse``. + """ + + _validate_comparison_inputs(inputs) + _validate_rope_inputs(inputs) + reference = run_unfused_rope_attention(inputs) + candidates = [run_fused_like_rope_attention(inputs)] + drifts = tuple(_compare_path(candidate, reference, inputs) for candidate in candidates) + return AttentionComparisonReport(reference_name=reference.name, drifts=drifts) + + +def compare_decode_kv_replay( + inputs: DecodeAttentionInputs, + *, + include_transformer_engine: bool = False, + strict_bitwise: bool = False, +) -> AttentionComparisonReport: + """Compare paged decode replay with a logical full-KV teacher-forcing view.""" + + _validate_decode_inputs(inputs) + if strict_bitwise: + reference = _run_decode_strict_shared_core(inputs, materialization="logical_prefill") + candidates = [ + _run_decode_strict_shared_core(inputs, materialization="paged_kv_layout") + ] + else: + reference = _run_decode_full_prefill_reference(inputs) + candidates = [_run_decode_kv_replay(inputs, merge_backend="rl_kernel")] + unavailable: list[str] = [] + if include_transformer_engine: + try: + candidates.append(_run_decode_kv_replay(inputs, merge_backend="transformer_engine")) + except TransformerEngineUnavailable as exc: + unavailable.append(f"transformer_engine_decode_kv_replay: {exc}") + drifts = tuple(_compare_decode_path(candidate, reference, inputs) for candidate in candidates) + return AttentionComparisonReport( + reference_name=reference.name, + drifts=drifts, + unavailable=tuple(unavailable), + ) + + +def run_decode_strict_shared_core( + inputs: DecodeAttentionInputs, +) -> AttentionPathResult: + """Run decode through the same fixed one-query arithmetic schedule.""" + + _validate_decode_inputs(inputs) + return _run_decode_strict_shared_core(inputs, materialization="paged_kv_layout") + + +def _run_decode_strict_shared_core( + inputs: DecodeAttentionInputs, + *, + materialization: str, +) -> AttentionPathResult: + outs: list[torch.Tensor] = [] + lses: list[torch.Tensor] = [] + for batch_index in range(inputs.q.size(0)): + q, k, v, logical_positions = _decode_logical_qkv(inputs, batch_index) + batch_out: list[torch.Tensor] = [] + batch_lse: list[torch.Tensor] = [] + for query_index in range(q.size(2)): + query_position = int(inputs.metadata.cache_position[batch_index, query_index].item()) + visible = logical_positions <= query_position + if not bool(visible.any()): + raise ValueError("each decode query must have at least one visible KV token") + out, lse = _strict_decode_attention_with_lse( + q[:, :, query_index : query_index + 1, :], + k[:, :, visible, :], + v[:, :, visible, :], + output_dtype=inputs.output_dtype, + ) + batch_out.append(out) + batch_lse.append(lse) + outs.append(torch.cat(batch_out, dim=2)) + lses.append(torch.cat(batch_lse, dim=2)) + return AttentionPathResult( + name=f"strict_shared_core_{materialization}", + out=torch.cat(outs, dim=0), + lse=torch.cat(lses, dim=0), + provenance={ + "attention_mode": "decode", + "materialization": materialization, + "execution_scope": "single_device_strict_shared_core", + "runtime_verified": False, + "actual_backend": "rlkernel.pytorch.strict_attention_reference", + "communication_backend": "none", + "production_ready": False, + "supports_backward": False, + "communication": "none", + "strict_mode": True, + "strict_core_id": "rlkernel.attention.deterministic_core.v1", + "strict_schedule": "single_batch_single_query_global_kv_blocks", + "native_attention_arithmetic": False, + "fallback": False, + "fallback_reason": None, + "split_kv_policy": "disabled", + "merge_order": "global_block_index", + "lse_domain": "attention", + "lse_exported": True, + "accum_dtype": "fp32", + "downcast_at": "final_write", + "cache_execution_fingerprint": _decode_cache_fingerprint( + inputs, + token_counts=tuple(int(length) for length in inputs.metadata.kv_seq_lens.tolist()), + domain="strict_shared_core_decode", + ), + "kv_cache_identity_fingerprint": _decode_cache_fingerprint( + inputs, + token_counts=tuple(int(length) for length in inputs.metadata.kv_seq_lens.tolist()), + domain="strict_shared_core_decode", + ), + "rope_identity_fingerprint": decode_rope_identity_fingerprint(inputs), + "query_position_ids": inputs.metadata.query_position_ids.tolist(), + "key_position_ids": inputs.metadata.key_position_ids.tolist(), + "block_table": inputs.metadata.block_table.tolist(), + "page_size": inputs.metadata.page_size, + }, + ) + + +def run_decode_full_prefill_reference(inputs: DecodeAttentionInputs) -> AttentionPathResult: + """Materialize the full logical KV sequence for each decode query. + + This is the teacher-forcing side of the PR6 comparison. It deliberately + ignores physical page boundaries after restoring logical token order. + """ + + _validate_decode_inputs(inputs) + return _run_decode_full_prefill_reference(inputs) + + +def _run_decode_full_prefill_reference(inputs: DecodeAttentionInputs) -> AttentionPathResult: + outs: list[torch.Tensor] = [] + lses: list[torch.Tensor] = [] + for batch_index in range(inputs.q.size(0)): + q, k, v, logical_positions = _decode_logical_qkv(inputs, batch_index) + batch_out: list[torch.Tensor] = [] + batch_lse: list[torch.Tensor] = [] + for query_index in range(q.size(2)): + query_position = int(inputs.metadata.cache_position[batch_index, query_index].item()) + visible = logical_positions <= query_position + out, lse = _attention_with_lse( + q[:, :, query_index : query_index + 1, :], + k[:, :, visible, :], + v[:, :, visible, :], + causal=False, + scale=inputs.scale, + key_padding_mask=None, + q_start=0, + k_start=0, + total_query_len=1, + total_kv_len=int(visible.sum().item()), + output_dtype=inputs.output_dtype, + ) + _require_finite_decode_result(out, lse, path="full logical-KV reference") + batch_out.append(out) + batch_lse.append(lse) + outs.append(torch.cat(batch_out, dim=2)) + lses.append(torch.cat(batch_lse, dim=2)) + return AttentionPathResult( + name="full_prefill_decode_reference", + out=torch.cat(outs, dim=0), + lse=torch.cat(lses, dim=0), + provenance={ + "attention_mode": "decode", + "materialization": "full_logical_kv", + "execution_scope": "single_device_logical_reference", + "runtime_verified": False, + "supports_backward": False, + "communication": "none", + "lse_domain": "attention", + "accum_dtype": "fp32", + "lse_dtype": "fp32", + "downcast_at": "final_write", + "output_dtype": str(inputs.output_dtype).replace("torch.", ""), + "scale": _decode_attention_scale(inputs), + "q_dtype": str(inputs.q.dtype).replace("torch.", ""), + "k_cache_dtype": str(inputs.k_cache.dtype).replace("torch.", ""), + "v_cache_dtype": str(inputs.v_cache.dtype).replace("torch.", ""), + }, + ) + + +def run_decode_kv_replay( + inputs: DecodeAttentionInputs, + *, + merge_backend: MergeBackend = "rl_kernel", +) -> AttentionPathResult: + """Replay decode over physical KV pages and merge by logical block index.""" + + _validate_decode_inputs(inputs) + return _run_decode_kv_replay(inputs, merge_backend=merge_backend) + + +def _run_decode_kv_replay( + inputs: DecodeAttentionInputs, + *, + merge_backend: MergeBackend, +) -> AttentionPathResult: + outs: list[torch.Tensor] = [] + lses: list[torch.Tensor] = [] + merge_orders: list[list[list[int]]] = [] + cp_block_owners: list[list[int]] = [] + for batch_index in range(inputs.q.size(0)): + q, k, v, logical_positions = _decode_logical_qkv(inputs, batch_index) + owners = _logical_block_owners(inputs, batch_index) + cp_block_owners.append(owners) + batch_out: list[torch.Tensor] = [] + batch_lse: list[torch.Tensor] = [] + batch_orders: list[list[int]] = [] + for query_index in range(q.size(2)): + query_position = int(inputs.metadata.cache_position[batch_index, query_index].item()) + states: list[_PartialAttentionState] = [] + order: list[int] = [] + for block_index, (block_start, block_end) in enumerate( + _chunk_bounds(k.size(2), inputs.metadata.page_size) + ): + block_positions = logical_positions[block_start:block_end] + visible = block_positions <= query_position + if not bool(visible.any()): + continue + visible_end = block_start + int(visible.sum().item()) + out, lse = _attention_with_lse( + q[:, :, query_index : query_index + 1, :], + k[:, :, block_start:visible_end, :], + v[:, :, block_start:visible_end, :], + causal=False, + scale=inputs.scale, + key_padding_mask=None, + q_start=0, + k_start=block_start, + total_query_len=1, + total_kv_len=visible_end, + output_dtype=torch.float32, + ) + states.append( + _PartialAttentionState( + out=out, + lse=lse, + block_start=block_index, + block_end=block_index + 1, + ) + ) + order.append(block_index) + if not states: + raise ValueError("each decode query must have at least one visible cached KV token") + out, lse = _merge_partial_states(states, backend=merge_backend) + _require_finite_decode_result(out, lse, path=f"{merge_backend} paged-KV replay") + batch_out.append(out.to(inputs.output_dtype)) + batch_lse.append(lse) + batch_orders.append(order) + merge_orders.append(batch_orders) + outs.append(torch.cat(batch_out, dim=2)) + lses.append(torch.cat(batch_lse, dim=2)) + + provenance: dict[str, Any] = { + "attention_mode": "decode", + "materialization": "paged_kv_replay", + "execution_scope": "single_device_logical_reference", + "runtime_verified": False, + "supports_backward": False, + "communication": "none", + "sq": inputs.q.size(2), + "page_size": inputs.metadata.page_size, + "cache_position": inputs.metadata.cache_position.tolist(), + "kv_seq_lens": inputs.metadata.kv_seq_lens.tolist(), + "block_table": inputs.metadata.block_table.tolist(), + "global_token_positions": inputs.metadata.global_token_positions.tolist(), + "query_position_ids": inputs.metadata.query_position_ids.tolist(), + "key_position_ids": inputs.metadata.key_position_ids.tolist(), + "prefix_cache_enabled": inputs.metadata.prefix_cache_enabled, + "prefix_cache_key": inputs.metadata.prefix_cache_key, + "prefix_length": inputs.metadata.prefix_length, + "prefix_cache_fingerprint": inputs.metadata.prefix_cache_fingerprint, + "q_rope_state": inputs.metadata.q_rope_state, + "k_cache_rope_state": inputs.metadata.k_cache_rope_state, + "rope_theta": float(inputs.rope_theta), + "rotary_dim": _decode_rope_rotary_dim(inputs), + "rope_cast_at": inputs.rope_cast_at, + "q_rope_output_dtype": str(_decode_q_rope_output_dtype(inputs)).replace("torch.", ""), + "k_cache_rope_output_dtype": str(_decode_k_rope_output_dtype(inputs)).replace("torch.", ""), + "cp_block_owners": cp_block_owners, + "merge_order": "global_block_index", + "logical_merge_orders": merge_orders, + "merge_backend": merge_backend, + "lse_domain": "attention", + "lse_exported": True, + "accum_dtype": "fp32", + "lse_dtype": "fp32", + "downcast_at": "final_write", + "output_dtype": str(inputs.output_dtype).replace("torch.", ""), + "scale": _decode_attention_scale(inputs), + "q_dtype": str(inputs.q.dtype).replace("torch.", ""), + "k_cache_dtype": str(inputs.k_cache.dtype).replace("torch.", ""), + "v_cache_dtype": str(inputs.v_cache.dtype).replace("torch.", ""), + "cache_execution_fingerprint": _decode_cache_fingerprint( + inputs, + token_counts=tuple(int(length) for length in inputs.metadata.kv_seq_lens.tolist()), + domain="full_logical_kv_cache", + ), + "kv_cache_identity_fingerprint": _decode_cache_fingerprint( + inputs, + token_counts=tuple(int(length) for length in inputs.metadata.kv_seq_lens.tolist()), + domain="full_logical_kv_cache", + ), + "rope_identity_fingerprint": decode_rope_identity_fingerprint(inputs), + "preprocess_policy": "reference_only_not_production", + "preprocess_backends": { + "qk_rmsnorm": "not_executed_projected_qk_input", + "rope": ( + "rlkernel.pytorch.rope_reference" + if inputs.metadata.q_rope_state == "pre_rope" + or inputs.metadata.k_cache_rope_state == "pre_rope" + else "external_post_rope_input" + ), + }, + "preprocess_fallback": True, + "preprocess_fallback_reason": ( + "decode logical replay is a single-device correctness reference, not native vLLM" + ), + } + if merge_backend == "transformer_engine": + provenance.update(_te_context_parallel_provenance()) + return AttentionPathResult( + name=f"{merge_backend}_decode_kv_replay", + out=torch.cat(outs, dim=0), + lse=torch.cat(lses, dim=0), + provenance=provenance, + ) + + +def run_full_attention(inputs: AttentionComparisonInputs) -> AttentionPathResult: + """Training-style full-sequence attention with exported attention-domain LSE.""" + + out, lse = _attention_with_lse( + inputs.q, + inputs.k, + inputs.v, + causal=inputs.causal, + scale=inputs.scale, + key_padding_mask=inputs.key_padding_mask, + q_start=0, + k_start=0, + total_query_len=inputs.q.size(2), + total_kv_len=inputs.k.size(2), + output_dtype=inputs.output_dtype, + ) + return AttentionPathResult( + name="full_prefill", + out=out, + lse=lse, + provenance={ + "attention_mode": "prefill", + "materialization": "full_sequence", + "lse_domain": "attention", + }, + ) + + +def run_unfused_rope_attention(inputs: AttentionComparisonInputs) -> AttentionPathResult: + """Canonical ``RoPE -> Attention`` reference materialization.""" + + post_rope_q, post_rope_k = _apply_rope_to_qk(inputs) + out, lse = _attention_with_lse( + post_rope_q, + post_rope_k, + inputs.v, + causal=inputs.causal, + scale=inputs.scale, + key_padding_mask=inputs.key_padding_mask, + q_start=0, + k_start=0, + total_query_len=post_rope_q.size(2), + total_kv_len=post_rope_k.size(2), + output_dtype=inputs.output_dtype, + ) + return AttentionPathResult( + name="unfused_rope_attention", + out=out, + lse=lse, + provenance=_rope_attention_provenance( + inputs, + materialization="rope_then_attention", + fusion_boundary="unfused_rope_attention", + ), + post_rope_q=post_rope_q, + post_rope_k=post_rope_k, + ) + + +def run_fused_like_rope_attention(inputs: AttentionComparisonInputs) -> AttentionPathResult: + """Semantic fused ``RoPE+Attention`` path using the same canonical RoPE rules.""" + + post_rope_q, post_rope_k = _apply_rope_to_qk(inputs) + out, lse = _attention_with_lse( + post_rope_q, + post_rope_k, + inputs.v, + causal=inputs.causal, + scale=inputs.scale, + key_padding_mask=inputs.key_padding_mask, + q_start=0, + k_start=0, + total_query_len=post_rope_q.size(2), + total_kv_len=post_rope_k.size(2), + output_dtype=inputs.output_dtype, + ) + return AttentionPathResult( + name="fused_like_rope_attention", + out=out, + lse=lse, + provenance=_rope_attention_provenance( + inputs, + materialization="fused_like_rope_attention", + fusion_boundary="fused_rope_attention", + ), + post_rope_q=post_rope_q, + post_rope_k=post_rope_k, + ) + + +def run_chunked_query_attention( + inputs: AttentionComparisonInputs, + *, + query_chunk_size: int | None, +) -> AttentionPathResult: + """Rollout-style chunked prefill replay over full KV on one device.""" + + sq = inputs.q.size(2) + chunk_size = ( + sq if query_chunk_size is None else _positive_int(query_chunk_size, "query_chunk_size") + ) + out_chunks: list[torch.Tensor] = [] + lse_chunks: list[torch.Tensor] = [] + chunk_bounds = _chunk_bounds(sq, chunk_size) + for q_start, q_end in chunk_bounds: + out, lse = _attention_with_lse( + inputs.q[:, :, q_start:q_end, :], + inputs.k, + inputs.v, + causal=inputs.causal, + scale=inputs.scale, + key_padding_mask=inputs.key_padding_mask, + q_start=q_start, + k_start=0, + total_query_len=sq, + total_kv_len=inputs.k.size(2), + output_dtype=inputs.output_dtype, + ) + out_chunks.append(out) + lse_chunks.append(lse) + + return AttentionPathResult( + name="chunked_prefill", + out=torch.cat(out_chunks, dim=2), + lse=torch.cat(lse_chunks, dim=2), + provenance={ + "attention_mode": "chunked_prefill", + "materialization": "query_chunks", + "query_chunk_size": chunk_size, + "chunk_bounds": [list(bound) for bound in chunk_bounds], + "lse_domain": "attention", + }, + ) + + +def run_paged_kv_attention( + inputs: AttentionComparisonInputs, + *, + kv_page_size: int | None, + merge_backend: MergeBackend = "rl_kernel", +) -> AttentionPathResult: + """Rollout-style paged-KV prefill replay with explicit LSE merge.""" + + skv = inputs.k.size(2) + page_size = skv if kv_page_size is None else _positive_int(kv_page_size, "kv_page_size") + states: list[_PartialAttentionState] = [] + page_bounds = _chunk_bounds(skv, page_size) + for k_start, k_end in page_bounds: + key_mask = ( + None if inputs.key_padding_mask is None else inputs.key_padding_mask[:, k_start:k_end] + ) + out, lse = _attention_with_lse( + inputs.q, + inputs.k[:, :, k_start:k_end, :], + inputs.v[:, :, k_start:k_end, :], + causal=inputs.causal, + scale=inputs.scale, + key_padding_mask=key_mask, + q_start=0, + k_start=k_start, + total_query_len=inputs.q.size(2), + total_kv_len=skv, + output_dtype=torch.float32, + ) + states.append( + _PartialAttentionState( + out=out, + lse=lse, + block_start=k_start, + block_end=k_end, + ) + ) + + out, lse = _merge_partial_states(states, backend=merge_backend) + provenance = { + "attention_mode": "prefill", + "materialization": "paged_kv", + "kv_page_size": page_size, + "kv_page_bounds": [list(bound) for bound in page_bounds], + "merge_backend": merge_backend, + "requested_backend": merge_backend, + "actual_backend": ( + "te_context_parallel_merge_helpers" + if merge_backend == "transformer_engine" + else "rl_kernel" + ), + "fallback": False, + "fallback_reason": None, + "merge_order": "global_block_index", + "lse_domain": "attention", + "lse_exported": True, + "accum_dtype": "fp32", + "downcast_at": "final_write", + } + if merge_backend == "transformer_engine": + provenance.update(_te_context_parallel_provenance()) + return AttentionPathResult( + name=f"{merge_backend}_paged_kv", + out=out.to(inputs.output_dtype), + lse=lse, + provenance=provenance, + ) + + +def transformer_engine_context_parallel_available() -> bool: + """Return whether the optional TE context-parallel helper module imports.""" + + try: + _load_te_context_parallel() + except TransformerEngineUnavailable: + return False + return True + + +def decode_prefix_cache_fingerprint( + inputs: DecodeAttentionInputs, + *, + prefix_length: int, +) -> str: + """Fingerprint logical prefix positions and cached K/V content. + + The fingerprint is invariant to physical page placement because cache slots + are first restored to logical token order. It includes every cache-side + RoPE materialization fact that can change the replayed K values, so a + prefix cannot be reused under a different rotary configuration. + """ + + prefix_length = _positive_int(prefix_length, "prefix_length") + if bool((inputs.metadata.kv_seq_lens < prefix_length).any()): + raise ValueError("prefix_length must not exceed any kv_seq_lens entry") + return _decode_cache_fingerprint( + inputs, + token_counts=(prefix_length,) * inputs.q.size(0), + domain="shared_prefix", + ) + + +def decode_kv_cache_fingerprint(inputs: DecodeAttentionInputs) -> str: + """Fingerprint complete logical cache execution identity after validation.""" + + _validate_decode_inputs(inputs) + return _decode_cache_fingerprint( + inputs, + token_counts=tuple(int(length) for length in inputs.metadata.kv_seq_lens.tolist()), + domain="full_logical_kv_cache", + ) + + +def decode_rope_identity_fingerprint(inputs: DecodeAttentionInputs) -> str: + """Fingerprint every RoPE fact that changes decode Q/K interpretation.""" + + _validate_decode_inputs(inputs) + digest = hashlib.sha256() + digest.update( + ( + f"q_rope_state={inputs.metadata.q_rope_state};" + f"k_cache_rope_state={inputs.metadata.k_cache_rope_state};" + f"rope_theta={float(inputs.rope_theta):.17g};" + f"rotary_dim={_decode_rope_rotary_dim(inputs)};" + f"rope_cast_at={inputs.rope_cast_at};" + f"q_rope_output_dtype={_decode_q_rope_output_dtype(inputs)};" + f"k_rope_output_dtype={_decode_k_rope_output_dtype(inputs)}\n" + ).encode() + ) + for tensor in (inputs.metadata.query_position_ids, inputs.metadata.key_position_ids): + digest.update(str(tuple(tensor.shape)).encode()) + digest.update(str(tensor.dtype).encode()) + digest.update(tensor.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes()) + return digest.hexdigest() + + +def _decode_cache_fingerprint( + inputs: DecodeAttentionInputs, + *, + token_counts: tuple[int, ...], + domain: str, +) -> str: + if len(token_counts) != inputs.q.size(0): + raise ValueError("cache fingerprint requires one token count per batch item") + digest = hashlib.sha256() + digest.update( + ( + f"domain={domain};page_size={inputs.metadata.page_size};" + f"prefix_cache_enabled={inputs.metadata.prefix_cache_enabled};" + f"prefix_cache_key={inputs.metadata.prefix_cache_key!r};" + f"prefix_length={inputs.metadata.prefix_length};" + f"k_rope_state={inputs.metadata.k_cache_rope_state};" + f"rope_theta={float(inputs.rope_theta):.17g};" + f"rotary_dim={_decode_rope_rotary_dim(inputs)};" + f"rope_cast_at={inputs.rope_cast_at};" + f"k_rope_output_dtype={_decode_k_rope_output_dtype(inputs)};" + f"k_dtype={inputs.k_cache.dtype};v_dtype={inputs.v_cache.dtype}\n" + ).encode() + ) + for batch_index, token_count in enumerate(token_counts): + if token_count <= 0 or token_count > int(inputs.metadata.kv_seq_lens[batch_index].item()): + raise ValueError("cache fingerprint token count is outside kv_seq_lens") + slots = _decode_logical_slot_index(inputs, batch_index)[:token_count] + owner_count = math.ceil(token_count / inputs.metadata.page_size) + owners = _logical_block_owners(inputs, batch_index)[:owner_count] + digest.update(f"batch={batch_index};tokens={token_count};owners={owners}\n".encode()) + for tensor in ( + inputs.metadata.global_token_positions[batch_index, slots], + inputs.metadata.key_position_ids[batch_index, slots], + inputs.k_cache[batch_index, :, slots, :], + inputs.v_cache[batch_index, :, slots, :], + ): + digest.update(str(tuple(tensor.shape)).encode()) + digest.update(str(tensor.dtype).encode()) + digest.update(tensor.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes()) + return digest.hexdigest() + + +def _compare_path( + candidate: AttentionPathResult, + reference: AttentionPathResult, + inputs: AttentionComparisonInputs, +) -> AttentionPathDrift: + dlogp = None + if inputs.lm_head_weight is not None and inputs.target_ids is not None: + candidate_logp = _selected_logps_from_attention(candidate.out, inputs) + reference_logp = _selected_logps_from_attention(reference.out, inputs) + dlogp = _drift_stats(candidate_logp, reference_logp, mask=inputs.active_token_mask) + + return AttentionPathDrift( + candidate_name=candidate.name, + out=_drift_stats(candidate.out, reference.out), + lse=_drift_stats(candidate.lse, reference.lse), + dlogp=dlogp, + provenance=candidate.provenance, + post_rope_q=( + None + if candidate.post_rope_q is None or reference.post_rope_q is None + else _drift_stats(candidate.post_rope_q, reference.post_rope_q) + ), + post_rope_k=( + None + if candidate.post_rope_k is None or reference.post_rope_k is None + else _drift_stats(candidate.post_rope_k, reference.post_rope_k) + ), + ) + + +def _compare_decode_path( + candidate: AttentionPathResult, + reference: AttentionPathResult, + inputs: DecodeAttentionInputs, +) -> AttentionPathDrift: + dlogp = None + if inputs.lm_head_weight is not None and inputs.target_ids is not None: + candidate_logp = _selected_logps_from_decode_attention(candidate.out, inputs) + reference_logp = _selected_logps_from_decode_attention(reference.out, inputs) + dlogp = _drift_stats(candidate_logp, reference_logp, mask=inputs.active_token_mask) + return AttentionPathDrift( + candidate_name=candidate.name, + out=_drift_stats(candidate.out, reference.out), + lse=_drift_stats(candidate.lse, reference.lse), + dlogp=dlogp, + provenance=candidate.provenance, + ) + + +def _apply_rope_to_qk(inputs: AttentionComparisonInputs) -> tuple[torch.Tensor, torch.Tensor]: + _validate_rope_inputs(inputs) + assert inputs.rope_positions is not None + rope = NativeRoPEOp() + output_dtype = _rope_output_dtype(inputs) + q = rope.forward_fp32(inputs.q, inputs.rope_positions, theta=inputs.rope_theta).to(output_dtype) + k = rope.forward_fp32(inputs.k, inputs.rope_positions, theta=inputs.rope_theta).to(output_dtype) + return q, k + + +def _decode_logical_qkv( + inputs: DecodeAttentionInputs, + batch_index: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Restore one batch's cache to logical order and materialize RoPE state.""" + + metadata = inputs.metadata + slot_index = _decode_logical_slot_index(inputs, batch_index) + logical_position_tensor = metadata.global_token_positions[batch_index, slot_index].long() + k = inputs.k_cache[batch_index : batch_index + 1, :, slot_index, :] + v = inputs.v_cache[batch_index : batch_index + 1, :, slot_index, :] + q = inputs.q[batch_index : batch_index + 1] + + rope = NativeRoPEOp() + if metadata.q_rope_state == "pre_rope": + q = rope.forward_fp32( + q, + metadata.query_position_ids[batch_index : batch_index + 1], + theta=inputs.rope_theta, + ).to(_decode_q_rope_output_dtype(inputs)) + if metadata.k_cache_rope_state == "pre_rope": + key_positions = metadata.key_position_ids[batch_index, slot_index].unsqueeze(0) + k = rope.forward_fp32(k, key_positions, theta=inputs.rope_theta).to( + _decode_k_rope_output_dtype(inputs) + ) + return q, k, v, logical_position_tensor + + +def _decode_logical_slot_index( + inputs: DecodeAttentionInputs, + batch_index: int, +) -> torch.Tensor: + metadata = inputs.metadata + sequence_length = int(metadata.kv_seq_lens[batch_index].item()) + logical_block_count = math.ceil(sequence_length / metadata.page_size) + logical_index = torch.arange( + sequence_length, + device=inputs.k_cache.device, + dtype=torch.long, + ) + pages = metadata.block_table[batch_index, :logical_block_count].long() + return ( + pages[logical_index // metadata.page_size] * metadata.page_size + + logical_index % metadata.page_size + ) + + +def _logical_block_owners(inputs: DecodeAttentionInputs, batch_index: int) -> list[int]: + block_count = math.ceil( + int(inputs.metadata.kv_seq_lens[batch_index].item()) / inputs.metadata.page_size + ) + if inputs.metadata.cp_block_owners is None: + return [0] * block_count + return [ + int(owner) for owner in inputs.metadata.cp_block_owners[batch_index, :block_count].tolist() + ] + + +def _decode_q_rope_output_dtype(inputs: DecodeAttentionInputs) -> torch.dtype: + return inputs.q.dtype if inputs.q_rope_output_dtype is None else inputs.q_rope_output_dtype + + +def _decode_k_rope_output_dtype(inputs: DecodeAttentionInputs) -> torch.dtype: + return ( + inputs.k_cache.dtype + if inputs.k_cache_rope_output_dtype is None + else inputs.k_cache_rope_output_dtype + ) + + +def _decode_rope_rotary_dim(inputs: DecodeAttentionInputs) -> int: + return inputs.q.size(-1) if inputs.rope_rotary_dim is None else inputs.rope_rotary_dim + + +def _decode_attention_scale(inputs: DecodeAttentionInputs) -> float: + return 1.0 / math.sqrt(inputs.q.size(-1)) if inputs.scale is None else float(inputs.scale) + + +def _rope_output_dtype(inputs: AttentionComparisonInputs) -> torch.dtype: + return inputs.q.dtype if inputs.rope_output_dtype is None else inputs.rope_output_dtype + + +def _rope_rotary_dim(inputs: AttentionComparisonInputs) -> int: + return inputs.q.size(-1) if inputs.rope_rotary_dim is None else inputs.rope_rotary_dim + + +def _rope_attention_provenance( + inputs: AttentionComparisonInputs, + *, + materialization: str, + fusion_boundary: str, +) -> dict[str, Any]: + assert inputs.rope_positions is not None + return { + "attention_mode": "prefill", + "materialization": materialization, + "rope_state": "post_rope", + "q_rope_state": "post_rope", + "k_rope_state": "post_rope", + "position_kind": "position_ids", + "position_ids_shape": list(inputs.rope_positions.shape), + "position_ids_min": int(inputs.rope_positions.min().item()), + "position_ids_max": int(inputs.rope_positions.max().item()), + "rope_theta": float(inputs.rope_theta), + "rotary_dim": _rope_rotary_dim(inputs), + "rope_cast_at": inputs.rope_cast_at, + "rope_output_dtype": str(_rope_output_dtype(inputs)).replace("torch.", ""), + "fusion_boundary": fusion_boundary, + "lse_domain": "attention", + } + + +def _strict_decode_attention_with_lse( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + output_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + """Canonical one-query score/softmax/value schedule for decode replay.""" + + qf, kf, vf = q.float(), k.float(), v.float() + hq, hkv = qf.size(1), kf.size(1) + if hq % hkv: + raise ValueError("strict decode Attention requires GQA-compatible head counts") + if hq != hkv: + kf = kf.repeat_interleave(hq // hkv, dim=1) + vf = vf.repeat_interleave(hq // hkv, dim=1) + scores = torch.matmul(qf, kf.transpose(-1, -2)) / math.sqrt(qf.size(-1)) + row_max = scores.amax(dim=-1, keepdim=True) + finite = torch.isfinite(row_max) + exp_scores = torch.where(finite, torch.exp(scores - row_max), torch.zeros_like(scores)) + row_sum = exp_scores.sum(dim=-1, keepdim=True) + lse = torch.where( + row_sum > 0, + row_max + torch.log(row_sum), + torch.full_like(row_sum, float("-inf")), + ) + weights = torch.where(row_sum > 0, exp_scores / row_sum, torch.zeros_like(exp_scores)) + return torch.matmul(weights, vf).to(output_dtype), lse.squeeze(-1) + + +def _attention_with_lse( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool, + scale: float | None, + key_padding_mask: torch.Tensor | None, + q_start: int, + k_start: int, + total_query_len: int, + total_kv_len: int, + output_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + _validate_qkv(q, k, v) + if key_padding_mask is not None: + if key_padding_mask.shape != (q.size(0), k.size(2)): + raise ValueError("key_padding_mask must have shape [B, local_skv]") + if key_padding_mask.dtype != torch.bool: + raise ValueError("key_padding_mask must be bool") + + qf, kf, vf = q.float(), k.float(), v.float() + hq, sq, dim = qf.shape[1], qf.shape[2], qf.shape[3] + hkv, skv = kf.shape[1], kf.shape[2] + if hkv != hq: + repeat = hq // hkv + kf = kf.repeat_interleave(repeat, dim=1) + vf = vf.repeat_interleave(repeat, dim=1) + + scale_value = scale if scale is not None else 1.0 / math.sqrt(dim) + scores = torch.matmul(qf, kf.transpose(-1, -2)) * scale_value + if causal: + query_offset = total_kv_len - total_query_len + q_pos = torch.arange(sq, device=q.device) + q_start + query_offset + k_pos = torch.arange(skv, device=q.device) + k_start + scores = scores.masked_fill(k_pos[None, :] > q_pos[:, None], float("-inf")) + if key_padding_mask is not None: + scores = scores.masked_fill(~key_padding_mask[:, None, None, :], float("-inf")) + + lse = torch.logsumexp(scores, dim=-1) + finite_lse = torch.isfinite(lse) + weights = torch.exp(scores - lse.unsqueeze(-1)) + weights = torch.where(finite_lse.unsqueeze(-1), weights, torch.zeros_like(weights)) + out = torch.matmul(weights, vf) + return out.to(output_dtype), lse + + +def _merge_partial_states( + states: list[_PartialAttentionState], + *, + backend: MergeBackend, +) -> tuple[torch.Tensor, torch.Tensor]: + if not states: + raise ValueError("at least one partial state is required") + ordered = sorted(states, key=lambda state: (state.block_start, state.block_end)) + _validate_partial_states(ordered) + if backend == "rl_kernel": + return _merge_partial_states_rl_kernel(ordered) + if backend == "transformer_engine": + return _merge_partial_states_transformer_engine(ordered) + raise ValueError(f"unsupported merge backend: {backend}") + + +def _merge_partial_states_rl_kernel( + states: list[_PartialAttentionState], +) -> tuple[torch.Tensor, torch.Tensor]: + merged_out = states[0].out.float() + merged_lse = states[0].lse.float() + for state in states[1:]: + next_lse = torch.logaddexp(merged_lse, state.lse.float()) + finite = torch.isfinite(next_lse) + weight_prev = torch.where( + finite, + torch.exp(merged_lse - next_lse), + torch.zeros_like(next_lse), + ) + weight_next = torch.where( + finite, + torch.exp(state.lse.float() - next_lse), + torch.zeros_like(next_lse), + ) + merged_out = ( + weight_prev.unsqueeze(-1) * merged_out + weight_next.unsqueeze(-1) * state.out.float() + ) + merged_lse = next_lse + return merged_out, merged_lse + + +def _merge_partial_states_transformer_engine( + states: list[_PartialAttentionState], +) -> tuple[torch.Tensor, torch.Tensor]: + te_cp = _load_te_context_parallel() + merged_out = states[0].out.float() + merged_lse = states[0].lse.float() + for state in states[1:]: + previous_lse = merged_lse + state_out = state.out.float() + state_lse = state.lse.float() + both_masked = torch.isneginf(previous_lse) & torch.isneginf(state_lse) + te_previous_lse = torch.where(both_masked, torch.zeros_like(previous_lse), previous_lse) + te_state_lse = torch.where(both_masked, torch.zeros_like(state_lse), state_lse) + merged_lse = te_previous_lse.clone() + te_cp.flash_attn_fwd_softmax_lse_correction(merged_lse, te_state_lse) + merged_out = te_cp.flash_attn_fwd_out_correction_init( + merged_out, + merged_lse, + te_previous_lse, + seq_dim=2, + ) + te_cp.flash_attn_fwd_out_correction( + merged_out, + state_out, + merged_lse, + te_state_lse, + seq_dim=2, + ) + if both_masked.any(): + merged_lse = torch.where(both_masked, previous_lse, merged_lse) + merged_out = torch.where( + both_masked.unsqueeze(-1), + torch.zeros_like(merged_out), + merged_out, + ) + return merged_out, merged_lse + + +def _load_te_context_parallel() -> Any: + try: + module = importlib.import_module(_TE_CONTEXT_PARALLEL_MODULE) + except (ImportError, OSError, RuntimeError) as exc: + raise TransformerEngineUnavailable(str(exc)) from exc + _probe_te_context_parallel(module) + return module + + +def _probe_te_context_parallel(module: Any) -> None: + missing = [ + name for name in _TE_CONTEXT_PARALLEL_HELPERS if not callable(getattr(module, name, None)) + ] + if missing: + raise TransformerEngineUnavailable( + f"{_TE_CONTEXT_PARALLEL_MODULE} missing required helpers: {', '.join(missing)}" + ) + + for name, expected in _TE_CONTEXT_PARALLEL_HELPERS.items(): + helper = getattr(module, name) + try: + parameters = tuple(inspect.signature(helper).parameters) + except (TypeError, ValueError) as exc: + raise TransformerEngineUnavailable( + f"{_TE_CONTEXT_PARALLEL_MODULE}.{name} signature is not inspectable" + ) from exc + if parameters[: len(expected)] != expected: + raise TransformerEngineUnavailable( + f"{_TE_CONTEXT_PARALLEL_MODULE}.{name} has incompatible signature " + f"{parameters}; expected prefix {expected}" + ) + + try: + lse_a = torch.tensor([[[0.0, -1.0]]], dtype=torch.float32) + lse_b = torch.tensor([[[1.0, -3.0]]], dtype=torch.float32) + out_a = torch.tensor([[[[1.0, -2.0], [0.5, 2.0]]]], dtype=torch.float32) + out_b = torch.tensor([[[[-1.0, 4.0], [3.0, -0.5]]]], dtype=torch.float32) + expected_lse = torch.logaddexp(lse_a, lse_b) + expected_out = ( + torch.exp(lse_a - expected_lse).unsqueeze(-1) * out_a + + torch.exp(lse_b - expected_lse).unsqueeze(-1) * out_b + ) + + probed_lse = lse_a.clone() + module.flash_attn_fwd_softmax_lse_correction(probed_lse, lse_b) + probed_out = module.flash_attn_fwd_out_correction_init( + out_a.clone(), + probed_lse, + lse_a, + seq_dim=2, + ) + module.flash_attn_fwd_out_correction( + probed_out, + out_b, + probed_lse, + lse_b, + seq_dim=2, + ) + except Exception as exc: + raise TransformerEngineUnavailable( + f"{_TE_CONTEXT_PARALLEL_MODULE} helper numeric self-test failed: {exc}" + ) from exc + + if not torch.allclose(probed_lse, expected_lse, atol=1.0e-6, rtol=0.0): + raise TransformerEngineUnavailable( + f"{_TE_CONTEXT_PARALLEL_MODULE} LSE helper numeric self-test failed" + ) + if not torch.allclose(probed_out, expected_out, atol=1.0e-6, rtol=0.0): + raise TransformerEngineUnavailable( + f"{_TE_CONTEXT_PARALLEL_MODULE} out helper numeric self-test failed" + ) + + +def _te_context_parallel_provenance() -> dict[str, Any]: + return { + "te_available": True, + "te_version": _te_version(), + "te_module": _TE_CONTEXT_PARALLEL_MODULE, + "te_symbols": list(_TE_CONTEXT_PARALLEL_HELPERS), + "te_capability_probe": "passed", + "te_signature_checked": True, + "te_numeric_selftest": "passed", + "actual_backend_source": "rl_kernel_te_context_parallel_adapter", + "deterministic_controls": "not_applicable_merge_only", + "dropout_policy": "not_applicable_merge_only", + } + + +def _te_version() -> str | None: + for package_name in ("transformer-engine", "transformer_engine"): + try: + return importlib_metadata.version(package_name) + except importlib_metadata.PackageNotFoundError: + continue + return None + + +def _selected_logps_from_attention( + out: torch.Tensor, + inputs: AttentionComparisonInputs, +) -> torch.Tensor: + if inputs.lm_head_weight is None or inputs.target_ids is None: + raise ValueError("lm_head_weight and target_ids are required for dlogp drift") + batch, heads, seq, dim = out.shape + hidden = out.transpose(1, 2).reshape(batch, seq, heads * dim) + if inputs.lm_head_weight.shape[1] != hidden.size(-1): + raise ValueError( + "lm_head_weight hidden dimension must equal Hq * D; " + f"got {inputs.lm_head_weight.shape[1]} and {hidden.size(-1)}" + ) + logits = torch.matmul(hidden.float(), inputs.lm_head_weight.float().transpose(0, 1)) + return selected_logprobs_reference( + logits, + inputs.target_ids, + mask=inputs.active_token_mask, + output_dtype=torch.float32, + ) + + +def _selected_logps_from_decode_attention( + out: torch.Tensor, + inputs: DecodeAttentionInputs, +) -> torch.Tensor: + if inputs.lm_head_weight is None or inputs.target_ids is None: + raise ValueError("lm_head_weight and target_ids are required for decode dlogp drift") + batch, heads, seq, dim = out.shape + hidden = out.transpose(1, 2).reshape(batch, seq, heads * dim) + if inputs.lm_head_weight.shape[1] != hidden.size(-1): + raise ValueError( + "lm_head_weight hidden dimension must equal Hq * D; " + f"got {inputs.lm_head_weight.shape[1]} and {hidden.size(-1)}" + ) + logits = torch.matmul(hidden.float(), inputs.lm_head_weight.float().transpose(0, 1)) + return selected_logprobs_reference( + logits, + inputs.target_ids, + mask=inputs.active_token_mask, + output_dtype=torch.float32, + ) + + +def _drift_stats( + candidate: torch.Tensor, + reference: torch.Tensor, + *, + mask: torch.Tensor | None = None, +) -> DriftStats: + if candidate.shape != reference.shape: + raise ValueError( + f"candidate shape {tuple(candidate.shape)} must match " + f"reference shape {tuple(reference.shape)}" + ) + candidate_fp32 = candidate.float() + reference_fp32 = reference.float() + raw_diff = (candidate_fp32 - reference_fp32).abs() + diff = torch.where( + candidate_fp32 == reference_fp32, + torch.zeros_like(raw_diff), + raw_diff, + ) + values = _active_values(diff, mask) + active_count = int(values.numel()) + if active_count == 0: + return DriftStats(0.0, 0.0, 0.0, 0.0, 0) + return DriftStats( + max_abs=float(values.max().item()), + mean_abs=float(values.mean().item()), + p95_abs=float(torch.quantile(values, 0.95).item()), + p99_abs=float(torch.quantile(values, 0.99).item()), + active_count=active_count, + ) + + +def _active_values(diff: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor: + if mask is None: + return diff.reshape(-1) + if mask.shape == diff.shape: + return diff[mask.to(device=diff.device, dtype=torch.bool)] + if mask.ndim == 2 and diff.ndim == 4 and mask.shape == (diff.size(0), diff.size(2)): + expanded = mask[:, None, :, None].expand_as(diff) + return diff[expanded.to(device=diff.device, dtype=torch.bool)] + if mask.ndim == 2 and diff.ndim == 3 and mask.shape == (diff.size(0), diff.size(2)): + expanded = mask[:, None, :].expand_as(diff) + return diff[expanded.to(device=diff.device, dtype=torch.bool)] + raise ValueError(f"mask shape {tuple(mask.shape)} cannot select diff shape {tuple(diff.shape)}") + + +def _validate_comparison_inputs(inputs: AttentionComparisonInputs) -> None: + _validate_qkv(inputs.q, inputs.k, inputs.v) + if inputs.key_padding_mask is not None: + if inputs.key_padding_mask.shape != (inputs.q.size(0), inputs.k.size(2)): + raise ValueError("key_padding_mask must have shape [B, Skv]") + if inputs.key_padding_mask.dtype != torch.bool: + raise ValueError("key_padding_mask must be bool") + if (inputs.lm_head_weight is None) != (inputs.target_ids is None): + raise ValueError("lm_head_weight and target_ids must be provided together") + if inputs.target_ids is not None and inputs.target_ids.shape != ( + inputs.q.size(0), + inputs.q.size(2), + ): + raise ValueError("target_ids must have shape [B, Sq]") + if inputs.active_token_mask is not None: + if inputs.active_token_mask.shape != (inputs.q.size(0), inputs.q.size(2)): + raise ValueError("active_token_mask must have shape [B, Sq]") + if inputs.active_token_mask.dtype != torch.bool: + raise ValueError("active_token_mask must be bool") + if ( + not isinstance(inputs.rope_theta, (float, int)) + or isinstance(inputs.rope_theta, bool) + or not math.isfinite(float(inputs.rope_theta)) + or float(inputs.rope_theta) <= 0 + ): + raise ValueError("rope_theta must be a positive number") + if inputs.rope_output_dtype is not None and not isinstance( + inputs.rope_output_dtype, torch.dtype + ): + raise ValueError("rope_output_dtype must be a torch.dtype when provided") + + +def _validate_rope_inputs(inputs: AttentionComparisonInputs) -> None: + if inputs.rope_positions is None: + raise ValueError("rope_positions are required for RoPE+Attention comparison") + if inputs.q.size(2) != inputs.k.size(2): + raise ValueError("RoPE+Attention comparison currently requires Sq == Skv") + if inputs.rope_rotary_dim is not None: + if isinstance(inputs.rope_rotary_dim, bool) or inputs.rope_rotary_dim <= 0: + raise ValueError("rope_rotary_dim must be a positive integer when provided") + if inputs.rope_rotary_dim != inputs.q.size(-1): + raise ValueError( + "rope_rotary_dim must equal head_dim until partial-rotary RoPE is supported" + ) + if inputs.rope_cast_at != "after_rope": + raise ValueError("rope_cast_at must be 'after_rope' for the current fp32 RoPE reference") + if ( + inputs.rope_positions.device != inputs.q.device + or inputs.rope_positions.device != inputs.k.device + ): + raise ValueError("rope_positions must be on the same device as q/k") + if inputs.rope_positions.dtype not in {torch.int32, torch.int64, torch.long}: + raise ValueError("rope_positions must contain integer token positions") + if inputs.rope_positions.ndim == 1: + if inputs.rope_positions.numel() != inputs.q.size(2): + raise ValueError("1D rope_positions must have length Sq") + elif inputs.rope_positions.ndim == 2: + if inputs.rope_positions.shape != (inputs.q.size(0), inputs.q.size(2)): + raise ValueError("2D rope_positions must have shape [B, Sq]") + else: + raise ValueError("rope_positions must have shape [Sq] or [B, Sq]") + + +def _validate_decode_inputs(inputs: DecodeAttentionInputs) -> None: + _validate_qkv(inputs.q, inputs.k_cache, inputs.v_cache) + if inputs.q.device != inputs.k_cache.device or inputs.q.device != inputs.v_cache.device: + raise ValueError("q, k_cache, and v_cache must be on the same device") + if ( + not inputs.q.is_floating_point() + or not inputs.k_cache.is_floating_point() + or not inputs.v_cache.is_floating_point() + ): + raise ValueError("q, k_cache, and v_cache must use floating-point dtypes") + if ( + not isinstance(inputs.output_dtype, torch.dtype) + or not inputs.output_dtype.is_floating_point + ): + raise ValueError("output_dtype must be a floating-point torch.dtype") + if inputs.scale is not None: + if ( + not isinstance(inputs.scale, (float, int)) + or isinstance(inputs.scale, bool) + or not math.isfinite(float(inputs.scale)) + or float(inputs.scale) <= 0 + ): + raise ValueError("scale must be a positive finite number") + metadata = inputs.metadata + batch, _, sq, head_dim = inputs.q.shape + cache_capacity = inputs.k_cache.size(2) + page_size = _positive_int(metadata.page_size, "page_size") + if cache_capacity % page_size != 0: + raise ValueError("physical KV cache capacity must be divisible by page_size") + physical_page_count = cache_capacity // page_size + if metadata.cache_position.shape != (batch, sq): + raise ValueError("cache_position must have shape [B, Sq]") + if metadata.query_position_ids.shape != (batch, sq): + raise ValueError("query_position_ids must have shape [B, Sq]") + if metadata.kv_seq_lens.shape != (batch,): + raise ValueError("kv_seq_lens must have shape [B]") + if metadata.block_table.ndim != 2 or metadata.block_table.size(0) != batch: + raise ValueError("block_table must have shape [B, max_blocks]") + expected_cache_shape = (batch, cache_capacity) + if metadata.global_token_positions.shape != expected_cache_shape: + raise ValueError("global_token_positions must have shape [B, cache_capacity]") + if metadata.key_position_ids.shape != expected_cache_shape: + raise ValueError("key_position_ids must have shape [B, cache_capacity]") + integer_tensors = { + "cache_position": metadata.cache_position, + "query_position_ids": metadata.query_position_ids, + "kv_seq_lens": metadata.kv_seq_lens, + "block_table": metadata.block_table, + "global_token_positions": metadata.global_token_positions, + "key_position_ids": metadata.key_position_ids, + } + if metadata.cp_block_owners is not None: + integer_tensors["cp_block_owners"] = metadata.cp_block_owners + for name, tensor in integer_tensors.items(): + if tensor.device != inputs.q.device: + raise ValueError(f"{name} must be on the same device as q/k/v") + if tensor.dtype not in {torch.int32, torch.int64, torch.long}: + raise ValueError(f"{name} must contain integers") + if metadata.cp_block_owners is not None: + if metadata.cp_block_owners.shape != metadata.block_table.shape: + raise ValueError("cp_block_owners must have the same shape as block_table") + if not torch.equal(metadata.cache_position, metadata.query_position_ids): + raise ValueError("cache_position and query_position_ids must identify the same positions") + if metadata.q_rope_state not in {"pre_rope", "post_rope"}: + raise ValueError("q_rope_state must be 'pre_rope' or 'post_rope'") + if metadata.k_cache_rope_state not in {"pre_rope", "post_rope"}: + raise ValueError("k_cache_rope_state must be 'pre_rope' or 'post_rope'") + if metadata.prefix_cache_enabled: + if not metadata.prefix_cache_key: + raise ValueError("prefix_cache_key is required when prefix cache is enabled") + _positive_int(metadata.prefix_length, "prefix_length") + if not metadata.prefix_cache_fingerprint: + raise ValueError("prefix_cache_fingerprint is required when prefix cache is enabled") + elif ( + metadata.prefix_cache_key is not None + or metadata.prefix_length != 0 + or metadata.prefix_cache_fingerprint is not None + ): + raise ValueError( + "prefix cache key/fingerprint must be None and prefix_length must be 0 " + "when prefix cache is disabled" + ) + if inputs.rope_cast_at != "after_rope": + raise ValueError("rope_cast_at must be 'after_rope' for the current fp32 RoPE reference") + if inputs.rope_rotary_dim is not None: + if inputs.rope_rotary_dim != head_dim: + raise ValueError("rope_rotary_dim must equal head_dim") + _positive_int(inputs.rope_rotary_dim, "rope_rotary_dim") + if ( + not isinstance(inputs.rope_theta, (float, int)) + or isinstance(inputs.rope_theta, bool) + or not math.isfinite(float(inputs.rope_theta)) + or float(inputs.rope_theta) <= 0 + ): + raise ValueError("rope_theta must be a positive number") + if inputs.q_rope_output_dtype is not None and not isinstance( + inputs.q_rope_output_dtype, torch.dtype + ): + raise ValueError("q_rope_output_dtype must be a torch.dtype when provided") + if inputs.k_cache_rope_output_dtype is not None and not isinstance( + inputs.k_cache_rope_output_dtype, torch.dtype + ): + raise ValueError("k_cache_rope_output_dtype must be a torch.dtype when provided") + for name, dtype in ( + ("q_rope_output_dtype", inputs.q_rope_output_dtype), + ("k_cache_rope_output_dtype", inputs.k_cache_rope_output_dtype), + ): + if dtype is not None and not dtype.is_floating_point: + raise ValueError(f"{name} must be a floating-point torch.dtype") + if ( + metadata.q_rope_state == "post_rope" + and inputs.q_rope_output_dtype is not None + and inputs.q.dtype != inputs.q_rope_output_dtype + ): + raise ValueError("post-RoPE q dtype must match q_rope_output_dtype") + if ( + metadata.k_cache_rope_state == "post_rope" + and inputs.k_cache_rope_output_dtype is not None + and inputs.k_cache.dtype != inputs.k_cache_rope_output_dtype + ): + raise ValueError("post-RoPE k_cache dtype must match k_cache_rope_output_dtype") + if (inputs.lm_head_weight is None) != (inputs.target_ids is None): + raise ValueError("lm_head_weight and target_ids must be provided together") + if inputs.target_ids is not None and inputs.target_ids.shape != (batch, sq): + raise ValueError("target_ids must have shape [B, Sq]") + if inputs.active_token_mask is not None: + if inputs.active_token_mask.shape != (batch, sq): + raise ValueError("active_token_mask must have shape [B, Sq]") + if inputs.active_token_mask.dtype != torch.bool: + raise ValueError("active_token_mask must be bool") + + for batch_index in range(batch): + sequence_length = int(metadata.kv_seq_lens[batch_index].item()) + if sequence_length <= 0 or sequence_length > cache_capacity: + raise ValueError("each kv_seq_lens entry must be in [1, cache_capacity]") + block_count = math.ceil(sequence_length / page_size) + if block_count > metadata.block_table.size(1): + raise ValueError("block_table does not contain enough logical KV blocks") + pages = metadata.block_table[batch_index, :block_count] + if bool(((pages < 0) | (pages >= physical_page_count)).any()): + raise ValueError("block_table contains an out-of-range physical page") + if torch.unique(pages).numel() != block_count: + raise ValueError("active block_table entries must not contain duplicate pages") + if bool((metadata.block_table[batch_index, block_count:] != -1).any()): + raise ValueError("unused block_table entries must be -1") + if metadata.cp_block_owners is not None: + active_owners = metadata.cp_block_owners[batch_index, :block_count] + inactive_owners = metadata.cp_block_owners[batch_index, block_count:] + if bool((active_owners < 0).any()): + raise ValueError("active cp_block_owners must be non-negative") + if bool((inactive_owners != -1).any()): + raise ValueError("unused cp_block_owners entries must be -1") + slot_index = _decode_logical_slot_index(inputs, batch_index) + if not torch.isfinite(inputs.q[batch_index]).all(): + raise ValueError("active decode q values must be finite") + if not torch.isfinite(inputs.k_cache[batch_index, :, slot_index, :]).all(): + raise ValueError("active logical k_cache values must be finite") + if not torch.isfinite(inputs.v_cache[batch_index, :, slot_index, :]).all(): + raise ValueError("active logical v_cache values must be finite") + active_slot_mask = torch.zeros(cache_capacity, device=inputs.q.device, dtype=torch.bool) + active_slot_mask[slot_index] = True + if bool((metadata.global_token_positions[batch_index, ~active_slot_mask] != -1).any()): + raise ValueError("unused global_token_positions entries must be -1") + if bool((metadata.key_position_ids[batch_index, ~active_slot_mask] != -1).any()): + raise ValueError("unused key_position_ids entries must be -1") + global_positions = metadata.global_token_positions[batch_index, slot_index] + if bool((global_positions < 0).any()) or ( + sequence_length > 1 and bool((global_positions[1:] <= global_positions[:-1]).any()) + ): + raise ValueError( + "block_table/global_token_positions must reconstruct logical positions " + "in strictly increasing global order" + ) + key_positions = metadata.key_position_ids[batch_index, slot_index] + if not torch.equal(key_positions, global_positions): + raise ValueError("key_position_ids must match cached global token positions") + cache_positions = metadata.cache_position[batch_index] + query_is_cached = (cache_positions[:, None] == global_positions[None, :]).any(dim=1) + if not bool(query_is_cached.all()): + raise ValueError("cache_position must refer to a token present in the KV cache") + if sq > 1 and bool((cache_positions[1:] <= cache_positions[:-1]).any()): + raise ValueError("few-query cache_position values must be strictly increasing") + + if metadata.prefix_cache_enabled: + actual_fingerprint = decode_prefix_cache_fingerprint( + inputs, + prefix_length=metadata.prefix_length, + ) + if actual_fingerprint != metadata.prefix_cache_fingerprint: + raise ValueError( + "prefix_cache_fingerprint does not match the logical prefix positions/content" + ) + + +def _validate_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError("q, k, and v must have shape [B, H, S, D]") + if k.shape != v.shape: + raise ValueError("k and v must have matching shape") + if q.size(0) != k.size(0) or q.size(3) != k.size(3): + raise ValueError("q, k, and v must share batch size and head dim") + if q.size(1) % k.size(1) != 0: + raise ValueError(f"Hq={q.size(1)} must be divisible by Hkv={k.size(1)}") + + +def _validate_partial_states(states: list[_PartialAttentionState]) -> None: + first = states[0] + if first.block_start < 0 or first.block_end <= first.block_start: + raise ValueError("partial state ranges must satisfy 0 <= block_start < block_end") + previous_end = first.block_end + for state in states[1:]: + if state.out.shape != first.out.shape or state.lse.shape != first.lse.shape: + raise ValueError("all partial states must have matching shapes") + if state.block_end <= state.block_start: + raise ValueError("partial state block ranges must have positive width") + if state.block_start != previous_end: + raise ValueError("partial state block ranges must be contiguous and non-overlapping") + previous_end = state.block_end + + +def _require_finite_decode_result( + out: torch.Tensor, + lse: torch.Tensor, + *, + path: str, +) -> None: + if not torch.isfinite(out).all() or not torch.isfinite(lse).all(): + raise ValueError(f"{path} produced non-finite attention output or LSE") + + +def _chunk_bounds(length: int, chunk_size: int) -> list[tuple[int, int]]: + if length <= 0: + raise ValueError("sequence length must be positive") + bounds: list[tuple[int, int]] = [] + cursor = 0 + while cursor < length: + end = min(cursor + chunk_size, length) + bounds.append((cursor, end)) + cursor = end + return bounds + + +def _positive_int(value: int, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, Integral) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + return int(value) + + +__all__ = [ + "AttentionComparisonInputs", + "AttentionComparisonReport", + "AttentionPathDrift", + "AttentionPathResult", + "DecodeAttentionInputs", + "DecodeKVCacheMetadata", + "DriftStats", + "TransformerEngineUnavailable", + "compare_single_gpu_rope_attention", + "compare_single_gpu_attention", + "compare_decode_kv_replay", + "decode_kv_cache_fingerprint", + "decode_rope_identity_fingerprint", + "decode_prefix_cache_fingerprint", + "run_chunked_query_attention", + "run_fused_like_rope_attention", + "run_full_attention", + "run_decode_full_prefill_reference", + "run_decode_kv_replay", + "run_paged_kv_attention", + "run_unfused_rope_attention", + "transformer_engine_context_parallel_available", +] diff --git a/tests/test_attention_comparison.py b/tests/test_attention_comparison.py new file mode 100644 index 00000000..f26f5f66 --- /dev/null +++ b/tests/test_attention_comparison.py @@ -0,0 +1,1004 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import argparse +import importlib +import json +import sys +import types +from dataclasses import replace +from typing import Literal + +import pytest +import torch + +from rl_engine.kernels.gtest import run_operator_suite +from rl_engine.kernels.gtest.operator_specs import make_candidate, make_operator_case +from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp +from rl_engine.testing.attention_comparison import ( + AttentionComparisonInputs, + DecodeAttentionInputs, + DecodeKVCacheMetadata, + compare_decode_kv_replay, + compare_single_gpu_attention, + compare_single_gpu_rope_attention, + decode_kv_cache_fingerprint, + decode_rope_identity_fingerprint, + decode_prefix_cache_fingerprint, + run_decode_kv_replay, + run_paged_kv_attention, +) + +_TE_CONTEXT_PARALLEL_MODULE = ( + "transformer_engine.pytorch.attention.dot_product_attention.context_parallel" +) + + +def _qkv(*, seed: int = 1): + gen = torch.Generator().manual_seed(seed) + q = torch.randn(2, 4, 6, 8, generator=gen) + k = torch.randn(2, 2, 6, 8, generator=gen) + v = torch.randn(2, 2, 6, 8, generator=gen) + return q, k, v + + +def _comparison_inputs() -> AttentionComparisonInputs: + q, k, v = _qkv() + gen = torch.Generator().manual_seed(2) + lm_head_weight = torch.randn(13, q.size(1) * q.size(3), generator=gen) + target_ids = torch.randint(0, 13, (q.size(0), q.size(2)), generator=gen) + active_mask = torch.tensor( + [ + [True, True, True, False, False, False], + [False, True, True, True, True, False], + ], + dtype=torch.bool, + ) + return AttentionComparisonInputs( + q=q, + k=k, + v=v, + causal=True, + lm_head_weight=lm_head_weight, + target_ids=target_ids, + active_token_mask=active_mask, + ) + + +def _decode_inputs( + *, + page_order: tuple[int, ...] = (0, 1, 2), + prefix_cache_enabled: bool = False, + q_rope_state: Literal["pre_rope", "post_rope"] = "post_rope", + k_cache_rope_state: Literal["pre_rope", "post_rope"] = "post_rope", + position_offset: int = 0, +) -> DecodeAttentionInputs: + q, logical_k, logical_v = _qkv(seed=17) + q = q[:, :, 4:6, :] + page_size = 2 + physical_k = torch.empty_like(logical_k) + physical_v = torch.empty_like(logical_v) + positions = torch.full((2, 6), -1, dtype=torch.long) + for logical_page, physical_page in enumerate(page_order): + logical_slice = slice(logical_page * page_size, (logical_page + 1) * page_size) + physical_slice = slice(physical_page * page_size, (physical_page + 1) * page_size) + physical_k[:, :, physical_slice, :] = logical_k[:, :, logical_slice, :] + physical_v[:, :, physical_slice, :] = logical_v[:, :, logical_slice, :] + positions[:, physical_slice] = torch.arange( + position_offset + logical_page * page_size, + position_offset + (logical_page + 1) * page_size, + ) + inputs = DecodeAttentionInputs( + q=q, + k_cache=physical_k, + v_cache=physical_v, + metadata=DecodeKVCacheMetadata( + cache_position=torch.tensor( + [[position_offset + 4, position_offset + 5]] * 2, + dtype=torch.long, + ), + kv_seq_lens=torch.tensor([6, 6], dtype=torch.long), + block_table=torch.tensor([page_order, page_order], dtype=torch.long), + global_token_positions=positions, + query_position_ids=torch.tensor( + [[position_offset + 4, position_offset + 5]] * 2, + dtype=torch.long, + ), + key_position_ids=positions.clone(), + page_size=page_size, + q_rope_state=q_rope_state, + k_cache_rope_state=k_cache_rope_state, + cp_block_owners=torch.tensor([[0, 1, 0], [0, 1, 0]], dtype=torch.long), + ), + lm_head_weight=torch.randn( + 11, q.size(1) * q.size(3), generator=torch.Generator().manual_seed(18) + ), + target_ids=torch.tensor([[1, 2], [3, 4]], dtype=torch.long), + active_token_mask=torch.tensor([[True, True], [False, True]], dtype=torch.bool), + ) + if not prefix_cache_enabled: + return inputs + prefix_length = 4 + with_prefix_identity = replace( + inputs, + metadata=replace( + inputs.metadata, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + prefix_length=prefix_length, + ), + ) + fingerprint = decode_prefix_cache_fingerprint( + with_prefix_identity, + prefix_length=prefix_length, + ) + return replace( + with_prefix_identity, + metadata=replace( + with_prefix_identity.metadata, + prefix_cache_fingerprint=fingerprint, + ), + ) + + +def test_single_gpu_attention_harness_reports_out_lse_and_dlogp_drift(): + report = compare_single_gpu_attention( + _comparison_inputs(), + query_chunk_size=2, + kv_page_size=3, + ) + + by_name = {drift.candidate_name: drift for drift in report.drifts} + assert set(by_name) == {"chunked_prefill", "rl_kernel_paged_kv"} + for drift in by_name.values(): + assert drift.out.max_abs <= 1.0e-6 + assert drift.lse.max_abs <= 1.0e-6 + assert drift.dlogp is not None + assert drift.dlogp.active_count == 7 + assert drift.dlogp.p95_abs <= 1.0e-6 + + payload = report.to_dict() + assert payload["reference_name"] == "full_prefill" + assert payload["drifts"][0]["out"]["p99_abs"] >= 0.0 + json.dumps(payload) + + +def test_single_gpu_attention_harness_preserves_key_padding_mask(): + q, k, v = _qkv(seed=3) + key_padding_mask = torch.tensor( + [ + [True, True, True, False, False, False], + [True, False, True, True, False, False], + ], + dtype=torch.bool, + ) + + report = compare_single_gpu_attention( + AttentionComparisonInputs( + q=q, + k=k, + v=v, + causal=True, + key_padding_mask=key_padding_mask, + ), + query_chunk_size=4, + kv_page_size=2, + ) + + assert report.unavailable == () + for drift in report.drifts: + assert drift.out.max_abs <= 1.0e-6 + assert drift.lse.max_abs <= 1.0e-6 + + +def test_single_gpu_rope_attention_harness_reports_rope_and_attention_drift(): + base = _comparison_inputs() + report = compare_single_gpu_rope_attention( + AttentionComparisonInputs( + q=base.q, + k=base.k, + v=base.v, + causal=True, + lm_head_weight=base.lm_head_weight, + target_ids=base.target_ids, + active_token_mask=base.active_token_mask, + rope_positions=torch.arange(base.q.size(2), dtype=torch.long), + rope_theta=1_000_000.0, + rope_rotary_dim=base.q.size(-1), + rope_output_dtype=torch.float32, + ) + ) + + assert report.reference_name == "unfused_rope_attention" + assert len(report.drifts) == 1 + drift = report.drifts[0] + assert drift.candidate_name == "fused_like_rope_attention" + assert drift.post_rope_q is not None + assert drift.post_rope_k is not None + assert drift.post_rope_q.max_abs <= 1.0e-6 + assert drift.post_rope_k.max_abs <= 1.0e-6 + assert drift.out.max_abs <= 1.0e-6 + assert drift.lse.max_abs <= 1.0e-6 + assert drift.dlogp is not None + assert drift.dlogp.max_abs <= 1.0e-6 + assert drift.provenance["position_kind"] == "position_ids" + assert drift.provenance["position_ids_shape"] == [base.q.size(2)] + assert drift.provenance["rotary_dim"] == base.q.size(-1) + assert drift.provenance["rope_cast_at"] == "after_rope" + assert drift.provenance["fusion_boundary"] == "fused_rope_attention" + + payload = report.to_dict() + assert payload["drifts"][0]["post_rope_q"]["active_count"] == base.q.numel() + json.dumps(payload) + + +def test_single_gpu_rope_attention_requires_position_metadata(): + base = _comparison_inputs() + + with pytest.raises(ValueError, match="rope_positions are required"): + compare_single_gpu_rope_attention(AttentionComparisonInputs(q=base.q, k=base.k, v=base.v)) + + +def test_decode_replay_matches_full_prefill_for_single_and_few_query(): + inputs = _decode_inputs() + report = compare_decode_kv_replay(inputs) + + assert report.reference_name == "full_prefill_decode_reference" + drift = report.drifts[0] + assert drift.candidate_name == "rl_kernel_decode_kv_replay" + assert drift.out.max_abs <= 1.0e-6 + assert drift.lse.max_abs <= 1.0e-6 + assert drift.dlogp is not None + assert drift.dlogp.max_abs <= 3.0e-6 + assert drift.dlogp.active_count == 3 + assert drift.provenance["attention_mode"] == "decode" + assert drift.provenance["cache_position"] == [[4, 5], [4, 5]] + assert drift.provenance["cp_block_owners"] == [[0, 1, 0], [0, 1, 0]] + assert drift.provenance["merge_order"] == "global_block_index" + assert drift.provenance["accum_dtype"] == "fp32" + assert drift.provenance["lse_dtype"] == "fp32" + assert drift.provenance["downcast_at"] == "final_write" + assert drift.provenance["output_dtype"] == "float32" + assert drift.provenance["scale"] == pytest.approx(1.0 / (8.0**0.5)) + assert drift.provenance["q_dtype"] == "float32" + assert drift.provenance["k_cache_dtype"] == "float32" + assert drift.provenance["v_cache_dtype"] == "float32" + assert drift.provenance["execution_scope"] == "single_device_logical_reference" + assert drift.provenance["runtime_verified"] is False + assert drift.provenance["supports_backward"] is False + assert drift.provenance["communication"] == "none" + assert drift.provenance["cache_execution_fingerprint"] == decode_kv_cache_fingerprint(inputs) + assert drift.provenance["kv_cache_identity_fingerprint"] == decode_kv_cache_fingerprint(inputs) + assert drift.provenance["rope_identity_fingerprint"] == decode_rope_identity_fingerprint(inputs) + assert drift.provenance["preprocess_policy"] == "reference_only_not_production" + assert drift.provenance["preprocess_fallback"] is True + assert drift.provenance["logical_merge_orders"] == [ + [[0, 1, 2], [0, 1, 2]], + [[0, 1, 2], [0, 1, 2]], + ] + json.dumps(report.to_dict()) + + single_query = DecodeAttentionInputs( + q=inputs.q[:, :, -1:, :], + k_cache=inputs.k_cache, + v_cache=inputs.v_cache, + metadata=DecodeKVCacheMetadata( + cache_position=inputs.metadata.cache_position[:, -1:], + kv_seq_lens=inputs.metadata.kv_seq_lens, + block_table=inputs.metadata.block_table, + global_token_positions=inputs.metadata.global_token_positions, + query_position_ids=inputs.metadata.query_position_ids[:, -1:], + key_position_ids=inputs.metadata.key_position_ids, + page_size=inputs.metadata.page_size, + cp_block_owners=inputs.metadata.cp_block_owners, + ), + ) + single_report = compare_decode_kv_replay(single_query) + assert single_report.drifts[0].out.max_abs <= 1.0e-6 + assert single_report.drifts[0].lse.max_abs <= 1.0e-6 + + +def test_strict_decode_replay_is_bitwise_across_logical_and_paged_materialization(): + report = compare_decode_kv_replay(_decode_inputs(page_order=(2, 0, 1)), strict_bitwise=True) + assert report.reference_name == "strict_shared_core_logical_prefill" + assert len(report.drifts) == 1 + drift = report.drifts[0] + assert drift.candidate_name == "strict_shared_core_paged_kv_layout" + assert drift.out.max_abs == 0.0 + assert drift.lse.max_abs == 0.0 + assert drift.provenance["strict_core_id"] == ( + "rlkernel.attention.deterministic_core.v1" + ) + assert drift.provenance["strict_schedule"] == ( + "single_batch_single_query_global_kv_blocks" + ) + assert drift.provenance["split_kv_policy"] == "disabled" + + +def test_decode_replay_is_invariant_to_physical_page_and_prefix_layout(): + contiguous = run_decode_kv_replay(_decode_inputs()) + permuted_inputs = _decode_inputs(page_order=(2, 0, 1), prefix_cache_enabled=True) + permuted = run_decode_kv_replay(permuted_inputs) + + torch.testing.assert_close(permuted.out, contiguous.out, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(permuted.lse, contiguous.lse, atol=1.0e-6, rtol=0.0) + assert permuted.provenance["prefix_cache_enabled"] is True + assert permuted.provenance["prefix_cache_key"] == "shared-prefix" + assert permuted.provenance["prefix_length"] == 4 + assert permuted.provenance["prefix_cache_fingerprint"] == decode_prefix_cache_fingerprint( + permuted_inputs, + prefix_length=permuted_inputs.metadata.prefix_length, + ) + + +def test_decode_replay_supports_nonzero_global_position_origin(): + base = compare_decode_kv_replay(_decode_inputs()) + offset = compare_decode_kv_replay(_decode_inputs(position_offset=128)) + + assert offset.drifts[0].out.max_abs <= 1.0e-6 + assert offset.drifts[0].lse.max_abs <= 1.0e-6 + assert base.drifts[0].out.max_abs <= 1.0e-6 + + +def test_decode_cache_execution_fingerprint_is_layout_invariant_and_content_bound(): + contiguous = _decode_inputs() + permuted = _decode_inputs(page_order=(2, 0, 1)) + + assert decode_kv_cache_fingerprint(contiguous) == decode_kv_cache_fingerprint(permuted) + assert run_decode_kv_replay(contiguous).provenance[ + "cache_execution_fingerprint" + ] == decode_kv_cache_fingerprint(contiguous) + + changed_v = contiguous.v_cache.clone() + first_slot = int(contiguous.metadata.block_table[0, 0].item()) * contiguous.metadata.page_size + changed_v[0, 0, first_slot, 0] += 1.0 + changed = replace(contiguous, v_cache=changed_v) + assert decode_kv_cache_fingerprint(changed) != decode_kv_cache_fingerprint(contiguous) + + contiguous_prefix = _decode_inputs(prefix_cache_enabled=True) + permuted_prefix = _decode_inputs(page_order=(2, 0, 1), prefix_cache_enabled=True) + assert contiguous_prefix.metadata.prefix_cache_fingerprint == ( + permuted_prefix.metadata.prefix_cache_fingerprint + ) + + +def test_decode_cache_fingerprint_ignores_inactive_physical_page_content(): + inputs = _decode_inputs() + generator = torch.Generator().manual_seed(91) + extra_k = torch.randn( + inputs.k_cache.size(0), + inputs.k_cache.size(1), + inputs.metadata.page_size, + inputs.k_cache.size(3), + generator=generator, + ) + extra_v = torch.randn(extra_k.shape, generator=generator) + extended = replace( + inputs, + k_cache=torch.cat((inputs.k_cache, extra_k), dim=2), + v_cache=torch.cat((inputs.v_cache, extra_v), dim=2), + metadata=replace( + inputs.metadata, + block_table=torch.cat( + (inputs.metadata.block_table, torch.full((2, 1), -1, dtype=torch.long)), + dim=1, + ), + global_token_positions=torch.cat( + ( + inputs.metadata.global_token_positions, + torch.full((2, inputs.metadata.page_size), -1, dtype=torch.long), + ), + dim=1, + ), + key_position_ids=torch.cat( + ( + inputs.metadata.key_position_ids, + torch.full((2, inputs.metadata.page_size), -1, dtype=torch.long), + ), + dim=1, + ), + cp_block_owners=torch.cat( + ( + inputs.metadata.cp_block_owners, + torch.full((2, 1), -1, dtype=torch.long), + ), + dim=1, + ), + ), + ) + changed_inactive = replace( + extended, + k_cache=extended.k_cache.clone(), + v_cache=extended.v_cache.clone(), + ) + changed_inactive.k_cache[:, :, -2:, :].fill_(123.0) + changed_inactive.v_cache[:, :, -2:, :].fill_(-456.0) + + assert decode_kv_cache_fingerprint(extended) == decode_kv_cache_fingerprint(changed_inactive) + extended_result = run_decode_kv_replay(extended) + changed_result = run_decode_kv_replay(changed_inactive) + torch.testing.assert_close(extended_result.out, changed_result.out) + torch.testing.assert_close(extended_result.lse, changed_result.lse) + + +def test_decode_replay_supports_per_batch_lengths_and_partial_final_page(): + inputs = _decode_inputs() + metadata = replace( + inputs.metadata, + cache_position=torch.tensor([[3, 4], [4, 5]], dtype=torch.long), + query_position_ids=torch.tensor([[3, 4], [4, 5]], dtype=torch.long), + kv_seq_lens=torch.tensor([5, 6], dtype=torch.long), + global_token_positions=inputs.metadata.global_token_positions.clone(), + key_position_ids=inputs.metadata.key_position_ids.clone(), + ) + first_unused_slot = int(metadata.block_table[0, 2].item()) * metadata.page_size + 1 + metadata.global_token_positions[0, first_unused_slot] = -1 + metadata.key_position_ids[0, first_unused_slot] = -1 + shortened = replace(inputs, metadata=metadata) + + report = compare_decode_kv_replay(shortened) + + assert report.drifts[0].out.max_abs <= 1.0e-6 + assert report.drifts[0].lse.max_abs <= 1.0e-6 + assert report.drifts[0].provenance["kv_seq_lens"] == [5, 6] + + +@pytest.mark.parametrize("field_name", ["q", "k_cache", "v_cache"]) +def test_decode_replay_rejects_non_finite_active_values(field_name): + inputs = _decode_inputs() + tensor = getattr(inputs, field_name).clone() + tensor[0, 0, 0, 0] = float("nan") + + with pytest.raises(ValueError, match="must be finite"): + run_decode_kv_replay(replace(inputs, **{field_name: tensor})) + + +def test_decode_replay_rejects_finite_inputs_that_overflow_attention(): + inputs = _decode_inputs() + huge = torch.full_like(inputs.q, torch.finfo(inputs.q.dtype).max) + + with pytest.raises(ValueError, match="produced non-finite"): + compare_decode_kv_replay(replace(inputs, q=huge)) + + +def test_decode_replay_rejects_prefix_cache_key_drift(): + inputs = _decode_inputs(prefix_cache_enabled=True) + changed_key = replace( + inputs, + metadata=replace(inputs.metadata, prefix_cache_key="different-prefix"), + ) + + with pytest.raises(ValueError, match="prefix_cache_fingerprint"): + run_decode_kv_replay(changed_key) + + +def test_decode_replay_is_invariant_to_equivalent_cp_block_ownership(): + cp2_inputs = _decode_inputs() + cp1_inputs = replace( + cp2_inputs, + metadata=replace( + cp2_inputs.metadata, + cp_block_owners=torch.zeros_like(cp2_inputs.metadata.cp_block_owners), + ), + ) + + cp1 = run_decode_kv_replay(cp1_inputs) + cp2 = run_decode_kv_replay(cp2_inputs) + torch.testing.assert_close(cp2.out, cp1.out, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(cp2.lse, cp1.lse, atol=1.0e-6, rtol=0.0) + assert cp1.provenance["cp_block_owners"] == [[0, 0, 0], [0, 0, 0]] + assert cp2.provenance["cp_block_owners"] == [[0, 1, 0], [0, 1, 0]] + + +def test_decode_replay_pre_rope_cache_matches_equivalent_post_rope_cache(): + pre_rope = _decode_inputs(q_rope_state="pre_rope", k_cache_rope_state="pre_rope") + rope = NativeRoPEOp() + post_q = rope.forward_fp32( + pre_rope.q, + pre_rope.metadata.query_position_ids, + theta=pre_rope.rope_theta, + ) + post_k = rope.forward_fp32( + pre_rope.k_cache, + pre_rope.metadata.key_position_ids, + theta=pre_rope.rope_theta, + ) + post_rope = replace( + pre_rope, + q=post_q, + k_cache=post_k, + metadata=replace( + pre_rope.metadata, + q_rope_state="post_rope", + k_cache_rope_state="post_rope", + ), + ) + + pre_result = run_decode_kv_replay(pre_rope) + post_result = run_decode_kv_replay(post_rope) + torch.testing.assert_close(post_result.out, pre_result.out, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(post_result.lse, pre_result.lse, atol=1.0e-6, rtol=0.0) + + +def test_decode_replay_preserves_separate_q_and_k_rope_output_dtypes(): + base = _decode_inputs(q_rope_state="pre_rope", k_cache_rope_state="pre_rope") + mixed = replace( + base, + k_cache=base.k_cache.to(torch.bfloat16), + v_cache=base.v_cache.to(torch.bfloat16), + ) + rope = NativeRoPEOp() + post = replace( + mixed, + q=rope.forward_fp32( + mixed.q, + mixed.metadata.query_position_ids, + theta=mixed.rope_theta, + ).to(torch.float32), + k_cache=rope.forward_fp32( + mixed.k_cache, + mixed.metadata.key_position_ids, + theta=mixed.rope_theta, + ).to(torch.bfloat16), + metadata=replace( + mixed.metadata, + q_rope_state="post_rope", + k_cache_rope_state="post_rope", + ), + ) + + pre_result = run_decode_kv_replay(mixed) + post_result = run_decode_kv_replay(post) + torch.testing.assert_close(post_result.out, pre_result.out, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(post_result.lse, pre_result.lse, atol=1.0e-6, rtol=0.0) + assert pre_result.provenance["q_rope_output_dtype"] == "float32" + assert pre_result.provenance["k_cache_rope_output_dtype"] == "bfloat16" + + +def test_decode_replay_rejects_stale_prefix_cache_content(): + inputs = _decode_inputs(prefix_cache_enabled=True) + stale_k = inputs.k_cache.clone() + first_prefix_slot = int(inputs.metadata.block_table[0, 0].item()) * inputs.metadata.page_size + stale_k[0, 0, first_prefix_slot, 0] += 1.0 + + with pytest.raises(ValueError, match="prefix_cache_fingerprint"): + run_decode_kv_replay(replace(inputs, k_cache=stale_k)) + + +def test_decode_replay_rejects_prefix_cache_theta_drift(): + inputs = _decode_inputs(prefix_cache_enabled=True) + + with pytest.raises(ValueError, match="prefix_cache_fingerprint"): + run_decode_kv_replay(replace(inputs, rope_theta=10_000.0)) + + +def test_decode_replay_rejects_prefix_cache_storage_dtype_drift(): + inputs = _decode_inputs(prefix_cache_enabled=True) + + with pytest.raises(ValueError, match="prefix_cache_fingerprint"): + run_decode_kv_replay(replace(inputs, k_cache=inputs.k_cache.to(torch.bfloat16))) + + +def test_decode_replay_rejects_unsupported_rope_cast_boundary(): + inputs = _decode_inputs(prefix_cache_enabled=True) + + with pytest.raises(ValueError, match="rope_cast_at"): + run_decode_kv_replay(replace(inputs, rope_cast_at="before_rope")) + + +@pytest.mark.parametrize("theta", [float("nan"), float("inf"), float("-inf")]) +def test_decode_replay_rejects_non_finite_rope_theta(theta): + with pytest.raises(ValueError, match="rope_theta"): + run_decode_kv_replay(replace(_decode_inputs(), rope_theta=theta)) + + +def test_decode_replay_rejects_noncanonical_inactive_page_metadata(): + inputs = _decode_inputs() + bad_metadata = replace( + inputs.metadata, + block_table=torch.cat( + [inputs.metadata.block_table, torch.zeros((2, 1), dtype=torch.long)], dim=1 + ), + cp_block_owners=torch.cat( + [inputs.metadata.cp_block_owners, torch.full((2, 1), -1, dtype=torch.long)], dim=1 + ), + ) + + with pytest.raises(ValueError, match="unused block_table"): + run_decode_kv_replay(replace(inputs, metadata=bad_metadata)) + + +def test_decode_replay_rejects_negative_active_cp_owner(): + inputs = _decode_inputs() + bad_owners = inputs.metadata.cp_block_owners.clone() + bad_owners[0, 1] = -1 + + with pytest.raises(ValueError, match="active cp_block_owners"): + run_decode_kv_replay( + replace(inputs, metadata=replace(inputs.metadata, cp_block_owners=bad_owners)) + ) + + +def test_decode_replay_rejects_noncanonical_inactive_cp_owner(): + inputs = _decode_inputs() + bad_metadata = replace( + inputs.metadata, + block_table=torch.cat( + [inputs.metadata.block_table, torch.full((2, 1), -1, dtype=torch.long)], dim=1 + ), + cp_block_owners=torch.cat( + [inputs.metadata.cp_block_owners, torch.zeros((2, 1), dtype=torch.long)], dim=1 + ), + ) + + with pytest.raises(ValueError, match="unused cp_block_owners"): + run_decode_kv_replay(replace(inputs, metadata=bad_metadata)) + + +@pytest.mark.parametrize("scale", [float("nan"), float("inf"), 0.0, -1.0]) +def test_decode_replay_rejects_invalid_scale(scale): + with pytest.raises(ValueError, match="scale"): + run_decode_kv_replay(replace(_decode_inputs(), scale=scale)) + + +def test_decode_replay_rejects_non_floating_cache_dtypes(): + inputs = _decode_inputs() + with pytest.raises(ValueError, match="floating-point dtypes"): + run_decode_kv_replay(replace(inputs, q=inputs.q.to(torch.int32))) + + +def test_decode_replay_fails_loudly_on_position_identity_mismatch(): + inputs = _decode_inputs() + bad_query_positions = inputs.metadata.query_position_ids.clone() + bad_query_positions[0, -1] = 4 + bad_metadata = DecodeKVCacheMetadata( + cache_position=inputs.metadata.cache_position, + kv_seq_lens=inputs.metadata.kv_seq_lens, + block_table=inputs.metadata.block_table, + global_token_positions=inputs.metadata.global_token_positions, + query_position_ids=bad_query_positions, + key_position_ids=inputs.metadata.key_position_ids, + page_size=inputs.metadata.page_size, + cp_block_owners=inputs.metadata.cp_block_owners, + ) + + with pytest.raises(ValueError, match="cache_position and query_position_ids"): + run_decode_kv_replay( + DecodeAttentionInputs( + q=inputs.q, + k_cache=inputs.k_cache, + v_cache=inputs.v_cache, + metadata=bad_metadata, + ) + ) + + +def test_decode_replay_fails_loudly_on_invalid_page_identity(): + inputs = _decode_inputs() + bad_positions = inputs.metadata.global_token_positions.clone() + bad_positions[:, 0] = 1 + bad_metadata = DecodeKVCacheMetadata( + cache_position=inputs.metadata.cache_position, + kv_seq_lens=inputs.metadata.kv_seq_lens, + block_table=inputs.metadata.block_table, + global_token_positions=bad_positions, + query_position_ids=inputs.metadata.query_position_ids, + key_position_ids=bad_positions.clone(), + page_size=inputs.metadata.page_size, + cp_block_owners=inputs.metadata.cp_block_owners, + ) + + with pytest.raises(ValueError, match="reconstruct logical positions"): + compare_decode_kv_replay( + DecodeAttentionInputs( + q=inputs.q, + k_cache=inputs.k_cache, + v_cache=inputs.v_cache, + metadata=bad_metadata, + ) + ) + + +def test_decode_replay_covers_qwen3_gqa_head_layout(): + generator = torch.Generator().manual_seed(23) + q = torch.randn(1, 32, 1, 128, generator=generator, dtype=torch.bfloat16) + k = torch.randn(1, 8, 4, 128, generator=generator, dtype=torch.bfloat16) + v = torch.randn(1, 8, 4, 128, generator=generator, dtype=torch.bfloat16) + positions = torch.arange(4, dtype=torch.long).unsqueeze(0) + inputs = DecodeAttentionInputs( + q=q, + k_cache=k, + v_cache=v, + metadata=DecodeKVCacheMetadata( + cache_position=torch.tensor([[3]], dtype=torch.long), + kv_seq_lens=torch.tensor([4], dtype=torch.long), + block_table=torch.tensor([[0, 1]], dtype=torch.long), + global_token_positions=positions, + query_position_ids=torch.tensor([[3]], dtype=torch.long), + key_position_ids=positions.clone(), + page_size=2, + cp_block_owners=torch.tensor([[0, 1]], dtype=torch.long), + ), + output_dtype=torch.bfloat16, + ) + + report = compare_decode_kv_replay(inputs) + assert report.drifts[0].out.max_abs <= 2 * torch.finfo(torch.bfloat16).eps + assert report.drifts[0].lse.max_abs <= 1.0e-6 + + +def test_decode_transformer_engine_oracle_reuses_sorted_partial_states(monkeypatch): + calls = {"lse": 0, "out": 0} + + def lse_correction(softmax_lse, softmax_lse_per_step): + calls["lse"] += 1 + softmax_lse.copy_(torch.logaddexp(softmax_lse, softmax_lse_per_step)) + + def out_correction_init(out_init_step, softmax_lse, softmax_lse_init_step, seq_dim): + scale = torch.exp(softmax_lse_init_step - softmax_lse).movedim(2, seq_dim) + return out_init_step * scale.unsqueeze(-1) + + def out_correction(out, out_per_step, softmax_lse, softmax_lse_per_step, seq_dim): + calls["out"] += 1 + scale = torch.exp(softmax_lse_per_step - softmax_lse).movedim(2, seq_dim) + out.add_(out_per_step * scale.unsqueeze(-1)) + + monkeypatch.setitem( + sys.modules, + _TE_CONTEXT_PARALLEL_MODULE, + types.SimpleNamespace( + flash_attn_fwd_softmax_lse_correction=lse_correction, + flash_attn_fwd_out_correction_init=out_correction_init, + flash_attn_fwd_out_correction=out_correction, + ), + ) + + report = compare_decode_kv_replay( + _decode_inputs(page_order=(2, 0, 1)), + include_transformer_engine=True, + ) + + by_name = {drift.candidate_name: drift for drift in report.drifts} + assert set(by_name) == { + "rl_kernel_decode_kv_replay", + "transformer_engine_decode_kv_replay", + } + assert by_name["transformer_engine_decode_kv_replay"].out.max_abs <= 1.0e-6 + assert by_name["transformer_engine_decode_kv_replay"].lse.max_abs <= 1.0e-6 + assert by_name["transformer_engine_decode_kv_replay"].provenance["logical_merge_orders"] == [ + [[0, 1, 2], [0, 1, 2]], + [[0, 1, 2], [0, 1, 2]], + ] + assert calls["lse"] > 0 + assert calls["out"] > 0 + assert report.unavailable == () + + +def test_decode_transformer_engine_unavailable_is_reported(monkeypatch): + real_import_module = importlib.import_module + + def fake_import_module(name, package=None): + if name == _TE_CONTEXT_PARALLEL_MODULE: + raise ImportError("decode TE unavailable") + return real_import_module(name, package) + + monkeypatch.setattr(importlib, "import_module", fake_import_module) + + report = compare_decode_kv_replay( + _decode_inputs(), + include_transformer_engine=True, + ) + + assert {drift.candidate_name for drift in report.drifts} == {"rl_kernel_decode_kv_replay"} + assert report.unavailable == ("transformer_engine_decode_kv_replay: decode TE unavailable",) + + +def test_transformer_engine_merge_oracle_can_be_reused_when_available(monkeypatch): + calls = {"lse": 0, "out": 0} + + def lse_correction(softmax_lse, softmax_lse_per_step): + calls["lse"] += 1 + softmax_lse.copy_(torch.logaddexp(softmax_lse, softmax_lse_per_step)) + + def out_correction_init(out_init_step, softmax_lse, softmax_lse_init_step, seq_dim): + scale = torch.exp(softmax_lse_init_step - softmax_lse).movedim(2, seq_dim) + return out_init_step * scale.unsqueeze(-1) + + def out_correction(out, out_per_step, softmax_lse, softmax_lse_per_step, seq_dim): + calls["out"] += 1 + scale = torch.exp(softmax_lse_per_step - softmax_lse).movedim(2, seq_dim) + out.add_(out_per_step * scale.unsqueeze(-1)) + + monkeypatch.setitem( + sys.modules, + _TE_CONTEXT_PARALLEL_MODULE, + types.SimpleNamespace( + flash_attn_fwd_softmax_lse_correction=lse_correction, + flash_attn_fwd_out_correction_init=out_correction_init, + flash_attn_fwd_out_correction=out_correction, + ), + ) + + report = compare_single_gpu_attention( + _comparison_inputs(), + query_chunk_size=3, + kv_page_size=2, + include_transformer_engine=True, + ) + + by_name = {drift.candidate_name: drift for drift in report.drifts} + assert "transformer_engine_paged_kv" in by_name + assert by_name["transformer_engine_paged_kv"].out.max_abs <= 1.0e-6 + assert by_name["transformer_engine_paged_kv"].lse.max_abs <= 1.0e-6 + provenance = by_name["transformer_engine_paged_kv"].provenance + assert provenance["te_available"] is True + assert provenance["te_module"] == _TE_CONTEXT_PARALLEL_MODULE + assert provenance["te_capability_probe"] == "passed" + assert provenance["te_signature_checked"] is True + assert provenance["te_numeric_selftest"] == "passed" + assert provenance["actual_backend"] == "te_context_parallel_merge_helpers" + assert provenance["actual_backend_source"] == "rl_kernel_te_context_parallel_adapter" + assert provenance["accum_dtype"] == "fp32" + assert provenance["downcast_at"] == "final_write" + assert calls["lse"] > 0 + assert calls["out"] > 0 + assert report.unavailable == () + + +def test_transformer_engine_merge_oracle_keeps_all_masked_rows_stable(monkeypatch): + def lse_correction(softmax_lse, softmax_lse_per_step): + max_scale = torch.max(softmax_lse, softmax_lse_per_step) + min_scale = torch.min(softmax_lse, softmax_lse_per_step) + softmax_lse.copy_(max_scale + torch.log1p(torch.exp(min_scale - max_scale))) + + def out_correction_init(out_init_step, softmax_lse, softmax_lse_init_step, seq_dim): + scale = torch.exp(softmax_lse_init_step - softmax_lse).movedim(2, seq_dim) + return out_init_step * scale.unsqueeze(-1) + + def out_correction(out, out_per_step, softmax_lse, softmax_lse_per_step, seq_dim): + scale = torch.exp(softmax_lse_per_step - softmax_lse).movedim(2, seq_dim) + out.add_(out_per_step * scale.unsqueeze(-1)) + + monkeypatch.setitem( + sys.modules, + _TE_CONTEXT_PARALLEL_MODULE, + types.SimpleNamespace( + flash_attn_fwd_softmax_lse_correction=lse_correction, + flash_attn_fwd_out_correction_init=out_correction_init, + flash_attn_fwd_out_correction=out_correction, + ), + ) + + q, k, v = _qkv(seed=11) + inputs = AttentionComparisonInputs( + q=q, + k=k, + v=v, + causal=False, + key_padding_mask=torch.zeros(q.size(0), k.size(2), dtype=torch.bool), + ) + + te_result = run_paged_kv_attention( + inputs, + kv_page_size=2, + merge_backend="transformer_engine", + ) + assert torch.equal(te_result.out, torch.zeros_like(te_result.out)) + assert torch.isneginf(te_result.lse).all() + + report = compare_single_gpu_attention( + inputs, + query_chunk_size=3, + kv_page_size=2, + include_transformer_engine=True, + ) + by_name = {drift.candidate_name: drift for drift in report.drifts} + assert by_name["transformer_engine_paged_kv"].out.max_abs == 0.0 + assert by_name["transformer_engine_paged_kv"].lse.max_abs == 0.0 + assert report.unavailable == () + + +def test_transformer_engine_path_reports_unavailable_without_failing(monkeypatch): + real_import_module = importlib.import_module + + def fake_import_module(name, package=None): + if name == _TE_CONTEXT_PARALLEL_MODULE: + raise ImportError("test TE unavailable") + return real_import_module(name, package) + + monkeypatch.setattr(importlib, "import_module", fake_import_module) + + report = compare_single_gpu_attention( + _comparison_inputs(), + query_chunk_size=3, + kv_page_size=2, + include_transformer_engine=True, + ) + + assert {drift.candidate_name for drift in report.drifts} == { + "chunked_prefill", + "rl_kernel_paged_kv", + } + assert report.unavailable == ("transformer_engine_paged_kv: test TE unavailable",) + + +def test_transformer_engine_path_reports_missing_helpers(monkeypatch): + monkeypatch.setitem( + sys.modules, + _TE_CONTEXT_PARALLEL_MODULE, + types.SimpleNamespace( + flash_attn_fwd_softmax_lse_correction=lambda softmax_lse, per_step: None, + ), + ) + + report = compare_single_gpu_attention( + _comparison_inputs(), + query_chunk_size=3, + kv_page_size=2, + include_transformer_engine=True, + ) + + assert len(report.unavailable) == 1 + assert "missing required helpers" in report.unavailable[0] + + +def test_transformer_engine_path_reports_incompatible_helper_signature(monkeypatch): + def lse_correction(wrong_name, softmax_lse_per_step): + wrong_name.copy_(torch.logaddexp(wrong_name, softmax_lse_per_step)) + + def out_correction_init(out_init_step, softmax_lse, softmax_lse_init_step, seq_dim): + scale = torch.exp(softmax_lse_init_step - softmax_lse).movedim(2, seq_dim) + return out_init_step * scale.unsqueeze(-1) + + def out_correction(out, out_per_step, softmax_lse, softmax_lse_per_step, seq_dim): + scale = torch.exp(softmax_lse_per_step - softmax_lse).movedim(2, seq_dim) + out.add_(out_per_step * scale.unsqueeze(-1)) + + monkeypatch.setitem( + sys.modules, + _TE_CONTEXT_PARALLEL_MODULE, + types.SimpleNamespace( + flash_attn_fwd_softmax_lse_correction=lse_correction, + flash_attn_fwd_out_correction_init=out_correction_init, + flash_attn_fwd_out_correction=out_correction, + ), + ) + + report = compare_single_gpu_attention( + _comparison_inputs(), + query_chunk_size=3, + kv_page_size=2, + include_transformer_engine=True, + ) + + assert len(report.unavailable) == 1 + assert "incompatible signature" in report.unavailable[0] + + +def test_operator_comparison_specs_register_attention(): + args = argparse.Namespace( + op="attention", + candidate="pytorch", + arch_key=None, + batch=1, + seq=3, + vocab=17, + seed=7, + input_mode="random", + constant_value=0.5, + token_value=3, + normalized_dim=128, + k_dim=16, + n_dim=32, + theta=1.0e6, + eps=1.0e-6, + ) + + case = make_operator_case(args, torch.float32, torch.device("cpu")) + candidate = make_candidate(args) + report = run_operator_suite("attention", candidates=[candidate], cases=[case]) + + assert report.passed + assert report.candidates[0].cases[0].op_class == "attention"