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 42be8c1b..792b36d0 100644 --- a/rl_engine/testing/__init__.py +++ b/rl_engine/testing/__init__.py @@ -3,6 +3,22 @@ """Testing helpers for RL-shaped kernel validation.""" +from .attention_comparison import ( + AttentionComparisonInputs, + AttentionComparisonReport, + AttentionPathDrift, + AttentionPathResult, + DriftStats, + TransformerEngineUnavailable, + compare_single_gpu_attention, + compare_single_gpu_rope_attention, + run_chunked_query_attention, + 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, @@ -15,13 +31,27 @@ from .rl_batch import SyntheticRLKernelBatch, make_synthetic_rl_kernel_batch __all__ = [ + "AttentionComparisonInputs", + "AttentionComparisonReport", + "AttentionPathDrift", + "AttentionPathResult", + "DriftStats", "SyntheticRLKernelBatch", + "TransformerEngineUnavailable", "active_token_count", + "compare_single_gpu_rope_attention", + "compare_single_gpu_attention", "compute_policy_ratio", "compute_reference_kl", "make_synthetic_rl_kernel_batch", "masked_mean", "masked_sum", + "run_chunked_query_attention", + "run_fused_like_rope_attention", + "run_full_attention", + "run_paged_kv_attention", + "run_unfused_rope_attention", "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..dbee69af --- /dev/null +++ b/rl_engine/testing/attention_comparison.py @@ -0,0 +1,907 @@ +# 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 importlib +import importlib.metadata as importlib_metadata +import inspect +import math +from dataclasses import dataclass +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"] + +_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 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 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 _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 _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 _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 _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 _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): + raise ValueError("rope_theta must be a positive number") + if 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_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] + 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_start < previous_end: + raise ValueError("partial state block ranges must not overlap") + previous_end = state.block_end + + +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 value <= 0: + raise ValueError(f"{name} must be a positive integer") + return int(value) + + +__all__ = [ + "AttentionComparisonInputs", + "AttentionComparisonReport", + "AttentionPathDrift", + "AttentionPathResult", + "DriftStats", + "TransformerEngineUnavailable", + "compare_single_gpu_rope_attention", + "compare_single_gpu_attention", + "run_chunked_query_attention", + "run_fused_like_rope_attention", + "run_full_attention", + "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..061e8edd --- /dev/null +++ b/tests/test_attention_comparison.py @@ -0,0 +1,364 @@ +# 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 + +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.testing.attention_comparison import ( + AttentionComparisonInputs, + compare_single_gpu_attention, + compare_single_gpu_rope_attention, + 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 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_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"