feat(attention): add single-gpu comparison harness

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.

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 语义契约后，才允许作为生产候选后端。

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",
 ]