diff --git a/ci/run_gpu_ci.sh b/ci/run_gpu_ci.sh index b8bdca1f..e3ec7724 100644 --- a/ci/run_gpu_ci.sh +++ b/ci/run_gpu_ci.sh @@ -197,12 +197,17 @@ TORCH_INDEX_URL="${TORCH_INDEX_URL:-https://download.pytorch.org/whl/cu124}" # --no-build-isolation: torch must be visible to setup.py, else the extension is silently skipped. # --no-deps: keep the pinned torch; do not let the editable install re-resolve it. "$PY" -m pip install --no-build-isolation --no-deps -e . -"$PY" -m pip install --no-cache-dir numpy tabulate accelerate "transformers==5.13.1" pytest triton +"$PY" -m pip install --no-cache-dir numpy tabulate accelerate "transformers==5.13.1" pytest triton "flashinfer-python>=0.6.0,<0.7" nvidia-smi # Fail fast if _C did not build or cannot launch, instead of silently using native fallbacks. "$PY" scripts/ci_smoke.py # Enforce _C in the pytest suite too (test_extension_smoke.py skips unless this is set). export RL_KERNEL_REQUIRE_EXT=1 +"$PY" -m pytest tests/test_flashinfer_pr7_attention.py -q +"$PY" scripts/ws2_pr7_flashinfer_attention_check.py --no-dry-run --device cuda --mode decode --split-kv-policy disabled --output artifacts/pr7-decode-disabled.json +"$PY" scripts/ws2_pr7_flashinfer_attention_check.py --no-dry-run --device cuda --mode decode --split-kv-policy fixed --fixed-split-size 4 --output artifacts/pr7-decode-fixed.json +"$PY" scripts/ws2_pr7_flashinfer_attention_check.py --no-dry-run --device cuda --mode prefill --query-len 4 --split-kv-policy disabled --output artifacts/pr7-prefill-disabled.json +"$PY" scripts/ws2_pr7_flashinfer_attention_check.py --no-dry-run --device cuda --mode prefill --query-len 4 --split-kv-policy fixed --fixed-split-size 4 --output artifacts/pr7-prefill-fixed.json export WS1_C8_JSON=/tmp/ws1-c8-ci.json export WS1_WEIGHTS_PATH="'"${WS1_WEIGHTS_PATH:-}"'" if [ "$WS1_TEST_SUITE" = "ws1-chain" ]; then diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda index 1617d184..51a4f4eb 100644 --- a/docker/Dockerfile.cuda +++ b/docker/Dockerfile.cuda @@ -15,6 +15,7 @@ COPY pyproject.toml setup.py* requirements*.txt ./ RUN pip install --no-cache-dir -U pip \ && pip install --no-cache-dir -r requirements.txt \ + && pip install --no-cache-dir "flashinfer-python>=0.6.0,<0.7" \ && pip install --no-cache-dir pytest WORKDIR /workspace diff --git a/docs/design/ws2-attention-pr7-flashinfer-rope-splitk.md b/docs/design/ws2-attention-pr7-flashinfer-rope-splitk.md new file mode 100644 index 00000000..ad7ea13f --- /dev/null +++ b/docs/design/ws2-attention-pr7-flashinfer-rope-splitk.md @@ -0,0 +1,405 @@ +# WS2 PR7 Fused Attention Backend Alignment + +Status: PR7 candidate plus deterministic communication integration for #235 + +## Goal + +PR7 is the production-backend alignment layer after PR2/PR3/PR6 have defined +the reference semantics. It evaluates fused prefill/decode candidates while +keeping RL-Kernel's reference contract as the source of truth. + +The full PR7 plan has two backend lanes: + +| Lane | Role | Status in this scaffold | +| --- | --- | --- | +| Training full-prefill candidate | Evaluate TE public `DotProductAttention` / fused attention for training-style full prefill | Planned; not implemented here | +| Rollout paged prefill/decode candidate | Evaluate FlashInfer paged attention with Qwen3 RoPE fusion, split-KV policy, LSE export, and batch-invariant sweep | Implemented as opt-in scaffold | + +This file therefore documents both the original PR7 requirements and the new +FlashInfer/RoPE/split-KV/batch-invariant additions. It should not be read as a +claim that FlashInfer or TE replaces the deterministic reference. + +## Source Of Truth + +The correctness source remains: + +```text +AttentionContract +PR2 single-GPU full/chunked/paged attention reference +PR3 deterministic CP reference and global_block_index merge +PR6 decode paged-KV replay vs full logical KV reference +attention-domain lse +dlogp drift report +``` + +PR7 backends are candidates. A candidate can be promoted only if it exports or +reconstructs the states required by the contract and passes the drift gates. + +## Original PR7 Requirements + +| Requirement | PR7 rule | +| --- | --- | +| Backend selection | Record `requested_backend`, `actual_backend`, `fallback`, and `fallback_reason`; no silent fallback. | +| Training path | Training-side forward is full-prefill / teacher-forcing. TE may be evaluated through its public `DotProductAttention` path, not by redefining RL-Kernel semantics. | +| Rollout path | Rollout-side forward is paged KV prefill/decode. FlashInfer is a rollout candidate, not a training backward backend. | +| LSE | Candidate must return attention-domain `lse` shaped as `[B, Hq, Sq]`. | +| CP/TP communication | PR7 preserves the P2P NCCL reference and maps the production `cuda_ag_rs` path to the deterministic CUDA AllGather/ReduceScatter operators from PR311/PR312. Missing compiled operators fail closed. | +| CP order | CP correctness remains PR3's `global_block_index` merge. A black-box backend that only returns final output must be validated against PR3/PR6; it cannot define CP merge order. | +| Paged KV | Page table, `cache_position`, `query_position_ids`, `key_position_ids`, and logical token order must be validated before execution. | +| Precision | Record accumulation/downcast policy. BF16 output is acceptable only after FP32/reference drift is reported. | +| Backward | Training backward belongs to PR8. PR7 forward results must not claim `dq/dk/dv` alignment. | + +## Current Implementation + +Implemented files: + +```text +rl_engine/kernels/ops/cuda/attention/cp_comm.py +rl_engine/kernels/ops/cuda/attention/flashinfer_paged_attention.py +tests/test_flashinfer_pr7_attention.py +scripts/ws2_pr7_flashinfer_attention_check.py +``` + +The FlashInfer adapter is opt-in and lazy-imported. Importing RL-Kernel does not +require FlashInfer. Real FlashInfer execution requires CUDA tensors; CPU tests +use a fake wrapper only to validate parameter binding and provenance. + +Current code path: + +```text +FlashInferQwen3PagedAttentionOp.forward(...) + validate q / k_cache / v_cache + validate Qwen3 RoPE fusion config + validate split-KV policy + validate CP=2 / TP=2 AG/RS communication interface contract + if strict_mode and require_cp_comm: + validate owner-local Q/K/V and real position IDs + deterministic AG(Q/K/V/position IDs) + apply WS1 RoPESM90Op row-by-row + run the shared WS1 no-Split-K deterministic Attention core + deterministic RS(Out, LSE) -> local query shard + return without constructing a FlashInfer arithmetic wrapper + build_flashinfer_paged_kv_plan(...) + validate page bounds + validate block_table/global_token_positions logical order + validate key_position_ids + reject duplicate pages and non-canonical inactive metadata + bind metadata q/k RoPE state to the fused-RoPE config + validate cache_position == query_position_ids and trailing query positions + validate prefix-cache fingerprint against K/V content and RoPE identity + materialize_flashinfer_paged_kv_cache(...) + if strict_mode: + materialize logical paged KV and run the shared WS1 deterministic core + else: + flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper + or flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper + wrapper.plan(..., pos_encoding_mode="ROPE_LLAMA", rope_theta=1e6, rope_scale=1.0, ...) + wrapper.run_return_lse(...) + restore out -> [B, Hq, Sq, D] + restore lse -> [B, Hq, Sq] +``` + +## CP/TP Communication Interface + +Issue #235 targets `Qwen3-8B, TP=2, CP=2, BF16`. PR7 surfaces the distributed +attention boundary and reuses the self-owned deterministic CUDA collectives. + +The exposed interface is: + +```text +AttentionParallelSpec(tp_world_size=2, cp_world_size=2) +AttentionCPCommunicationPlan(backend="cuda_ag_rs", status="implemented") +AttentionCPPartialState(out, lse, AttentionCPBlockMetadata(...)) +CUDAAGRSAttentionCPCommunication.all_gather_query(...) +CUDAAGRSAttentionCPCommunication.all_gather_kv(...) +CUDAAGRSAttentionCPCommunication.all_gather_position_ids(...) +CUDAAGRSAttentionCPCommunication.all_gather_partial_states(...) +CUDAAGRSAttentionCPCommunication.reduce_scatter_merged_state(...) +CUDAAGRSAttentionCPCommunication.reduce_scatter_strict_result(...) +sort_attention_cp_partial_states(..., plan=...) +``` + +The CP execution order is explicit and keeps compute and communication +decoupled: + +```text +strict production: + local Q/K/V + position IDs + -> custom CUDA AG(Q/K/V/position IDs) + -> shared full-logical-QKV WS1 deterministic Attention core + -> custom CUDA RS(Out, LSE) + -> local query shard + +reference compatibility path: + local Q shard -> custom AG(Q) + -> owner-local partial AttentionCPState(out, lse, global_block_index) + -> custom AG(partial states) -> PR3 FP32 ordered merge + -> custom RS(merged state) +``` + +`CUDAAGRSAttentionCPCommunication` uses the self-owned deterministic CUDA +collectives from PR311/PR312 for strict Q/K/V/position AG and the final Out/LSE +RS. `P2PNCCLAttentionCPCommunication` implements the same strict boundary as +an independent NCCL reference. The old partial-state methods remain available +for PR3/reference validation, but strict mode never labels a full-KV result as +an owner-local partial or enters native FlashInfer Attention arithmetic. + +Ordering is part of the interface: + +```text +merge_key = AttentionCPBlockMetadata.global_block_index +accum_dtype = fp32 +required_state = (out, lse) +communication_pattern = ag_rs +compute_communication = decoupled +duplicate global_block_index -> error +``` + +Strict production disables Split-KV in the shared CUDA core. This gives CP=1 +and CP=2 the same full-QKV kernel grid and reduction graph. Fixed/auto +Split-KV remains available only on the diagnostic/reference FlashInfer lane; +it cannot claim strict bitwise identity. + +## FlashInfer RoPE Fusion + +The implemented fused boundary is: + +```text +pre-RoPE Q, pre-RoPE K cache, V + -> FlashInfer paged attention with pos_encoding_mode="ROPE_LLAMA" + -> out, attention-domain lse +``` + +The accepted Qwen3 settings are locked to: + +```text +pos_encoding_mode = "ROPE_LLAMA" +rope_theta = 1_000_000.0 +rope_scale = 1.0 +rotary_dim = head_dim +layout = Qwen3 rotate-half / non-interleaved +``` + +The adapter rejects post-RoPE Q/K in both the config and the actual runtime +metadata. It also requires `cache_position == query_position_ids` and trailing +contiguous query positions, because this adapter currently relies on the +wrapper's implicit RoPE positions. That avoids silent double rotation and +position drift. If a later rollout path stores post-RoPE K or supports arbitrary +query positions, it must be represented as a separate capability with explicit +position tensors. + +Prefix-cache identity binds logical K/V content, storage dtypes, cached-K RoPE +state, theta, rotary dim, cast boundary, and output dtype. Equivalent physical +page placement produces the same logical identity; stale content or a changed +RoPE configuration fails before execution. + +## Split-KV Policy + +PR7 exposes split-KV as a contract/provenance knob instead of inheriting +backend defaults: + +| Policy | FlashInfer plan kwargs | Batch-invariant status | +| --- | --- | --- | +| `disabled` | `disable_split_kv=True` | strict candidate | +| `fixed:` | `fixed_split_size=N`, `disable_split_kv=False` | candidate; must pass drift sweep | +| `auto` | `disable_split_kv=False` | rejected when batch invariance is required | + +RL-Kernel's fixed size is measured in logical KV tokens. FlashInfer 0.6 names +`fixed_split_size` in physical pages, so PR7 accepts a fixed token size only +when it is divisible by `page_size`, passes `fixed_split_size / page_size` to +the backend, and converts page boundaries back to logical token boundaries in +the report. Runtime callbacks must declare `split_size_unit="pages"` and +`boundary_unit="pages"`; otherwise strict provenance fails closed. + +The shared WS2 contract now treats Split-KV as a first-class semantic field. PR7 still +must export the backend's actual token boundaries; a requested FlashInfer knob alone is +not sufficient evidence of train/rollout equivalence. +Once PR1/PR4 add a first-class field, the PR7 provenance can be wired into that +field without changing the backend adapter. + +## Batch-Invariant Validation + +PR7 distinguishes "configured for batch invariance" from "proven batch +invariant": + +```text +configured: + split-KV disabled or fixed + no auto backend scheduling in the contract + page/order metadata validated + +proven: + same sample alone vs same sample inside a batch has zero or tolerated drift + out and lse both reported + real CUDA/H-card run, not fake wrapper +``` + +The validation script reports: + +```text +batch_invariant_sweep.method = single_row_vs_same_row_inside_batch +batch_invariant_sweep.out_max_abs +batch_invariant_sweep.lse_max_abs +page_layout_invariant_sweep.out.max_abs +page_layout_invariant_sweep.lse.max_abs +drift.out.{max_abs,mean_abs,p95_abs,p99_abs} +drift.lse.{max_abs,mean_abs,p95_abs,p99_abs} +drift.dlogp.{max_abs,mean_abs,p95_abs,p99_abs} +``` + +FlashInfer is not declared batch-invariant by default. It becomes an accepted +candidate only if the real-hardware sweep passes under the selected split-KV +policy. + +## TE Relationship + +The TE lane is still part of the full PR7 plan, but it is not implemented in +this scaffold. + +| TE use | PR7 position | +| --- | --- | +| TE CP correction helpers | Already used by PR2/PR3/PR6 as optional merge oracle. | +| TE full fused attention | Future PR7 training full-prefill candidate through public `DotProductAttention`. | +| TE internal CP black box | Not a source of RL-Kernel `global_block_index` order unless it exposes compatible partial states or passes reference drift gates. | +| TE backward | PR8 only, and only if compatible saved forward/backward state is available. | + +FlashInfer and TE can coexist in PR7: + +```text +training candidate: + TE DotProductAttention full prefill + +rollout candidate: + FlashInfer ROPE_LLAMA paged prefill/decode + +shared gate: + compare both against RL-Kernel reference states and selected-logprob drift +``` + +## Relation To Existing PRs + +| PR | Relationship | +| --- | --- | +| PR1 | Carries the strict Split-KV policy and capability requirements; runtime adapters still export actual boundaries and fallback provenance. | +| PR2 | Provides full/chunked/paged single-GPU references and fused-like RoPE attribution. PR7 uses the same semantic boundary, but calls a real backend candidate. | +| PR3 | Owns CP=1/2 deterministic semantics and `global_block_index` merge. PR7 does not replace this. | +| PR4 / #263 | Records actual backend and split-KV provenance. PR7 matches that policy and can later wire into PR4 fields. | +| PR5 | Should run the cross matrix once PR7 backends exist: full/chunked/paged, split-KV disabled/fixed, batch shape, TP/CP, dtype. | +| PR6 / #260 | Provides decode paged-KV replay and full logical KV reference. PR7 consumes PR6-style metadata and validates the same page/position identity before FlashInfer execution. | +| PR8 | Owns training backward alignment. PR7 forward-only results are not backward claims. | + +## Conflict Check + +| Topic | Possible conflict | PR7 resolution | +| --- | --- | --- | +| RoPE fusion vs PR2/PR6 post-RoPE references | A fused backend might hide post-RoPE Q/K state. | PR7 records `rope_fusion_boundary`, rejects post-RoPE inputs for this path, and compares final `out/lse` against references. | +| FlashInfer split-KV vs PR3 CP merge order | FlashInfer internal split-KV order is not PR3 `global_block_index`. | Treat split-KV as backend-local reduction. CP merge remains PR3; FlashInfer must pass drift gates and cannot define CP order. | +| CP=2/TP=2 target vs communication ordering | The issue requires actual distributed semantics, not an interface claim. | `CUDAAGRSAttentionCPCommunication` executes PR311/PR312, keeps FP32 partial states, sorts the authoritative manifest by `global_block_index`, and fails closed if the compiled collective is absent. | +| Batch-invariant claim vs backend heuristics | Auto split-KV may depend on batch composition/runtime scheduling. | Reject `auto` when `require_batch_invariant=True`; report disabled/fixed policy explicitly. | +| TE training lane vs FlashInfer rollout lane | Two libraries may have different materialization boundaries. | Both bind to the same RL-Kernel semantic contract and are judged by shared `out/lse/dlogp` reports. | +| Requested policy vs actual plan | A requested fixed/max-splits knob may not equal the runtime token boundaries. | Require actual runtime plan callbacks in strict fixed mode; fail closed when unavailable. | +| Decode vs training full prefill | Training does not own a persistent paged KV cache. | PR7 compares rollout paged KV decode/prefill to training-style full logical KV reference, not to a nonexistent training decode cache. | + +No direct conflict was found between the three new points and the original PR7 +plan. The important restriction is that FlashInfer split-KV/RoPE fusion remains +a candidate path until real CUDA/H-card drift reports prove it. + +## Validation + +CPU-safe structural tests: + +```bash +pytest tests/test_flashinfer_pr7_attention.py -q +``` + +Dry-run plan/provenance checks without CUDA or FlashInfer: + +```bash +python scripts/ws2_pr7_flashinfer_attention_check.py --dry-run --json +python scripts/ws2_pr7_flashinfer_attention_check.py --dry-run --mode prefill --query-len 16 --json +``` + +CUDA/H-card validation once hardware is available: + +```bash +python scripts/ws2_pr7_flashinfer_attention_check.py \ + --no-dry-run \ + --device cuda \ + --mode decode \ + --split-kv-policy disabled \ + --output artifacts/pr7-decode-disabled.json \ + --json + +python scripts/ws2_pr7_flashinfer_attention_check.py \ + --no-dry-run \ + --device cuda \ + --mode prefill \ + --query-len 16 \ + --split-kv-policy fixed \ + --fixed-split-size 4 \ + --output artifacts/pr7-prefill-fixed.json \ + --json +``` + +The CUDA report must include: + +```text +passed / errors +drift.out / drift.lse / drift.dlogp +batch_invariant_sweep +page_layout_invariant_sweep +rope_fusion_boundary +split_kv_policy +actual_split_kv_plans +actual_split_kv_plan_set +arithmetic_semantics_verified +actual_backend +fallback / fallback_reason +``` + +## Non-Claims + +This candidate does not enable FlashInfer by default, does not implement the TE +training lane, does not replace PR3 CP merge, does not prove real H-card +batch-invariance locally, and does not implement training backward. +# P2P NCCL reference and strict arithmetic provenance + +The existing `CUDAAGRSAttentionCPCommunication` interface is preserved and now +uses the self-owned deterministic CUDA AG/RS implementation. It remains +fail-closed when PR311/PR312 or their compiled symbols are absent. For GPU +reference validation, `P2PNCCLAttentionCPCommunication` implements the same partial-state +and Q-AG protocol with `torch.distributed.batch_isend_irecv` on an NCCL process +group. +It requires an authoritative manifest containing every logical KV block, +token range, CP owner, TP owner, and query scatter range. Missing blocks, +gaps, overlaps, wrong owners, incomplete gathers, non-NCCL groups, and CPU +tensors are rejected before merge. + +Strict FlashInfer validation also requires runtime callbacks for both: + +- the actual Split-KV token boundaries for every request; and +- arithmetic provenance declaring FP32 accumulation, FP32 LSE, and downcast + only at the final output write. + +Requested knobs, maximum split counts, or labels such as +`flashinfer_internal` are not accepted as proof. + +The repository pins the PR7 lane to `flashinfer-python>=0.6.0,<0.7`; older +versions do not expose both `fixed_split_size` and `disable_split_kv` plan +knobs. Upstream 0.6 wrappers still do not expose RL-Kernel's three strict +provenance callbacks, so an unpatched stock wrapper is expected to fail closed. +Passing strict acceptance requires a small wrapper/adapter that reads the actual +plan information returned by FlashInfer and reports it through the documented +callbacks; the requested plan arguments alone are deliberately insufficient. + +Run the two-GPU transport check with: + +```bash +torchrun --standalone --nproc-per-node=2 \ + scripts/ws2_p2p_nccl_attention_reference_check.py +``` + +The validation CLI materializes the TP-local Qwen3-8B shard. For the default +`TP=2` target this is `Hq=16`, `Hkv=4`, `D=128`; `32/8` are global model head +counts and are rejected as a mislabeled local execution. Any non-finite drift +statistic also fails strict acceptance. diff --git a/pyproject.toml b/pyproject.toml index 3aa5ded1..e69de29b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,52 +0,0 @@ -[build-system] -requires = ["setuptools>=64", "wheel"] -build-backend = "setuptools.build_meta" - -[project] -name = "RL-Kernel" -version = "0.1.0" -description = "High-performance RL training engine focused on kernel fusion and memory efficiency." -readme = "README.md" -requires-python = ">=3.10" -license = {text = "Apache-2.0"} -authors = [ - {name = "RL-Kernel Contributors"} -] -dependencies = [ - "torch>=2.4.1", - "tabulate", - "numpy", - "accelerate", - "transformers==5.13.1", -] - -[project.optional-dependencies] -cuda = ["flashinfer-python>=0.1.6", "nvidia-ml-py"] -rocm = ["aiter"] -vllm = ["vllm>=0.6.0"] -drift-viewer = ["Pillow>=10", "PySide6>=6.6"] -dev = ["pytest", "black", "isort", "ruff", "mypy", "pre-commit"] - -[tool.setuptools.packages.find] -where = ["."] -include = ["rl_engine*"] - -[tool.ruff] -line-length = 100 - -[tool.ruff.lint] -select = ["E", "F", "B"] -ignore = [] - -[tool.ruff.lint.per-file-ignores] -"__init__.py" = ["F401"] - -[tool.mypy] -ignore_missing_imports = true -follow_imports = "silent" - -[tool.pytest.ini_options] -markers = [ - "smoke_operator: temporary smoke-only operator plumbing tests", - "unit: CPU-safe unit tests", -] diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index 29faedda..79c24823 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -18,9 +18,14 @@ _EnumT = TypeVar("_EnumT", bound=Enum) -# Stable identity shared by the training and rollout deterministic Attention -# core. Backend adapters may differ, but strict mode must report this ID. -STRICT_ATTENTION_CORE_ID = "rlkernel.attention.deterministic_core.v1" +# Stable identities for Attention arithmetic shared by training and rollout. +# The FA4 core is the strict production path. The materializing RL-Kernel core +# remains available as an explicit reference and capability-gap fallback. +STRICT_ATTENTION_PRODUCTION_CORE_ID = "rlkernel.attention.flash_attention4.num_splits1.v1" +STRICT_ATTENTION_REFERENCE_CORE_ID = "rlkernel.attention.deterministic_core.v1" +# Compatibility alias for callers that explicitly select the original core. +STRICT_ATTENTION_CORE_ID = STRICT_ATTENTION_REFERENCE_CORE_ID +STRICT_ATTENTION_FA4_SCHEDULE_ID = "single_batch_flash_attention4_num_splits1" STRICT_ATTENTION_SCHEDULE_ID = "single_batch_single_query_global_kv_blocks" @@ -1672,6 +1677,9 @@ class AttentionDispatchResult: "SplitKVRuntimePlanSet", "SplitKVSpec", "STRICT_ATTENTION_CORE_ID", + "STRICT_ATTENTION_FA4_SCHEDULE_ID", + "STRICT_ATTENTION_PRODUCTION_CORE_ID", + "STRICT_ATTENTION_REFERENCE_CORE_ID", "STRICT_ATTENTION_SCHEDULE_ID", "validate_split_kv_alignment", "validate_split_kv_plan_set_alignment", diff --git a/rl_engine/kernels/ops/cuda/attention/__init__.py b/rl_engine/kernels/ops/cuda/attention/__init__.py index 8ca8df7b..91fa6f99 100644 --- a/rl_engine/kernels/ops/cuda/attention/__init__.py +++ b/rl_engine/kernels/ops/cuda/attention/__init__.py @@ -1,8 +1,14 @@ from .deterministic_attn import DeterministicAttentionOp -from .flash_attn import FlashAttentionOp +from .flash_attn import FlashAttentionOp, StrictFlashAttention4Core, StrictFlashAttentionUnavailable from .prefix_shared_attn import PrefixSharedAttentionOp -__all__ = ["DeterministicAttentionOp", "FlashAttentionOp", "PrefixSharedAttentionOp"] +__all__ = [ + "DeterministicAttentionOp", + "FlashAttentionOp", + "PrefixSharedAttentionOp", + "StrictFlashAttention4Core", + "StrictFlashAttentionUnavailable", +] # CP communication and FlashInfer are optional layers owned by later WS2 PRs. # Keep the base Attention package importable while those PRs are developed or @@ -14,6 +20,7 @@ AttentionCPCommunicationPlan, AttentionCPCommunicationUnavailable, AttentionCPMergedState, + AttentionCPOutputShard, AttentionCPPartialState, AttentionParallelSpec, CPCommunicationBackend, @@ -32,6 +39,7 @@ "AttentionCPCommunicationPlan", "AttentionCPCommunicationUnavailable", "AttentionCPMergedState", + "AttentionCPOutputShard", "AttentionCPPartialState", "AttentionParallelSpec", "CPCommunicationBackend", diff --git a/rl_engine/kernels/ops/cuda/attention/cp_comm.py b/rl_engine/kernels/ops/cuda/attention/cp_comm.py new file mode 100644 index 00000000..d1470bfd --- /dev/null +++ b/rl_engine/kernels/ops/cuda/attention/cp_comm.py @@ -0,0 +1,1293 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CP/TP Attention communication and a P2P NCCL reference. + +PR7 evaluates fused attention backends for the #235 target +``Qwen3-8B, TP=2, CP=2, BF16``. The self-owned CUDA communication operators are +AG/RS and compute-communication decoupled. This module adapts the deterministic +collectives from PR311/PR312 to the Attention partial-state contract. + +The strict production path uses one arithmetic graph at every CP size: + +```text +owner-local Q/K/V and position IDs + -> deterministic AG(Q/K/V/positions) + -> shared no-Split-K CUDA Attention core on full logical Q/K/V + -> deterministic RS(Out, LSE) + -> owner-local query result +``` + +The older partial-state path remains available as a reference interface: + +```text +owner-local deterministic fallback or TE attention over rank-owned KV blocks + -> AttentionCPPartialState(out, lse, global_block_index, tp/cp rank metadata) + -> custom CUDA AG communication operator + -> sort by global_block_index + -> PR3 FP32 online-softmax merge + -> custom CUDA RS communication operator +``` +The CUDA backend uses the self-owned deterministic CUDA collectives when PR311 / +PR312 are present. It keeps the same manifest and FP32 merge contract as the +P2P NCCL reference, and fails closed when those compiled operators are absent. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Literal, Protocol, Sequence + +import torch + +CPCommunicationBackend = Literal["cuda_ag_rs", "p2p_nccl_reference", "local_debug"] +CPCommunicationStatus = Literal["interface_only", "implemented"] + + +class AttentionCPCommunicationUnavailable(RuntimeError): + """Raised when a requested CP communication backend is not implemented.""" + + +@dataclass(frozen=True) +class AttentionParallelSpec: + """TP/CP identity carried by PR7 attention backend reports.""" + + tp_world_size: int = 2 + tp_rank: int = 0 + cp_world_size: int = 2 + cp_rank: int = 0 + + def validate(self) -> None: + _positive_int(self.tp_world_size, "tp_world_size") + _positive_int(self.cp_world_size, "cp_world_size") + _rank_in_world(self.tp_rank, self.tp_world_size, "tp_rank") + _rank_in_world(self.cp_rank, self.cp_world_size, "cp_rank") + + def provenance(self) -> dict[str, int]: + self.validate() + return { + "tp_world_size": int(self.tp_world_size), + "tp_rank": int(self.tp_rank), + "cp_world_size": int(self.cp_world_size), + "cp_rank": int(self.cp_rank), + } + + +@dataclass(frozen=True) +class AttentionCPBlockMetadata: + """Logical identity for one attention partial state.""" + + global_block_index: int + kv_block_start: int + kv_block_end: int + owner_cp_rank: int + owner_tp_rank: int + + def validate(self, parallel: AttentionParallelSpec) -> None: + parallel.validate() + if ( + isinstance(self.global_block_index, bool) + or not isinstance(self.global_block_index, int) + or self.global_block_index < 0 + ): + raise ValueError("global_block_index must be non-negative") + if ( + isinstance(self.kv_block_start, bool) + or isinstance(self.kv_block_end, bool) + or not isinstance(self.kv_block_start, int) + or not isinstance(self.kv_block_end, int) + or self.kv_block_start < 0 + or self.kv_block_end <= self.kv_block_start + ): + raise ValueError("KV block bounds must satisfy 0 <= start < end") + _rank_in_world(self.owner_cp_rank, parallel.cp_world_size, "owner_cp_rank") + _rank_in_world(self.owner_tp_rank, parallel.tp_world_size, "owner_tp_rank") + + def provenance(self) -> dict[str, int]: + return { + "global_block_index": int(self.global_block_index), + "kv_block_start": int(self.kv_block_start), + "kv_block_end": int(self.kv_block_end), + "owner_cp_rank": int(self.owner_cp_rank), + "owner_tp_rank": int(self.owner_tp_rank), + } + + +@dataclass(frozen=True) +class AttentionCPPartialState: + """One local or received ``(out, lse)`` state before CP merge.""" + + out: torch.Tensor + lse: torch.Tensor + block: AttentionCPBlockMetadata + + def validate(self, parallel: AttentionParallelSpec) -> None: + self.block.validate(parallel) + if self.out.ndim != 4: + raise ValueError("partial out must have shape [B, Hq, Sq, D]") + if self.lse.ndim != 3: + raise ValueError("partial lse must have shape [B, Hq, Sq]") + if self.out.shape[:3] != self.lse.shape: + raise ValueError("partial out and lse must share [B, Hq, Sq]") + if self.out.device != self.lse.device: + raise ValueError("partial out and lse must be on the same device") + if self.lse.dtype != torch.float32: + raise ValueError("partial lse must be attention-domain FP32") + if self.out.dtype != torch.float32: + raise ValueError("partial out must remain FP32 until the final write") + + +@dataclass(frozen=True) +class AttentionCPMergedState: + """Merged attention state before the CUDA RS communication operator.""" + + out: torch.Tensor + lse: torch.Tensor + + def validate(self) -> None: + if self.out.ndim != 4: + raise ValueError("merged out must have shape [B, Hq, Sq, D]") + if self.lse.ndim != 3: + raise ValueError("merged lse must have shape [B, Hq, Sq]") + if self.out.shape[:3] != self.lse.shape: + raise ValueError("merged out and lse must share [B, Hq, Sq]") + if self.out.device != self.lse.device: + raise ValueError("merged out and lse must be on the same device") + if self.lse.dtype != torch.float32: + raise ValueError("merged lse must be attention-domain FP32") + if self.out.dtype != torch.float32: + raise ValueError("merged out must remain FP32 until the final write") + + +@dataclass(frozen=True) +class AttentionCPOutputShard: + """Final strict-core output shard after RS.""" + + out: torch.Tensor + lse: torch.Tensor + + def validate(self) -> None: + if self.out.ndim != 4 or self.lse.ndim != 3: + raise ValueError("strict output must have out [B,Hq,Sq,D] and lse [B,Hq,Sq]") + if self.out.shape[:3] != self.lse.shape: + raise ValueError("strict output and lse must share [B,Hq,Sq]") + if self.out.device != self.lse.device: + raise ValueError("strict output and lse must be on the same device") + if self.out.dtype not in (torch.float16, torch.bfloat16): + raise ValueError("strict Attention output must be FP16 or BF16") + if self.lse.dtype is not torch.float32: + raise ValueError("strict Attention LSE must be FP32") + + +@dataclass(frozen=True) +class AttentionCPCommunicationPlan: + """Requested AG/RS communication contract for CP attention partial states.""" + + parallel: AttentionParallelSpec + backend: CPCommunicationBackend = "cuda_ag_rs" + status: CPCommunicationStatus = "interface_only" + pattern: str = "ag_rs" + compute_communication: str = "decoupled" + merge_order: str = "global_block_index" + accum_dtype: torch.dtype = torch.float32 + return_lse: bool = True + expected_blocks: tuple[AttentionCPBlockMetadata, ...] = () + expected_kv_token_range: tuple[int, int] | None = None + query_token_ranges: tuple[tuple[int, int], ...] = () + merge_root_cp_rank: int = 0 + + def validate(self) -> None: + self.parallel.validate() + if self.backend not in {"cuda_ag_rs", "p2p_nccl_reference", "local_debug"}: + raise ValueError(f"unsupported CP communication backend: {self.backend}") + if self.status not in {"interface_only", "implemented"}: + raise ValueError(f"unsupported CP communication status: {self.status}") + if self.pattern != "ag_rs": + raise ValueError("PR7 CP communication must use the custom CUDA AG/RS interface") + if self.compute_communication != "decoupled": + raise ValueError("PR7 CP communication must keep compute and communication decoupled") + if self.merge_order != "global_block_index": + raise ValueError("PR7 CP communication must preserve global_block_index merge order") + if self.accum_dtype is not torch.float32: + raise ValueError("PR7 CP merge accumulation must be FP32") + if not self.return_lse: + raise ValueError("PR7 CP communication requires LSE-carrying partial states") + _rank_in_world( + self.merge_root_cp_rank, + self.parallel.cp_world_size, + "merge_root_cp_rank", + ) + _validate_expected_block_manifest(self) + _validate_query_token_ranges(self) + if self.backend == "p2p_nccl_reference": + if self.status != "implemented": + raise ValueError("P2P NCCL reference plans must use status='implemented'") + if not self.expected_blocks or self.expected_kv_token_range is None: + raise ValueError("P2P NCCL reference requires a complete expected block manifest") + if not self.query_token_ranges: + raise ValueError("P2P NCCL reference requires one query range per CP rank") + + def provenance(self) -> dict[str, object]: + self.validate() + return { + "cp_comm_backend": self.backend, + "cp_comm_status": self.status, + "cp_comm_pattern": self.pattern, + "cp_comm_compute_communication": self.compute_communication, + "cp_comm_merge_order": self.merge_order, + "cp_comm_accum_dtype": "fp32", + "cp_comm_return_lse": self.return_lse, + "cp_comm_contract": "partial_out_lse_global_block_index", + "cp_comm_strict_contract": "ag_qkv_positions_shared_core_rs_out_lse", + "cp_comm_strict_kv_communication": "all_gather", + "cp_comm_strict_position_communication": "all_gather", + "cp_comm_strict_backward": "rs_out_backward_ag_then_ag_qkv_backward_rs", + "cp_comm_expected_kv_token_range": ( + None if self.expected_kv_token_range is None else list(self.expected_kv_token_range) + ), + "cp_comm_expected_blocks": [block.provenance() for block in self.expected_blocks], + "cp_comm_query_token_ranges": [list(bounds) for bounds in self.query_token_ranges], + "cp_comm_merge_root_cp_rank": int(self.merge_root_cp_rank), + **self.parallel.provenance(), + } + + +class AttentionCPCommunication(Protocol): + """Protocol implemented by custom CUDA AG/RS communication operators.""" + + def all_gather_query( + self, + local_q: torch.Tensor, + plan: AttentionCPCommunicationPlan, + ) -> torch.Tensor: + """Gather the query sequence shards in logical CP-rank order.""" + + def all_gather_kv( + self, + local_k: torch.Tensor, + local_v: torch.Tensor, + plan: AttentionCPCommunicationPlan, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Gather owner-local K/V in logical CP-rank order.""" + + def all_gather_position_ids( + self, + local_query_positions: torch.Tensor, + local_key_positions: torch.Tensor, + plan: AttentionCPCommunicationPlan, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Gather the position IDs paired with owner-local Q/K.""" + + def all_gather_partial_states( + self, + local_states: tuple[AttentionCPPartialState, ...], + plan: AttentionCPCommunicationPlan, + ) -> tuple[AttentionCPPartialState, ...]: + """Run the custom CUDA AG operator and return gathered partial states.""" + + def reduce_scatter_merged_state( + self, + merged_state: AttentionCPMergedState, + plan: AttentionCPCommunicationPlan, + ) -> AttentionCPMergedState: + """Run the custom CUDA RS operator and return this rank's output shard.""" + + def reduce_scatter_strict_result( + self, + out: torch.Tensor, + lse: torch.Tensor, + plan: AttentionCPCommunicationPlan, + ) -> AttentionCPOutputShard: + """RS a shared-core result without changing its output dtype.""" + + +class _AllGatherSequence(torch.autograd.Function): + """Autograd bridge for the self-owned rank-ordered sequence AllGather.""" + + @staticmethod + def forward(ctx, local: torch.Tensor, collective: Any, sequence_dim: int) -> torch.Tensor: + ctx.collective = collective + ctx.sequence_dim = int(sequence_dim) + packed = local.movedim(ctx.sequence_dim, 0).contiguous() + gathered = collective.all_gather(packed) + return gathered.movedim(0, ctx.sequence_dim).contiguous() + + @staticmethod + def backward(ctx, grad_global: torch.Tensor) -> tuple[torch.Tensor, None, None]: + packed = grad_global.movedim(ctx.sequence_dim, 0).contiguous() + grad_local = ctx.collective.reduce_scatter(packed) + return grad_local.movedim(0, ctx.sequence_dim).contiguous(), None, None + + +class _RootReduceScatterSequence(torch.autograd.Function): + """Scatter one authoritative full result and gather its gradient back to the root.""" + + @staticmethod + def forward( + ctx, + full: torch.Tensor, + collective: Any, + sequence_dim: int, + rank: int, + root: int, + ) -> torch.Tensor: + ctx.collective = collective + ctx.sequence_dim = int(sequence_dim) + ctx.rank = int(rank) + ctx.root = int(root) + packed = full.movedim(ctx.sequence_dim, 0).contiguous() + if ctx.rank != ctx.root: + packed = torch.zeros_like(packed) + local = collective.reduce_scatter(packed) + return local.movedim(0, ctx.sequence_dim).contiguous() + + @staticmethod + def backward(ctx, grad_local: torch.Tensor) -> tuple[torch.Tensor, None, None, None, None]: + packed = grad_local.movedim(ctx.sequence_dim, 0).contiguous() + grad_full = ctx.collective.all_gather(packed).movedim(0, ctx.sequence_dim).contiguous() + if ctx.rank != ctx.root: + grad_full.zero_() + return grad_full, None, None, None, None + + +class CUDAAGRSAttentionCPCommunication: + """Deterministic CUDA AG/RS adapter backed by PR311/PR312.""" + + backend_id = "cuda_ag_rs" + supports_autograd = True + + def __init__(self, *, process_group: Any = None, collective: Any = None) -> None: + self._process_group = process_group + self._collective = collective + + def _get_collective(self, plan: AttentionCPCommunicationPlan): + if self._collective is not None: + return self._collective + try: + from rl_engine.distributed import DeterministicCollective + except ImportError as exc: + raise AttentionCPCommunicationUnavailable( + "self-owned CUDA AG/RS requires PR311/PR312 DeterministicCollective" + ) from exc + try: + self._collective = DeterministicCollective( + group=self._process_group, + device=torch.device("cuda", torch.cuda.current_device()), + ) + except (RuntimeError, ValueError, TypeError) as exc: + raise AttentionCPCommunicationUnavailable( + f"self-owned CUDA AG/RS is unavailable: {exc}" + ) from exc + if self._collective.world_size != plan.parallel.cp_world_size: + raise AttentionCPCommunicationUnavailable( + "self-owned collective world size does not match CP communication plan" + ) + return self._collective + + def all_gather_query( + self, + local_q: torch.Tensor, + plan: AttentionCPCommunicationPlan, + ) -> torch.Tensor: + """Gather Q with the self-owned CUDA AG operator. + + The collective works on the leading dimension, so Q is temporarily + laid out as ``[S_local, B, H, D]``. The returned tensor is restored to + the Attention layout ``[B, H, S_global, D]`` in CP-rank order. + """ + + self._validate_cuda_plan(plan) + _validate_query_shard(local_q, plan) + ranges = plan.query_token_ranges + if len({end - start for start, end in ranges}) != 1: + raise AttentionCPCommunicationUnavailable( + "self-owned CUDA AG requires equal query shard lengths" + ) + collective = self._get_collective(plan) + return self._all_gather_sequence_tensor(local_q, collective, sequence_dim=2) + + def all_gather_kv( + self, + local_k: torch.Tensor, + local_v: torch.Tensor, + plan: AttentionCPCommunicationPlan, + ) -> tuple[torch.Tensor, torch.Tensor]: + self._validate_cuda_plan(plan) + _validate_local_kv_shard(local_k, local_v, plan) + _require_equal_kv_owner_widths(plan, "self-owned CUDA AG") + collective = self._get_collective(plan) + global_k = self._all_gather_sequence_tensor(local_k, collective, sequence_dim=2) + global_v = self._all_gather_sequence_tensor(local_v, collective, sequence_dim=2) + return global_k, global_v + + def all_gather_position_ids( + self, + local_query_positions: torch.Tensor, + local_key_positions: torch.Tensor, + plan: AttentionCPCommunicationPlan, + ) -> tuple[torch.Tensor, torch.Tensor]: + self._validate_cuda_plan(plan) + _validate_local_position_ids(local_query_positions, local_key_positions, plan) + _require_equal_kv_owner_widths(plan, "self-owned CUDA AG") + collective = self._get_collective(plan) + query_positions = self._all_gather_sequence_tensor( + local_query_positions, collective, sequence_dim=1 + ) + key_positions = self._all_gather_sequence_tensor( + local_key_positions, collective, sequence_dim=1 + ) + return query_positions, key_positions + + @staticmethod + def _all_gather_sequence_tensor( + local: torch.Tensor, + collective: Any, + *, + sequence_dim: int, + ) -> torch.Tensor: + return _AllGatherSequence.apply(local, collective, sequence_dim) + + def all_gather_partial_states( + self, + local_states: tuple[AttentionCPPartialState, ...], + plan: AttentionCPCommunicationPlan, + ) -> tuple[AttentionCPPartialState, ...]: + self._validate_cuda_plan(plan) + _validate_local_partial_states(local_states, plan) + ordered = tuple(sorted(local_states, key=lambda state: state.block.global_block_index)) + block_count = len(ordered) + counts = [None] * plan.parallel.cp_world_size + self._dist().all_gather_object(counts, block_count, group=self._process_group) + if any(count != block_count for count in counts): + raise AttentionCPCommunicationUnavailable( + "self-owned CUDA AG requires equal block counts on all CP ranks" + ) + collective = self._get_collective(plan) + packed_out = torch.stack([state.out for state in ordered], dim=0).contiguous() + packed_lse = torch.stack([state.lse for state in ordered], dim=0).contiguous() + gathered_out = collective.all_gather(packed_out) + gathered_lse = collective.all_gather(packed_lse) + received: list[AttentionCPPartialState] = [] + blocks_by_rank = [ + _expected_blocks_for_cp_rank(plan, cp_rank) + for cp_rank in range(plan.parallel.cp_world_size) + ] + for cp_rank, blocks in enumerate(blocks_by_rank): + for block_index, block in enumerate(blocks): + row = cp_rank * block_count + block_index + received.append( + AttentionCPPartialState( + out=gathered_out[row], + lse=gathered_lse[row], + block=block, + ) + ) + return sort_attention_cp_partial_states(tuple(received), plan=plan) + + def reduce_scatter_merged_state( + self, + merged_state: AttentionCPMergedState, + plan: AttentionCPCommunicationPlan, + ) -> AttentionCPMergedState: + self._validate_cuda_plan(plan) + merged_state.validate() + ranges = plan.query_token_ranges + if not ranges or len({end - start for start, end in ranges}) != 1: + raise AttentionCPCommunicationUnavailable( + "self-owned CUDA RS currently requires equal contiguous query ranges" + ) + collective = self._get_collective(plan) + rank = plan.parallel.cp_rank + root = plan.merge_root_cp_rank + out_local = _RootReduceScatterSequence.apply(merged_state.out, collective, 2, rank, root) + lse_local = _RootReduceScatterSequence.apply(merged_state.lse, collective, 2, rank, root) + result = AttentionCPMergedState(out=out_local, lse=lse_local) + result.validate() + return result + + def reduce_scatter_strict_result( + self, + out: torch.Tensor, + lse: torch.Tensor, + plan: AttentionCPCommunicationPlan, + ) -> AttentionCPOutputShard: + self._validate_cuda_plan(plan) + _validate_strict_full_result(out, lse, plan) + ranges = plan.query_token_ranges + if len({end - start for start, end in ranges}) != 1: + raise AttentionCPCommunicationUnavailable( + "self-owned CUDA RS requires equal contiguous query ranges" + ) + collective = self._get_collective(plan) + result = AttentionCPOutputShard( + out=_RootReduceScatterSequence.apply( + out, + collective, + 2, + plan.parallel.cp_rank, + plan.merge_root_cp_rank, + ), + lse=_RootReduceScatterSequence.apply( + lse, + collective, + 2, + plan.parallel.cp_rank, + plan.merge_root_cp_rank, + ), + ) + result.validate() + return result + + def _dist(self): + import torch.distributed as dist + + if not dist.is_available() or not dist.is_initialized(): + raise AttentionCPCommunicationUnavailable( + "self-owned CUDA AG/RS requires initialized torch.distributed" + ) + return dist + + def _validate_cuda_plan(self, plan: AttentionCPCommunicationPlan) -> None: + plan.validate() + if plan.backend != "cuda_ag_rs" or plan.status != "implemented": + raise AttentionCPCommunicationUnavailable( + "self-owned CUDA AG/RS requires an implemented cuda_ag_rs plan" + ) + if not torch.cuda.is_available(): + raise AttentionCPCommunicationUnavailable("self-owned CUDA AG/RS requires CUDA") + + +class P2PNCCLAttentionCPCommunication: + """Correctness-first P2P NCCL implementation of the CP protocol. + + The block manifest is authoritative. Only ``out`` and ``lse`` tensors are + transported, in deterministic peer/block order; received metadata is + reconstructed from the manifest and then validated as a complete set. + ``reduce_scatter_merged_state`` uses a designated root and explicit P2P + sends so its numerical behavior is easy to compare with a future CUDA RS. + """ + + backend_id = "p2p_nccl_reference" + supports_autograd = False + + def __init__( + self, + *, + process_group: Any = None, + dist_module: Any = None, + validate_cuda_tensors: bool = True, + ) -> None: + if dist_module is None: + import torch.distributed as dist + + dist_module = dist + self._dist = dist_module + self._group = process_group + self._validate_cuda_tensors = validate_cuda_tensors + + def all_gather_query( + self, + local_q: torch.Tensor, + plan: AttentionCPCommunicationPlan, + ) -> torch.Tensor: + """Reference query AG implemented with explicit NCCL P2P traffic.""" + + self._validate_runtime(plan) + _validate_query_shard(local_q, plan) + self._require_cuda(local_q) + ranges = plan.query_token_ranges + if len({end - start for start, end in ranges}) != 1: + raise AttentionCPCommunicationUnavailable( + "P2P NCCL query AG currently requires equal query shard lengths" + ) + received: dict[int, torch.Tensor] = {plan.parallel.cp_rank: local_q.contiguous()} + operations: list[Any] = [] + for peer_cp_rank in range(plan.parallel.cp_world_size): + if peer_cp_rank == plan.parallel.cp_rank: + continue + peer = self._global_peer(peer_cp_rank) + remote = torch.empty_like(local_q) + received[peer_cp_rank] = remote + operations.extend( + ( + self._dist.P2POp(self._dist.irecv, remote, peer, group=self._group), + self._dist.P2POp( + self._dist.isend, + local_q.contiguous(), + peer, + group=self._group, + ), + ) + ) + self._run_operations(operations) + return torch.cat( + [received[rank] for rank in range(plan.parallel.cp_world_size)], + dim=2, + ) + + def all_gather_kv( + self, + local_k: torch.Tensor, + local_v: torch.Tensor, + plan: AttentionCPCommunicationPlan, + ) -> tuple[torch.Tensor, torch.Tensor]: + self._validate_runtime(plan) + _validate_local_kv_shard(local_k, local_v, plan) + self._require_cuda(local_k, local_v) + ranges = _kv_owner_ranges(plan) + local_rank = plan.parallel.cp_rank + gathered_k: dict[int, torch.Tensor] = {local_rank: local_k.contiguous()} + gathered_v: dict[int, torch.Tensor] = {local_rank: local_v.contiguous()} + operations: list[Any] = [] + for peer_rank, (start, end) in enumerate(ranges): + if peer_rank == local_rank: + continue + peer = self._global_peer(peer_rank) + shape = (*local_k.shape[:2], end - start, local_k.size(3)) + peer_k = torch.empty(shape, dtype=local_k.dtype, device=local_k.device) + peer_v = torch.empty(shape, dtype=local_v.dtype, device=local_v.device) + gathered_k[peer_rank] = peer_k + gathered_v[peer_rank] = peer_v + operations.extend( + ( + self._dist.P2POp(self._dist.irecv, peer_k, peer, group=self._group), + self._dist.P2POp(self._dist.irecv, peer_v, peer, group=self._group), + self._dist.P2POp( + self._dist.isend, local_k.contiguous(), peer, group=self._group + ), + self._dist.P2POp( + self._dist.isend, local_v.contiguous(), peer, group=self._group + ), + ) + ) + self._run_operations(operations) + return ( + torch.cat([gathered_k[rank] for rank in range(len(ranges))], dim=2), + torch.cat([gathered_v[rank] for rank in range(len(ranges))], dim=2), + ) + + def all_gather_position_ids( + self, + local_query_positions: torch.Tensor, + local_key_positions: torch.Tensor, + plan: AttentionCPCommunicationPlan, + ) -> tuple[torch.Tensor, torch.Tensor]: + self._validate_runtime(plan) + _validate_local_position_ids(local_query_positions, local_key_positions, plan) + self._require_cuda(local_query_positions, local_key_positions) + query_ranges = plan.query_token_ranges + key_ranges = _kv_owner_ranges(plan) + local_rank = plan.parallel.cp_rank + gathered_q = {local_rank: local_query_positions.contiguous()} + gathered_k = {local_rank: local_key_positions.contiguous()} + operations: list[Any] = [] + for peer_rank, ((q_start, q_end), (k_start, k_end)) in enumerate( + zip(query_ranges, key_ranges, strict=True) + ): + if peer_rank == local_rank: + continue + peer = self._global_peer(peer_rank) + peer_q = torch.empty( + (local_query_positions.size(0), q_end - q_start), + dtype=local_query_positions.dtype, + device=local_query_positions.device, + ) + peer_k = torch.empty( + (local_key_positions.size(0), k_end - k_start), + dtype=local_key_positions.dtype, + device=local_key_positions.device, + ) + gathered_q[peer_rank] = peer_q + gathered_k[peer_rank] = peer_k + operations.extend( + ( + self._dist.P2POp(self._dist.irecv, peer_q, peer, group=self._group), + self._dist.P2POp(self._dist.irecv, peer_k, peer, group=self._group), + self._dist.P2POp( + self._dist.isend, + local_query_positions.contiguous(), + peer, + group=self._group, + ), + self._dist.P2POp( + self._dist.isend, + local_key_positions.contiguous(), + peer, + group=self._group, + ), + ) + ) + self._run_operations(operations) + return ( + torch.cat([gathered_q[rank] for rank in range(len(query_ranges))], dim=1), + torch.cat([gathered_k[rank] for rank in range(len(key_ranges))], dim=1), + ) + + def all_gather_partial_states( + self, + local_states: tuple[AttentionCPPartialState, ...], + plan: AttentionCPCommunicationPlan, + ) -> tuple[AttentionCPPartialState, ...]: + self._validate_runtime(plan) + _validate_local_partial_states(local_states, plan) + self._require_cuda(local_states[0].out, local_states[0].lse) + ordered_local_states = tuple( + sorted(local_states, key=lambda state: state.block.global_block_index) + ) + template = ordered_local_states[0] + if template.out.size(2) != plan.query_token_ranges[-1][1]: + raise ValueError( + "local partial states must cover the complete query range before gather" + ) + received: list[AttentionCPPartialState] = [] + operations: list[Any] = [] + receive_tensors: list[tuple[AttentionCPBlockMetadata, torch.Tensor, torch.Tensor]] = [] + + for peer_cp_rank in range(plan.parallel.cp_world_size): + if peer_cp_rank == plan.parallel.cp_rank: + continue + peer = self._global_peer(peer_cp_rank) + peer_blocks = _expected_blocks_for_cp_rank(plan, peer_cp_rank) + for block in peer_blocks: + out = torch.empty_like(template.out) + lse = torch.empty_like(template.lse) + operations.extend( + ( + self._dist.P2POp( + self._dist.irecv, + out, + peer, + group=self._group, + ), + self._dist.P2POp( + self._dist.irecv, + lse, + peer, + group=self._group, + ), + ) + ) + receive_tensors.append((block, out, lse)) + for state in ordered_local_states: + operations.extend( + ( + self._dist.P2POp( + self._dist.isend, + state.out.contiguous(), + peer, + group=self._group, + ), + self._dist.P2POp( + self._dist.isend, + state.lse.contiguous(), + peer, + group=self._group, + ), + ) + ) + + self._run_operations(operations) + for block, out, lse in receive_tensors: + received.append(AttentionCPPartialState(out=out, lse=lse, block=block)) + return sort_attention_cp_partial_states( + (*ordered_local_states, *received), + plan=plan, + ) + + def reduce_scatter_merged_state( + self, + merged_state: AttentionCPMergedState, + plan: AttentionCPCommunicationPlan, + ) -> AttentionCPMergedState: + self._validate_runtime(plan) + merged_state.validate() + self._require_cuda(merged_state.out, merged_state.lse) + ranges = plan.query_token_ranges + full_query_tokens = ranges[-1][1] + if merged_state.out.size(2) != full_query_tokens: + raise ValueError("merged state query length does not match query_token_ranges coverage") + + rank = plan.parallel.cp_rank + root = plan.merge_root_cp_rank + local_start, local_end = ranges[rank] + if rank == root: + operations: list[Any] = [] + for peer_cp_rank, (start, end) in enumerate(ranges): + if peer_cp_rank == root or start == end: + continue + peer = self._global_peer(peer_cp_rank) + operations.extend( + ( + self._dist.P2POp( + self._dist.isend, + merged_state.out[:, :, start:end, :].contiguous(), + peer, + group=self._group, + ), + self._dist.P2POp( + self._dist.isend, + merged_state.lse[:, :, start:end].contiguous(), + peer, + group=self._group, + ), + ) + ) + self._run_operations(operations) + result = AttentionCPMergedState( + out=merged_state.out[:, :, local_start:local_end, :].contiguous(), + lse=merged_state.lse[:, :, local_start:local_end].contiguous(), + ) + else: + local_query_tokens = local_end - local_start + if local_query_tokens == 0: + result = AttentionCPMergedState( + out=merged_state.out[:, :, 0:0, :].contiguous(), + lse=merged_state.lse[:, :, 0:0].contiguous(), + ) + result.validate() + return result + out = torch.empty( + ( + *merged_state.out.shape[:2], + local_query_tokens, + merged_state.out.size(3), + ), + dtype=merged_state.out.dtype, + device=merged_state.out.device, + ) + lse = torch.empty( + (*merged_state.lse.shape[:2], local_query_tokens), + dtype=merged_state.lse.dtype, + device=merged_state.lse.device, + ) + peer = self._global_peer(root) + self._run_operations( + [ + self._dist.P2POp( + self._dist.irecv, + out, + peer, + group=self._group, + ), + self._dist.P2POp( + self._dist.irecv, + lse, + peer, + group=self._group, + ), + ] + ) + result = AttentionCPMergedState(out=out, lse=lse) + result.validate() + return result + + def reduce_scatter_strict_result( + self, + out: torch.Tensor, + lse: torch.Tensor, + plan: AttentionCPCommunicationPlan, + ) -> AttentionCPOutputShard: + self._validate_runtime(plan) + _validate_strict_full_result(out, lse, plan) + self._require_cuda(out, lse) + rank = plan.parallel.cp_rank + root = plan.merge_root_cp_rank + local_start, local_end = plan.query_token_ranges[rank] + if rank == root: + operations: list[Any] = [] + for peer_rank, (start, end) in enumerate(plan.query_token_ranges): + if peer_rank == root or start == end: + continue + peer = self._global_peer(peer_rank) + operations.extend( + ( + self._dist.P2POp( + self._dist.isend, + out[:, :, start:end, :].contiguous(), + peer, + group=self._group, + ), + self._dist.P2POp( + self._dist.isend, + lse[:, :, start:end].contiguous(), + peer, + group=self._group, + ), + ) + ) + self._run_operations(operations) + result = AttentionCPOutputShard( + out=out[:, :, local_start:local_end, :].contiguous(), + lse=lse[:, :, local_start:local_end].contiguous(), + ) + else: + out_local = torch.empty( + (*out.shape[:2], local_end - local_start, out.size(3)), + dtype=out.dtype, + device=out.device, + ) + lse_local = torch.empty( + (*lse.shape[:2], local_end - local_start), + dtype=lse.dtype, + device=lse.device, + ) + peer = self._global_peer(root) + self._run_operations( + ( + self._dist.P2POp(self._dist.irecv, out_local, peer, group=self._group), + self._dist.P2POp(self._dist.irecv, lse_local, peer, group=self._group), + ) + ) + result = AttentionCPOutputShard(out=out_local, lse=lse_local) + result.validate() + return result + + def _validate_runtime(self, plan: AttentionCPCommunicationPlan) -> None: + plan.validate() + if plan.backend != "p2p_nccl_reference" or plan.status != "implemented": + raise AttentionCPCommunicationUnavailable( + "P2P NCCL communication requires an implemented p2p_nccl_reference plan" + ) + dist = self._dist + if not dist.is_available() or not dist.is_initialized(): + raise AttentionCPCommunicationUnavailable( + "P2P NCCL communication requires initialized torch.distributed" + ) + backend = str(dist.get_backend(self._group)).lower() + if "nccl" not in backend: + raise AttentionCPCommunicationUnavailable( + f"P2P NCCL communication requires the NCCL backend; got {backend}" + ) + world_size = int(dist.get_world_size(self._group)) + rank = int(dist.get_rank(self._group)) + if world_size != plan.parallel.cp_world_size: + raise AttentionCPCommunicationUnavailable( + "process-group world size does not match cp_world_size" + ) + if rank != plan.parallel.cp_rank: + raise AttentionCPCommunicationUnavailable( + "process-group rank does not match the communication plan cp_rank" + ) + local_blocks = _expected_blocks_for_cp_rank(plan, plan.parallel.cp_rank) + if not local_blocks: + raise AttentionCPCommunicationUnavailable( + "P2P NCCL communication requires every CP rank to own at least one block" + ) + + def _global_peer(self, peer_cp_rank: int) -> int: + get_global_rank = getattr(self._dist, "get_global_rank", None) + if self._group is not None and callable(get_global_rank): + return int(get_global_rank(self._group, peer_cp_rank)) + return peer_cp_rank + + def _run_operations(self, operations: Sequence[Any]) -> None: + if not operations: + return + requests = self._dist.batch_isend_irecv(list(operations)) + for request in requests: + request.wait() + + def _require_cuda(self, *tensors: torch.Tensor) -> None: + if self._validate_cuda_tensors and any(tensor.device.type != "cuda" for tensor in tensors): + raise AttentionCPCommunicationUnavailable( + "P2P NCCL communication requires CUDA tensors" + ) + + +def sort_attention_cp_partial_states( + states: tuple[AttentionCPPartialState, ...], + *, + plan: AttentionCPCommunicationPlan, +) -> tuple[AttentionCPPartialState, ...]: + """Validate and sort partial states by ``global_block_index``.""" + + plan.validate() + if not states: + raise ValueError("at least one CP attention partial state is required") + for state in states: + state.validate(plan.parallel) + ordered = tuple(sorted(states, key=lambda state: state.block.global_block_index)) + indices = [state.block.global_block_index for state in ordered] + if len(set(indices)) != len(indices): + raise ValueError("duplicate global_block_index values are not allowed") + if not plan.expected_blocks and indices != list(range(len(ordered))): + raise ValueError( + "partial states without a manifest must cover global_block_index " + "values [0, block_count)" + ) + if not plan.expected_blocks and ordered[0].block.kv_block_start != 0: + raise ValueError("partial states without a manifest must start at KV token 0") + if plan.expected_kv_token_range is not None: + expected_start, expected_end = plan.expected_kv_token_range + if ( + ordered[0].block.kv_block_start != expected_start + or ordered[-1].block.kv_block_end != expected_end + ): + raise ValueError("partial states do not cover the declared expected KV token range") + _validate_partial_state_set(ordered, plan) + return ordered + + +def _validate_expected_block_manifest(plan: AttentionCPCommunicationPlan) -> None: + blocks = plan.expected_blocks + if not blocks: + if plan.expected_kv_token_range is not None: + raise ValueError("expected_kv_token_range requires expected_blocks") + return + for block in blocks: + block.validate(plan.parallel) + if block.owner_tp_rank != plan.parallel.tp_rank: + raise ValueError("expected block owner_tp_rank must match the plan TP shard") + ordered = tuple(sorted(blocks, key=lambda block: block.global_block_index)) + indices = tuple(block.global_block_index for block in ordered) + if len(set(indices)) != len(indices): + raise ValueError("expected block manifest contains duplicate global_block_index values") + if len(set(blocks)) != len(blocks): + raise ValueError("expected block manifest contains duplicate metadata") + owners = {block.owner_cp_rank for block in blocks} + if owners != set(range(plan.parallel.cp_world_size)): + raise ValueError("expected block manifest must assign work to every CP rank") + expected_range = plan.expected_kv_token_range + if expected_range is None: + raise ValueError("expected block manifest requires expected_kv_token_range") + try: + start, end = expected_range + except (TypeError, ValueError) as exc: + raise ValueError("expected KV range must contain exactly (start, end)") from exc + if ( + isinstance(start, bool) + or isinstance(end, bool) + or not isinstance(start, int) + or not isinstance(end, int) + or start < 0 + or end <= start + ): + raise ValueError("expected KV range must satisfy 0 <= start < end") + if ordered[0].kv_block_start != start or ordered[-1].kv_block_end != end: + raise ValueError("expected block manifest does not cover the declared KV range") + previous_end = start + for block in ordered: + if block.kv_block_start != previous_end: + raise ValueError("expected block manifest must be gap-free and non-overlapping") + previous_end = block.kv_block_end + + +def _validate_query_token_ranges(plan: AttentionCPCommunicationPlan) -> None: + ranges = plan.query_token_ranges + if not ranges: + return + if len(ranges) != plan.parallel.cp_world_size: + raise ValueError("query_token_ranges must contain one range per CP rank") + previous_end = 0 + for bounds in ranges: + try: + start, end = bounds + except (TypeError, ValueError) as exc: + raise ValueError("query token ranges must contain (start, end) pairs") from exc + if ( + isinstance(start, bool) + or isinstance(end, bool) + or not isinstance(start, int) + or not isinstance(end, int) + ): + raise ValueError("query token ranges must contain (start, end) pairs") + if start != previous_end or end < start: + raise ValueError("query token ranges must be non-negative, contiguous, and start at 0") + previous_end = end + if previous_end == 0: + raise ValueError("query token ranges must cover at least one query token") + + +def _validate_query_shard( + local_q: torch.Tensor, + plan: AttentionCPCommunicationPlan, +) -> None: + if local_q.ndim != 4: + raise ValueError("local Q must have shape [B, Hq, Sq_local, D]") + if not plan.query_token_ranges: + raise ValueError("query AG requires query_token_ranges") + start, end = plan.query_token_ranges[plan.parallel.cp_rank] + if local_q.size(2) != end - start: + raise ValueError("local Q sequence length does not match its query_token_range") + + +def _kv_owner_ranges( + plan: AttentionCPCommunicationPlan, +) -> tuple[tuple[int, int], ...]: + """Return one gap-free logical KV range for each CP owner.""" + + ranges: list[tuple[int, int]] = [] + for owner_rank in range(plan.parallel.cp_world_size): + blocks = _expected_blocks_for_cp_rank(plan, owner_rank) + if not blocks: + raise ValueError(f"CP block manifest has no blocks for owner rank {owner_rank}") + ordered = sorted(blocks, key=lambda block: block.kv_block_start) + start = ordered[0].kv_block_start + cursor = start + for block in ordered: + if block.kv_block_start != cursor: + raise ValueError("CP owner KV blocks must form a contiguous range") + cursor = block.kv_block_end + ranges.append((start, cursor)) + expected = plan.expected_kv_token_range + if expected is None or ranges[0][0] != expected[0] or ranges[-1][1] != expected[1]: + raise ValueError("CP owner ranges do not cover expected_kv_token_range") + for left, right in zip(ranges, ranges[1:], strict=False): + if left[1] != right[0]: + raise ValueError("CP owner KV ranges must be gap-free and rank ordered") + return tuple(ranges) + + +def _require_equal_kv_owner_widths( + plan: AttentionCPCommunicationPlan, + backend: str, +) -> None: + widths = {end - start for start, end in _kv_owner_ranges(plan)} + if len(widths) != 1: + raise AttentionCPCommunicationUnavailable(f"{backend} requires equal KV shard lengths") + + +def _validate_local_kv_shard( + local_k: torch.Tensor, + local_v: torch.Tensor, + plan: AttentionCPCommunicationPlan, +) -> None: + if local_k.ndim != 4 or local_v.ndim != 4: + raise ValueError("local K/V must have shape [B,Hkv,Skv_local,D]") + if local_k.shape != local_v.shape: + raise ValueError("local K/V must have matching shapes") + if local_k.dtype != local_v.dtype or local_k.device != local_v.device: + raise ValueError("local K/V must have matching dtype and device") + start, end = _kv_owner_ranges(plan)[plan.parallel.cp_rank] + if local_k.size(2) != end - start: + raise ValueError("local K/V width does not match the CP owner range") + + +def _validate_local_position_ids( + local_query_positions: torch.Tensor, + local_key_positions: torch.Tensor, + plan: AttentionCPCommunicationPlan, +) -> None: + integer_dtypes = (torch.int32, torch.int64) + if local_query_positions.ndim != 2 or local_key_positions.ndim != 2: + raise ValueError("local Q/K position IDs must have shape [B,S_local]") + if local_query_positions.dtype not in integer_dtypes or ( + local_key_positions.dtype not in integer_dtypes + ): + raise ValueError("local Q/K position IDs must contain integers") + if local_query_positions.device != local_key_positions.device: + raise ValueError("local Q/K position IDs must be on the same device") + query_start, query_end = plan.query_token_ranges[plan.parallel.cp_rank] + key_start, key_end = _kv_owner_ranges(plan)[plan.parallel.cp_rank] + if local_query_positions.size(1) != query_end - query_start: + raise ValueError("local query position width does not match query ownership") + if local_key_positions.size(1) != key_end - key_start: + raise ValueError("local key position width does not match KV ownership") + if local_query_positions.size(0) != local_key_positions.size(0): + raise ValueError("local Q/K position IDs must share batch size") + + +def _validate_strict_full_result( + out: torch.Tensor, + lse: torch.Tensor, + plan: AttentionCPCommunicationPlan, +) -> None: + result = AttentionCPOutputShard(out=out, lse=lse) + result.validate() + total_queries = plan.query_token_ranges[-1][1] + if out.size(2) != total_queries: + raise ValueError("strict full result does not cover all query_token_ranges") + + +def _validate_local_partial_states( + states: tuple[AttentionCPPartialState, ...], + plan: AttentionCPCommunicationPlan, +) -> None: + expected = _expected_blocks_for_cp_rank(plan, plan.parallel.cp_rank) + if not states: + raise ValueError("each CP rank must provide at least one local partial state") + for state in states: + state.validate(plan.parallel) + if state.block.owner_cp_rank != plan.parallel.cp_rank: + raise ValueError("local partial state has the wrong CP owner") + if state.block.owner_tp_rank != plan.parallel.tp_rank: + raise ValueError("local partial state has the wrong TP owner") + actual = tuple( + state.block for state in sorted(states, key=lambda item: item.block.global_block_index) + ) + if actual != expected: + raise ValueError("local partial states do not exactly match the rank manifest") + _validate_common_state_shapes(states) + + +def _expected_blocks_for_cp_rank( + plan: AttentionCPCommunicationPlan, + cp_rank: int, +) -> tuple[AttentionCPBlockMetadata, ...]: + return tuple( + block + for block in sorted( + plan.expected_blocks, + key=lambda item: item.global_block_index, + ) + if block.owner_cp_rank == cp_rank + ) + + +def _validate_partial_state_set( + states: tuple[AttentionCPPartialState, ...], + plan: AttentionCPCommunicationPlan, +) -> None: + _validate_common_state_shapes(states) + previous_end = states[0].block.kv_block_start + for state in states: + if state.block.owner_tp_rank != plan.parallel.tp_rank: + raise ValueError("partial state has the wrong TP owner for this CP group") + if state.block.kv_block_start != previous_end: + raise ValueError("partial state KV ranges must be gap-free and non-overlapping") + previous_end = state.block.kv_block_end + if plan.expected_blocks: + actual = tuple(state.block for state in states) + expected = tuple(sorted(plan.expected_blocks, key=lambda block: block.global_block_index)) + if actual != expected: + raise ValueError( + "gathered partial states do not exactly match the complete block manifest" + ) + + +def _validate_common_state_shapes(states: Sequence[AttentionCPPartialState]) -> None: + first = states[0] + for state in states[1:]: + if state.out.shape != first.out.shape or state.lse.shape != first.lse.shape: + raise ValueError("all CP partial states must have matching out/lse shapes") + if state.out.dtype != first.out.dtype or state.lse.dtype != first.lse.dtype: + raise ValueError("all CP partial states must have matching out/lse dtypes") + if state.out.device != first.out.device or state.lse.device != first.lse.device: + raise ValueError("all CP partial states must be on the same device") + + +def _positive_int(value: int, name: str) -> None: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + + +def _rank_in_world(rank: int, world_size: int, name: str) -> None: + if isinstance(rank, bool) or not isinstance(rank, int) or rank < 0 or rank >= world_size: + raise ValueError(f"{name} must be in [0, world_size)") + + +__all__ = [ + "AttentionCPBlockMetadata", + "AttentionCPCommunication", + "AttentionCPCommunicationPlan", + "AttentionCPCommunicationUnavailable", + "AttentionCPMergedState", + "AttentionCPOutputShard", + "AttentionCPPartialState", + "AttentionParallelSpec", + "CPCommunicationBackend", + "CPCommunicationStatus", + "CUDAAGRSAttentionCPCommunication", + "P2PNCCLAttentionCPCommunication", + "sort_attention_cp_partial_states", +] diff --git a/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py b/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py index 1e3d01e2..c3ef6aa3 100644 --- a/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py +++ b/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py @@ -11,6 +11,7 @@ from __future__ import annotations import math +from dataclasses import dataclass from typing import Optional import torch @@ -20,6 +21,8 @@ from rl_engine.kernels.attention_contract import ( STRICT_ATTENTION_CORE_ID, STRICT_ATTENTION_SCHEDULE_ID, + SplitKVMode, + SplitKVSpec, ) from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE from rl_engine.utils.logger import logger @@ -27,6 +30,15 @@ _HEAD_DIM = 128 +@dataclass(frozen=True) +class DeterministicAttentionCoreResult: + """Output and auditable arithmetic identity of the shared strict core.""" + + out: torch.Tensor + lse: torch.Tensor + provenance: dict[str, object] + + class _DeterministicAttentionFn(Function): @staticmethod def forward( @@ -207,3 +219,136 @@ def _validate_inputs( ) if sq < 1 or skv < 1: raise ValueError(f"Sq and Skv must be positive, got Sq={sq}, Skv={skv}") + + +class RLKernelDeterministicAttentionCore: + """Materializing CUDA reference core shared by training and rollout. + + This remains useful for correctness and capability-gap diagnosis. The + production default is the shared FA4 CuTe core with ``num_splits=1``. + """ + + core_id = STRICT_ATTENTION_CORE_ID + strict_schedule = STRICT_ATTENTION_SCHEDULE_ID + backend_id = "rlkernel.cuda.deterministic_attention" + merge_order = "global_block_index" + accum_dtype = "fp32" + downcast_at = "final_write" + fallback = False + native_attention_arithmetic = False + production_ready = False + reference_only = True + + def __init__( + self, + *, + split_kv: SplitKVSpec | None = None, + ) -> None: + requested = SplitKVSpec.disabled() if split_kv is None else split_kv + if not isinstance(requested, SplitKVSpec): + raise TypeError("split_kv must be a SplitKVSpec") + if requested.mode is not SplitKVMode.DISABLED: + raise ValueError("the strict CUDA Attention core requires Split-KV to be disabled") + self.split_kv = requested + self._op = DeterministicAttentionOp() + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + **kwargs, + ) -> DeterministicAttentionCoreResult: + return self.forward_with_lse(q, k, v, **kwargs) + + def forward_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: float | None = None, + key_padding_mask: torch.Tensor | None = None, + query_position_ids: torch.Tensor | None = None, + key_position_ids: torch.Tensor | None = None, + output_dtype: torch.dtype | None = None, + ) -> DeterministicAttentionCoreResult: + self._validate_positions( + q, + k, + causal=causal, + query_position_ids=query_position_ids, + key_position_ids=key_position_ids, + ) + resolved_dtype = q.dtype if output_dtype is None else output_dtype + if resolved_dtype != q.dtype: + raise ValueError("strict Attention output_dtype must match the Q/K/V input dtype") + out, lse = self._op.forward_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + ) + return DeterministicAttentionCoreResult( + out=out, + lse=lse, + provenance={ + "strict_core_id": self.core_id, + "attention_backend": self.backend_id, + "split_kv": self.split_kv.resolve(k.size(2), backend=self.backend_id).to_dict(), + "merge_order": self.merge_order, + "accum_dtype": self.accum_dtype, + "downcast_at": self.downcast_at, + "fallback": self.fallback, + "fallback_reason": None, + "native_attention_arithmetic": self.native_attention_arithmetic, + "production_ready": self.production_ready, + "reference_only": self.reference_only, + "strict_schedule": self.strict_schedule, + }, + ) + + @staticmethod + def _validate_positions( + q: torch.Tensor, + k: torch.Tensor, + *, + causal: bool, + query_position_ids: torch.Tensor | None, + key_position_ids: torch.Tensor | None, + ) -> None: + if not causal: + return + if query_position_ids is None or key_position_ids is None: + raise ValueError( + "strict CUDA Attention requires query_position_ids and " "key_position_ids" + ) + expected_q_shape = (q.size(0), q.size(2)) + expected_k_shape = (k.size(0), k.size(2)) + if tuple(query_position_ids.shape) != expected_q_shape: + raise ValueError(f"query_position_ids must have shape {expected_q_shape}") + if tuple(key_position_ids.shape) != expected_k_shape: + raise ValueError(f"key_position_ids must have shape {expected_k_shape}") + if query_position_ids.device != q.device or key_position_ids.device != k.device: + raise ValueError("strict Attention position IDs must be on the Q/K device") + integer_dtypes = (torch.int32, torch.int64) + if query_position_ids.dtype not in integer_dtypes or ( + key_position_ids.dtype not in integer_dtypes + ): + raise ValueError("strict Attention position IDs must contain integers") + if q.size(2) > k.size(2): + raise ValueError("causal strict Attention requires Sq <= Skv") + if q.size(2) > 1 and bool( + (query_position_ids[:, 1:] - query_position_ids[:, :-1] != 1).any() + ): + raise ValueError("query_position_ids must be contiguous and increasing") + if k.size(2) > 1 and bool((key_position_ids[:, 1:] - key_position_ids[:, :-1] != 1).any()): + raise ValueError("key_position_ids must be contiguous and increasing") + if not torch.equal(query_position_ids, key_position_ids[:, -q.size(2) :]): + raise ValueError( + "strict CUDA Attention requires queries to be the trailing " + "contiguous positions of the logical KV sequence" + ) diff --git a/rl_engine/kernels/ops/cuda/attention/flash_attn.py b/rl_engine/kernels/ops/cuda/attention/flash_attn.py index 0a068af6..24afdfc5 100644 --- a/rl_engine/kernels/ops/cuda/attention/flash_attn.py +++ b/rl_engine/kernels/ops/cuda/attention/flash_attn.py @@ -1,11 +1,235 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors +from __future__ import annotations + +import importlib +import importlib.metadata +import inspect +from typing import Any, Callable + import torch +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_FA4_SCHEDULE_ID, + STRICT_ATTENTION_PRODUCTION_CORE_ID, + SplitKVMode, + SplitKVSpec, +) from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.cuda.attention.deterministic_attn import ( + DeterministicAttentionCoreResult, + RLKernelDeterministicAttentionCore, +) from rl_engine.utils.logger import logger +_FA4_API_SOURCE = "flash_attn.cute.interface" +_FA4_REQUIRED_PARAMETERS = frozenset( + {"softmax_scale", "causal", "num_splits", "pack_gqa", "deterministic", "return_lse"} +) + + +class StrictFlashAttentionUnavailable(RuntimeError): + """Raised when the exact FlashAttention production contract is unavailable.""" + + +def _load_fa4_cute_op() -> tuple[Callable[..., Any], str]: + try: + module = importlib.import_module(_FA4_API_SOURCE) + op = getattr(module, "flash_attn_func") + except (AttributeError, ImportError, OSError, RuntimeError) as exc: + raise StrictFlashAttentionUnavailable( + "strict CUDA Attention requires flash_attn.cute.interface.flash_attn_func" + ) from exc + try: + package_version = importlib.metadata.version("flash-attn") + except importlib.metadata.PackageNotFoundError: + package_version = "unknown" + return op, package_version + + +class StrictFlashAttention4Core: + """Shared CUDA production core with the reduction-affecting FA4 knobs fixed.""" + + core_id = STRICT_ATTENTION_PRODUCTION_CORE_ID + strict_schedule = STRICT_ATTENTION_FA4_SCHEDULE_ID + backend_id = "flash_attention_4.cute" + api_source = _FA4_API_SOURCE + merge_order = "global_block_index" + accum_dtype = "fp32" + downcast_at = "final_write" + fallback = False + native_attention_arithmetic = True + production_ready = True + reference_only = False + num_splits = 1 + deterministic_backward = True + + def __init__( + self, + *, + split_kv: SplitKVSpec | None = None, + _op: Callable[..., Any] | None = None, + _package_version: str | None = None, + ) -> None: + requested = SplitKVSpec.disabled() if split_kv is None else split_kv + if not isinstance(requested, SplitKVSpec): + raise TypeError("split_kv must be a SplitKVSpec") + if requested.mode is not SplitKVMode.DISABLED: + raise ValueError("strict FA4 Attention requires Split-KV to be disabled") + if _op is None: + op, package_version = _load_fa4_cute_op() + else: + op = _op + package_version = "test-double" if _package_version is None else _package_version + self._validate_api(op) + self.split_kv = requested + self.package_version = package_version + self._op = op + + @staticmethod + def _validate_api(op: Callable[..., Any]) -> None: + try: + parameters = inspect.signature(op).parameters + except (TypeError, ValueError) as exc: + raise StrictFlashAttentionUnavailable( + "cannot inspect the FlashAttention CuTe API signature" + ) from exc + missing = sorted(_FA4_REQUIRED_PARAMETERS.difference(parameters)) + if missing: + raise StrictFlashAttentionUnavailable( + "FlashAttention CuTe API is missing strict controls: " + ", ".join(missing) + ) + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + **kwargs: Any, + ) -> DeterministicAttentionCoreResult: + return self.forward_with_lse(q, k, v, **kwargs) + + def forward_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: float | None = None, + key_padding_mask: torch.Tensor | None = None, + query_position_ids: torch.Tensor | None = None, + key_position_ids: torch.Tensor | None = None, + output_dtype: torch.dtype | None = None, + ) -> DeterministicAttentionCoreResult: + self._validate_inputs(q, k, v, key_padding_mask) + RLKernelDeterministicAttentionCore._validate_positions( + q, + k, + causal=causal, + query_position_ids=query_position_ids, + key_position_ids=key_position_ids, + ) + resolved_dtype = q.dtype if output_dtype is None else output_dtype + if resolved_dtype != q.dtype: + raise ValueError("strict Attention output_dtype must match the Q/K/V input dtype") + + q_fa = q.transpose(1, 2).contiguous() + k_fa = k.transpose(1, 2).contiguous() + v_fa = v.transpose(1, 2).contiguous() + result = self._op( + q_fa, + k_fa, + v_fa, + softmax_scale=scale, + causal=causal, + num_splits=self.num_splits, + pack_gqa=q.size(1) > k.size(1), + deterministic=self.deterministic_backward, + return_lse=True, + ) + if not isinstance(result, tuple) or len(result) != 2: + raise StrictFlashAttentionUnavailable( + "FlashAttention CuTe must return exactly (out, lse) when return_lse=True" + ) + out_fa, lse = result + if not isinstance(out_fa, torch.Tensor) or not isinstance(lse, torch.Tensor): + raise StrictFlashAttentionUnavailable("FlashAttention CuTe returned non-tensor output") + expected_out_shape = q_fa.shape + expected_lse_shape = (q.size(0), q.size(1), q.size(2)) + if tuple(out_fa.shape) != tuple(expected_out_shape): + raise StrictFlashAttentionUnavailable( + f"FlashAttention output must have shape {tuple(expected_out_shape)}" + ) + if tuple(lse.shape) != expected_lse_shape: + raise StrictFlashAttentionUnavailable( + f"FlashAttention LSE must have shape {expected_lse_shape}" + ) + if out_fa.dtype != resolved_dtype: + raise StrictFlashAttentionUnavailable( + "FlashAttention output dtype does not match the requested output dtype" + ) + if lse.dtype != torch.float32: + raise StrictFlashAttentionUnavailable( + "FlashAttention attention-domain LSE must be FP32" + ) + + out = out_fa.transpose(1, 2).contiguous() + return DeterministicAttentionCoreResult( + out=out, + lse=lse.contiguous(), + provenance={ + "strict_core_id": self.core_id, + "strict_schedule": self.strict_schedule, + "attention_backend": self.backend_id, + "fa_api_source": self.api_source, + "fa_package_version": self.package_version, + "num_splits": self.num_splits, + "deterministic_backward": self.deterministic_backward, + "dropout_p": 0.0, + "split_kv": self.split_kv.resolve(k.size(2), backend=self.backend_id).to_dict(), + "merge_order": self.merge_order, + "accum_dtype": self.accum_dtype, + "downcast_at": self.downcast_at, + "fallback": self.fallback, + "fallback_reason": None, + "native_attention_arithmetic": self.native_attention_arithmetic, + "production_ready": self.production_ready, + "reference_only": self.reference_only, + }, + ) + + @staticmethod + def _validate_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + key_padding_mask: torch.Tensor | None, + ) -> None: + if key_padding_mask is not None: + raise ValueError( + "strict FA4 core does not accept padding masks; materialize each logical row" + ) + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError("q/k/v must be 4-D [B, H, S, D]") + if q.size(0) != 1: + raise ValueError("strict FA4 core executes one logical batch row at a time") + if k.size(0) != 1 or v.size(0) != 1: + raise ValueError("q/k/v batch sizes must match") + if k.shape != v.shape or q.size(3) != k.size(3): + raise ValueError("k/v shapes and q/k/v head dimensions must match") + if q.size(1) % k.size(1) != 0: + raise ValueError("Q heads must be divisible by KV heads for GQA") + if q.dtype not in (torch.float16, torch.bfloat16): + raise ValueError("strict FA4 core supports FP16/BF16 only") + if k.dtype != q.dtype or v.dtype != q.dtype: + raise ValueError("q/k/v must share one dtype") + if not (q.is_cuda and k.is_cuda and v.is_cuda): + raise ValueError("strict FA4 core requires CUDA tensors") + if not (q.device == k.device == v.device): + raise ValueError("q/k/v must be on one CUDA device") + class FlashAttentionOp: """ diff --git a/rl_engine/kernels/ops/cuda/attention/flashinfer_paged_attention.py b/rl_engine/kernels/ops/cuda/attention/flashinfer_paged_attention.py new file mode 100644 index 00000000..1735a19e --- /dev/null +++ b/rl_engine/kernels/ops/cuda/attention/flashinfer_paged_attention.py @@ -0,0 +1,2173 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""FlashInfer paged-attention candidate for WS2 PR7. + +This module is intentionally opt-in. It adapts RL-Kernel's +``[B, H, S, D]`` attention tensors and PR6-style paged-KV metadata to +FlashInfer's paged attention wrappers, while recording the three PR7 contract +choices that affect rollout/training alignment: + +* Qwen3-exact RoPE fused into attention through ``ROPE_LLAMA``; +* split-KV policy, with auto split rejected when batch invariance is required; +* LSE export and provenance for downstream drift reports. +""" + +from __future__ import annotations + +import hashlib +import importlib +import inspect +import math +from dataclasses import dataclass, field +from typing import Any, Literal + +import torch + +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_FA4_SCHEDULE_ID, + STRICT_ATTENTION_PRODUCTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, + AttentionContractError, + SplitKVExecutionPlan, + SplitKVMode, + SplitKVRuntimeCoordinate, + SplitKVRuntimePlanEntry, + SplitKVRuntimePlanSet, + SplitKVSpec, + validate_split_kv_alignment, +) +from rl_engine.kernels.ops.cuda.attention.cp_comm import ( + AttentionCPCommunication, + AttentionCPCommunicationPlan, + AttentionCPMergedState, + AttentionCPPartialState, + AttentionParallelSpec, +) +from rl_engine.kernels.ops.cuda.attention.deterministic_attn import DeterministicAttentionCoreResult +from rl_engine.kernels.ops.cuda.attention.flash_attn import StrictFlashAttention4Core +from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + AttentionPartialState, + AttentionRingSchedule, + DeterministicCPAttentionReferenceOp, + build_reference_split_kv_runtime_plan_set, + merge_attention_partial_states, +) +from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp + +RoPEState = Literal["pre_rope", "post_rope"] +FlashInferAttentionMode = Literal["prefill", "decode"] +_FLASHINFER_MODULE = "flashinfer" + + +class FlashInferUnavailable(RuntimeError): + """Raised when FlashInfer cannot be imported or lacks required symbols.""" + + +@dataclass(frozen=True) +class FlashInferRoPEFusionConfig: + """Qwen3 RoPE settings used when FlashInfer performs RoPE inside attention.""" + + pos_encoding_mode: str = "ROPE_LLAMA" + rope_theta: float = 1_000_000.0 + rope_scale: float = 1.0 + rotary_dim: int | None = None + q_rope_state: RoPEState = "pre_rope" + k_cache_rope_state: RoPEState = "pre_rope" + + def validate(self, head_dim: int) -> None: + if self.pos_encoding_mode != "ROPE_LLAMA": + raise ValueError("PR7 RoPE fusion requires FlashInfer pos_encoding_mode='ROPE_LLAMA'") + if float(self.rope_theta) != 1_000_000.0: + raise ValueError("Qwen3-8B RoPE fusion requires rope_theta=1_000_000.0") + if float(self.rope_scale) != 1.0: + raise ValueError("Qwen3-8B RoPE fusion requires rope_scale=1.0") + rotary_dim = head_dim if self.rotary_dim is None else int(self.rotary_dim) + if rotary_dim != head_dim: + raise ValueError("FlashInfer PR7 candidate supports full-head Qwen3 RoPE only") + if self.q_rope_state != "pre_rope" or self.k_cache_rope_state != "pre_rope": + raise ValueError( + "FlashInfer ROPE_LLAMA attention fusion expects pre-RoPE Q and pre-RoPE K cache; " + "post-RoPE tensors would be rotated twice" + ) + + def provenance(self, head_dim: int) -> dict[str, Any]: + rotary_dim = head_dim if self.rotary_dim is None else int(self.rotary_dim) + return { + "rope_fusion": True, + "rope_fusion_boundary": "flashinfer_attention_kernel", + "pos_encoding_mode": self.pos_encoding_mode, + "rope_backend": "flashinfer", + "rope_theta": float(self.rope_theta), + "rope_scale": float(self.rope_scale), + "rotary_dim": rotary_dim, + "rope_layout": "qwen3_rotate_half_non_interleaved", + "q_rope_state": self.q_rope_state, + "k_cache_rope_state": self.k_cache_rope_state, + } + + +FlashInferSplitKVPolicy = SplitKVSpec + + +@dataclass(frozen=True) +class FlashInferPagedAttentionConfig: + """Runtime knobs for the opt-in FlashInfer paged attention candidate.""" + + mode: FlashInferAttentionMode = "prefill" + causal: bool = True + kv_layout: str = "NHD" + softmax_scale: float | None = None + return_lse: bool = True + require_batch_invariant: bool = True + workspace_size_bytes: int = 128 * 1024 * 1024 + rope: FlashInferRoPEFusionConfig = field(default_factory=FlashInferRoPEFusionConfig) + split_kv: SplitKVSpec = field(default_factory=SplitKVSpec.disabled) + cp_comm_plan: AttentionCPCommunicationPlan = field( + default_factory=lambda: AttentionCPCommunicationPlan( + parallel=AttentionParallelSpec(tp_world_size=2, cp_world_size=2), + ) + ) + require_cp_comm: bool = False + require_verified_arithmetic: bool = True + cp_communication: AttentionCPCommunication | None = None + strict_mode: bool = False + deterministic_core: Any | None = None + strict_rope_op: Any | None = None + + def validate(self, *, head_dim: int, query_len: int) -> None: + if self.mode not in {"prefill", "decode"}: + raise ValueError("mode must be 'prefill' or 'decode'") + if self.mode == "decode" and query_len != 1: + raise ValueError("BatchDecodeWithPagedKVCacheWrapper requires Sq == 1") + if self.kv_layout != "NHD": + raise ValueError("PR7 FlashInfer adapter currently supports kv_layout='NHD' only") + if not self.return_lse: + raise ValueError("PR7 requires attention-domain LSE export") + if self.workspace_size_bytes <= 0: + raise ValueError("workspace_size_bytes must be positive") + self.rope.validate(head_dim) + if not isinstance(self.split_kv, SplitKVSpec): + raise ValueError("split_kv must be a SplitKVSpec") + if self.require_batch_invariant and self.split_kv.mode is SplitKVMode.AUTO: + raise ValueError( + "FlashInfer auto split-KV is not a batch-invariant candidate; " + "use disabled split-KV or a fixed split size" + ) + self.cp_comm_plan.validate() + if self.require_cp_comm: + if self.cp_comm_plan.status != "implemented": + raise ValueError( + "require_cp_comm=True needs an implemented CP communication plan; " + "interface-only plans cannot produce owner-local partial states" + ) + elif self.cp_comm_plan.status != "interface_only": + raise ValueError("implemented CP communication plans require require_cp_comm=True") + if not isinstance(self.require_verified_arithmetic, bool): + raise ValueError("require_verified_arithmetic must be a bool") + if not isinstance(self.strict_mode, bool): + raise ValueError("strict_mode must be a bool") + if self.strict_mode: + if not self.require_batch_invariant: + raise ValueError("strict Attention requires batch invariance") + if self.split_kv.mode is not SplitKVMode.DISABLED: + raise ValueError("strict Attention requires Split-KV to be disabled") + if self.deterministic_core is not None: + _validate_strict_core(self.deterministic_core) + if self.require_cp_comm: + if self.cp_communication is None: + raise ValueError("strict CP Attention requires a communication adapter") + for method_name in ( + "all_gather_query", + "all_gather_kv", + "all_gather_position_ids", + "reduce_scatter_strict_result", + ): + if not callable(getattr(self.cp_communication, method_name, None)): + raise ValueError( + "strict CP communication adapter must implement " f"{method_name}" + ) + + +@dataclass(frozen=True) +class FlashInferPagedKVPlan: + """FlashInfer paged-KV tensors derived from PR6-style metadata.""" + + qo_indptr: torch.Tensor + paged_kv_indptr: torch.Tensor + paged_kv_indices: torch.Tensor + paged_kv_last_page_len: torch.Tensor + kv_seq_lens: torch.Tensor + seq_lens_q: torch.Tensor + page_size: int + physical_page_count_per_batch: int + logical_block_counts: tuple[int, ...] + + def provenance(self) -> dict[str, Any]: + return { + "page_size": self.page_size, + "physical_page_count_per_batch": self.physical_page_count_per_batch, + "logical_block_counts": list(self.logical_block_counts), + "qo_indptr": self.qo_indptr.detach().cpu().tolist(), + "paged_kv_indptr": self.paged_kv_indptr.detach().cpu().tolist(), + "paged_kv_indices": self.paged_kv_indices.detach().cpu().tolist(), + "paged_kv_last_page_len": self.paged_kv_last_page_len.detach().cpu().tolist(), + "kv_seq_lens": self.kv_seq_lens.detach().cpu().tolist(), + "seq_lens_q": self.seq_lens_q.detach().cpu().tolist(), + } + + +@dataclass(frozen=True) +class FlashInferAttentionResult: + """Output of the FlashInfer PR7 candidate.""" + + out: torch.Tensor + lse: torch.Tensor + provenance: dict[str, Any] + + +def build_flashinfer_paged_kv_plan( + metadata: Any, + *, + batch_size: int, + query_len: int, + cache_capacity: int, + device: torch.device, +) -> FlashInferPagedKVPlan: + """Convert PR6-style paged metadata to FlashInfer page table tensors.""" + + page_size = _positive_int(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.kv_seq_lens.shape != (batch_size,): + raise ValueError("kv_seq_lens must have shape [B]") + if metadata.block_table.ndim != 2 or metadata.block_table.size(0) != batch_size: + raise ValueError("block_table must have shape [B, max_blocks]") + + qo_indptr = [0] + paged_kv_indptr = [0] + paged_kv_indices: list[int] = [] + paged_kv_last_page_len: list[int] = [] + kv_seq_lens: list[int] = [] + seq_lens_q: list[int] = [] + logical_block_counts: list[int] = [] + for batch_index in range(batch_size): + seq_len = _positive_int(int(metadata.kv_seq_lens[batch_index].item()), "kv_seq_len") + if seq_len > cache_capacity: + raise ValueError("kv_seq_len must not exceed cache capacity") + block_count = (seq_len + page_size - 1) // page_size + if block_count > metadata.block_table.size(1): + raise ValueError("block_table does not contain enough logical KV blocks") + kv_seq_lens.append(seq_len) + seq_lens_q.append(query_len) + logical_block_counts.append(block_count) + qo_indptr.append(qo_indptr[-1] + query_len) + paged_kv_indptr.append(paged_kv_indptr[-1] + block_count) + last_len = ((seq_len - 1) % page_size) + 1 + paged_kv_last_page_len.append(last_len) + for logical_block in range(block_count): + local_page = int(metadata.block_table[batch_index, logical_block].item()) + if local_page < 0 or local_page >= physical_page_count: + raise ValueError("block_table contains an out-of-range physical page") + paged_kv_indices.append(batch_index * physical_page_count + local_page) + active_pages = metadata.block_table[batch_index, :block_count] + if torch.unique(active_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") + _validate_metadata_logical_positions( + metadata, + batch_index=batch_index, + seq_len=seq_len, + page_size=page_size, + block_count=block_count, + device=device, + ) + + return FlashInferPagedKVPlan( + qo_indptr=torch.tensor(qo_indptr, device=device, dtype=torch.int32), + paged_kv_indptr=torch.tensor(paged_kv_indptr, device=device, dtype=torch.int32), + paged_kv_indices=torch.tensor(paged_kv_indices, device=device, dtype=torch.int32), + paged_kv_last_page_len=torch.tensor( + paged_kv_last_page_len, + device=device, + dtype=torch.int32, + ), + kv_seq_lens=torch.tensor(kv_seq_lens, device=device, dtype=torch.int32), + seq_lens_q=torch.tensor(seq_lens_q, device=device, dtype=torch.int32), + page_size=page_size, + physical_page_count_per_batch=physical_page_count, + logical_block_counts=tuple(logical_block_counts), + ) + + +def materialize_flashinfer_paged_kv_cache( + k_cache: torch.Tensor, + v_cache: torch.Tensor, + *, + page_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Flatten ``[B, Hkv, P*page, D]`` caches to FlashInfer NHD pages.""" + + if k_cache.shape != v_cache.shape: + raise ValueError("k_cache and v_cache must have matching shape") + if k_cache.ndim != 4: + raise ValueError("k_cache and v_cache must have shape [B, Hkv, cache_capacity, D]") + batch, heads, cache_capacity, head_dim = k_cache.shape + if cache_capacity % page_size != 0: + raise ValueError("cache capacity must be divisible by page_size") + page_count = cache_capacity // page_size + k_pages = ( + k_cache.contiguous() + .reshape(batch, heads, page_count, page_size, head_dim) + .permute(0, 2, 3, 1, 4) + .reshape(batch * page_count, page_size, heads, head_dim) + .contiguous() + ) + v_pages = ( + v_cache.contiguous() + .reshape(batch, heads, page_count, page_size, head_dim) + .permute(0, 2, 3, 1, 4) + .reshape(batch * page_count, page_size, heads, head_dim) + .contiguous() + ) + return k_pages, v_pages + + +class _NativeFlashInferRuntimeAdapter: + """Expose strict provenance from FlashInfer's materialized FA2 plan. + + Upstream FlashInfer does not provide the RL-Kernel provenance callbacks. + Its FA2 scheduler does, however, materialize the request/tile schedule and + the token chunk size into the wrapper's caller-owned workspace. Read that + schedule after plan() so strict acceptance describes the kernel plan that + will actually run instead of merely echoing requested knobs. + """ + + def __init__(self, wrapper: Any, cfg: FlashInferPagedAttentionConfig) -> None: + self._wrapper = wrapper + self._cfg = cfg + self._plan_kwargs: dict[str, Any] | None = None + + def __getattr__(self, name: str) -> Any: + return getattr(self._wrapper, name) + + def plan(self, *args: Any, **kwargs: Any) -> Any: + result = self._wrapper.plan(*args, **kwargs) + self._plan_kwargs = dict(kwargs) + self._validate_materialized_plan() + return result + + def get_actual_split_kv_plan(self) -> list[dict[str, Any]]: + seq_lens, page_size, fixed_split_pages = self._runtime_split_layout() + result = [] + for seq_len in seq_lens: + if fixed_split_pages is None: + boundaries = [(0, (seq_len + page_size - 1) // page_size)] + mode = SplitKVMode.DISABLED.value + else: + page_count = (seq_len + page_size - 1) // page_size + boundaries = [ + (start, min(start + fixed_split_pages, page_count)) + for start in range(0, page_count, fixed_split_pages) + ] + mode = SplitKVMode.FIXED.value + result.append( + { + "mode": mode, + "split_size": fixed_split_pages, + "split_size_unit": "pages", + "boundary_unit": "pages", + "boundaries": boundaries, + "fallback": False, + "fallback_reason": None, + } + ) + return result + + def get_actual_split_kv_plan_set(self) -> dict[str, Any]: + seq_lens, page_size, fixed_split_pages = self._runtime_split_layout() + parallel = self._cfg.cp_comm_plan.parallel + entries = [] + for batch_index, total in enumerate(seq_lens): + owner_ranges = _balanced_token_ranges(total, parallel.cp_world_size) + for tp_rank in range(parallel.tp_world_size): + for cp_rank in range(parallel.cp_world_size): + for owner_cp_rank, (owner_start, owner_end) in enumerate(owner_ranges): + if owner_start % page_size or owner_end % page_size: + raise FlashInferUnavailable( + "strict FlashInfer CP owner ranges must align to KV pages" + ) + owner_pages = (owner_end - owner_start) // page_size + if fixed_split_pages is None: + boundaries = [(0, owner_pages)] + mode = SplitKVMode.DISABLED.value + else: + boundaries = [ + (start, min(start + fixed_split_pages, owner_pages)) + for start in range(0, owner_pages, fixed_split_pages) + ] + mode = SplitKVMode.FIXED.value + entries.append( + { + "batch_index": batch_index, + "tp_rank": tp_rank, + "cp_rank": cp_rank, + "owner_cp_rank": owner_cp_rank, + "expected_kv_range": [owner_start, owner_end], + "mode": mode, + "split_size": fixed_split_pages, + "split_size_unit": "pages", + "boundary_unit": "pages", + "boundaries": boundaries, + "merge_order": "global_block_index", + "accum_dtype": "fp32", + "downcast_at": "final_write", + "fallback": False, + "fallback_reason": None, + } + ) + return { + "batch_size": len(seq_lens), + "tp_world_size": parallel.tp_world_size, + "cp_world_size": parallel.cp_world_size, + "total_kv_tokens": list(seq_lens), + "entries": entries, + } + + def get_attention_arithmetic_provenance(self) -> dict[str, str]: + self._validate_materialized_plan() + return { + "accum_dtype": "fp32", + "downcast_at": "final_write", + "lse_dtype": "fp32", + "backend_lse_log_base": "2", + "export_lse_log_base": "e", + "source": "flashinfer_fa2_materialized_plan", + } + + @staticmethod + def normalize_lse(lse: torch.Tensor) -> torch.Tensor: + """Convert FlashInfer FA2's log2 LSE to the natural-log contract.""" + + return lse * math.log(2.0) + + def _runtime_split_layout(self) -> tuple[tuple[int, ...], int, int | None]: + self._validate_materialized_plan() + assert self._plan_kwargs is not None + seq_lens_raw = self._plan_kwargs.get("seq_lens") + if not isinstance(seq_lens_raw, torch.Tensor): + raise FlashInferUnavailable("native FlashInfer provenance requires explicit seq_lens") + seq_lens = tuple(int(value) for value in seq_lens_raw.detach().cpu().tolist()) + page_size = int(self._plan_kwargs["page_size"]) + disabled = bool(self._plan_kwargs.get("disable_split_kv", False)) + fixed_split_pages_raw = self._plan_kwargs.get("fixed_split_size") + fixed_split_pages = ( + None if disabled or fixed_split_pages_raw is None else int(fixed_split_pages_raw) + ) + return seq_lens, page_size, fixed_split_pages + + def _validate_materialized_plan(self) -> None: + if self._plan_kwargs is None: + raise FlashInferUnavailable("FlashInfer plan() has not materialized a runtime plan") + if getattr(self._wrapper, "_backend", None) != "fa2": + raise FlashInferUnavailable( + "strict native FlashInfer provenance currently requires the FA2 backend" + ) + plan_info = getattr(self._wrapper, "_plan_info", None) + if plan_info is None or not hasattr(plan_info, "__getitem__") or len(plan_info) != 15: + raise FlashInferUnavailable( + "native FlashInfer FA2 did not expose the expected PrefillPlanInfo" + ) + seq_lens_raw = self._plan_kwargs.get("seq_lens") + if not isinstance(seq_lens_raw, torch.Tensor): + raise FlashInferUnavailable("native FlashInfer provenance requires explicit seq_lens") + seq_lens = tuple(int(value) for value in seq_lens_raw.detach().cpu().tolist()) + page_size = int(self._plan_kwargs["page_size"]) + disabled = bool(self._plan_kwargs.get("disable_split_kv", False)) + fixed_split_pages_raw = self._plan_kwargs.get("fixed_split_size") + fixed_split_pages = ( + None if disabled or fixed_split_pages_raw is None else int(fixed_split_pages_raw) + ) + runtime_split = bool(plan_info[14]) + if disabled and runtime_split: + raise FlashInferUnavailable( + "FlashInfer materialized split-KV despite disable_split_kv=True" + ) + if fixed_split_pages is not None: + expected_chunk_tokens = fixed_split_pages * page_size + actual_chunk_tokens = self._workspace_i32(int(plan_info[9]), 1)[0] + if actual_chunk_tokens != expected_chunk_tokens: + raise FlashInferUnavailable( + "FlashInfer materialized KV chunk differs from fixed_split_size" + ) + expected_runtime_split = any(seq_len > expected_chunk_tokens for seq_len in seq_lens) + if runtime_split != expected_runtime_split: + raise FlashInferUnavailable( + "FlashInfer materialized split flag differs from fixed Split-KV plan" + ) + padded_batch_size = int(plan_info[0]) + request_indices = self._workspace_i32(int(plan_info[4]), padded_batch_size) + kv_tile_indices = self._workspace_i32(int(plan_info[6]), padded_batch_size) + actual_tiles = set(zip(request_indices, kv_tile_indices, strict=True)) + expected_tiles = { + (batch_index, tile_index) + for batch_index, seq_len in enumerate(seq_lens) + for tile_index in range( + 1 + if fixed_split_pages is None + else ( + (seq_len + fixed_split_pages * page_size - 1) // (fixed_split_pages * page_size) + ) + ) + } + if actual_tiles != expected_tiles: + raise FlashInferUnavailable( + "FlashInfer materialized request/KV-tile schedule differs from the strict plan" + ) + + def _workspace_i32(self, byte_offset: int, count: int) -> tuple[int, ...]: + workspace = getattr(self._wrapper, "_pin_memory_int_workspace_buffer", None) + if not isinstance(workspace, torch.Tensor) or workspace.device.type != "cpu": + raise FlashInferUnavailable( + "native FlashInfer did not expose its materialized host plan workspace" + ) + byte_count = count * torch.tensor([], dtype=torch.int32).element_size() + values = workspace.narrow(0, byte_offset, byte_count).view(torch.int32) + return tuple(int(value) for value in values.tolist()) + + +def _balanced_token_ranges(total: int, parts: int) -> tuple[tuple[int, int], ...]: + base, extra = divmod(total, parts) + ranges = [] + start = 0 + for index in range(parts): + end = start + base + (1 if index < extra else 0) + ranges.append((start, end)) + start = end + return tuple(ranges) + + +class FlashInferQwen3PagedAttentionOp: + """Opt-in FlashInfer paged attention backend candidate for #235 PR7.""" + + op_class = "attention" + + def __init__(self, *, flashinfer_module: Any | None = None) -> None: + self._flashinfer_module = flashinfer_module + + def __call__( + self, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + metadata: Any, + *, + config: FlashInferPagedAttentionConfig | None = None, + ) -> FlashInferAttentionResult: + return self.forward(q, k_cache, v_cache, metadata, config=config) + + def forward( + self, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + metadata: Any, + *, + config: FlashInferPagedAttentionConfig | None = None, + ) -> FlashInferAttentionResult: + """Run FlashInfer paged attention and return RL-Kernel shaped tensors. + + Args: + q: pre-RoPE query tensor, ``[B, Hq, Sq, D]``. + k_cache: pre-RoPE paged key cache, ``[B, Hkv, cache_capacity, D]``. + v_cache: paged value cache, ``[B, Hkv, cache_capacity, D]``. + metadata: PR6 ``DecodeKVCacheMetadata``-compatible object. + config: PR7 FlashInfer backend knobs. + """ + + _validate_qkv_cache(q, k_cache, v_cache) + cfg = FlashInferPagedAttentionConfig() if config is None else config + batch_size, q_heads, query_len, head_dim = q.shape + kv_heads = k_cache.size(1) + cfg.validate(head_dim=head_dim, query_len=query_len) + if cfg.require_cp_comm and cfg.strict_mode: + return self._run_strict_cp(q, k_cache, v_cache, metadata, cfg) + if self._flashinfer_module is None and q.device.type != "cuda" and not cfg.require_cp_comm: + raise FlashInferUnavailable("FlashInfer PR7 candidate requires CUDA tensors") + + plan = build_flashinfer_paged_kv_plan( + metadata, + batch_size=batch_size, + query_len=query_len, + cache_capacity=k_cache.size(2), + device=q.device, + ) + _validate_flashinfer_rope_metadata(metadata, cfg, q) + _validate_flashinfer_prefix_cache(q, k_cache, v_cache, metadata, cfg) + if cfg.require_cp_comm: + return self._forward_deterministic_cp_fallback( + q, + k_cache, + v_cache, + metadata, + cfg, + plan, + ) + if cfg.strict_mode: + return self._run_strict_core(q, k_cache, v_cache, metadata, cfg, plan) + q_flat = q.transpose(1, 2).reshape(batch_size * query_len, q_heads, head_dim).contiguous() + k_pages, v_pages = materialize_flashinfer_paged_kv_cache( + k_cache, + v_cache, + page_size=plan.page_size, + ) + wrapper = self._make_wrapper(cfg, q) + applied_plan_kwargs = self._plan_wrapper( + wrapper, + cfg, + plan, + q_dtype=q.dtype, + q_heads=q_heads, + kv_heads=kv_heads, + head_dim=head_dim, + query_len=query_len, + ) + actual_split_plans = self._actual_split_kv_plans(wrapper, cfg, plan) + actual_split_plan_set = self._actual_split_kv_plan_set( + wrapper, + cfg, + plan, + ) + arithmetic = self._actual_arithmetic_semantics(wrapper, cfg) + out_flat, lse_flat = self._run_wrapper(wrapper, q_flat, (k_pages, v_pages), cfg) + self._validate_runtime_outputs( + out_flat, + lse_flat, + q, + require_fp32_output=cfg.require_cp_comm, + ) + out = _restore_out(out_flat, batch_size=batch_size, query_len=query_len) + lse = _restore_lse( + lse_flat, + batch_size=batch_size, + query_len=query_len, + q_heads=q_heads, + ) + provenance = { + "attention_backend": "flashinfer", + "requested_backend": "flashinfer_qwen3_rope_paged_attention", + "actual_backend": f"flashinfer_batch_{cfg.mode}_paged_kv", + "attention_mode": cfg.mode, + "materialization": "flashinfer_rope_llama_paged_kv", + "kv_layout": cfg.kv_layout, + "causal": cfg.causal, + "softmax_scale": cfg.softmax_scale, + "lse_domain": "attention", + "lse_exported": True, + **arithmetic, + "fallback": False, + "fallback_reason": None, + "paged_kv_policy": "flashinfer_page_table", + } + provenance.update(cfg.rope.provenance(head_dim)) + provenance.update( + _split_kv_provenance( + cfg.split_kv, + actual_split_plans, + applied_plan_kwargs=applied_plan_kwargs, + require_batch_invariant=cfg.require_batch_invariant, + ) + ) + provenance["actual_split_kv_plan_set"] = ( + None if actual_split_plan_set is None else actual_split_plan_set.to_dict() + ) + provenance.update(cfg.cp_comm_plan.provenance()) + provenance["cp_comm_required"] = cfg.require_cp_comm + provenance.update(plan.provenance()) + return FlashInferAttentionResult( + out=out.to(dtype=q.dtype), + lse=lse, + provenance=provenance, + ) + + @staticmethod + def _run_strict_core( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + metadata: Any, + cfg: FlashInferPagedAttentionConfig, + paged_plan: FlashInferPagedKVPlan, + ) -> FlashInferAttentionResult: + """Use FlashInfer only for paged-KV layout, never Attention arithmetic.""" + + core = _resolve_strict_core(cfg) + _validate_strict_core(core) + rope = _resolve_strict_rope(cfg) + logical_k, logical_v, key_positions = _materialize_strict_logical_kv( + k_cache, v_cache, metadata, paged_plan + ) + query_positions = getattr(metadata, "query_position_ids", None) + if not isinstance(query_positions, torch.Tensor): + raise FlashInferUnavailable( + "strict Attention requires query_position_ids runtime metadata" + ) + + outputs: list[torch.Tensor] = [] + lses: list[torch.Tensor] = [] + row_provenance: list[dict[str, object]] = [] + for batch_index, seq_len_value in enumerate(paged_plan.kv_seq_lens.tolist()): + seq_len = int(seq_len_value) + q_row = q[batch_index : batch_index + 1] + k_row = logical_k[batch_index : batch_index + 1, :, :seq_len, :] + v_row = logical_v[batch_index : batch_index + 1, :, :seq_len, :] + q_pos = query_positions[batch_index : batch_index + 1] + k_pos = key_positions[batch_index : batch_index + 1, :seq_len] + q_ready = _apply_strict_rope(rope, q_row, q_pos, cfg.rope.rope_theta) + k_ready = _apply_strict_rope(rope, k_row, k_pos, cfg.rope.rope_theta) + result = core.forward_with_lse( + q_ready, + k_ready, + v_row, + causal=cfg.causal, + scale=cfg.softmax_scale, + query_position_ids=q_pos, + key_position_ids=k_pos, + output_dtype=q.dtype, + ) + _validate_strict_core_result(result, core) + outputs.append(result.out) + lses.append(result.lse) + row_provenance.append(dict(result.provenance)) + + provenance = _strict_attention_provenance( + cfg, + core_provenance=row_provenance[0], + materialization="flashinfer_paged_kv_layout_shared_core", + cp_required=False, + rope=rope, + ) + provenance["strict_core_row_plans"] = [item["split_kv"] for item in row_provenance] + provenance.update(cfg.cp_comm_plan.provenance()) + provenance.update(paged_plan.provenance()) + return FlashInferAttentionResult( + out=torch.cat(outputs, dim=0), + lse=torch.cat(lses, dim=0), + provenance=provenance, + ) + + @staticmethod + def _run_strict_cp( + q_local: torch.Tensor, + k_local: torch.Tensor, + v_local: torch.Tensor, + metadata: Any, + cfg: FlashInferPagedAttentionConfig, + ) -> FlashInferAttentionResult: + """AG full Q/K/V, execute the shared core, then RS final Out/LSE.""" + + plan = cfg.cp_comm_plan + communication = cfg.cp_communication + assert communication is not None + if any(tensor.requires_grad for tensor in (q_local, k_local, v_local)) and not getattr( + communication, "supports_autograd", False + ): + raise FlashInferUnavailable( + "strict training requires the autograd-capable self-owned CUDA AG/RS backend" + ) + query_start, query_end = plan.query_token_ranges[plan.parallel.cp_rank] + key_start, key_end = _cp_owner_ranges(plan)[plan.parallel.cp_rank] + if q_local.size(2) != query_end - query_start: + raise ValueError("strict CP Q must contain only the owner-local query range") + if k_local.size(2) != key_end - key_start: + raise ValueError("strict CP K/V must contain only the owner-local KV range") + if getattr(metadata, "q_rope_state", None) != "pre_rope" or ( + getattr(metadata, "k_cache_rope_state", None) != "pre_rope" + ): + raise ValueError("strict CP Attention requires pre-RoPE Q and K") + local_query_positions = _strict_local_query_positions( + metadata, q_local, (query_start, query_end) + ) + local_key_positions = _strict_local_key_positions(metadata, k_local, (key_start, key_end)) + + global_q = communication.all_gather_query(q_local, plan) + global_k, global_v = communication.all_gather_kv(k_local, v_local, plan) + global_q_positions, global_k_positions = communication.all_gather_position_ids( + local_query_positions, + local_key_positions, + plan, + ) + expected_q_tokens = plan.query_token_ranges[-1][1] + expected_kv_range = plan.expected_kv_token_range + if expected_kv_range is None: + raise FlashInferUnavailable("strict CP Attention requires an expected KV range") + expected_k_tokens = expected_kv_range[1] - expected_kv_range[0] + ring_schedule = AttentionRingSchedule.build( + expected_k_tokens, + cp_world_size=plan.parallel.cp_world_size, + kv_chunk_size=None, + ) + if global_q.size(2) != expected_q_tokens: + raise FlashInferUnavailable("strict AG(Q) returned the wrong global width") + if global_k.size(2) != expected_k_tokens or global_v.shape != global_k.shape: + raise FlashInferUnavailable("strict AG(K/V) returned the wrong global shape") + + rope = _resolve_strict_rope(cfg) + q_ready = _apply_strict_rope(rope, global_q, global_q_positions, cfg.rope.rope_theta) + k_ready = _apply_strict_rope(rope, global_k, global_k_positions, cfg.rope.rope_theta) + core = _resolve_strict_core(cfg) + _validate_strict_core(core) + + outputs: list[torch.Tensor] = [] + lses: list[torch.Tensor] = [] + row_provenance: list[dict[str, object]] = [] + for batch_index in range(global_q.size(0)): + result = core.forward_with_lse( + q_ready[batch_index : batch_index + 1], + k_ready[batch_index : batch_index + 1], + global_v[batch_index : batch_index + 1], + causal=cfg.causal, + scale=cfg.softmax_scale, + query_position_ids=global_q_positions[batch_index : batch_index + 1], + key_position_ids=global_k_positions[batch_index : batch_index + 1], + output_dtype=q_local.dtype, + ) + _validate_strict_core_result(result, core) + outputs.append(result.out) + lses.append(result.lse) + row_provenance.append(dict(result.provenance)) + full_out = torch.cat(outputs, dim=0) + full_lse = torch.cat(lses, dim=0) + local_result = communication.reduce_scatter_strict_result(full_out, full_lse, plan) + + provenance = _strict_attention_provenance( + cfg, + core_provenance=row_provenance[0], + materialization="ag_qkv_positions_shared_core_rs", + cp_required=True, + rope=rope, + ) + provenance.update(plan.provenance()) + provenance.update( + { + "strict_core_row_plans": [item["split_kv"] for item in row_provenance], + "strict_full_qkv_all_gather": True, + "strict_position_ids_all_gather": True, + "strict_split_kv": "disabled", + "strict_comm_autograd": bool(getattr(communication, "supports_autograd", False)), + "strict_local_query_range": [query_start, query_end], + "strict_local_kv_range": [key_start, key_end], + **ring_schedule.provenance(), + "ring_schedule_default": True, + "ring_partial_arithmetic": False, + } + ) + return FlashInferAttentionResult( + out=local_result.out, + lse=local_result.lse, + provenance=provenance, + ) + + @staticmethod + def _forward_deterministic_cp_fallback( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + metadata: Any, + cfg: FlashInferPagedAttentionConfig, + paged_plan: FlashInferPagedKVPlan, + ) -> FlashInferAttentionResult: + """Execute the owner-local CP contract without relabeling full-KV output. + + FlashInfer's public paged wrapper does not expose one FP32 partial state + per CP-owned KV block. Until it does, strict CP execution uses the + deterministic reference arithmetic while preserving the production + communication boundary: AG Q, owner-local partials, ordered FP32 merge, + and RS of the merged ``(Out, LSE)`` state. + """ + + communication = cfg.cp_communication + if communication is None: + raise FlashInferUnavailable( + "strict CP fallback requires an AttentionCPCommunication implementation" + ) + cp_plan = cfg.cp_comm_plan + total_query_tokens = cp_plan.query_token_ranges[-1][1] + if q.size(2) != total_query_tokens: + raise ValueError("strict CP fallback expects the complete logical Q sequence before AG") + kv_seq_lens = tuple(int(value) for value in paged_plan.kv_seq_lens.tolist()) + if len(set(kv_seq_lens)) != 1: + raise ValueError( + "strict CP fallback currently requires equal KV lengths across the batch" + ) + total_kv_tokens = kv_seq_lens[0] + if cp_plan.expected_kv_token_range != (0, total_kv_tokens): + raise ValueError("CP block manifest must cover the complete logical KV token range") + + query_start, query_end = cp_plan.query_token_ranges[cp_plan.parallel.cp_rank] + local_q = q[:, :, query_start:query_end, :].contiguous() + try: + gathered_q = communication.all_gather_query(local_q, cp_plan) + except AttributeError as exc: + raise FlashInferUnavailable( + "strict CP communication must implement the query AllGather boundary" + ) from exc + if gathered_q.shape != q.shape: + raise FlashInferUnavailable( + "query AllGather did not reconstruct the complete logical Q tensor" + ) + + logical_k, logical_v, logical_key_positions = _materialize_logical_kv_cache( + k_cache, + v_cache, + metadata, + total_kv_tokens=total_kv_tokens, + ) + rope = NativeRoPEOp() + q_ready = gathered_q + if cfg.rope.q_rope_state == "pre_rope": + q_ready = rope.forward_fp32( + gathered_q, + metadata.query_position_ids, + theta=cfg.rope.rope_theta, + ).to(gathered_q.dtype) + k_ready = logical_k + if cfg.rope.k_cache_rope_state == "pre_rope": + k_ready = rope.forward_fp32( + logical_k, + logical_key_positions, + theta=cfg.rope.rope_theta, + ).to(logical_k.dtype) + + reference = DeterministicCPAttentionReferenceOp() + local_states = [] + for block in cp_plan.expected_blocks: + if block.owner_cp_rank != cp_plan.parallel.cp_rank: + continue + partial = reference.local_partial_state( + q_ready, + k_ready[:, :, block.kv_block_start : block.kv_block_end, :], + logical_v[:, :, block.kv_block_start : block.kv_block_end, :], + q_start=0, + k_start=block.kv_block_start, + total_kv_len=total_kv_tokens, + total_query_len=q_ready.size(2), + causal=cfg.causal, + scale=cfg.softmax_scale, + query_position_offsets=metadata.query_position_ids[:, 0], + key_position_offsets=logical_key_positions[:, 0], + ) + local_states.append( + AttentionCPPartialState( + out=partial.out, + lse=partial.lse, + block=block, + ) + ) + gathered_states = communication.all_gather_partial_states( + tuple(local_states), + cp_plan, + ) + merged = merge_attention_partial_states( + [ + AttentionPartialState( + out=state.out, + lse=state.lse, + block_start=state.block.kv_block_start, + block_end=state.block.kv_block_end, + ) + for state in gathered_states + ] + ) + local = communication.reduce_scatter_merged_state( + AttentionCPMergedState(out=merged.out, lse=merged.lse), + cp_plan, + ) + + kv_chunk_size = ( + cfg.split_kv.fixed_split_size if cfg.split_kv.mode is SplitKVMode.FIXED else None + ) + runtime_plan_set = build_reference_split_kv_runtime_plan_set( + kv_seq_lens, + tp_world_size=cp_plan.parallel.tp_world_size, + cp_world_size=cp_plan.parallel.cp_world_size, + kv_chunk_size=kv_chunk_size, + backend="deterministic_cp_fallback", + ) + actual_plans = tuple( + cfg.split_kv.resolve(length, backend="deterministic_cp_fallback") + for length in kv_seq_lens + ) + applied_split_knobs = ( + {"fixed_split_size": cfg.split_kv.fixed_split_size} + if cfg.split_kv.mode is SplitKVMode.FIXED + else {"disable_split_kv": True} + ) + provenance = { + "attention_backend": "deterministic_cp_fallback", + "requested_backend": "flashinfer_qwen3_rope_paged_attention", + "actual_backend": "rlkernel_deterministic_cp_reference", + "attention_mode": cfg.mode, + "materialization": "logical_paged_kv_owner_local", + "kv_layout": cfg.kv_layout, + "causal": cfg.causal, + "softmax_scale": cfg.softmax_scale, + "lse_domain": "attention", + "lse_exported": True, + "accum_dtype": "fp32", + "downcast_at": "final_write", + "lse_dtype": "fp32", + "arithmetic_plan_source": "rlkernel_deterministic_cp_reference", + "arithmetic_semantics_verified": True, + "fallback": True, + "fallback_reason": "flashinfer_owner_local_cp_partial_api_unavailable", + "paged_kv_policy": "validated_logical_page_table", + "cp_comm_required": True, + "query_ag": "cp_rank_order", + "query_range": [query_start, query_end], + } + provenance.update(cfg.rope.provenance(q.size(-1))) + provenance.update( + _split_kv_provenance( + cfg.split_kv, + actual_plans, + applied_plan_kwargs=applied_split_knobs, + require_batch_invariant=cfg.require_batch_invariant, + ) + ) + provenance["actual_split_kv_plan_set"] = runtime_plan_set.to_dict() + provenance.update(cp_plan.provenance()) + provenance.update(paged_plan.provenance()) + return FlashInferAttentionResult( + out=local.out.to(dtype=q.dtype), + lse=local.lse, + provenance=provenance, + ) + + def _load_flashinfer(self) -> Any: + if self._flashinfer_module is not None: + return self._flashinfer_module + try: + self._flashinfer_module = importlib.import_module(_FLASHINFER_MODULE) + except (ImportError, OSError, RuntimeError) as exc: + raise FlashInferUnavailable(str(exc)) from exc + return self._flashinfer_module + + def _make_wrapper(self, cfg: FlashInferPagedAttentionConfig, q: torch.Tensor) -> Any: + module = self._load_flashinfer() + namespace_name = "decode" if cfg.mode == "decode" else "prefill" + class_name = ( + "BatchDecodeWithPagedKVCacheWrapper" + if cfg.mode == "decode" + else "BatchPrefillWithPagedKVCacheWrapper" + ) + namespace = getattr(module, namespace_name, None) + wrapper_cls = getattr(namespace, class_name, None) if namespace is not None else None + if wrapper_cls is None: + raise FlashInferUnavailable(f"flashinfer.{namespace_name}.{class_name} is unavailable") + + workspace = torch.zeros(cfg.workspace_size_bytes, dtype=torch.uint8, device=q.device) + constructor_kwargs: dict[str, Any] = {"kv_layout": cfg.kv_layout} + if cfg.mode == "decode": + constructor_kwargs["use_tensor_cores"] = True + try: + wrapper = wrapper_cls(workspace, **constructor_kwargs) + except TypeError: + try: + wrapper = wrapper_cls( + float_workspace_buffer=workspace, + **constructor_kwargs, + ) + except TypeError: + constructor_kwargs.pop("use_tensor_cores", None) + try: + wrapper = wrapper_cls(workspace, **constructor_kwargs) + except TypeError as exc: + raise FlashInferUnavailable( + f"could not instantiate flashinfer.{namespace_name}.{class_name}" + ) from exc + if type(wrapper).__module__.startswith("flashinfer."): + return _NativeFlashInferRuntimeAdapter(wrapper, cfg) + return wrapper + + @staticmethod + def _plan_wrapper( + wrapper: Any, + cfg: FlashInferPagedAttentionConfig, + plan: FlashInferPagedKVPlan, + *, + q_dtype: torch.dtype, + q_heads: int, + kv_heads: int, + head_dim: int, + query_len: int, + ) -> dict[str, Any]: + plan_kwargs = { + "num_qo_heads": q_heads, + "num_kv_heads": kv_heads, + "page_size": plan.page_size, + "pos_encoding_mode": cfg.rope.pos_encoding_mode, + "rope_scale": float(cfg.rope.rope_scale), + "rope_theta": float(cfg.rope.rope_theta), + "q_data_type": q_dtype, + "kv_data_type": q_dtype, + "o_data_type": torch.float32 if cfg.require_cp_comm else q_dtype, + "seq_lens": plan.kv_seq_lens, + } + if cfg.mode == "decode": + plan_kwargs.update( + { + "indptr": plan.paged_kv_indptr, + "indices": plan.paged_kv_indices, + "last_page_len": plan.paged_kv_last_page_len, + "head_dim": head_dim, + "q_len_per_req": query_len, + } + ) + else: + plan_kwargs.update( + { + "qo_indptr": plan.qo_indptr, + "paged_kv_indptr": plan.paged_kv_indptr, + "paged_kv_indices": plan.paged_kv_indices, + "paged_kv_last_page_len": plan.paged_kv_last_page_len, + "head_dim_qk": head_dim, + "causal": cfg.causal, + "seq_lens_q": plan.seq_lens_q, + } + ) + scale = cfg.softmax_scale + if scale is not None: + plan_kwargs["sm_scale"] = float(scale) + plan_kwargs.update( + _flashinfer_split_kv_plan_kwargs( + cfg.split_kv, + page_size=plan.page_size, + ) + ) + applied = _call_with_supported_kwargs(wrapper.plan, plan_kwargs, return_applied=True) + assert isinstance(applied, dict) + if cfg.split_kv.mode is SplitKVMode.FIXED and "fixed_split_size" not in applied: + raise FlashInferUnavailable( + "FlashInfer plan() did not accept required Split-KV knob 'fixed_split_size'" + ) + if cfg.split_kv.mode is SplitKVMode.DISABLED and "disable_split_kv" not in applied: + raise FlashInferUnavailable( + "FlashInfer plan() did not accept required Split-KV knob 'disable_split_kv'" + ) + return applied + + @staticmethod + def _actual_split_kv_plans( + wrapper: Any, + cfg: FlashInferPagedAttentionConfig, + plan: FlashInferPagedKVPlan, + ) -> tuple[SplitKVExecutionPlan, ...]: + getter = getattr(wrapper, "get_actual_split_kv_plan", None) + if not callable(getter): + if cfg.split_kv.mode is SplitKVMode.DISABLED: + return tuple( + cfg.split_kv.resolve(int(seq_len), backend="flashinfer_disabled_verified") + for seq_len in plan.kv_seq_lens.tolist() + ) + if cfg.require_batch_invariant: + raise FlashInferUnavailable( + "strict fixed Split-KV consistency requires runtime actual-plan provenance; " + "FlashInfer wrapper has no get_actual_split_kv_plan() callback. A requested " + "max-splits/count knob is not proof of token boundaries" + ) + return tuple( + cfg.split_kv.resolve(int(seq_len), backend="flashinfer_requested_only") + for seq_len in plan.kv_seq_lens.tolist() + ) + raw_plans = getter() + if not isinstance(raw_plans, (list, tuple)) or len(raw_plans) != len(plan.kv_seq_lens): + raise FlashInferUnavailable( + "get_actual_split_kv_plan() must return one plan per batch request" + ) + result: list[SplitKVExecutionPlan] = [] + for batch_index, (raw, seq_len) in enumerate( + zip(raw_plans, plan.kv_seq_lens.tolist(), strict=True) + ): + if not isinstance(raw, dict): + raise FlashInferUnavailable("actual Split-KV runtime plan entries must be dicts") + try: + required_keys = { + "mode", + "split_size", + "boundaries", + "fallback", + "fallback_reason", + } + missing_keys = sorted(required_keys.difference(raw)) + if missing_keys: + raise FlashInferUnavailable( + "actual Split-KV runtime plan is missing required fields: " + + ", ".join(missing_keys) + ) + execution = SplitKVExecutionPlan( + requested_mode=cfg.split_kv.mode, + requested_split_size=cfg.split_kv.fixed_split_size, + actual_mode=raw.get("mode"), + actual_split_size=( + None + if raw.get("split_size") is None + else _flashinfer_split_size_tokens( + raw.get("split_size"), + page_size=plan.page_size, + unit=raw.get("split_size_unit"), + ) + ), + boundaries=_normalize_flashinfer_split_boundaries( + raw.get("boundaries", ()), + page_size=plan.page_size, + seq_len=int(seq_len), + unit=raw.get("boundary_unit"), + ), + backend="flashinfer", + source="runtime_callback", + fallback=raw["fallback"], + fallback_reason=raw["fallback_reason"], + ) + except AttentionContractError as exc: + raise FlashInferUnavailable( + f"invalid actual Split-KV plan for batch {batch_index}: {exc}" + ) from exc + expected = cfg.split_kv.resolve(int(seq_len), backend="flashinfer_contract") + if cfg.require_batch_invariant: + try: + validate_split_kv_alignment(expected, execution) + except AttentionContractError as exc: + raise FlashInferUnavailable( + f"FlashInfer actual Split-KV plan for batch {batch_index} " + f"does not match the requested strict logical plan: {exc}" + ) from exc + result.append(execution) + return tuple(result) + + @staticmethod + def _actual_split_kv_plan_set( + wrapper: Any, + cfg: FlashInferPagedAttentionConfig, + plan: FlashInferPagedKVPlan, + ) -> SplitKVRuntimePlanSet | None: + getter = getattr(wrapper, "get_actual_split_kv_plan_set", None) + if not callable(getter): + if cfg.require_batch_invariant: + raise FlashInferUnavailable( + "strict Split-KV consistency requires a complete " + "batch/TP/CP/owner " + "runtime plan set; FlashInfer wrapper has no " + "get_actual_split_kv_plan_set() callback" + ) + return None + raw = getter() + if not isinstance(raw, dict): + raise FlashInferUnavailable("get_actual_split_kv_plan_set() must return a dict") + try: + entries = tuple( + SplitKVRuntimePlanEntry( + coordinate=SplitKVRuntimeCoordinate( + batch_index=int(entry["batch_index"]), + tp_rank=int(entry["tp_rank"]), + cp_rank=int(entry["cp_rank"]), + owner_cp_rank=int(entry["owner_cp_rank"]), + ), + expected_kv_range=tuple(entry["expected_kv_range"]), + execution=SplitKVExecutionPlan( + requested_mode=cfg.split_kv.mode, + requested_split_size=cfg.split_kv.fixed_split_size, + actual_mode=entry["mode"], + actual_split_size=( + None + if entry["split_size"] is None + else _flashinfer_split_size_tokens( + entry["split_size"], + page_size=plan.page_size, + unit=entry.get("split_size_unit"), + ) + ), + boundaries=_normalize_flashinfer_split_boundaries( + entry["boundaries"], + page_size=plan.page_size, + seq_len=int( + entry["expected_kv_range"][1] - entry["expected_kv_range"][0] + ), + unit=entry.get("boundary_unit"), + offset=int(entry["expected_kv_range"][0]), + ), + merge_order=entry["merge_order"], + acc_dtype=entry["accum_dtype"], + downcast_at=entry["downcast_at"], + backend="flashinfer", + source="runtime_plan_set_callback", + fallback=entry["fallback"], + fallback_reason=entry["fallback_reason"], + ), + ) + for entry in raw["entries"] + ) + plan_set = SplitKVRuntimePlanSet( + batch_size=int(raw["batch_size"]), + tp_world_size=int(raw["tp_world_size"]), + cp_world_size=int(raw["cp_world_size"]), + total_kv_tokens=tuple(int(value) for value in raw["total_kv_tokens"]), + entries=entries, + ) + except (KeyError, TypeError, ValueError) as exc: + raise FlashInferUnavailable( + f"invalid FlashInfer runtime Split-KV plan set: {exc}" + ) from exc + parallel = cfg.cp_comm_plan.parallel + expected_topology = ( + len(plan.kv_seq_lens), + parallel.tp_world_size, + parallel.cp_world_size, + tuple(int(value) for value in plan.kv_seq_lens.tolist()), + ) + actual_topology = ( + plan_set.batch_size, + plan_set.tp_world_size, + plan_set.cp_world_size, + plan_set.total_kv_tokens, + ) + if actual_topology != expected_topology: + raise FlashInferUnavailable( + "FlashInfer runtime Split-KV plan-set topology does not match the request" + ) + if cfg.require_batch_invariant: + for entry in plan_set.entries: + start, end = entry.expected_kv_range + expected_local = cfg.split_kv.resolve( + end - start, + backend="flashinfer_contract", + ) + expected = SplitKVExecutionPlan( + requested_mode=expected_local.requested_mode, + requested_split_size=expected_local.requested_split_size, + actual_mode=expected_local.actual_mode, + actual_split_size=expected_local.actual_split_size, + boundaries=tuple( + (start + local_start, start + local_end) + for local_start, local_end in expected_local.boundaries + ), + backend="flashinfer_contract", + source="contract_exact", + ) + try: + validate_split_kv_alignment(expected, entry.execution) + except AttentionContractError as exc: + raise FlashInferUnavailable( + "FlashInfer runtime Split-KV plan-set entry does not match " + f"the strict owner-local plan at {entry.coordinate}: {exc}" + ) from exc + return plan_set + + @staticmethod + def _actual_arithmetic_semantics( + wrapper: Any, + cfg: FlashInferPagedAttentionConfig, + ) -> dict[str, Any]: + getter = getattr(wrapper, "get_attention_arithmetic_provenance", None) + if not callable(getter): + if cfg.require_verified_arithmetic: + raise FlashInferUnavailable( + "strict attention consistency requires runtime arithmetic provenance; " + "FlashInfer wrapper has no get_attention_arithmetic_provenance() callback" + ) + return { + "accum_dtype": None, + "downcast_at": None, + "lse_dtype": None, + "arithmetic_plan_source": "unverified_backend_internal", + "arithmetic_semantics_verified": False, + } + raw = getter() + if not isinstance(raw, dict): + raise FlashInferUnavailable("get_attention_arithmetic_provenance() must return a dict") + required = { + "accum_dtype": "fp32", + "downcast_at": "final_write", + "lse_dtype": "fp32", + } + mismatches = [key for key, expected in required.items() if raw.get(key) != expected] + source = raw.get("source") + if not isinstance(source, str) or not source.strip(): + mismatches.append("source") + if mismatches: + raise FlashInferUnavailable( + "FlashInfer runtime arithmetic semantics do not satisfy the strict " + "attention contract: " + ", ".join(mismatches) + ) + result = { + **required, + "arithmetic_plan_source": source, + "arithmetic_semantics_verified": True, + } + for key in ("backend_lse_log_base", "export_lse_log_base"): + if key in raw: + result[key] = raw[key] + return result + + @staticmethod + def _run_wrapper( + wrapper: Any, + q_flat: torch.Tensor, + paged_kv_cache: tuple[torch.Tensor, torch.Tensor], + cfg: FlashInferPagedAttentionConfig, + ) -> tuple[torch.Tensor, torch.Tensor]: + if hasattr(wrapper, "run_return_lse"): + result = wrapper.run_return_lse(q_flat, paged_kv_cache) + else: + result = _call_with_supported_kwargs( + wrapper.run, + { + "q": q_flat, + "paged_kv_cache": paged_kv_cache, + "return_lse": cfg.return_lse, + }, + ) + if not isinstance(result, tuple) or len(result) != 2: + raise FlashInferUnavailable("FlashInfer PR7 candidate must return (out, lse)") + out_flat, lse_flat = result + normalize_lse = getattr(wrapper, "normalize_lse", None) + if callable(normalize_lse): + lse_flat = normalize_lse(lse_flat) + return out_flat, lse_flat + + @staticmethod + def _validate_runtime_outputs( + out_flat: torch.Tensor, + lse_flat: torch.Tensor, + q: torch.Tensor, + *, + require_fp32_output: bool, + ) -> None: + if not isinstance(out_flat, torch.Tensor) or not isinstance(lse_flat, torch.Tensor): + raise FlashInferUnavailable("FlashInfer output and LSE must be tensors") + if out_flat.device != q.device or lse_flat.device != q.device: + raise FlashInferUnavailable("FlashInfer output and LSE must remain on the query device") + expected_out_dtype = torch.float32 if require_fp32_output else q.dtype + if out_flat.dtype != expected_out_dtype: + raise FlashInferUnavailable( + "FlashInfer final output dtype does not match the requested output dtype" + ) + if lse_flat.dtype != torch.float32: + raise FlashInferUnavailable("FlashInfer attention-domain LSE must be FP32") + + +def _validate_strict_core(core: Any) -> None: + if not callable(getattr(core, "forward_with_lse", None)): + raise ValueError("strict Attention core must implement forward_with_lse") + expected_schedules = { + STRICT_ATTENTION_PRODUCTION_CORE_ID: STRICT_ATTENTION_FA4_SCHEDULE_ID, + STRICT_ATTENTION_CORE_ID: STRICT_ATTENTION_SCHEDULE_ID, + } + core_id = getattr(core, "core_id", None) + if core_id not in expected_schedules: + raise ValueError( + "strict Attention core ID must identify the FA4 production core or explicit reference" + ) + if getattr(core, "strict_schedule", None) != expected_schedules[core_id]: + raise ValueError("strict Attention core schedule does not match its exact core identity") + required = { + "merge_order": "global_block_index", + "accum_dtype": "fp32", + "downcast_at": "final_write", + "fallback": False, + } + mismatches = [ + name for name, expected in required.items() if getattr(core, name, None) != expected + ] + if mismatches: + raise ValueError( + "strict Attention core has incompatible arithmetic identity: " + ", ".join(mismatches) + ) + if core_id == STRICT_ATTENTION_PRODUCTION_CORE_ID: + production_required = { + "backend_id": "flash_attention_4.cute", + "native_attention_arithmetic": True, + "num_splits": 1, + "deterministic_backward": True, + "production_ready": True, + "reference_only": False, + } + production_mismatches = [ + name + for name, expected in production_required.items() + if getattr(core, name, None) != expected + ] + if production_mismatches: + raise ValueError( + "strict FA4 production core has incompatible controls: " + + ", ".join(production_mismatches) + ) + + +def _resolve_strict_core(cfg: FlashInferPagedAttentionConfig) -> Any: + if cfg.deterministic_core is not None: + return cfg.deterministic_core + return StrictFlashAttention4Core(split_kv=cfg.split_kv) + + +def _validate_strict_core_result( + result: Any, + core: Any, +) -> None: + if not isinstance(result, DeterministicAttentionCoreResult): + raise FlashInferUnavailable( + "strict Attention core must return DeterministicAttentionCoreResult" + ) + if result.out.dtype not in (torch.float16, torch.bfloat16): + raise FlashInferUnavailable("strict Attention core output must be FP16/BF16") + if result.lse.dtype is not torch.float32: + raise FlashInferUnavailable("strict Attention core LSE must be FP32") + expected = { + "strict_core_id": core.core_id, + "strict_schedule": core.strict_schedule, + "attention_backend": core.backend_id, + "merge_order": core.merge_order, + "accum_dtype": core.accum_dtype, + "downcast_at": core.downcast_at, + "fallback": False, + "native_attention_arithmetic": core.native_attention_arithmetic, + } + mismatches = [name for name, value in expected.items() if result.provenance.get(name) != value] + if mismatches: + raise FlashInferUnavailable( + "strict core result changed its declared arithmetic identity: " + ", ".join(mismatches) + ) + split_plan = result.provenance.get("split_kv") + if not isinstance(split_plan, dict) or split_plan.get("actual_split_kv_policy") != ( + SplitKVMode.DISABLED.value + ): + raise FlashInferUnavailable("strict core result did not prove Split-KV disabled") + + +def _resolve_strict_rope(cfg: FlashInferPagedAttentionConfig) -> Any: + if cfg.strict_rope_op is not None: + return cfg.strict_rope_op + try: + from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RoPESM90Op + + return RoPESM90Op() + except (ImportError, RuntimeError) as exc: + raise FlashInferUnavailable( + "strict Attention requires the RL-Kernel WS1 RoPE CUDA operator" + ) from exc + + +def _apply_strict_rope( + rope: Any, + x: torch.Tensor, + position_ids: torch.Tensor, + theta: float, +) -> torch.Tensor: + """RoPESM90Op accepts shared 1-D positions, so execute one batch row at a time.""" + + if position_ids.shape != (x.size(0), x.size(2)): + raise ValueError("strict RoPE position IDs must have shape [B,S]") + rows = [] + for batch_index in range(x.size(0)): + row = x[batch_index : batch_index + 1] + positions = position_ids[batch_index] + try: + rotated = rope(row, positions, theta=float(theta)) + except TypeError: + rotated = rope(row, positions) + if not isinstance(rotated, torch.Tensor) or rotated.shape != row.shape: + raise FlashInferUnavailable("strict RoPE returned an invalid tensor") + if rotated.dtype != row.dtype or rotated.device != row.device: + raise FlashInferUnavailable("strict RoPE changed dtype or device") + rows.append(rotated) + return torch.cat(rows, dim=0) + + +def _materialize_strict_logical_kv( + k_cache: torch.Tensor, + v_cache: torch.Tensor, + metadata: Any, + plan: FlashInferPagedKVPlan, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Restore each paged cache row to logical order without padding arithmetic.""" + + max_len = int(plan.kv_seq_lens.max().item()) + logical_k = torch.zeros( + (k_cache.size(0), k_cache.size(1), max_len, k_cache.size(3)), + dtype=k_cache.dtype, + device=k_cache.device, + ) + logical_v = torch.zeros_like(logical_k) + positions = torch.full( + (k_cache.size(0), max_len), + -1, + dtype=torch.long, + device=k_cache.device, + ) + page_size = int(metadata.page_size) + for batch_index, seq_len_value in enumerate(plan.kv_seq_lens.tolist()): + seq_len = int(seq_len_value) + block_count = (seq_len + page_size - 1) // page_size + token_index = torch.arange(seq_len, device=k_cache.device, dtype=torch.long) + pages = metadata.block_table[batch_index, :block_count].long() + slots = pages[token_index // page_size] * page_size + token_index % page_size + logical_k[batch_index, :, :seq_len, :] = k_cache[batch_index, :, slots, :] + logical_v[batch_index, :, :seq_len, :] = v_cache[batch_index, :, slots, :] + source_positions = getattr(metadata, "key_position_ids", None) + if not isinstance(source_positions, torch.Tensor): + source_positions = getattr(metadata, "global_token_positions", None) + if not isinstance(source_positions, torch.Tensor): + raise FlashInferUnavailable( + "strict Attention requires key_position_ids runtime metadata" + ) + positions[batch_index, :seq_len] = source_positions[batch_index, slots].long() + return logical_k, logical_v, positions + + +def _cp_owner_ranges( + plan: AttentionCPCommunicationPlan, +) -> tuple[tuple[int, int], ...]: + ranges = [] + for owner_rank in range(plan.parallel.cp_world_size): + blocks = sorted( + ( + block + for block in plan.expected_blocks + if block.owner_cp_rank == owner_rank + and block.owner_tp_rank == plan.parallel.tp_rank + ), + key=lambda block: block.kv_block_start, + ) + if not blocks: + raise ValueError(f"strict CP manifest has no blocks for owner {owner_rank}") + start = blocks[0].kv_block_start + cursor = start + for block in blocks: + if block.kv_block_start != cursor: + raise ValueError("strict CP owner blocks must form a contiguous range") + cursor = block.kv_block_end + ranges.append((start, cursor)) + expected = plan.expected_kv_token_range + if expected is None or ranges[0][0] != expected[0] or ranges[-1][1] != expected[1]: + raise ValueError("strict CP owner ranges do not cover expected_kv_token_range") + for left, right in zip(ranges, ranges[1:], strict=False): + if left[1] != right[0]: + raise ValueError("strict CP owner ranges must be gap-free and rank ordered") + return tuple(ranges) + + +def _strict_local_query_positions( + metadata: Any, + q_local: torch.Tensor, + query_range: tuple[int, int], +) -> torch.Tensor: + positions = getattr(metadata, "query_position_ids", None) + if not isinstance(positions, torch.Tensor): + raise ValueError("strict CP requires query_position_ids") + if positions.shape == (q_local.size(0), q_local.size(2)): + return positions.contiguous() + start, end = query_range + if positions.ndim == 2 and positions.size(0) == q_local.size(0) and positions.size(1) >= end: + return positions[:, start:end].contiguous() + raise ValueError("strict CP query_position_ids do not match local query ownership") + + +def _strict_local_key_positions( + metadata: Any, + k_local: torch.Tensor, + key_range: tuple[int, int], +) -> torch.Tensor: + positions = getattr(metadata, "key_position_ids", None) + if not isinstance(positions, torch.Tensor): + positions = getattr(metadata, "global_token_positions", None) + if not isinstance(positions, torch.Tensor): + raise ValueError("strict CP requires key_position_ids") + if positions.shape == (k_local.size(0), k_local.size(2)): + return positions.contiguous() + start, end = key_range + if positions.ndim == 2 and positions.size(0) == k_local.size(0) and positions.size(1) >= end: + return positions[:, start:end].contiguous() + raise ValueError("strict CP key_position_ids do not match local KV ownership") + + +def _strict_attention_provenance( + cfg: FlashInferPagedAttentionConfig, + *, + core_provenance: dict[str, object], + materialization: str, + cp_required: bool, + rope: Any, +) -> dict[str, Any]: + communication_id = getattr(cfg.cp_communication, "backend_id", "unknown") + communication_backend = ( + "self_owned_cuda_ag_rs" if communication_id == "cuda_ag_rs" else communication_id + ) + return { + "attention_backend": core_provenance["attention_backend"], + "requested_backend": "flashinfer_layout_adapter", + "actual_backend": core_provenance["attention_backend"], + "adapter_backend": "flashinfer", + "attention_mode": cfg.mode, + "materialization": materialization, + "causal": cfg.causal, + "softmax_scale": cfg.softmax_scale, + "lse_domain": "attention", + "lse_exported": True, + "lse_dtype": "fp32", + "strict_mode": True, + "strict_core_id": core_provenance["strict_core_id"], + "strict_schedule": core_provenance["strict_schedule"], + "accum_dtype": core_provenance["accum_dtype"], + "downcast_at": core_provenance["downcast_at"], + "arithmetic_plan_source": core_provenance.get("fa_api_source", "rlkernel_reference_core"), + "arithmetic_semantics_verified": True, + "native_attention_arithmetic": core_provenance["native_attention_arithmetic"], + "fallback": False, + "fallback_reason": None, + "rope_backend": getattr(rope, "backend_id", "rlkernel.cuda.rope_sm90"), + "rope_theta": float(cfg.rope.rope_theta), + "rotary_dim": cfg.rope.rotary_dim, + "rope_fusion": False, + "rope_fusion_boundary": "rlkernel_rope_then_attention", + "q_rope_state": "post_rope", + "k_cache_rope_state": "post_rope", + "batch_invariant_claim": "strict_runtime_verified", + "cp_comm_required": cp_required, + "communication_backend": (communication_backend if cp_required else "none"), + "num_splits": core_provenance.get("num_splits"), + "deterministic_backward": core_provenance.get("deterministic_backward"), + "fa_api_source": core_provenance.get("fa_api_source"), + "fa_package_version": core_provenance.get("fa_package_version"), + "reference_only": bool(core_provenance.get("reference_only", False)), + "production_ready": bool( + core_provenance.get("production_ready", False) + and ( + not cp_required or getattr(cfg.cp_communication, "backend_id", None) == "cuda_ag_rs" + ) + ), + } + + +def flashinfer_qwen3_paged_attention_available() -> bool: + """Return whether the FlashInfer paged attention wrappers are importable.""" + + try: + module = FlashInferQwen3PagedAttentionOp()._load_flashinfer() + prefill = getattr( + getattr(module, "prefill", None), + "BatchPrefillWithPagedKVCacheWrapper", + None, + ) + decode = getattr( + getattr(module, "decode", None), + "BatchDecodeWithPagedKVCacheWrapper", + None, + ) + if not callable(prefill) or not callable(decode): + return False + except FlashInferUnavailable: + return False + return True + + +def _validate_metadata_logical_positions( + metadata: Any, + *, + batch_index: int, + seq_len: int, + page_size: int, + block_count: int, + device: torch.device, +) -> None: + if not hasattr(metadata, "global_token_positions"): + return + global_token_positions = metadata.global_token_positions + if global_token_positions.ndim != 2 or global_token_positions.size(0) <= batch_index: + raise ValueError("global_token_positions must have shape [B, cache_capacity]") + physical_slots: list[int] = [] + for logical_block in range(block_count): + local_page = int(metadata.block_table[batch_index, logical_block].item()) + token_count = min(page_size, seq_len - logical_block * page_size) + for page_offset in range(token_count): + physical_slots.append(local_page * page_size + page_offset) + slot_index = torch.tensor(physical_slots, device=device, dtype=torch.long) + actual = global_token_positions[batch_index, slot_index] + expected = torch.arange( + 0, + seq_len, + device=device, + dtype=global_token_positions.dtype, + ) + if not torch.equal(actual, expected): + raise ValueError( + "block_table/global_token_positions must reconstruct logical positions " + "as one contiguous global range" + ) + if hasattr(metadata, "key_position_ids"): + key_positions = metadata.key_position_ids[batch_index, slot_index] + if not torch.equal(key_positions, expected.to(dtype=key_positions.dtype)): + raise ValueError("key_position_ids must match cached global token positions") + active_slot_mask = torch.zeros( + global_token_positions.size(1), + device=device, + dtype=torch.bool, + ) + active_slot_mask[slot_index] = True + if bool((global_token_positions[batch_index, ~active_slot_mask] != -1).any()): + raise ValueError("unused global_token_positions entries must be -1") + if hasattr(metadata, "key_position_ids") and bool( + (metadata.key_position_ids[batch_index, ~active_slot_mask] != -1).any() + ): + raise ValueError("unused key_position_ids entries must be -1") + + +def _validate_flashinfer_rope_metadata( + metadata: Any, + cfg: FlashInferPagedAttentionConfig, + q: torch.Tensor, +) -> None: + """Bind the configured fused-RoPE boundary to the actual cache metadata.""" + + for name, expected in ( + ("q_rope_state", cfg.rope.q_rope_state), + ("k_cache_rope_state", cfg.rope.k_cache_rope_state), + ): + actual = getattr(metadata, name, None) + if actual != expected: + raise ValueError( + f"metadata.{name}={actual!r} does not match the FlashInfer fused-RoPE " + f"contract {expected!r}" + ) + cache_position = getattr(metadata, "cache_position", None) + query_position_ids = getattr(metadata, "query_position_ids", None) + if not isinstance(cache_position, torch.Tensor) or not isinstance( + query_position_ids, torch.Tensor + ): + raise ValueError("cache_position and query_position_ids are required tensors") + if cache_position.device != q.device or query_position_ids.device != q.device: + raise ValueError("query position metadata must be on the query device") + if cache_position.dtype not in {torch.int32, torch.int64, torch.long} or ( + query_position_ids.dtype not in {torch.int32, torch.int64, torch.long} + ): + raise ValueError("query position metadata must contain integers") + if cache_position.shape != (q.size(0), q.size(2)): + raise ValueError("cache_position must have shape [B, Sq]") + if query_position_ids.shape != cache_position.shape: + raise ValueError("query_position_ids must have shape [B, Sq]") + if not torch.equal(cache_position, query_position_ids): + raise ValueError("cache_position and query_position_ids must match exactly") + kv_seq_lens = getattr(metadata, "kv_seq_lens", None) + if not isinstance(kv_seq_lens, torch.Tensor) or kv_seq_lens.shape != (q.size(0),): + raise ValueError("kv_seq_lens must have shape [B]") + if bool((cache_position < 0).any()) or bool((cache_position >= kv_seq_lens[:, None]).any()): + raise ValueError("cache_position must identify tokens present in each KV sequence") + if q.size(2) > 1 and bool((cache_position[:, 1:] <= cache_position[:, :-1]).any()): + raise ValueError("few-query cache_position values must be strictly increasing") + expected_query_positions = torch.stack( + [ + torch.arange( + int(seq_len.item()) - q.size(2), + int(seq_len.item()), + device=q.device, + dtype=cache_position.dtype, + ) + for seq_len in kv_seq_lens + ] + ) + if not torch.equal(cache_position, expected_query_positions): + raise ValueError( + "FlashInfer implicit RoPE positions require queries to be the trailing " + "contiguous positions of each KV sequence" + ) + + +def _materialize_logical_kv_cache( + k_cache: torch.Tensor, + v_cache: torch.Tensor, + metadata: Any, + *, + total_kv_tokens: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Restore paged KV to logical order for the deterministic CP fallback.""" + + page_size = int(metadata.page_size) + block_count = (total_kv_tokens + page_size - 1) // page_size + logical_k: list[torch.Tensor] = [] + logical_v: list[torch.Tensor] = [] + logical_positions: list[torch.Tensor] = [] + for batch_index in range(k_cache.size(0)): + pages = metadata.block_table[batch_index, :block_count].long() + token_index = torch.arange(total_kv_tokens, device=k_cache.device, dtype=torch.long) + slots = pages[token_index // page_size] * page_size + token_index % page_size + logical_k.append(k_cache[batch_index, :, slots, :]) + logical_v.append(v_cache[batch_index, :, slots, :]) + if hasattr(metadata, "key_position_ids"): + logical_positions.append(metadata.key_position_ids[batch_index, slots].long()) + else: + logical_positions.append(token_index) + return ( + torch.stack(logical_k, dim=0).contiguous(), + torch.stack(logical_v, dim=0).contiguous(), + torch.stack(logical_positions, dim=0).contiguous(), + ) + + +def flashinfer_prefix_cache_fingerprint( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + metadata: Any, + cfg: FlashInferPagedAttentionConfig, + *, + prefix_length: int, +) -> str: + """Hash the logical prefix and the fused-RoPE materialization identity.""" + + prefix_length = _positive_int(prefix_length, "prefix_length") + if bool((metadata.kv_seq_lens < prefix_length).any()): + raise ValueError("prefix_length must not exceed any kv_seq_lens entry") + digest = hashlib.sha256() + rotary_dim = q.size(-1) if cfg.rope.rotary_dim is None else cfg.rope.rotary_dim + digest.update( + ( + f"k_rope_state={metadata.k_cache_rope_state};" + f"rope_theta={float(cfg.rope.rope_theta):.17g};" + f"rotary_dim={rotary_dim};" + "rope_cast_at=after_rope;" + f"k_rope_output_dtype={k_cache.dtype}\n" + ).encode() + ) + digest.update(f"k_dtype={k_cache.dtype};v_dtype={v_cache.dtype}\n".encode()) + page_size = int(metadata.page_size) + for batch_index in range(q.size(0)): + logical_index = torch.arange(prefix_length, device=q.device, dtype=torch.long) + pages = metadata.block_table[batch_index].long() + slots = pages[logical_index // page_size] * page_size + logical_index % page_size + for tensor in ( + metadata.global_token_positions[batch_index, slots], + metadata.key_position_ids[batch_index, slots], + k_cache[batch_index, :, slots, :], + 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 _validate_flashinfer_prefix_cache( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + metadata: Any, + cfg: FlashInferPagedAttentionConfig, +) -> None: + enabled = bool(getattr(metadata, "prefix_cache_enabled", False)) + key = getattr(metadata, "prefix_cache_key", None) + length = getattr(metadata, "prefix_length", 0) + fingerprint = getattr(metadata, "prefix_cache_fingerprint", None) + if not enabled: + if key is not None or length != 0 or fingerprint is not None: + raise ValueError( + "prefix cache key/fingerprint must be None and prefix_length must be 0 " + "when prefix cache is disabled" + ) + return + if not isinstance(key, str) or not key: + raise ValueError("prefix_cache_key is required when prefix cache is enabled") + if not isinstance(fingerprint, str) or not fingerprint: + raise ValueError("prefix_cache_fingerprint is required when prefix cache is enabled") + actual = flashinfer_prefix_cache_fingerprint( + q, + k_cache, + v_cache, + metadata, + cfg, + prefix_length=length, + ) + if actual != fingerprint: + raise ValueError( + "prefix_cache_fingerprint does not match logical prefix content or fused-RoPE identity" + ) + + +def _validate_qkv_cache(q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor) -> None: + if q.ndim != 4 or k_cache.ndim != 4 or v_cache.ndim != 4: + raise ValueError("q, k_cache, and v_cache must have shape [B, H, S, D]") + if k_cache.shape != v_cache.shape: + raise ValueError("k_cache and v_cache must have matching shape") + if q.size(0) != k_cache.size(0) or q.size(3) != k_cache.size(3): + raise ValueError("q and KV cache must share batch size and head_dim") + if q.size(1) % k_cache.size(1) != 0: + raise ValueError("Q head count must be divisible by KV head count") + + +def _restore_out(out_flat: torch.Tensor, *, batch_size: int, query_len: int) -> torch.Tensor: + if out_flat.ndim != 3: + raise FlashInferUnavailable("FlashInfer output must have shape [B*Sq, Hq, D]") + _, q_heads, head_dim = out_flat.shape + expected = batch_size * query_len + if out_flat.size(0) != expected: + raise FlashInferUnavailable( + f"FlashInfer output first dim must be B*Sq={expected}, got {out_flat.size(0)}" + ) + return out_flat.reshape(batch_size, query_len, q_heads, head_dim).transpose(1, 2).contiguous() + + +def _restore_lse( + lse_flat: torch.Tensor, + *, + batch_size: int, + query_len: int, + q_heads: int, +) -> torch.Tensor: + expected_tokens = batch_size * query_len + if lse_flat.shape == (expected_tokens, q_heads): + return lse_flat.reshape(batch_size, query_len, q_heads).transpose(1, 2).contiguous() + if lse_flat.shape == (q_heads, expected_tokens): + return lse_flat.transpose(0, 1).reshape(batch_size, query_len, q_heads).transpose(1, 2) + raise FlashInferUnavailable( + f"FlashInfer LSE must have shape [B*Sq, Hq] or [Hq, B*Sq]; got {tuple(lse_flat.shape)}" + ) + + +def _call_with_supported_kwargs( + fn: Any, + kwargs: dict[str, Any], + *, + return_applied: bool = False, +) -> Any: + try: + signature = inspect.signature(fn) + except (TypeError, ValueError): + result = fn(**kwargs) + return dict(kwargs) if return_applied else result + parameters = signature.parameters + if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in parameters.values()): + result = fn(**kwargs) + return dict(kwargs) if return_applied else result + supported = {name: value for name, value in kwargs.items() if name in parameters} + missing_required = [ + name + for name, param in parameters.items() + if param.default is inspect.Parameter.empty + and param.kind + in { + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + } + and name not in supported + ] + if missing_required: + raise FlashInferUnavailable( + f"{getattr(fn, '__qualname__', fn)} missing supported arguments: " + f"{', '.join(missing_required)}" + ) + result = fn(**supported) + return supported if return_applied else result + + +def _flashinfer_split_kv_plan_kwargs( + spec: SplitKVSpec, + *, + page_size: int, +) -> dict[str, Any]: + if spec.mode is SplitKVMode.DISABLED: + return {"disable_split_kv": True} + if spec.mode is SplitKVMode.FIXED: + assert spec.fixed_split_size is not None + if spec.fixed_split_size % page_size != 0: + raise FlashInferUnavailable( + "FlashInfer fixed Split-K size is expressed in pages; the WS2 token " + "split size must be divisible by page_size" + ) + return { + # FlashInfer names this in pages, while the WS2 contract is tokens. + "fixed_split_size": int(spec.fixed_split_size) // page_size, + "disable_split_kv": False, + } + return {"disable_split_kv": False} + + +def _flashinfer_split_size_tokens(value: Any, *, page_size: int, unit: Any) -> int: + if unit != "pages": + raise FlashInferUnavailable( + "FlashInfer fixed Split-K provenance must declare split_size_unit='pages'" + ) + return _positive_int(int(value), "split_size") * page_size + + +def _normalize_flashinfer_split_boundaries( + boundaries: Any, + *, + page_size: int, + seq_len: int, + unit: Any, + offset: int = 0, +) -> tuple[tuple[int, int], ...]: + if unit != "pages": + raise FlashInferUnavailable( + "FlashInfer Split-K boundaries must declare boundary_unit='pages'" + ) + normalized = [] + for boundary in boundaries: + start_page, end_page = boundary + normalized.append( + ( + offset + int(start_page) * page_size, + offset + min(int(end_page) * page_size, seq_len), + ) + ) + return tuple(normalized) + + +def _split_kv_provenance( + spec: SplitKVSpec, + plans: tuple[SplitKVExecutionPlan, ...], + *, + applied_plan_kwargs: dict[str, Any], + require_batch_invariant: bool, +) -> dict[str, Any]: + policy = spec.mode.value + if spec.mode is SplitKVMode.FIXED: + policy = f"fixed:{spec.fixed_split_size}" + return { + "split_kv_policy": policy, + "requested_split_kv_policy": spec.mode.value, + "requested_split_kv_size": spec.fixed_split_size, + "actual_split_kv_plans": [plan.to_dict() for plan in plans], + "backend_native_split_kv_knobs": { + key: value + for key, value in applied_plan_kwargs.items() + if key in {"fixed_split_size", "disable_split_kv"} + }, + "batch_invariant_required": bool(require_batch_invariant), + "batch_invariant_claim": ( + "strict_runtime_verified" if require_batch_invariant else "diagnostic_only" + ), + } + + +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__ = [ + "FlashInferAttentionMode", + "FlashInferAttentionResult", + "FlashInferPagedAttentionConfig", + "FlashInferPagedKVPlan", + "FlashInferQwen3PagedAttentionOp", + "FlashInferRoPEFusionConfig", + "FlashInferSplitKVPolicy", + "FlashInferUnavailable", + "build_flashinfer_paged_kv_plan", + "flashinfer_qwen3_paged_attention_available", + "flashinfer_prefix_cache_fingerprint", + "materialize_flashinfer_paged_kv_cache", +] diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index d68dbb18..8f952c12 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -60,6 +60,104 @@ def __post_init__(self) -> None: raise ValueError("block_end must be >= block_start") +@dataclass(frozen=True) +class AttentionRingBlock: + """One logical KV block in the decoupled Ring Attention schedule.""" + + global_block_index: int + block_start: int + block_end: int + owner_cp_rank: int + + +@dataclass(frozen=True) +class AttentionRingSchedule: + """Static compute order separated from the fixed arithmetic merge order.""" + + schedule_id: str + total_kv_tokens: int + cp_world_size: int + kv_chunk_size: Optional[int] + blocks: tuple[AttentionRingBlock, ...] + compute_order: tuple[int, ...] + merge_order: tuple[int, ...] + compute_communication: str = "decoupled" + overlap: str = "disabled" + + @classmethod + def build( + cls, + total_kv_tokens: int, + *, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> "AttentionRingSchedule": + if ( + isinstance(cp_world_size, bool) + or not isinstance(cp_world_size, int) + or cp_world_size < 1 + ): + raise ValueError("cp_world_size must be >= 1") + if kv_chunk_size is not None and ( + isinstance(kv_chunk_size, bool) + or not isinstance(kv_chunk_size, int) + or kv_chunk_size < 1 + ): + raise ValueError("kv_chunk_size must be >= 1 when provided") + if ( + isinstance(total_kv_tokens, bool) + or not isinstance(total_kv_tokens, int) + or total_kv_tokens < 1 + ): + raise ValueError("total_kv_tokens must be an integer >= 1") + if total_kv_tokens < cp_world_size: + raise ValueError("Ring Attention requires at least one KV token per CP rank") + blocks: list[AttentionRingBlock] = [] + for owner_cp_rank, (owner_start, owner_end) in enumerate( + _split_bounds(total_kv_tokens, cp_world_size) + ): + cursor = owner_start + while cursor < owner_end: + block_end = ( + owner_end if kv_chunk_size is None else min(cursor + kv_chunk_size, owner_end) + ) + blocks.append( + AttentionRingBlock( + global_block_index=len(blocks), + block_start=cursor, + block_end=block_end, + owner_cp_rank=owner_cp_rank, + ) + ) + cursor = block_end + compute_order: list[int] = [] + left, right = 0, len(blocks) - 1 + while left <= right: + compute_order.append(left) + if left != right: + compute_order.append(right) + left += 1 + right -= 1 + return cls( + schedule_id="rlkernel.attention.strict_ring_state.v1", + total_kv_tokens=total_kv_tokens, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + blocks=tuple(blocks), + compute_order=tuple(compute_order), + merge_order=tuple(range(len(blocks))), + ) + + def provenance(self) -> dict[str, object]: + return { + "compute_communication": self.compute_communication, + "compute_schedule": self.schedule_id, + "compute_order": list(self.compute_order), + "merge_order_indices": list(self.merge_order), + "communication_overlap": self.overlap, + } + + @dataclass(frozen=True) class DeterministicAttentionCoreResult: """Strict-core output and the exact arithmetic plan used for it.""" @@ -359,6 +457,21 @@ def split_kv_execution_plans( backend="deterministic_cp_reference", ) + @staticmethod + def ring_schedule( + total_kv_tokens: int, + *, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> AttentionRingSchedule: + """Build the production-default pre-overlap Ring schedule.""" + + return AttentionRingSchedule.build( + total_kv_tokens, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + @staticmethod def execution_provenance( total_kv_tokens: int, @@ -732,6 +845,11 @@ def backward_reference( ) out = state.out.to(resolved_output_dtype) lse = state.lse + ring_schedule = self.ring_schedule( + k.size(2), + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) return AttentionBackwardPathResult( name=name @@ -779,6 +897,9 @@ def backward_reference( "merge_order": "global_block_index", "accum_dtype": "fp32", "downcast_at": "final_write", + **ring_schedule.provenance(), + "ring_schedule_default": True, + "ring_partial_arithmetic": False, "strict_bitwise": self.strict_bitwise, "strict_core_id": (STRICT_ATTENTION_CORE_ID if self.strict_bitwise else None), "strict_schedule": (STRICT_ATTENTION_SCHEDULE_ID if self.strict_bitwise else None), @@ -1803,6 +1924,8 @@ def build_reference_split_kv_runtime_plan_set( "AttentionBackwardPathResult", "AttentionBackwardRankDrift", "AttentionPartialState", + "AttentionRingBlock", + "AttentionRingSchedule", "AttentionSavedForwardState", "build_reference_split_kv_runtime_plan_set", "CPAttentionReferenceOp", diff --git a/scripts/ws2_p2p_nccl_attention_reference_check.py b/scripts/ws2_p2p_nccl_attention_reference_check.py index 5eed0ec5..7b799d74 100644 --- a/scripts/ws2_p2p_nccl_attention_reference_check.py +++ b/scripts/ws2_p2p_nccl_attention_reference_check.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Two-GPU P2P NCCL correctness reference for issue #235. +"""NCCL correctness reference for issue #235 CP Attention communication. -Run with: +Use two ranks for a CP-only diagnostic, four ranks for the formal TP=2/CP=2 +target, or eight ranks for two independent TP=2/CP=2 replicas:: - torchrun --standalone --nproc-per-node=2 \ + torchrun --standalone --nproc-per-node=4 \ scripts/ws2_p2p_nccl_attention_reference_check.py """ @@ -16,7 +17,8 @@ import os import sys from pathlib import Path -from typing import Sequence +from types import SimpleNamespace +from typing import Any, Sequence import torch import torch.distributed as dist @@ -25,14 +27,26 @@ if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) +from rl_engine.kernels.attention_contract import ( # noqa: E402 + STRICT_ATTENTION_FA4_SCHEDULE_ID, + STRICT_ATTENTION_PRODUCTION_CORE_ID, +) from rl_engine.kernels.ops.cuda.attention.cp_comm import ( # noqa: E402 AttentionCPBlockMetadata, AttentionCPCommunicationPlan, AttentionCPMergedState, AttentionCPPartialState, AttentionParallelSpec, + CUDAAGRSAttentionCPCommunication, P2PNCCLAttentionCPCommunication, ) +from rl_engine.kernels.ops.cuda.attention.flash_attn import StrictFlashAttention4Core # noqa: E402 +from rl_engine.kernels.ops.cuda.attention.flashinfer_paged_attention import ( # noqa: E402 + FlashInferPagedAttentionConfig, + FlashInferQwen3PagedAttentionOp, + _apply_strict_rope, +) +from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RoPESM90Op # noqa: E402 from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( # noqa: E402 AttentionPartialState, DeterministicCPAttentionReferenceOp, @@ -49,9 +63,25 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser.add_argument("--head-dim", type=int, default=128) parser.add_argument("--chunk-size", type=int, default=4) parser.add_argument("--seed", type=int, default=2357) + parser.add_argument("--repeats", type=int, default=3) parser.add_argument("--atol", type=float, default=2.0e-4) parser.add_argument("--final-write-atol", type=float, default=2.0e-2) - return parser.parse_args(argv) + parser.add_argument( + "--transport", + choices=("p2p_nccl_reference", "cuda_ag_rs"), + default="p2p_nccl_reference", + help="P2P is the correctness reference; cuda_ag_rs selects PR311/PR312", + ) + parser.add_argument( + "--strict-shared-core", + action="store_true", + help="run AG(Q/K/V/positions) -> shared CUDA core -> RS(Out/LSE) with backward", + ) + parser.add_argument("--output", type=Path) + args = parser.parse_args(argv) + if args.strict_shared_core and args.transport != "cuda_ag_rs": + parser.error("--strict-shared-core requires --transport cuda_ag_rs") + return args def main(argv: Sequence[str] | None = None) -> int: @@ -61,13 +91,33 @@ def main(argv: Sequence[str] | None = None) -> int: dist.init_process_group("nccl", init_method="env://") try: world_size = dist.get_world_size() - rank = dist.get_rank() - if world_size != 2: - raise RuntimeError("this reference check requires exactly two NCCL ranks") - local_rank = int(os.environ.get("LOCAL_RANK", str(rank))) + global_rank = dist.get_rank() + if world_size not in {2, 4, 8}: + raise RuntimeError("this check requires 2, 4, or 8 NCCL ranks") + if torch.cuda.device_count() < world_size: + raise RuntimeError("this single-node check requires one visible GPU per NCCL rank") + local_rank = int(os.environ.get("LOCAL_RANK", str(global_rank))) torch.cuda.set_device(local_rank) device = torch.device("cuda", local_rank) - result = run_check(args, rank=rank, device=device) + + cp_groups = [ + dist.new_group(ranks=[pair_start, pair_start + 1]) + for pair_start in range(0, world_size, 2) + ] + rank_in_replica = global_rank if world_size < 8 else global_rank % 4 + replica_index = 0 if world_size < 8 else global_rank // 4 + tp_rank = 0 if world_size == 2 else rank_in_replica // 2 + cp_rank = rank_in_replica % 2 + cp_group = cp_groups[global_rank // 2] + result = run_check( + args, + global_rank=global_rank, + tp_rank=tp_rank, + cp_rank=cp_rank, + replica_index=replica_index, + cp_group=cp_group, + device=device, + ) failures = torch.tensor( [0 if result["passed"] else 1], dtype=torch.int32, @@ -77,20 +127,27 @@ def main(argv: Sequence[str] | None = None) -> int: result["global_failure_count"] = int(failures.item()) reports: list[dict[str, object] | None] = [None] * world_size dist.all_gather_object(reports, result) - if rank == 0: - print( - json.dumps( - { - "schema_version": "ws2_p2p_nccl_attention_reference/v1", - "backend": str(dist.get_backend()), - "world_size": world_size, - "global_failure_count": int(failures.item()), - "ranks": reports, - }, - indent=2, - sort_keys=True, - ) - ) + if global_rank == 0: + report = { + "schema_version": ( + "ws2_p2p_nccl_attention_reference/v1" + if args.transport == "p2p_nccl_reference" + else "ws2_cuda_ag_rs_attention/v1" + ), + "backend": str(dist.get_backend()), + "transport": args.transport, + "world_size": world_size, + "tp_world_size": 1 if world_size == 2 else 2, + "cp_world_size": 2, + "replica_count": 2 if world_size == 8 else 1, + "global_failure_count": int(failures.item()), + "ranks": reports, + } + serialized = json.dumps(report, indent=2, sort_keys=True) + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(serialized + "\n", encoding="utf-8") + print(serialized) return 0 if int(failures.item()) == 0 else 1 finally: dist.destroy_process_group() @@ -99,7 +156,11 @@ def main(argv: Sequence[str] | None = None) -> int: def run_check( args: argparse.Namespace, *, - rank: int, + global_rank: int, + tp_rank: int, + cp_rank: int, + replica_index: int, + cp_group: Any, device: torch.device, ) -> dict[str, object]: if args.batch < 1: @@ -108,6 +169,8 @@ def run_check( raise ValueError("seq_len must be positive and divisible by CP=2") if args.chunk_size < 1: raise ValueError("chunk_size must be positive") + if args.repeats < 2: + raise ValueError("repeats must be at least 2 for a bitwise stability check") if args.q_heads != 16 or args.kv_heads != 4 or args.head_dim != 128: raise ValueError("TP=2 Qwen3-8B local heads must be Hq=16, Hkv=4, D=128") for name in ("atol", "final_write_atol"): @@ -115,7 +178,7 @@ def run_check( if not math.isfinite(value) or value < 0: raise ValueError(f"{name} must be finite and non-negative") - generator = torch.Generator(device="cpu").manual_seed(args.seed) + generator = torch.Generator(device="cpu").manual_seed(args.seed + tp_rank + 100 * replica_index) shape_q = (args.batch, args.q_heads, args.seq_len, args.head_dim) shape_kv = (args.batch, args.kv_heads, args.seq_len, args.head_dim) q = torch.randn(shape_q, generator=generator, dtype=torch.bfloat16).to(device) @@ -131,29 +194,39 @@ def run_check( kv_block_start=start, kv_block_end=min(start + args.chunk_size, owner_end), owner_cp_rank=owner, - owner_tp_rank=0, + owner_tp_rank=tp_rank, ) ) plan = AttentionCPCommunicationPlan( parallel=AttentionParallelSpec( tp_world_size=2, - tp_rank=0, + tp_rank=tp_rank, cp_world_size=2, - cp_rank=rank, + cp_rank=cp_rank, ), - backend="p2p_nccl_reference", + backend=args.transport, status="implemented", expected_blocks=tuple(blocks), expected_kv_token_range=(0, args.seq_len), query_token_ranges=owner_ranges, ) + communication: P2PNCCLAttentionCPCommunication | CUDAAGRSAttentionCPCommunication + if args.transport == "p2p_nccl_reference": + communication = P2PNCCLAttentionCPCommunication(process_group=cp_group) + else: + communication = CUDAAGRSAttentionCPCommunication(process_group=cp_group) + + query_start, query_end = owner_ranges[cp_rank] + q_local = q[:, :, query_start:query_end, :].contiguous() + q_gathered = communication.all_gather_query(q_local, plan) + query_ag_max_abs = float((q_gathered - q).abs().max().item()) reference = DeterministicCPAttentionReferenceOp() local_states: list[AttentionCPPartialState] = [] for block in reversed(blocks): - if block.owner_cp_rank != rank: + if block.owner_cp_rank != cp_rank: continue state = reference.local_partial_state( - q, + q_gathered, k[:, :, block.kv_block_start : block.kv_block_end, :], v[:, :, block.kv_block_start : block.kv_block_end, :], q_start=0, @@ -164,58 +237,238 @@ def run_check( ) local_states.append(AttentionCPPartialState(state.out, state.lse, block)) - communication = P2PNCCLAttentionCPCommunication() - gathered = communication.all_gather_partial_states(tuple(local_states), plan) - merged = merge_attention_partial_states( - [ - AttentionPartialState( - state.out, - state.lse, - state.block.kv_block_start, - state.block.kv_block_end, - ) - for state in gathered - ] - ) - local = communication.reduce_scatter_merged_state( - AttentionCPMergedState(merged.out, merged.lse), - plan, - ) + def communicate() -> tuple[tuple[AttentionCPPartialState, ...], AttentionCPMergedState]: + gathered_states = communication.all_gather_partial_states(tuple(local_states), plan) + merged_state = merge_attention_partial_states( + [ + AttentionPartialState( + state.out, + state.lse, + state.block.kv_block_start, + state.block.kv_block_end, + ) + for state in gathered_states + ] + ) + local_state = communication.reduce_scatter_merged_state( + AttentionCPMergedState(merged_state.out, merged_state.lse), + plan, + ) + return gathered_states, local_state + + gathered, local = communicate() + gathered_indices = [state.block.global_block_index for state in gathered] + repeat_query_bitwise = True + repeat_out_bitwise = True + repeat_lse_bitwise = True + repeat_manifest_bitwise = True + for _ in range(args.repeats - 1): + repeated_q = communication.all_gather_query(q_local, plan) + repeated_gathered, repeated_local = communicate() + repeat_query_bitwise = repeat_query_bitwise and torch.equal(repeated_q, q_gathered) + repeat_out_bitwise = repeat_out_bitwise and torch.equal(repeated_local.out, local.out) + repeat_lse_bitwise = repeat_lse_bitwise and torch.equal(repeated_local.lse, local.lse) + repeat_manifest_bitwise = ( + repeat_manifest_bitwise + and [state.block.global_block_index for state in repeated_gathered] == gathered_indices + ) + full_out, full_lse = reference.forward_fp32_with_lse(q, k, v, causal=True) - start, end = owner_ranges[rank] + start, end = owner_ranges[cp_rank] out_max_abs = float((local.out - full_out[:, :, start:end, :]).abs().max().item()) lse_max_abs = float((local.lse - full_lse[:, :, start:end]).abs().max().item()) final_out = local.out.to(q.dtype) expected_final_out = full_out[:, :, start:end, :].to(q.dtype) final_out_max_abs = float((final_out.float() - expected_final_out.float()).abs().max().item()) - gathered_indices = [state.block.global_block_index for state in gathered] + strict_shared_core = ( + _run_strict_shared_core_check( + args, + plan=plan, + communication=communication, + q=q, + k=k, + v=v, + owner_ranges=owner_ranges, + ) + if args.strict_shared_core + else {"executed": False, "passed": False} + ) passed = ( gathered_indices == list(range(len(blocks))) + and query_ag_max_abs == 0.0 + and repeat_query_bitwise + and repeat_out_bitwise + and repeat_lse_bitwise + and repeat_manifest_bitwise and out_max_abs <= args.atol and lse_max_abs <= args.atol and final_out.dtype == q.dtype and final_out_max_abs <= args.final_write_atol + and (not args.strict_shared_core or strict_shared_core["passed"] is True) ) return { - "rank": rank, + "rank": global_rank, + "global_world_size": dist.get_world_size() if dist.is_initialized() else 1, + "tp_rank": tp_rank, + "tp_world_size": 1 if dist.is_initialized() and dist.get_world_size() == 2 else 2, + "cp_rank": cp_rank, + "cp_world_size": 2, + "replica_index": replica_index, + "replica_count": 2 if dist.is_initialized() and dist.get_world_size() == 8 else 1, "device": str(device), "dtype": "bf16", "accum_dtype": "fp32", "downcast_at": "final_write", "final_output_dtype": str(final_out.dtype).removeprefix("torch."), - "transport": "p2p_nccl_reference", + "transport": args.transport, + "protocol": "ag_query_local_kv_rs_out_lse", + "strict_protocol": "ag_qkv_positions_shared_core_rs_out_lse", + "query_ag": args.transport, + "query_ag_max_abs": query_ag_max_abs, "query_range": [start, end], - "world_size": 2, "expected_block_manifest": [block.provenance() for block in blocks], + "local_block_indices": sorted(state.block.global_block_index for state in local_states), "gathered_block_indices": gathered_indices, + "repeat_count": args.repeats, + "repeat_query_bitwise": repeat_query_bitwise, + "repeat_out_bitwise": repeat_out_bitwise, + "repeat_lse_bitwise": repeat_lse_bitwise, + "repeat_manifest_bitwise": repeat_manifest_bitwise, "out_max_abs": out_max_abs, "lse_max_abs": lse_max_abs, "final_out_max_abs": final_out_max_abs, "atol": args.atol, "final_write_atol": args.final_write_atol, + "strict_shared_core": strict_shared_core, "passed": passed, } +def _run_strict_shared_core_check( + args: argparse.Namespace, + *, + plan: AttentionCPCommunicationPlan, + communication: CUDAAGRSAttentionCPCommunication | P2PNCCLAttentionCPCommunication, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + owner_ranges: tuple[tuple[int, int], ...], +) -> dict[str, object]: + """Exercise the complete differentiable self-owned communication path.""" + + cp_rank = plan.parallel.cp_rank + start, end = owner_ranges[cp_rank] + q_local = q[:, :, start:end, :].detach().clone().requires_grad_() + k_local = k[:, :, start:end, :].detach().clone().requires_grad_() + v_local = v[:, :, start:end, :].detach().clone().requires_grad_() + positions = torch.arange(args.seq_len, dtype=torch.long, device=q.device).expand(args.batch, -1) + metadata = SimpleNamespace( + q_rope_state="pre_rope", + k_cache_rope_state="pre_rope", + query_position_ids=positions[:, start:end].contiguous(), + key_position_ids=positions[:, start:end].contiguous(), + ) + config = FlashInferPagedAttentionConfig( + mode="prefill", + cp_comm_plan=plan, + require_cp_comm=True, + strict_mode=True, + cp_communication=communication, + ) + op = FlashInferQwen3PagedAttentionOp(flashinfer_module=None) + distributed = op(q_local, k_local, v_local, metadata, config=config) + + q_ref = q.detach().clone().requires_grad_() + k_ref = k.detach().clone().requires_grad_() + v_ref = v.detach().clone().requires_grad_() + rope = RoPESM90Op() + q_ready = _apply_strict_rope(rope, q_ref, positions, config.rope.rope_theta) + k_ready = _apply_strict_rope(rope, k_ref, positions, config.rope.rope_theta) + reference = StrictFlashAttention4Core().forward_with_lse( + q_ready, + k_ready, + v_ref, + causal=True, + query_position_ids=positions, + key_position_ids=positions, + output_dtype=q.dtype, + ) + out_ref = reference.out[:, :, start:end, :] + lse_ref = reference.lse[:, :, start:end] + + dout = torch.randn( + q.shape, + generator=torch.Generator(device="cpu").manual_seed(args.seed + 10_000), + dtype=q.dtype, + ).to(q.device) + (distributed.out.float() * dout[:, :, start:end, :].float()).sum().backward() + (reference.out.float() * dout.float()).sum().backward() + + comparisons = { + "out": (distributed.out, out_ref), + "lse": (distributed.lse, lse_ref), + "dq": (q_local.grad, q_ref.grad[:, :, start:end, :]), + "dk": (k_local.grad, k_ref.grad[:, :, start:end, :]), + "dv": (v_local.grad, v_ref.grad[:, :, start:end, :]), + } + bitwise = { + name: left is not None and right is not None and torch.equal(left, right) + for name, (left, right) in comparisons.items() + } + max_abs = { + name: float((left.float() - right.float()).abs().max().item()) + for name, (left, right) in comparisons.items() + if left is not None and right is not None + } + + repeat_out_bitwise = True + repeat_lse_bitwise = True + for _ in range(args.repeats - 1): + repeated = op( + q_local.detach(), + k_local.detach(), + v_local.detach(), + metadata, + config=config, + ) + repeat_out_bitwise = repeat_out_bitwise and torch.equal(repeated.out, distributed.out) + repeat_lse_bitwise = repeat_lse_bitwise and torch.equal(repeated.lse, distributed.lse) + + provenance = distributed.provenance + identity_valid = ( + provenance.get("strict_core_id") == STRICT_ATTENTION_PRODUCTION_CORE_ID + and provenance.get("strict_schedule") == STRICT_ATTENTION_FA4_SCHEDULE_ID + and provenance.get("strict_mode") is True + and provenance.get("native_attention_arithmetic") is True + and provenance.get("num_splits") == 1 + and provenance.get("deterministic_backward") is True + and provenance.get("fa_api_source") == "flash_attn.cute.interface" + and provenance.get("fallback") is False + and provenance.get("strict_split_kv") == "disabled" + and provenance.get("strict_comm_autograd") is True + and provenance.get("production_ready") is True + ) + return { + "executed": True, + "passed": ( + all(bitwise.values()) and repeat_out_bitwise and repeat_lse_bitwise and identity_valid + ), + "strict_core_id": provenance.get("strict_core_id"), + "strict_schedule": provenance.get("strict_schedule"), + "actual_backend": provenance.get("actual_backend"), + "communication_backend": provenance.get("communication_backend"), + "production_ready": provenance.get("production_ready"), + "strict_mode": provenance.get("strict_mode"), + "native_attention_arithmetic": provenance.get("native_attention_arithmetic"), + "fallback": provenance.get("fallback"), + "split_kv_policy": provenance.get("strict_split_kv"), + "communication_autograd": provenance.get("strict_comm_autograd"), + "bitwise": bitwise, + "max_abs": max_abs, + "repeat_out_bitwise": repeat_out_bitwise, + "repeat_lse_bitwise": repeat_lse_bitwise, + } + + if __name__ == "__main__": raise SystemExit(main()) diff --git a/scripts/ws2_pr7_flashinfer_attention_check.py b/scripts/ws2_pr7_flashinfer_attention_check.py new file mode 100644 index 00000000..7ab64842 --- /dev/null +++ b/scripts/ws2_pr7_flashinfer_attention_check.py @@ -0,0 +1,748 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""PR7 FlashInfer RoPE-fused paged attention validation entry point. + +The default dry-run mode is CI/local friendly: it builds the FlashInfer page plan +and provenance without importing FlashInfer or requiring CUDA. On a CUDA host +with FlashInfer installed, omit ``--dry-run`` to run the opt-in PR7 candidate and +compare it with the PR6 full logical KV reference. +""" + +from __future__ import annotations + +import argparse +import json +import math +import sys +from dataclasses import replace +from pathlib import Path +from typing import Any, Sequence + +import torch + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from rl_engine.kernels.attention_contract import ( # noqa: E402 + STRICT_ATTENTION_FA4_SCHEDULE_ID, + STRICT_ATTENTION_PRODUCTION_CORE_ID, +) +from rl_engine.kernels.ops.cuda.attention.cp_comm import ( # noqa: E402 + AttentionCPCommunicationPlan, + AttentionParallelSpec, +) +from rl_engine.kernels.ops.cuda.attention.flash_attn import StrictFlashAttention4Core # noqa: E402 +from rl_engine.kernels.ops.cuda.attention.flashinfer_paged_attention import ( # noqa: E402 + FlashInferPagedAttentionConfig, + FlashInferQwen3PagedAttentionOp, + FlashInferRoPEFusionConfig, + FlashInferSplitKVPolicy, + FlashInferUnavailable, + _apply_strict_rope, + _materialize_strict_logical_kv, + build_flashinfer_paged_kv_plan, +) +from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RoPESM90Op # noqa: E402 +from rl_engine.testing.attention_comparison import ( # noqa: E402 + AttentionPathResult, + DecodeAttentionInputs, + DecodeKVCacheMetadata, + run_decode_full_prefill_reference, +) + + +def main(argv: Sequence[str] | None = None) -> int: + args = _parse_args(argv) + device = torch.device("cuda" if args.device == "cuda" else "cpu") + inputs = _make_inputs(args, device) + config = _make_config(args) + config.validate(head_dim=args.head_dim, query_len=args.query_len) + plan = build_flashinfer_paged_kv_plan( + inputs.metadata, + batch_size=args.batch_size, + query_len=args.query_len, + cache_capacity=inputs.k_cache.size(2), + device=device, + ) + report: dict[str, Any] = { + "status": "dry_run" if args.dry_run else "executed", + "pr": "PR7", + "target": "Qwen3-8B TP-local FlashInfer candidate; CP transport validated separately", + "mode": config.mode, + "device": str(device), + "shape": { + "batch_size": args.batch_size, + "query_len": args.query_len, + "kv_seq_len": args.kv_seq_len, + "page_size": args.page_size, + "q_heads": args.q_heads, + "kv_heads": args.kv_heads, + "head_dim": args.head_dim, + }, + "rope": config.rope.provenance(args.head_dim), + "split_kv": { + **config.split_kv.to_dict(), + "provenance_status": "requested_only_dry_run", + "actual_plan_required_for_strict_pass": config.require_batch_invariant, + "requested_execution_plans": [ + config.split_kv.resolve( + int(seq_len), + backend="flashinfer_dry_run_requested_only", + ).to_dict() + for seq_len in plan.kv_seq_lens.tolist() + ], + }, + "communication": config.cp_comm_plan.provenance() + | {"cp_comm_required": config.require_cp_comm}, + "paged_kv_plan": plan.provenance(), + "tests_expected": [ + "FlashInfer ROPE_LLAMA vs NativeRoPEOp + full logical KV reference", + "split-K disabled/fixed policy drift", + "batch composition/position invariant sweep", + "attention-domain LSE export drift", + "strict shared CUDA core with separate multi-rank AG/RS forward/backward evidence", + ], + "thresholds": { + "out_max_abs": args.out_atol, + "lse_max_abs": args.lse_atol, + "dlogp_max_abs": args.dlogp_atol, + }, + } + if args.dry_run: + report["passed"] = False + report["acceptance_eligible"] = False + report["errors"] = ["dry-run does not execute the FlashInfer candidate"] + _emit(report, json_output=args.json, output=args.output) + return 0 + + if device.type != "cuda": + raise SystemExit("non-dry-run PR7 validation requires --device cuda") + op = FlashInferQwen3PagedAttentionOp() + try: + candidate = op( + inputs.q, + inputs.k_cache, + inputs.v_cache, + inputs.metadata, + config=config, + ) + except (FlashInferUnavailable, RuntimeError) as exc: + # Keep optional dependency failures machine-readable while preserving + # a non-zero exit so aggregate acceptance cannot pass closed. + report.update( + { + "status": "not_available", + "passed": False, + "acceptance_eligible": False, + "errors": [f"Attention backend unavailable: {exc}"], + "unavailable_reason": str(exc), + } + ) + _emit(report, json_output=args.json, output=args.output) + return 1 + reference_inputs = replace( + inputs, + metadata=replace( + inputs.metadata, + q_rope_state="pre_rope", + k_cache_rope_state="pre_rope", + ), + ) + pytorch_reference = run_decode_full_prefill_reference(reference_inputs) + reference = ( + _run_strict_cuda_reference(inputs, config, plan) + if config.strict_mode + else pytorch_reference + ) + out_stats = _drift_stats(candidate.out, reference.out) + lse_stats = _drift_stats(candidate.lse, reference.lse) + dlogp_stats = _selected_logprob_drift( + candidate.out, + reference.out, + seed=args.seed + 101, + vocab_size=args.vocab_size, + ) + report["candidate_provenance"] = candidate.provenance + report["split_kv"].update( + { + "provenance_status": "runtime_verified", + "actual_execution_plans": candidate.provenance.get( + "actual_split_kv_plans", + candidate.provenance.get("strict_core_row_plans"), + ), + "actual_plan_set": candidate.provenance.get("actual_split_kv_plan_set"), + } + ) + report["drift"] = { + "out": out_stats, + "lse": lse_stats, + "dlogp": dlogp_stats, + } + report["reference_backend"] = ( + "rlkernel.cuda.deterministic_attention" + if config.strict_mode + else "rlkernel.pytorch.full_logical_kv_reference" + ) + if config.strict_mode: + report["diagnostic_drift_vs_pytorch"] = { + "out": _drift_stats(candidate.out, pytorch_reference.out), + "lse": _drift_stats(candidate.lse, pytorch_reference.lse), + "dlogp": _selected_logprob_drift( + candidate.out, + pytorch_reference.out, + seed=args.seed + 101, + vocab_size=args.vocab_size, + ), + } + if config.require_batch_invariant: + report["batch_invariant_sweep"] = _run_batch_invariance_sweep( + op, + inputs, + candidate, + config, + ) + report["page_layout_invariant_sweep"] = _run_page_layout_invariance_sweep( + op, + inputs, + candidate, + config, + ) + errors = _acceptance_errors(report, args) + report["errors"] = errors + report["passed"] = not errors + report["acceptance_eligible"] = not errors + report["status"] = "passed" if not errors else "failed" + _emit(report, json_output=args.json, output=args.output) + return 0 if not errors else 1 + + +def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--mode", choices=["prefill", "decode"], default="decode") + parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu") + parser.add_argument("--dry-run", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--json", action="store_true", help="emit machine-readable JSON") + parser.add_argument("--output", type=Path, help="write the JSON report to this path") + parser.add_argument("--batch-size", type=int, default=2) + parser.add_argument("--query-len", type=int, default=1) + parser.add_argument("--kv-seq-len", type=int, default=16) + parser.add_argument("--page-size", type=int, default=4) + parser.add_argument("--q-heads", type=int, default=16) + parser.add_argument("--kv-heads", type=int, default=4) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--dtype", choices=["bf16", "fp16", "fp32"], default="bf16") + parser.add_argument("--seed", type=int, default=2357) + parser.add_argument("--vocab-size", type=int, default=257) + parser.add_argument("--out-atol", type=float, default=1.0e-2) + parser.add_argument("--lse-atol", type=float, default=2.0e-3) + parser.add_argument("--dlogp-atol", type=float, default=2.0e-3) + parser.add_argument("--tp-world-size", type=int, default=2) + parser.add_argument("--tp-rank", type=int, default=0) + parser.add_argument("--cp-world-size", type=int, default=2) + parser.add_argument("--cp-rank", type=int, default=0) + parser.add_argument( + "--cp-comm-backend", + choices=["cuda_ag_rs", "local_debug"], + default="cuda_ag_rs", + ) + parser.add_argument("--require-cp-comm", action="store_true") + parser.add_argument( + "--strict", + action="store_true", + help="use FlashInfer only for paged layout and execute the RL-Kernel shared core", + ) + parser.add_argument("--fixed-split-size", type=int, default=None) + parser.add_argument( + "--split-kv-policy", + choices=["disabled", "fixed", "auto"], + default="disabled", + ) + parser.add_argument( + "--require-batch-invariant", + action=argparse.BooleanOptionalAction, + default=True, + ) + args = parser.parse_args(argv) + for name in ("out_atol", "lse_atol", "dlogp_atol"): + if getattr(args, name) < 0: + parser.error(f"--{name.replace('_', '-')} must be non-negative") + if args.vocab_size < 2: + parser.error("--vocab-size must be >= 2") + if 32 % args.tp_world_size != 0 or 8 % args.tp_world_size != 0: + parser.error("--tp-world-size must divide Qwen3-8B's 32 query and 8 KV heads") + expected_q_heads = 32 // args.tp_world_size + expected_kv_heads = 8 // args.tp_world_size + if args.q_heads != expected_q_heads or args.kv_heads != expected_kv_heads: + parser.error( + "--q-heads/--kv-heads must describe the TP-local Qwen3-8B shard: " + f"expected {expected_q_heads}/{expected_kv_heads} for TP={args.tp_world_size}" + ) + if args.head_dim != 128: + parser.error("--head-dim must be 128 for the Qwen3-8B acceptance target") + if args.strict and args.split_kv_policy != "disabled": + parser.error("--strict requires --split-kv-policy disabled") + return args + + +def _make_config(args: argparse.Namespace) -> FlashInferPagedAttentionConfig: + if args.split_kv_policy == "disabled": + split_kv = FlashInferSplitKVPolicy.disabled() + elif args.split_kv_policy == "fixed": + if args.fixed_split_size is None: + raise SystemExit("--split-kv-policy fixed requires --fixed-split-size") + split_kv = FlashInferSplitKVPolicy.fixed(args.fixed_split_size) + else: + split_kv = FlashInferSplitKVPolicy.auto() + return FlashInferPagedAttentionConfig( + mode=args.mode, + workspace_size_bytes=128 * 1024 * 1024, + require_batch_invariant=args.require_batch_invariant, + rope=FlashInferRoPEFusionConfig( + rope_theta=1_000_000.0, + rope_scale=1.0, + rotary_dim=args.head_dim, + q_rope_state="pre_rope", + k_cache_rope_state="pre_rope", + ), + split_kv=split_kv, + cp_comm_plan=AttentionCPCommunicationPlan( + parallel=AttentionParallelSpec( + tp_world_size=args.tp_world_size, + tp_rank=args.tp_rank, + cp_world_size=args.cp_world_size, + cp_rank=args.cp_rank, + ), + backend=args.cp_comm_backend, + status="interface_only", + ), + require_cp_comm=args.require_cp_comm, + strict_mode=args.strict, + ) + + +def _make_inputs(args: argparse.Namespace, device: torch.device) -> DecodeAttentionInputs: + dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[args.dtype] + if args.kv_seq_len % args.page_size != 0: + raise SystemExit("--kv-seq-len must be divisible by --page-size for this scaffold") + if args.mode == "decode" and args.query_len != 1: + raise SystemExit("--mode decode requires --query-len 1") + generator = torch.Generator(device=device).manual_seed(args.seed) + q = torch.randn( + args.batch_size, + args.q_heads, + args.query_len, + args.head_dim, + generator=generator, + device=device, + dtype=dtype, + ) + k_cache = torch.randn( + args.batch_size, + args.kv_heads, + args.kv_seq_len, + args.head_dim, + generator=generator, + device=device, + dtype=dtype, + ) + v_cache = torch.randn( + args.batch_size, + args.kv_heads, + args.kv_seq_len, + args.head_dim, + generator=generator, + device=device, + dtype=dtype, + ) + page_count = args.kv_seq_len // args.page_size + block_table = torch.arange(page_count, device=device, dtype=torch.long).repeat( + args.batch_size, + 1, + ) + positions = torch.arange(args.kv_seq_len, device=device, dtype=torch.long).repeat( + args.batch_size, + 1, + ) + query_start = args.kv_seq_len - args.query_len + query_positions = torch.arange( + query_start, + args.kv_seq_len, + device=device, + dtype=torch.long, + ).repeat(args.batch_size, 1) + metadata = DecodeKVCacheMetadata( + cache_position=query_positions.clone(), + kv_seq_lens=torch.full( + (args.batch_size,), + args.kv_seq_len, + device=device, + dtype=torch.long, + ), + block_table=block_table, + global_token_positions=positions, + query_position_ids=query_positions.clone(), + key_position_ids=positions.clone(), + page_size=args.page_size, + q_rope_state="pre_rope", + k_cache_rope_state="pre_rope", + ) + return DecodeAttentionInputs(q=q, k_cache=k_cache, v_cache=v_cache, metadata=metadata) + + +def _run_strict_cuda_reference( + inputs: DecodeAttentionInputs, + config: FlashInferPagedAttentionConfig, + paged_plan: Any, +) -> AttentionPathResult: + """Call the production FA4 core directly on the logical KV sequence.""" + + core = config.deterministic_core or StrictFlashAttention4Core(split_kv=config.split_kv) + rope = config.strict_rope_op or RoPESM90Op() + logical_k, logical_v, key_positions = _materialize_strict_logical_kv( + inputs.k_cache, + inputs.v_cache, + inputs.metadata, + paged_plan, + ) + query_positions = inputs.metadata.query_position_ids + outputs: list[torch.Tensor] = [] + lses: list[torch.Tensor] = [] + for batch_index, seq_len_value in enumerate(paged_plan.kv_seq_lens.tolist()): + seq_len = int(seq_len_value) + q_row = inputs.q[batch_index : batch_index + 1] + k_row = logical_k[batch_index : batch_index + 1, :, :seq_len, :] + v_row = logical_v[batch_index : batch_index + 1, :, :seq_len, :] + q_pos = query_positions[batch_index : batch_index + 1] + k_pos = key_positions[batch_index : batch_index + 1, :seq_len] + result = core.forward_with_lse( + _apply_strict_rope(rope, q_row, q_pos, config.rope.rope_theta), + _apply_strict_rope(rope, k_row, k_pos, config.rope.rope_theta), + v_row, + causal=config.causal, + scale=config.softmax_scale, + query_position_ids=q_pos, + key_position_ids=k_pos, + output_dtype=inputs.q.dtype, + ) + outputs.append(result.out) + lses.append(result.lse) + return AttentionPathResult( + name="direct_flash_attention4_num_splits1", + out=torch.cat(outputs, dim=0), + lse=torch.cat(lses, dim=0), + provenance={ + "strict_core_id": core.core_id, + "strict_schedule": core.strict_schedule, + "split_kv_policy": "disabled", + }, + ) + + +def _run_batch_invariance_sweep( + op: FlashInferQwen3PagedAttentionOp, + inputs: DecodeAttentionInputs, + batch_result: Any, + config: FlashInferPagedAttentionConfig, +) -> dict[str, Any]: + rows = [] + max_out = 0.0 + max_lse = 0.0 + for batch_index in range(inputs.q.size(0)): + single = _select_batch_row(inputs, batch_index) + single_result = op( + single.q, + single.k_cache, + single.v_cache, + single.metadata, + config=config, + ) + out_diff = ( + single_result.out.float() - batch_result.out[batch_index : batch_index + 1].float() + ).abs() + lse_diff = ( + single_result.lse.float() - batch_result.lse[batch_index : batch_index + 1].float() + ).abs() + row_out = float(out_diff.max().item()) + row_lse = float(lse_diff.max().item()) + max_out = max(max_out, row_out) + max_lse = max(max_lse, row_lse) + rows.append({"batch_index": batch_index, "out_max_abs": row_out, "lse_max_abs": row_lse}) + return { + "method": "single_row_vs_same_row_inside_batch", + "row_count": len(rows), + "out_max_abs": max_out, + "lse_max_abs": max_lse, + "rows": rows, + "passed": max_out == 0.0 and max_lse == 0.0, + } + + +def _run_page_layout_invariance_sweep( + op: FlashInferQwen3PagedAttentionOp, + inputs: DecodeAttentionInputs, + base_result: Any, + config: FlashInferPagedAttentionConfig, +) -> dict[str, Any]: + page_size = inputs.metadata.page_size + page_count = inputs.k_cache.size(2) // page_size + if page_count < 2: + return { + "method": "logical_page_table_permutation", + "status": "not_applicable", + "passed": True, + "reason": "fewer than two physical pages", + } + permutation = torch.arange( + page_count - 1, + -1, + -1, + device=inputs.k_cache.device, + dtype=torch.long, + ) + inverse = torch.empty_like(permutation) + inverse[permutation] = torch.arange(page_count, device=permutation.device) + + def permute_cache(cache: torch.Tensor) -> torch.Tensor: + pages = cache.reshape( + cache.size(0), + cache.size(1), + page_count, + page_size, + cache.size(3), + ) + return pages[:, :, permutation].reshape_as(cache).contiguous() + + block_table = inverse[inputs.metadata.block_table.long()] + positions = inputs.metadata.global_token_positions.reshape( + inputs.q.size(0), page_count, page_size + )[:, permutation].reshape_as(inputs.metadata.global_token_positions) + key_positions = inputs.metadata.key_position_ids.reshape( + inputs.q.size(0), page_count, page_size + )[:, permutation].reshape_as(inputs.metadata.key_position_ids) + permuted = DecodeAttentionInputs( + q=inputs.q, + k_cache=permute_cache(inputs.k_cache), + v_cache=permute_cache(inputs.v_cache), + metadata=replace( + inputs.metadata, + block_table=block_table, + global_token_positions=positions, + key_position_ids=key_positions, + ), + scale=inputs.scale, + output_dtype=inputs.output_dtype, + rope_theta=inputs.rope_theta, + rope_rotary_dim=inputs.rope_rotary_dim, + rope_cast_at=inputs.rope_cast_at, + q_rope_output_dtype=inputs.q_rope_output_dtype, + k_cache_rope_output_dtype=inputs.k_cache_rope_output_dtype, + ) + candidate = op( + permuted.q, + permuted.k_cache, + permuted.v_cache, + permuted.metadata, + config=config, + ) + out = _drift_stats(candidate.out, base_result.out) + lse = _drift_stats(candidate.lse, base_result.lse) + return { + "method": "logical_page_table_permutation", + "permutation": permutation.detach().cpu().tolist(), + "out": out, + "lse": lse, + "passed": out["max_abs"] == 0.0 and lse["max_abs"] == 0.0, + } + + +def _drift_stats(candidate: torch.Tensor, reference: torch.Tensor) -> dict[str, Any]: + if candidate.shape != reference.shape: + raise ValueError("candidate and reference tensors must have matching shapes") + candidate_fp32 = candidate.float() + reference_fp32 = reference.float() + raw = (candidate_fp32 - reference_fp32).abs() + diff = torch.where(candidate_fp32 == reference_fp32, torch.zeros_like(raw), raw).reshape(-1) + if diff.numel() == 0: + return { + "max_abs": 0.0, + "mean_abs": 0.0, + "p95_abs": 0.0, + "p99_abs": 0.0, + "active_count": 0, + } + return { + "max_abs": float(diff.max().item()), + "mean_abs": float(diff.mean().item()), + "p95_abs": float(torch.quantile(diff, 0.95).item()), + "p99_abs": float(torch.quantile(diff, 0.99).item()), + "active_count": int(diff.numel()), + } + + +def _selected_logprob_drift( + candidate_out: torch.Tensor, + reference_out: torch.Tensor, + *, + seed: int, + vocab_size: int, +) -> dict[str, Any]: + batch, heads, seq_len, head_dim = candidate_out.shape + generator = torch.Generator(device="cpu").manual_seed(seed) + weight = torch.randn( + vocab_size, + heads * head_dim, + generator=generator, + dtype=torch.float32, + ).to(candidate_out.device) + weight.mul_(1.0 / math.sqrt(heads * head_dim)) + target_ids = ( + torch.arange( + batch * seq_len, + device=candidate_out.device, + dtype=torch.long, + ).reshape(batch, seq_len) + % vocab_size + ) + + def selected(out: torch.Tensor) -> torch.Tensor: + hidden = out.float().transpose(1, 2).reshape(batch, seq_len, heads * head_dim) + logits = torch.matmul(hidden, weight.transpose(0, 1)) + return torch.log_softmax(logits, dim=-1).gather(-1, target_ids.unsqueeze(-1)).squeeze(-1) + + return _drift_stats(selected(candidate_out), selected(reference_out)) + + +def _acceptance_errors(report: dict[str, Any], args: argparse.Namespace) -> list[str]: + errors = [] + drift = report["drift"] + for name, threshold in ( + ("out", args.out_atol), + ("lse", args.lse_atol), + ("dlogp", args.dlogp_atol), + ): + value = drift.get(name, {}).get("max_abs") + if not isinstance(value, (int, float)) or not math.isfinite(float(value)) or value < 0: + errors.append(f"{name} max_abs must be finite and non-negative") + elif value > threshold: + errors.append(f"{name} max_abs={value} exceeds {threshold}") + provenance = report.get("candidate_provenance", {}) + if not str(report.get("device", "")).startswith("cuda"): + errors.append("strict PR7 acceptance requires a CUDA execution") + shape = report.get("shape", {}) + expected_shape = { + "q_heads": 32 // args.tp_world_size, + "kv_heads": 8 // args.tp_world_size, + "head_dim": 128, + } + if not isinstance(shape, dict) or any( + shape.get(key) != expected for key, expected in expected_shape.items() + ): + errors.append("runtime shape is not the Qwen3-8B TP-local head shard") + if provenance.get("attention_mode") != args.mode: + errors.append("runtime attention mode differs from the requested mode") + if provenance.get("fallback") is not False: + errors.append("FlashInfer execution used or omitted fallback provenance") + if args.strict: + if provenance.get("strict_mode") is not True: + errors.append("strict runtime did not execute the shared Attention core") + if provenance.get("strict_core_id") != STRICT_ATTENTION_PRODUCTION_CORE_ID: + errors.append("strict runtime did not execute the FA4 production core") + if provenance.get("strict_schedule") != STRICT_ATTENTION_FA4_SCHEDULE_ID: + errors.append("strict runtime arithmetic schedule is invalid") + if provenance.get("actual_backend") != "flash_attention_4.cute": + errors.append("strict runtime backend is not FlashAttention-4 CuTe") + if provenance.get("native_attention_arithmetic") is not True: + errors.append("strict runtime did not execute native FA4 Attention arithmetic") + if provenance.get("num_splits") != 1: + errors.append("strict runtime did not fix FA4 num_splits=1") + if provenance.get("deterministic_backward") is not True: + errors.append("strict runtime did not request deterministic FA4 backward") + if provenance.get("fa_api_source") != "flash_attn.cute.interface": + errors.append("strict runtime did not prove the FA4 CuTe API source") + if provenance.get("reference_only") is not False: + errors.append("strict runtime selected the reference core") + strict_plans = provenance.get("strict_core_row_plans") + if not isinstance(strict_plans, list) or not strict_plans: + errors.append("strict no-Split-K execution plans are missing") + elif any(plan.get("actual_split_kv_policy") != "disabled" for plan in strict_plans): + errors.append("strict runtime did not keep Split-KV disabled") + if provenance.get("rope_backend") not in { + "rlkernel.cuda.rope_sm90", + "rlkernel.cuda.rope_sm90_op", + }: + errors.append("strict runtime did not use the RL-Kernel WS1 RoPE operator") + elif provenance.get("pos_encoding_mode") != "ROPE_LLAMA": + errors.append("FlashInfer runtime did not use ROPE_LLAMA") + if provenance.get("rope_theta") != 1_000_000.0 or provenance.get("rotary_dim") != 128: + errors.append("FlashInfer runtime RoPE identity does not match Qwen3-8B") + if provenance.get("arithmetic_semantics_verified") is not True: + errors.append("runtime arithmetic semantics were not verified") + if not args.strict: + if not provenance.get("actual_split_kv_plans"): + errors.append("actual Split-KV runtime plans are missing") + plan_set = provenance.get("actual_split_kv_plan_set") + if not isinstance(plan_set, dict) or plan_set.get("coverage") != ( + "complete_batch_tp_cp_owner_cartesian_product" + ): + errors.append("complete batch/TP/CP/owner Split-KV plan set is missing") + for sweep_name in ("batch_invariant_sweep", "page_layout_invariant_sweep"): + if report.get(sweep_name, {}).get("passed") is not True: + errors.append(f"{sweep_name} failed") + return errors + + +def _select_batch_row(inputs: DecodeAttentionInputs, batch_index: int) -> DecodeAttentionInputs: + metadata = inputs.metadata + cp_block_owners = ( + None + if metadata.cp_block_owners is None + else metadata.cp_block_owners[batch_index : batch_index + 1] + ) + selected_metadata = DecodeKVCacheMetadata( + cache_position=metadata.cache_position[batch_index : batch_index + 1], + kv_seq_lens=metadata.kv_seq_lens[batch_index : batch_index + 1], + block_table=metadata.block_table[batch_index : batch_index + 1], + global_token_positions=metadata.global_token_positions[batch_index : batch_index + 1], + query_position_ids=metadata.query_position_ids[batch_index : batch_index + 1], + key_position_ids=metadata.key_position_ids[batch_index : batch_index + 1], + page_size=metadata.page_size, + prefix_cache_key=metadata.prefix_cache_key, + prefix_cache_enabled=metadata.prefix_cache_enabled, + prefix_length=metadata.prefix_length, + prefix_cache_fingerprint=metadata.prefix_cache_fingerprint, + q_rope_state=metadata.q_rope_state, + k_cache_rope_state=metadata.k_cache_rope_state, + cp_block_owners=cp_block_owners, + ) + return replace( + inputs, + q=inputs.q[batch_index : batch_index + 1], + k_cache=inputs.k_cache[batch_index : batch_index + 1], + v_cache=inputs.v_cache[batch_index : batch_index + 1], + metadata=selected_metadata, + ) + + +def _emit( + report: dict[str, Any], + *, + json_output: bool, + output: Path | None, +) -> None: + if output is not None: + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text( + json.dumps(report, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + if json_output: + print(json.dumps(report, indent=2, sort_keys=True)) + return + print(f"PR7 FlashInfer check: {report['status']}") + print(json.dumps(report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/setup.py b/setup.py index 79f882d9..e69de29b 100644 --- a/setup.py +++ b/setup.py @@ -1,313 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -import importlib.util -import os -import warnings -from pathlib import Path - -from setuptools import find_packages, setup - - -def _load_envs_module(): - envs_path = Path(__file__).with_name("envs.py") - spec = importlib.util.spec_from_file_location("_rl_kernel_envs", envs_path) - if spec is None or spec.loader is None: - raise RuntimeError(f"failed to load environment helpers from {envs_path}") - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -envs = _load_envs_module() - - -def _load_torch_extension_tools(): - try: - import torch - except ModuleNotFoundError as exc: - if exc.name != "torch": - raise - return None, None, None - - from torch.utils.cpp_extension import BuildExtension, CUDAExtension - - # CUDAExtension is also the supported extension entry point for ROCm - # PyTorch builds. BuildExtension dispatches .cu/.hip sources to hipcc when - # torch.version.hip is set. - return torch, BuildExtension, CUDAExtension - - -def _native_extension_required() -> bool: - """Whether the caller explicitly requested a native extension build.""" - return ( - envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) - or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) - or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) - or envs.env_flag("FORCE_CUDA") - ) - - -def _cuda_define_from_env(name: str, macro: str) -> list[str]: - value = os.environ.get(name) - if value is None: - return [] - parsed = int(value) - if parsed <= 0: - raise ValueError(f"{name} must be positive, got {value!r}") - return [f"-D{macro}={parsed}"] - - -_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES = ( - "-Xfatbin", - "-compress-all", - "-gencode", - "--generate-code", - "--expt-", - "-lineinfo", - "-allow-unsupported-compiler", - "-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH", -) -_ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE = { - "-Xfatbin", - "-gencode", - "--generate-code", -} - - -def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: - """Remove CUDA-only device compiler flags before BuildExtension calls hipcc.""" - filtered_flags = [] - skip_next = False - for flag in flags: - if skip_next: - skip_next = False - continue - if flag in _ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE: - skip_next = True - continue - if flag.startswith(_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES): - continue - filtered_flags.append(flag) - return filtered_flags - - -def get_extensions(): - torch, _, CUDAExtension = _load_torch_extension_tools() - if torch is None: - message = ( - "PyTorch is unavailable, so rl_engine._C cannot be built. Install a matching " - "CUDA/ROCm PyTorch build first, then run " - "`RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e .`." - ) - if _native_extension_required(): - raise RuntimeError(message) - warnings.warn( - f"{message} Continuing with the pure-Python fallback because no native extension " - "was explicitly requested.", - RuntimeWarning, - stacklevel=2, - ) - return [] - - extensions = [] - torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib") - torch_rpath = ["-Wl,-rpath,$ORIGIN/../torch/lib"] - if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": - torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") - is_rocm = getattr(torch.version, "hip", None) is not None - - # CUDAExtension is intentionally used for both CUDA and ROCm. On ROCm, - # PyTorch's BuildExtension hipifies CUDA sources and invokes hipcc; it also - # consumes PYTORCH_ROCM_ARCH (one or more ';'-separated gfx targets) to add - # --offload-arch. Do not require a visible GPU when a ROCm target was - # explicitly selected. - no_rocm_arch = not os.environ.get("PYTORCH_ROCM_ARCH", "").strip() - if is_rocm and no_rocm_arch and torch.cuda.device_count() == 0: - raise RuntimeError( - "ROCm builds without a visible GPU require PYTORCH_ROCM_ARCH. " - "Set one or more ';'-separated targets, for example " - "PYTORCH_ROCM_ARCH='gfx942;gfx950'." - ) - - if is_rocm or torch.cuda.is_available(): - cuda_sources = [ - "csrc/ops.cpp", - "csrc/fused_logp_kernel.cu", - "csrc/deterministic_logp_kernel.cu", - "csrc/cuda/gemm/det_gemm_kernel.cu", - "csrc/cuda/rmsnorm.cu", - "csrc/cuda/activation.cu", - "csrc/cuda/attention/deterministic_attention.cu", - "csrc/cuda/distributed/deterministic_collective.cu", - ] - if not is_rocm: - # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). - # The ROCm dispatcher falls back to PyTorch SDPA for this operator. - cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") - - nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] - if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): - nvcc_flags.append("--use_fast_math") - if not is_rocm: - cc_major, cc_minor = torch.cuda.get_device_capability() - enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" - if not enable_sm90: - # SM90 build emits 90a below; mixing plain compute_90 breaks TMA ptxas. - nvcc_flags.append( - f"-gencode=arch=compute_{cc_major}{cc_minor},code=sm_{cc_major}{cc_minor}" - ) - nvcc_flags.append("--expt-relaxed-constexpr") - nvcc_flags.append("--expt-extended-lambda") - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_TWOPASS_BLOCK_SIZE", - "FUSED_LOGP_TWOPASS_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_BLOCK_SIZE", - "FUSED_LOGP_ONLINE_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", - "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", - "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", - "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", - ) - ) - if not is_rocm and envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): - nvcc_flags.append("-lineinfo") - if ( - not is_rocm - and os.name == "nt" - and envs.env_flag(envs.KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC) - ): - nvcc_flags.append("-allow-unsupported-compiler") - nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") - - cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] - extra_link_args = list(torch_rpath) - if os.name != "nt": - # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). - extra_link_args.append("-lcuda") - - if not is_rocm: - sm90_srcs = [ - "csrc/cuda/fused_logp_sm90.cu", - "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob - "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp - "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build - # Single-card batch-invariant embedding/lm-head. - "csrc/cuda/embedding_lm_head_sm90.cu", - ] - enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) - present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] - if enable_sm90 and present_sm90: - tma_arch = f"{cc_major}{cc_minor}a" # WGMMA/TMA require the arch-native 'a' variant - cuda_sources.extend(present_sm90) - nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") - cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") - if "-lcuda" not in extra_link_args: - extra_link_args.append("-lcuda") - - # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp - # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in - # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. - enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" - if enable_det_gemm_sm90: - tma_arch = f"{cc_major}{cc_minor}a" - arch_flag = f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}" - if arch_flag not in nvcc_flags: - nvcc_flags.append(arch_flag) - if "-lcuda" not in extra_link_args: - extra_link_args.append("-lcuda") - nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") - cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") - - if is_rocm: - nvcc_flags = _filter_rocm_incompatible_nvcc_flags(nvcc_flags) - - extensions.append( - CUDAExtension( - name="rl_engine._C", - sources=cuda_sources, - include_dirs=[], - extra_compile_args={ - "cxx": cxx_flags, - "nvcc": nvcc_flags, - }, - extra_link_args=extra_link_args, - ) - ) - - if _native_extension_required() and not extensions: - raise RuntimeError( - "rl_engine._C was requested but no CUDA/ROCm build environment is available. " - "Use a matching GPU-enabled PyTorch build; for a GPU-less ROCm build, set " - "PYTORCH_ROCM_ARCH to the target architecture." - ) - - return extensions - - -def get_cmdclass(): - _, BuildExtension, _ = _load_torch_extension_tools() - if BuildExtension is None: - return {} - return {"build_ext": BuildExtension} - - -setup( - name="rl-engine", - version="0.1.0", - packages=find_packages(include=["rl_engine", "rl_engine.*"]), - install_requires=[ - "torch>=2.4.1", - "tabulate", - "numpy", - "accelerate", - "transformers==5.13.1", - ], - ext_modules=get_extensions(), - cmdclass=get_cmdclass(), - extras_require={ - "cuda": ["flashinfer"], - "rocm": ["aiter"], - "vllm": ["vllm>=0.6.0"], - "drift-viewer": ["Pillow>=10", "PySide6>=6.6"], - }, - entry_points={ - "console_scripts": [ - "rlk-drift-view=rl_engine.alignment.cross_config.drift_viewer:main", - ], - }, - python_requires=">=3.10", - include_package_data=True, - zip_safe=False, -) diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index c3aec008..4025acd2 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -19,8 +19,13 @@ STRICT_ATTENTION_SCHEDULE_ID, SplitKVSpec, ) +from rl_engine.kernels.ops.cuda.attention import deterministic_attn as deterministic_attn_module +from rl_engine.kernels.ops.cuda.attention.deterministic_attn import ( + RLKernelDeterministicAttentionCore, +) from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( AttentionPartialState, + AttentionRingSchedule, AttentionSavedForwardState, DeterministicAttentionCore, DeterministicCPAttentionReferenceOp, @@ -39,6 +44,17 @@ _GRAD_ATOL = 1.0e-5 +def test_ring_schedule_separates_compute_and_merge_order(): + schedule = AttentionRingSchedule.build(12, cp_world_size=2, kv_chunk_size=2) + + assert schedule.schedule_id == "rlkernel.attention.strict_ring_state.v1" + assert schedule.compute_communication == "decoupled" + assert schedule.overlap == "disabled" + assert schedule.merge_order == tuple(range(6)) + assert schedule.compute_order == (0, 5, 1, 4, 2, 3) + assert [block.owner_cp_rank for block in schedule.blocks] == [0, 0, 0, 1, 1, 1] + + @contextlib.contextmanager def _single_thread(): prev = torch.get_num_threads() @@ -923,6 +939,35 @@ def test_registry_dispatches_cp_attention_reference(): assert isinstance(kernel_registry.get_op("cp_attention"), DeterministicCPAttentionReferenceOp) +def test_shared_strict_core_reports_canonical_cuda_schedule(monkeypatch): + class FakeCUDAAttentionOp: + @staticmethod + def forward_with_lse(q, k, v, **_kwargs): + del k, v + return ( + torch.zeros_like(q), + torch.zeros(q.shape[:3], dtype=torch.float32, device=q.device), + ) + + monkeypatch.setattr( + deterministic_attn_module, + "DeterministicAttentionOp", + FakeCUDAAttentionOp, + ) + q, k, v = _qkv(1, 3, 4, seed=41, heads=4, kv_heads=2, dim=8) + q_positions = torch.arange(1, 4).view(1, -1) + k_positions = torch.arange(4).view(1, -1) + result = RLKernelDeterministicAttentionCore().forward_with_lse( + q, + k, + v, + query_position_ids=q_positions, + key_position_ids=k_positions, + ) + assert result.provenance["strict_core_id"] == STRICT_ATTENTION_CORE_ID + assert result.provenance["strict_schedule"] == STRICT_ATTENTION_SCHEDULE_ID + + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) def test_strict_forward_is_bitwise_invariant_to_batch_cp_and_chunk(dtype): q, k, v = _qkv(2, 5, 9, seed=41, dtype=dtype, heads=4, kv_heads=2, dim=8) diff --git a/tests/test_flashinfer_pr7_attention.py b/tests/test_flashinfer_pr7_attention.py new file mode 100644 index 00000000..ce09a3b6 --- /dev/null +++ b/tests/test_flashinfer_pr7_attention.py @@ -0,0 +1,2005 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import json +import types +from dataclasses import replace + +import pytest +import torch + +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_FA4_SCHEDULE_ID, + STRICT_ATTENTION_PRODUCTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, + SplitKVSpec, +) +from rl_engine.kernels.ops.cuda.attention import ( + flashinfer_paged_attention as paged_attention_module, +) +from rl_engine.kernels.ops.cuda.attention.cp_comm import ( + AttentionCPBlockMetadata, + AttentionCPCommunicationPlan, + AttentionCPCommunicationUnavailable, + AttentionCPMergedState, + AttentionCPOutputShard, + AttentionCPPartialState, + AttentionParallelSpec, + CUDAAGRSAttentionCPCommunication, + P2PNCCLAttentionCPCommunication, + sort_attention_cp_partial_states, +) +from rl_engine.kernels.ops.cuda.attention.deterministic_attn import DeterministicAttentionCoreResult +from rl_engine.kernels.ops.cuda.attention.flash_attn import ( + StrictFlashAttention4Core, + StrictFlashAttentionUnavailable, +) +from rl_engine.kernels.ops.cuda.attention.flashinfer_paged_attention import ( + FlashInferPagedAttentionConfig, + FlashInferQwen3PagedAttentionOp, + FlashInferRoPEFusionConfig, + FlashInferUnavailable, + _NativeFlashInferRuntimeAdapter, + build_flashinfer_paged_kv_plan, + flashinfer_prefix_cache_fingerprint, + materialize_flashinfer_paged_kv_cache, +) +from rl_engine.testing.attention_comparison import DecodeKVCacheMetadata +from scripts import ws2_p2p_nccl_attention_reference_check as p2p_check_script +from scripts import ws2_pr7_flashinfer_attention_check as check_script + + +class _FakeFlashInferWrapper: + instances: list["_FakeFlashInferWrapper"] = [] + + def __init__(self, workspace_buffer, *, kv_layout): + self.workspace_buffer = workspace_buffer + self.kv_layout = kv_layout + self.plan_kwargs = None + self.run_q = None + self.run_cache = None + self.instances.append(self) + + def plan(self, **kwargs): + self.plan_kwargs = kwargs + + def run_return_lse(self, q, paged_kv_cache): + self.run_q = q + self.run_cache = paged_kv_cache + out = torch.zeros( + q.shape, + dtype=self.plan_kwargs.get("o_data_type", q.dtype), + device=q.device, + ) + lse = torch.zeros(q.size(0), q.size(1), dtype=torch.float32, device=q.device) + return out, lse + + def get_actual_split_kv_plan(self): + seq_lens = self.plan_kwargs["seq_lens"].tolist() + disabled = bool(self.plan_kwargs.get("disable_split_kv", False)) + split_size = self.plan_kwargs.get("fixed_split_size") + plans = [] + for seq_len in seq_lens: + if disabled: + boundaries = [(0, (seq_len + 1) // 2)] + mode = "disabled" + actual_size = None + else: + assert split_size is not None + boundaries = [ + (start, min(start + split_size, (seq_len + 1) // 2)) + for start in range(0, (seq_len + 1) // 2, split_size) + ] + mode = "fixed" + actual_size = split_size + plans.append( + { + "mode": mode, + "split_size": actual_size, + "split_size_unit": "pages", + "boundary_unit": "pages", + "boundaries": boundaries, + "fallback": False, + "fallback_reason": None, + } + ) + return plans + + def get_attention_arithmetic_provenance(self): + return { + "accum_dtype": "fp32", + "downcast_at": "final_write", + "lse_dtype": "fp32", + "source": "fake_runtime_capability", + } + + def get_actual_split_kv_plan_set(self): + seq_lens = self.plan_kwargs["seq_lens"].tolist() + disabled = bool(self.plan_kwargs.get("disable_split_kv", False)) + split_size = self.plan_kwargs.get("fixed_split_size") + entries = [] + for batch_index, total in enumerate(seq_lens): + owner_ranges = ((0, total - 2), (total - 2, total)) + for tp_rank in range(2): + for cp_rank in range(2): + for owner_cp_rank, (owner_start, owner_end) in enumerate(owner_ranges): + if disabled: + mode = "disabled" + actual_size = None + boundaries = [(0, (owner_end - owner_start + 1) // 2)] + else: + mode = "fixed" + actual_size = split_size + boundaries = [ + ( + (start - owner_start) // 2, + min( + (start - owner_start) // 2 + split_size, + (owner_end - owner_start + 1) // 2, + ), + ) + for start in range(owner_start, owner_end, split_size * 2) + ] + entries.append( + { + "batch_index": batch_index, + "tp_rank": tp_rank, + "cp_rank": cp_rank, + "owner_cp_rank": owner_cp_rank, + "expected_kv_range": [owner_start, owner_end], + "mode": mode, + "split_size": actual_size, + "split_size_unit": "pages", + "boundary_unit": "pages", + "boundaries": boundaries, + "merge_order": "global_block_index", + "accum_dtype": "fp32", + "downcast_at": "final_write", + "fallback": False, + "fallback_reason": None, + } + ) + return { + "batch_size": len(seq_lens), + "tp_world_size": 2, + "cp_world_size": 2, + "total_kv_tokens": seq_lens, + "entries": entries, + } + + +def test_native_flashinfer_adapter_reads_materialized_plan_and_normalizes_lse(): + class _NativeWrapper: + def __init__(self): + self._backend = "fa2" + self._plan_info = (2, 2, 0, 16, 0, 16, 32, 0, 48, 60, 0, 0, 0, 0, 0) + self._pin_memory_int_workspace_buffer = torch.zeros(64, dtype=torch.uint8) + self._pin_memory_int_workspace_buffer[0:8].view(torch.int32).copy_( + torch.tensor([0, 1], dtype=torch.int32) + ) + self._pin_memory_int_workspace_buffer[32:40].view(torch.int32).zero_() + + def plan(self, **kwargs): + return None + + adapter = _NativeFlashInferRuntimeAdapter( + _NativeWrapper(), + FlashInferPagedAttentionConfig(workspace_size_bytes=1024), + ) + adapter.plan( + seq_lens=torch.tensor([16, 16], dtype=torch.int32), + page_size=4, + disable_split_kv=True, + ) + + assert adapter.get_actual_split_kv_plan()[0]["boundaries"] == [(0, 4)] + plan_set = adapter.get_actual_split_kv_plan_set() + assert len(plan_set["entries"]) == 16 + assert plan_set["entries"][1]["expected_kv_range"] == [8, 16] + normalized = adapter.normalize_lse(torch.ones(1)) + assert torch.allclose(normalized, torch.log(torch.tensor([2.0]))) + + +def _fake_flashinfer(): + _FakeFlashInferWrapper.instances = [] + return types.SimpleNamespace( + prefill=types.SimpleNamespace( + BatchPrefillWithPagedKVCacheWrapper=_FakeFlashInferWrapper, + ), + decode=types.SimpleNamespace( + BatchDecodeWithPagedKVCacheWrapper=_FakeFlashInferWrapper, + ), + ) + + +def _metadata(*, batch: int = 2, query_len: int = 1) -> DecodeKVCacheMetadata: + page_size = 2 + cache_capacity = 6 + positions = torch.arange(cache_capacity, dtype=torch.long).repeat(batch, 1) + query_positions = torch.arange( + cache_capacity - query_len, + cache_capacity, + dtype=torch.long, + ).repeat(batch, 1) + return DecodeKVCacheMetadata( + cache_position=query_positions.clone(), + kv_seq_lens=torch.full((batch,), cache_capacity, dtype=torch.long), + block_table=torch.tensor([[0, 1, 2]] * batch, dtype=torch.long), + global_token_positions=positions, + query_position_ids=query_positions.clone(), + key_position_ids=positions.clone(), + page_size=page_size, + q_rope_state="pre_rope", + k_cache_rope_state="pre_rope", + ) + + +def _qkv(*, batch: int = 2, query_len: int = 1): + gen = torch.Generator().manual_seed(7) + q = torch.randn(batch, 4, query_len, 8, generator=gen) + k = torch.randn(batch, 2, 6, 8, generator=gen) + v = torch.randn(batch, 2, 6, 8, generator=gen) + return q, k, v + + +def _partial_state(global_block_index: int) -> AttentionCPPartialState: + return AttentionCPPartialState( + out=torch.full((1, 2, 1, 4), float(global_block_index)), + lse=torch.full((1, 2, 1), float(global_block_index), dtype=torch.float32), + block=AttentionCPBlockMetadata( + global_block_index=global_block_index, + kv_block_start=global_block_index * 2, + kv_block_end=global_block_index * 2 + 2, + owner_cp_rank=global_block_index % 2, + owner_tp_rank=0, + ), + ) + + +def _p2p_plan(*, cp_rank: int = 0) -> AttentionCPCommunicationPlan: + return AttentionCPCommunicationPlan( + parallel=AttentionParallelSpec( + tp_world_size=2, + tp_rank=0, + cp_world_size=2, + cp_rank=cp_rank, + ), + backend="p2p_nccl_reference", + status="implemented", + expected_blocks=( + AttentionCPBlockMetadata(0, 0, 2, 0, 0), + AttentionCPBlockMetadata(1, 2, 4, 1, 0), + ), + expected_kv_token_range=(0, 4), + query_token_ranges=((0, 1), (1, 2)), + ) + + +class _CompletedRequest: + def wait(self): + return True + + +class _FakeP2POp: + def __init__(self, op, tensor, peer, *, group=None): + self.op = op + self.tensor = tensor + self.peer = peer + self.group = group + + +class _FakeNCCLDistributed: + def __init__(self, *, rank: int, receive_payloads=()): + self.rank = rank + self.receive_payloads = list(receive_payloads) + + @staticmethod + def is_available(): + return True + + @staticmethod + def is_initialized(): + return True + + @staticmethod + def get_backend(group=None): + return "nccl" + + @staticmethod + def get_world_size(group=None): + return 2 + + def get_rank(self, group=None): + return self.rank + + @staticmethod + def get_global_rank(group, rank): + return rank + + @staticmethod + def isend(tensor, dst, group=None): + raise AssertionError("P2POp should defer isend") + + @staticmethod + def irecv(tensor, src, group=None): + raise AssertionError("P2POp should defer irecv") + + P2POp = _FakeP2POp + + def batch_isend_irecv(self, operations): + for operation in operations: + if getattr(operation.op, "__name__", None) == "irecv": + operation.tensor.copy_(self.receive_payloads.pop(0)) + return [_CompletedRequest() for _ in operations] + + +class _FakeCPCommunication: + def all_gather_query(self, local_q, plan): + return torch.cat([local_q] * plan.parallel.cp_world_size, dim=2) + + def all_gather_partial_states(self, local_states, plan): + local = local_states[0] + remote_block = next( + block for block in plan.expected_blocks if block.owner_cp_rank != plan.parallel.cp_rank + ) + remote = AttentionCPPartialState( + out=torch.ones_like(local.out), + lse=torch.ones_like(local.lse), + block=remote_block, + ) + return tuple(sorted((local, remote), key=lambda state: state.block.global_block_index)) + + def reduce_scatter_merged_state(self, merged_state, plan): + start, end = plan.query_token_ranges[plan.parallel.cp_rank] + return AttentionCPMergedState( + out=merged_state.out[:, :, start:end, :], + lse=merged_state.lse[:, :, start:end], + ) + + +class _IdentityStrictRoPE: + backend_id = "rlkernel.cuda.rope_sm90" + + def __call__(self, x, positions, *, theta=1_000_000.0): + assert positions.ndim == 1 + assert float(theta) == 1_000_000.0 + return x + + +class _RecordingStrictCore: + core_id = STRICT_ATTENTION_CORE_ID + strict_schedule = STRICT_ATTENTION_SCHEDULE_ID + backend_id = "test.strict_cuda_core" + merge_order = "global_block_index" + accum_dtype = "fp32" + downcast_at = "final_write" + fallback = False + native_attention_arithmetic = False + + def __init__(self): + self.calls = [] + + def forward_with_lse( + self, + q, + k, + v, + *, + query_position_ids, + key_position_ids, + **_kwargs, + ): + self.calls.append( + { + "q": q.clone(), + "k": k.clone(), + "v": v.clone(), + "query_position_ids": query_position_ids.clone(), + "key_position_ids": key_position_ids.clone(), + } + ) + return DeterministicAttentionCoreResult( + out=torch.zeros_like(q), + lse=torch.zeros(q.shape[:3], dtype=torch.float32, device=q.device), + provenance={ + "strict_core_id": self.core_id, + "strict_schedule": self.strict_schedule, + "attention_backend": self.backend_id, + "split_kv": { + "actual_split_kv_policy": "disabled", + "actual_split_boundaries": [[0, k.size(2)]], + }, + "merge_order": self.merge_order, + "accum_dtype": self.accum_dtype, + "downcast_at": self.downcast_at, + "fallback": False, + "fallback_reason": None, + "native_attention_arithmetic": False, + }, + ) + + +class _StrictCPCommunication: + backend_id = "p2p_nccl_reference" + + def all_gather_query(self, local_q, plan): + return torch.cat((local_q, local_q + 1), dim=2) + + def all_gather_kv(self, local_k, local_v, plan): + return ( + torch.cat((local_k, local_k + 2), dim=2), + torch.cat((local_v, local_v + 3), dim=2), + ) + + def all_gather_position_ids(self, local_q_positions, local_k_positions, plan): + return ( + torch.cat((local_q_positions, local_q_positions + 1), dim=1), + torch.cat((local_k_positions, local_k_positions + 2), dim=1), + ) + + def reduce_scatter_strict_result(self, out, lse, plan): + start, end = plan.query_token_ranges[plan.parallel.cp_rank] + return AttentionCPOutputShard( + out=out[:, :, start:end, :].contiguous(), + lse=lse[:, :, start:end].contiguous(), + ) + + +class _FakeDeterministicCollective: + world_size = 2 + + @staticmethod + def all_gather(tensor): + return torch.cat((tensor, tensor + 10), dim=0) + + @staticmethod + def reduce_scatter(tensor): + return tensor.chunk(2, dim=0)[0].contiguous() + + +class _FakeAutogradCollective: + world_size = 2 + + def __init__(self): + self.all_gather_calls = 0 + self.reduce_scatter_calls = 0 + + def all_gather(self, tensor): + self.all_gather_calls += 1 + return torch.cat((tensor, tensor), dim=0) + + def reduce_scatter(self, tensor): + self.reduce_scatter_calls += 1 + first, second = tensor.chunk(2, dim=0) + return (first + second).contiguous() + + +class _FakeCollectiveDist: + @staticmethod + def all_gather_object(outputs, value, group=None): + outputs[:] = [value, value] + + +def test_flashinfer_pr7_prefill_adapter_passes_qwen3_rope_and_splitk_policy(): + q, k, v = _qkv(query_len=2) + metadata = _metadata(query_len=2) + op = FlashInferQwen3PagedAttentionOp(flashinfer_module=_fake_flashinfer()) + + result = op( + q, + k, + v, + metadata, + config=FlashInferPagedAttentionConfig( + mode="prefill", + workspace_size_bytes=1024, + split_kv=SplitKVSpec.fixed(4), + ), + ) + + wrapper = _FakeFlashInferWrapper.instances[-1] + assert wrapper.kv_layout == "NHD" + assert wrapper.run_q.shape == (q.size(0) * q.size(2), q.size(1), q.size(3)) + assert result.out.shape == q.shape + assert result.lse.shape == q.shape[:3] + assert result.provenance["actual_backend"] == "flashinfer_batch_prefill_paged_kv" + assert result.provenance["rope_fusion_boundary"] == "flashinfer_attention_kernel" + assert result.provenance["pos_encoding_mode"] == "ROPE_LLAMA" + assert result.provenance["rope_theta"] == 1_000_000.0 + assert result.provenance["rope_scale"] == 1.0 + assert result.provenance["split_kv_policy"] == "fixed:4" + assert result.provenance["batch_invariant_claim"] == "strict_runtime_verified" + assert result.provenance["requested_split_kv_policy"] == "fixed" + assert result.provenance["requested_split_kv_size"] == 4 + assert result.provenance["actual_split_kv_plans"][0]["actual_split_boundaries"] == [ + [0, 4], + [4, 6], + ] + assert result.provenance["tp_world_size"] == 2 + assert result.provenance["cp_world_size"] == 2 + assert result.provenance["cp_comm_backend"] == "cuda_ag_rs" + assert result.provenance["cp_comm_status"] == "interface_only" + assert result.provenance["cp_comm_pattern"] == "ag_rs" + assert result.provenance["cp_comm_compute_communication"] == "decoupled" + assert result.provenance["cp_comm_merge_order"] == "global_block_index" + assert result.provenance["cp_comm_accum_dtype"] == "fp32" + assert result.provenance["cp_comm_return_lse"] is True + assert result.provenance["cp_comm_contract"] == "partial_out_lse_global_block_index" + assert result.provenance["cp_comm_required"] is False + assert result.provenance["accum_dtype"] == "fp32" + assert result.provenance["downcast_at"] == "final_write" + assert result.provenance["arithmetic_semantics_verified"] is True + assert result.provenance["actual_split_kv_plan_set"]["coverage"] == ( + "complete_batch_tp_cp_owner_cartesian_product" + ) + + plan = wrapper.plan_kwargs + assert plan["qo_indptr"].tolist() == [0, 2, 4] + assert plan["paged_kv_indptr"].tolist() == [0, 3, 6] + assert plan["paged_kv_indices"].tolist() == [0, 1, 2, 3, 4, 5] + assert plan["paged_kv_last_page_len"].tolist() == [2, 2] + assert plan["pos_encoding_mode"] == "ROPE_LLAMA" + assert plan["rope_theta"] == 1_000_000.0 + assert plan["rope_scale"] == 1.0 + assert plan["q_data_type"] == q.dtype + assert plan["kv_data_type"] == q.dtype + assert plan["fixed_split_size"] == 2 + assert plan["disable_split_kv"] is False + + +def test_flashinfer_pr7_required_cp_comm_uses_explicit_deterministic_fallback(): + q, k, v = _qkv(query_len=2) + metadata = _metadata(query_len=2) + plan = AttentionCPCommunicationPlan( + parallel=AttentionParallelSpec(tp_world_size=2, cp_world_size=2), + backend="p2p_nccl_reference", + status="implemented", + expected_blocks=( + AttentionCPBlockMetadata(0, 0, 3, 0, 0), + AttentionCPBlockMetadata(1, 3, 6, 1, 0), + ), + expected_kv_token_range=(0, 6), + query_token_ranges=((0, 1), (1, 2)), + ) + config = FlashInferPagedAttentionConfig( + mode="prefill", + workspace_size_bytes=1024, + cp_comm_plan=plan, + require_cp_comm=True, + cp_communication=_FakeCPCommunication(), + ) + result = FlashInferQwen3PagedAttentionOp(flashinfer_module=_fake_flashinfer())( + q, + k, + v, + metadata, + config=config, + ) + + assert result.out.shape == (q.size(0), q.size(1), 1, q.size(3)) + assert result.lse.shape == (q.size(0), q.size(1), 1) + assert result.provenance["actual_backend"] == "rlkernel_deterministic_cp_reference" + assert result.provenance["fallback"] is True + assert result.provenance["fallback_reason"] == ( + "flashinfer_owner_local_cp_partial_api_unavailable" + ) + assert result.provenance["cp_comm_required"] is True + assert result.provenance["query_ag"] == "cp_rank_order" + assert result.provenance["actual_split_kv_plan_set"]["coverage"] == ( + "complete_batch_tp_cp_owner_cartesian_product" + ) + + +def test_strict_paged_path_uses_shared_core_and_never_runs_flashinfer_arithmetic(): + q, k, v = (tensor.to(torch.bfloat16) for tensor in _qkv(query_len=1)) + core = _RecordingStrictCore() + _FakeFlashInferWrapper.instances = [] + + result = FlashInferQwen3PagedAttentionOp(flashinfer_module=_fake_flashinfer())( + q, + k, + v, + _metadata(query_len=1), + config=FlashInferPagedAttentionConfig( + mode="decode", + workspace_size_bytes=1024, + strict_mode=True, + deterministic_core=core, + strict_rope_op=_IdentityStrictRoPE(), + ), + ) + + assert _FakeFlashInferWrapper.instances == [] + assert len(core.calls) == q.size(0) + assert core.calls[0]["query_position_ids"].tolist() == [[5]] + assert core.calls[0]["key_position_ids"].tolist() == [[0, 1, 2, 3, 4, 5]] + assert result.provenance["strict_core_id"] == STRICT_ATTENTION_CORE_ID + assert result.provenance["native_attention_arithmetic"] is False + assert result.provenance["fallback"] is False + assert result.provenance["materialization"] == ("flashinfer_paged_kv_layout_shared_core") + assert result.out.dtype is torch.bfloat16 + assert result.lse.dtype is torch.float32 + + +def test_strict_cp_path_gathers_qkv_and_real_position_ids_before_shared_core(): + generator = torch.Generator().manual_seed(19) + q = torch.randn(1, 4, 1, 8, generator=generator, dtype=torch.bfloat16) + k = torch.randn(1, 2, 2, 8, generator=generator, dtype=torch.bfloat16) + v = torch.randn(1, 2, 2, 8, generator=generator, dtype=torch.bfloat16) + core = _RecordingStrictCore() + metadata = types.SimpleNamespace( + q_rope_state="pre_rope", + k_cache_rope_state="pre_rope", + query_position_ids=torch.tensor([[2]], dtype=torch.long), + key_position_ids=torch.tensor([[0, 1]], dtype=torch.long), + ) + + result = FlashInferQwen3PagedAttentionOp(flashinfer_module=_fake_flashinfer())( + q, + k, + v, + metadata, + config=FlashInferPagedAttentionConfig( + mode="prefill", + workspace_size_bytes=1024, + cp_comm_plan=_p2p_plan(), + require_cp_comm=True, + strict_mode=True, + cp_communication=_StrictCPCommunication(), + deterministic_core=core, + strict_rope_op=_IdentityStrictRoPE(), + ), + ) + + assert len(core.calls) == 1 + assert core.calls[0]["query_position_ids"].tolist() == [[2, 3]] + assert core.calls[0]["key_position_ids"].tolist() == [[0, 1, 2, 3]] + assert core.calls[0]["q"].shape[2] == 2 + assert core.calls[0]["k"].shape[2] == 4 + assert result.out.shape == q.shape + assert result.provenance["materialization"] == ("ag_qkv_positions_shared_core_rs") + assert result.provenance["strict_full_qkv_all_gather"] is True + assert result.provenance["strict_position_ids_all_gather"] is True + assert result.provenance["compute_communication"] == "decoupled" + assert result.provenance["compute_schedule"] == ("rlkernel.attention.strict_ring_state.v1") + assert result.provenance["communication_overlap"] == "disabled" + assert result.provenance["ring_schedule_default"] is True + assert result.provenance["ring_partial_arithmetic"] is False + assert result.provenance["fallback"] is False + + +def test_strict_cp_training_rejects_non_autograd_p2p_reference(): + generator = torch.Generator().manual_seed(29) + q = torch.randn(1, 4, 1, 8, generator=generator, dtype=torch.bfloat16).requires_grad_() + k = torch.randn(1, 2, 2, 8, generator=generator, dtype=torch.bfloat16).requires_grad_() + v = torch.randn(1, 2, 2, 8, generator=generator, dtype=torch.bfloat16).requires_grad_() + metadata = types.SimpleNamespace( + q_rope_state="pre_rope", + k_cache_rope_state="pre_rope", + query_position_ids=torch.tensor([[2]], dtype=torch.long), + key_position_ids=torch.tensor([[0, 1]], dtype=torch.long), + ) + + with pytest.raises(FlashInferUnavailable, match="autograd-capable self-owned CUDA AG/RS"): + FlashInferQwen3PagedAttentionOp(flashinfer_module=_fake_flashinfer())( + q, + k, + v, + metadata, + config=FlashInferPagedAttentionConfig( + mode="prefill", + workspace_size_bytes=1024, + cp_comm_plan=_p2p_plan(), + require_cp_comm=True, + strict_mode=True, + cp_communication=_StrictCPCommunication(), + deterministic_core=_RecordingStrictCore(), + strict_rope_op=_IdentityStrictRoPE(), + ), + ) + + +def test_strict_mode_rejects_split_k_and_unverified_core(): + with pytest.raises(ValueError, match="Split-KV to be disabled"): + FlashInferPagedAttentionConfig( + strict_mode=True, + split_kv=SplitKVSpec.fixed(2), + deterministic_core=_RecordingStrictCore(), + ).validate(head_dim=8, query_len=1) + + bad_core = types.SimpleNamespace( + core_id="different", + forward_with_lse=lambda *_args, **_kwargs: None, + ) + with pytest.raises(ValueError, match="core ID"): + FlashInferPagedAttentionConfig( + strict_mode=True, + deterministic_core=bad_core, + ).validate(head_dim=8, query_len=1) + + +def test_flashinfer_pr7_decode_adapter_can_disable_splitk_for_strict_candidate(): + q, k, v = _qkv(query_len=1) + metadata = _metadata(query_len=1) + op = FlashInferQwen3PagedAttentionOp(flashinfer_module=_fake_flashinfer()) + + result = op( + q, + k, + v, + metadata, + config=FlashInferPagedAttentionConfig( + mode="decode", + workspace_size_bytes=1024, + split_kv=SplitKVSpec.disabled(), + ), + ) + + wrapper = _FakeFlashInferWrapper.instances[-1] + assert result.provenance["actual_backend"] == "flashinfer_batch_decode_paged_kv" + assert result.provenance["split_kv_policy"] == "disabled" + assert result.provenance["batch_invariant_claim"] == "strict_runtime_verified" + assert wrapper.plan_kwargs["disable_split_kv"] is True + + +def test_flashinfer_pr7_rejects_auto_splitk_when_batch_invariance_is_required(): + q, k, v = _qkv(query_len=1) + metadata = _metadata(query_len=1) + op = FlashInferQwen3PagedAttentionOp(flashinfer_module=_fake_flashinfer()) + + with pytest.raises(ValueError, match="auto split-KV"): + op( + q, + k, + v, + metadata, + config=FlashInferPagedAttentionConfig( + mode="decode", + workspace_size_bytes=1024, + split_kv=SplitKVSpec.auto(), + ), + ) + + +def test_flashinfer_pr7_strict_fixed_mode_requires_actual_runtime_split_plan(): + class _NoRuntimePlanWrapper(_FakeFlashInferWrapper): + get_actual_split_kv_plan = None + + fake = types.SimpleNamespace( + prefill=types.SimpleNamespace(BatchPrefillWithPagedKVCacheWrapper=_NoRuntimePlanWrapper), + decode=types.SimpleNamespace(BatchDecodeWithPagedKVCacheWrapper=_NoRuntimePlanWrapper), + ) + + q, k, v = _qkv(query_len=1) + + with pytest.raises(FlashInferUnavailable, match="actual-plan provenance"): + FlashInferQwen3PagedAttentionOp(flashinfer_module=fake)( + q, + k, + v, + _metadata(query_len=1), + config=FlashInferPagedAttentionConfig( + mode="decode", + workspace_size_bytes=1024, + split_kv=SplitKVSpec.fixed(2), + ), + ) + + +def test_flashinfer_pr7_disabled_plan_is_exact_when_disable_knob_is_accepted(): + class _NoRuntimePlanWrapper(_FakeFlashInferWrapper): + get_actual_split_kv_plan = None + + fake = types.SimpleNamespace( + prefill=types.SimpleNamespace(BatchPrefillWithPagedKVCacheWrapper=_NoRuntimePlanWrapper), + decode=types.SimpleNamespace(BatchDecodeWithPagedKVCacheWrapper=_NoRuntimePlanWrapper), + ) + q, k, v = _qkv(query_len=1) + result = FlashInferQwen3PagedAttentionOp(flashinfer_module=fake)( + q, + k, + v, + _metadata(query_len=1), + config=FlashInferPagedAttentionConfig( + mode="decode", + workspace_size_bytes=1024, + split_kv=SplitKVSpec.disabled(), + ), + ) + + assert result.provenance["actual_split_kv_plans"][0]["actual_split_boundaries"] == [[0, 6]] + assert result.provenance["actual_split_kv_plans"][0]["split_kv_backend"] == ( + "flashinfer_disabled_verified" + ) + + +def test_flashinfer_pr7_rejects_actual_split_plan_mismatch(): + class _MismatchedRuntimePlanWrapper(_FakeFlashInferWrapper): + def get_actual_split_kv_plan(self): + return [ + {"mode": "fixed", "split_size": 3, "boundaries": [(0, 3), (3, 6)]}, + {"mode": "fixed", "split_size": 3, "boundaries": [(0, 3), (3, 6)]}, + ] + + fake = types.SimpleNamespace( + prefill=types.SimpleNamespace( + BatchPrefillWithPagedKVCacheWrapper=_MismatchedRuntimePlanWrapper + ), + decode=types.SimpleNamespace( + BatchDecodeWithPagedKVCacheWrapper=_MismatchedRuntimePlanWrapper + ), + ) + q, k, v = _qkv(query_len=1) + + with pytest.raises( + FlashInferUnavailable, + match="does not match|invalid actual|missing required fields", + ): + FlashInferQwen3PagedAttentionOp(flashinfer_module=fake)( + q, + k, + v, + _metadata(query_len=1), + config=FlashInferPagedAttentionConfig( + mode="decode", + workspace_size_bytes=1024, + split_kv=SplitKVSpec.fixed(2), + ), + ) + + +def test_flashinfer_pr7_strict_mode_requires_runtime_arithmetic_provenance(): + class _NoArithmeticProvenanceWrapper(_FakeFlashInferWrapper): + get_attention_arithmetic_provenance = None + + fake = types.SimpleNamespace( + prefill=types.SimpleNamespace( + BatchPrefillWithPagedKVCacheWrapper=_NoArithmeticProvenanceWrapper + ), + decode=types.SimpleNamespace( + BatchDecodeWithPagedKVCacheWrapper=_NoArithmeticProvenanceWrapper + ), + ) + q, k, v = _qkv(query_len=1) + + with pytest.raises(FlashInferUnavailable, match="arithmetic provenance"): + FlashInferQwen3PagedAttentionOp(flashinfer_module=fake)( + q, + k, + v, + _metadata(query_len=1), + config=FlashInferPagedAttentionConfig( + mode="decode", + workspace_size_bytes=1024, + ), + ) + + +def test_flashinfer_pr7_strict_mode_requires_complete_runtime_plan_set(): + class _NoRuntimePlanSetWrapper(_FakeFlashInferWrapper): + get_actual_split_kv_plan_set = None + + fake = types.SimpleNamespace( + prefill=types.SimpleNamespace(BatchPrefillWithPagedKVCacheWrapper=_NoRuntimePlanSetWrapper), + decode=types.SimpleNamespace(BatchDecodeWithPagedKVCacheWrapper=_NoRuntimePlanSetWrapper), + ) + q, k, v = _qkv(query_len=1) + + with pytest.raises(FlashInferUnavailable, match="complete batch/TP/CP/owner"): + FlashInferQwen3PagedAttentionOp(flashinfer_module=fake)( + q, + k, + v, + _metadata(query_len=1), + config=FlashInferPagedAttentionConfig( + mode="decode", + workspace_size_bytes=1024, + ), + ) + + +def test_flashinfer_pr7_runtime_plan_set_requires_explicit_reduction_semantics(): + class _MissingReductionSemanticsWrapper(_FakeFlashInferWrapper): + def get_actual_split_kv_plan_set(self): + plan_set = super().get_actual_split_kv_plan_set() + del plan_set["entries"][0]["merge_order"] + return plan_set + + fake = types.SimpleNamespace( + prefill=types.SimpleNamespace( + BatchPrefillWithPagedKVCacheWrapper=_MissingReductionSemanticsWrapper + ), + decode=types.SimpleNamespace( + BatchDecodeWithPagedKVCacheWrapper=_MissingReductionSemanticsWrapper + ), + ) + q, k, v = _qkv(query_len=1) + + with pytest.raises(FlashInferUnavailable, match="merge_order"): + FlashInferQwen3PagedAttentionOp(flashinfer_module=fake)( + q, + k, + v, + _metadata(query_len=1), + config=FlashInferPagedAttentionConfig( + mode="decode", + workspace_size_bytes=1024, + ), + ) + + +def test_flashinfer_pr7_runtime_split_plan_requires_explicit_fallback_fields(): + class _MissingFallbackFieldsWrapper(_FakeFlashInferWrapper): + def get_actual_split_kv_plan(self): + plans = super().get_actual_split_kv_plan() + del plans[0]["fallback"] + return plans + + fake = types.SimpleNamespace( + prefill=types.SimpleNamespace( + BatchPrefillWithPagedKVCacheWrapper=_MissingFallbackFieldsWrapper + ), + decode=types.SimpleNamespace( + BatchDecodeWithPagedKVCacheWrapper=_MissingFallbackFieldsWrapper + ), + ) + q, k, v = _qkv(query_len=1) + + with pytest.raises(FlashInferUnavailable, match="fallback"): + FlashInferQwen3PagedAttentionOp(flashinfer_module=fake)( + q, + k, + v, + _metadata(query_len=1), + config=FlashInferPagedAttentionConfig( + mode="decode", + workspace_size_bytes=1024, + ), + ) + + +def test_flashinfer_pr7_rejects_non_fp32_runtime_accumulation(): + class _WrongArithmeticWrapper(_FakeFlashInferWrapper): + def get_attention_arithmetic_provenance(self): + return { + "accum_dtype": "bf16", + "downcast_at": "per_split", + "lse_dtype": "fp32", + "source": "fake_runtime_capability", + } + + fake = types.SimpleNamespace( + prefill=types.SimpleNamespace(BatchPrefillWithPagedKVCacheWrapper=_WrongArithmeticWrapper), + decode=types.SimpleNamespace(BatchDecodeWithPagedKVCacheWrapper=_WrongArithmeticWrapper), + ) + q, k, v = _qkv(query_len=1) + + with pytest.raises(FlashInferUnavailable, match="accum_dtype, downcast_at"): + FlashInferQwen3PagedAttentionOp(flashinfer_module=fake)( + q, + k, + v, + _metadata(query_len=1), + config=FlashInferPagedAttentionConfig( + mode="decode", + workspace_size_bytes=1024, + ), + ) + + +def test_flashinfer_pr7_rejects_runtime_output_dtype_boundary_mismatch(): + class _WrongOutputDTypeWrapper(_FakeFlashInferWrapper): + def run_return_lse(self, q, paged_kv_cache): + out, lse = super().run_return_lse(q, paged_kv_cache) + return out.double(), lse + + fake = types.SimpleNamespace( + prefill=types.SimpleNamespace(BatchPrefillWithPagedKVCacheWrapper=_WrongOutputDTypeWrapper), + decode=types.SimpleNamespace(BatchDecodeWithPagedKVCacheWrapper=_WrongOutputDTypeWrapper), + ) + q, k, v = _qkv(query_len=1) + + with pytest.raises(FlashInferUnavailable, match="final output dtype"): + FlashInferQwen3PagedAttentionOp(flashinfer_module=fake)( + q, + k, + v, + _metadata(query_len=1), + config=FlashInferPagedAttentionConfig( + mode="decode", + workspace_size_bytes=1024, + ), + ) + + +def test_flashinfer_pr7_rejects_non_fp32_runtime_lse(): + class _WrongLSEDTypeWrapper(_FakeFlashInferWrapper): + def run_return_lse(self, q, paged_kv_cache): + out, lse = super().run_return_lse(q, paged_kv_cache) + return out, lse.to(torch.bfloat16) + + fake = types.SimpleNamespace( + prefill=types.SimpleNamespace(BatchPrefillWithPagedKVCacheWrapper=_WrongLSEDTypeWrapper), + decode=types.SimpleNamespace(BatchDecodeWithPagedKVCacheWrapper=_WrongLSEDTypeWrapper), + ) + q, k, v = _qkv(query_len=1) + + with pytest.raises(FlashInferUnavailable, match="LSE must be FP32"): + FlashInferQwen3PagedAttentionOp(flashinfer_module=fake)( + q, + k, + v, + _metadata(query_len=1), + config=FlashInferPagedAttentionConfig( + mode="decode", + workspace_size_bytes=1024, + ), + ) + + +def test_flashinfer_pr7_required_cp_comm_needs_implemented_plan(): + q, k, v = _qkv(query_len=1) + metadata = _metadata(query_len=1) + op = FlashInferQwen3PagedAttentionOp(flashinfer_module=_fake_flashinfer()) + + with pytest.raises(ValueError, match="implemented CP communication plan"): + op( + q, + k, + v, + metadata, + config=FlashInferPagedAttentionConfig( + mode="decode", + workspace_size_bytes=1024, + require_cp_comm=True, + ), + ) + + +def test_flashinfer_pr7_implemented_cp_comm_status_requires_execution(): + config = FlashInferPagedAttentionConfig( + cp_comm_plan=AttentionCPCommunicationPlan( + parallel=AttentionParallelSpec(tp_world_size=2, cp_world_size=2), + status="implemented", + ) + ) + + with pytest.raises(ValueError, match="require_cp_comm"): + config.validate(head_dim=8, query_len=1) + + +def test_flashinfer_pr7_rejects_post_rope_inputs_for_rope_llama_fusion(): + q, k, v = _qkv(query_len=1) + metadata = _metadata(query_len=1) + op = FlashInferQwen3PagedAttentionOp(flashinfer_module=_fake_flashinfer()) + + with pytest.raises(ValueError, match="rotated twice"): + op( + q, + k, + v, + metadata, + config=FlashInferPagedAttentionConfig( + mode="decode", + workspace_size_bytes=1024, + rope=FlashInferRoPEFusionConfig(q_rope_state="post_rope"), + ), + ) + + +def test_flashinfer_pr7_rejects_metadata_rope_state_mismatch(): + q, k, v = _qkv(query_len=1) + metadata = _metadata(query_len=1) + metadata = DecodeKVCacheMetadata(**{**metadata.__dict__, "k_cache_rope_state": "post_rope"}) + + with pytest.raises(ValueError, match="metadata.k_cache_rope_state"): + FlashInferQwen3PagedAttentionOp(flashinfer_module=_fake_flashinfer())( + q, + k, + v, + metadata, + config=FlashInferPagedAttentionConfig( + mode="decode", + workspace_size_bytes=1024, + ), + ) + + +def test_flashinfer_pr7_rejects_query_position_identity_mismatch(): + q, k, v = _qkv(query_len=1) + metadata = _metadata(query_len=1) + metadata = DecodeKVCacheMetadata( + **{ + **metadata.__dict__, + "query_position_ids": torch.zeros_like(metadata.query_position_ids), + } + ) + + with pytest.raises(ValueError, match="must match exactly"): + FlashInferQwen3PagedAttentionOp(flashinfer_module=_fake_flashinfer())( + q, + k, + v, + metadata, + config=FlashInferPagedAttentionConfig( + mode="decode", + workspace_size_bytes=1024, + ), + ) + + +def test_flashinfer_pr7_rejects_nontrailing_query_positions(): + q, k, v = _qkv(query_len=1) + metadata = _metadata(query_len=1) + metadata = DecodeKVCacheMetadata( + **{ + **metadata.__dict__, + "cache_position": torch.full_like(metadata.cache_position, 4), + "query_position_ids": torch.full_like(metadata.query_position_ids, 4), + } + ) + + with pytest.raises(ValueError, match="trailing contiguous positions"): + FlashInferQwen3PagedAttentionOp(flashinfer_module=_fake_flashinfer())( + q, + k, + v, + metadata, + config=FlashInferPagedAttentionConfig( + mode="decode", + workspace_size_bytes=1024, + ), + ) + + +def test_flashinfer_pr7_prefix_cache_fingerprint_binds_rope_identity(): + q, k, v = _qkv(batch=1, query_len=1) + metadata = _metadata(batch=1, query_len=1) + config = FlashInferPagedAttentionConfig(mode="decode", workspace_size_bytes=1024) + fingerprint = flashinfer_prefix_cache_fingerprint( + q, + k, + v, + metadata, + config, + prefix_length=4, + ) + cached = DecodeKVCacheMetadata( + **{ + **metadata.__dict__, + "prefix_cache_enabled": True, + "prefix_cache_key": "prefix-0", + "prefix_length": 4, + "prefix_cache_fingerprint": fingerprint, + } + ) + + FlashInferQwen3PagedAttentionOp(flashinfer_module=_fake_flashinfer())( + q, + k, + v, + cached, + config=config, + ) + drifted = replace( + config, + rope=replace(config.rope, rope_theta=10_000.0), + ) + with pytest.raises(ValueError, match="rope_theta"): + FlashInferQwen3PagedAttentionOp(flashinfer_module=_fake_flashinfer())( + q, + k, + v, + cached, + config=drifted, + ) + + +def test_flashinfer_pr7_rejects_stale_prefix_cache_content(): + q, k, v = _qkv(batch=1, query_len=1) + metadata = _metadata(batch=1, query_len=1) + config = FlashInferPagedAttentionConfig(mode="decode", workspace_size_bytes=1024) + fingerprint = flashinfer_prefix_cache_fingerprint( + q, + k, + v, + metadata, + config, + prefix_length=4, + ) + cached = DecodeKVCacheMetadata( + **{ + **metadata.__dict__, + "prefix_cache_enabled": True, + "prefix_cache_key": "prefix-0", + "prefix_length": 4, + "prefix_cache_fingerprint": fingerprint, + } + ) + stale_k = k.clone() + stale_k[0, 0, 0, 0] += 1.0 + + with pytest.raises(ValueError, match="prefix_cache_fingerprint"): + FlashInferQwen3PagedAttentionOp(flashinfer_module=_fake_flashinfer())( + q, + stale_k, + v, + cached, + config=config, + ) + + +def test_attention_cp_partial_states_sort_by_global_block_index(): + plan = AttentionCPCommunicationPlan( + parallel=AttentionParallelSpec(tp_world_size=2, cp_world_size=2), + status="implemented", + ) + + ordered = sort_attention_cp_partial_states( + (_partial_state(2), _partial_state(0), _partial_state(1)), + plan=plan, + ) + + assert [state.block.global_block_index for state in ordered] == [0, 1, 2] + + +def test_attention_cp_partial_states_reject_duplicate_global_block_index(): + plan = AttentionCPCommunicationPlan( + parallel=AttentionParallelSpec(tp_world_size=2, cp_world_size=2) + ) + + with pytest.raises(ValueError, match="duplicate global_block_index"): + sort_attention_cp_partial_states( + (_partial_state(1), _partial_state(1)), + plan=plan, + ) + + +def test_cuda_ag_rs_attention_cp_comm_requires_compiled_collective(): + plan = AttentionCPCommunicationPlan( + parallel=AttentionParallelSpec(tp_world_size=2, cp_world_size=2), + expected_blocks=( + AttentionCPBlockMetadata(0, 0, 2, 0, 0), + AttentionCPBlockMetadata(1, 2, 4, 1, 0), + ), + expected_kv_token_range=(0, 4), + query_token_ranges=((0, 1), (1, 2)), + status="implemented", + ) + communication = CUDAAGRSAttentionCPCommunication() + + with pytest.raises( + AttentionCPCommunicationUnavailable, + match="requires CUDA|requires initialized|unavailable|extension|DeterministicCollective", + ): + communication.all_gather_partial_states((_partial_state(0),), plan) + + merged = AttentionCPMergedState( + out=torch.zeros(1, 2, 1, 4), + lse=torch.zeros(1, 2, 1, dtype=torch.float32), + ) + with pytest.raises( + AttentionCPCommunicationUnavailable, + match="requires CUDA|requires initialized|unavailable|extension|DeterministicCollective", + ): + communication.reduce_scatter_merged_state(merged, plan) + + +def test_cuda_ag_rs_attention_cp_comm_executes_injected_collective(monkeypatch): + plan = AttentionCPCommunicationPlan( + parallel=AttentionParallelSpec(tp_world_size=2, cp_world_size=2), + backend="cuda_ag_rs", + status="implemented", + expected_blocks=( + AttentionCPBlockMetadata(0, 0, 2, 0, 0), + AttentionCPBlockMetadata(1, 2, 4, 1, 0), + ), + expected_kv_token_range=(0, 4), + query_token_ranges=((0, 1), (1, 2)), + ) + communication = CUDAAGRSAttentionCPCommunication(collective=_FakeDeterministicCollective()) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(communication, "_dist", lambda: _FakeCollectiveDist()) + + local_q = torch.zeros(1, 2, 1, 4) + gathered_q = communication.all_gather_query(local_q, plan) + assert gathered_q.shape == (1, 2, 2, 4) + assert torch.equal(gathered_q[:, :, :1, :], local_q) + assert torch.equal(gathered_q[:, :, 1:, :], local_q + 10) + + local_state = AttentionCPPartialState( + out=torch.zeros(1, 2, 2, 4), + lse=torch.zeros(1, 2, 2, dtype=torch.float32), + block=plan.expected_blocks[0], + ) + gathered = communication.all_gather_partial_states((local_state,), plan) + assert [state.block.global_block_index for state in gathered] == [0, 1] + assert torch.equal(gathered[1].out, local_state.out + 10) + + merged = AttentionCPMergedState( + out=torch.arange(16, dtype=torch.float32).reshape(1, 2, 2, 4), + lse=torch.arange(4, dtype=torch.float32).reshape(1, 2, 2), + ) + shard = communication.reduce_scatter_merged_state(merged, plan) + assert torch.equal(shard.out, merged.out[:, :, :1, :]) + assert torch.equal(shard.lse, merged.lse[:, :, :1]) + + +def test_cuda_ag_rs_sequence_collectives_preserve_attention_gradients(monkeypatch): + plan = AttentionCPCommunicationPlan( + parallel=AttentionParallelSpec(tp_world_size=2, cp_world_size=2), + backend="cuda_ag_rs", + status="implemented", + expected_blocks=( + AttentionCPBlockMetadata(0, 0, 2, 0, 0), + AttentionCPBlockMetadata(1, 2, 4, 1, 0), + ), + expected_kv_token_range=(0, 4), + query_token_ranges=((0, 1), (1, 2)), + ) + collective = _FakeAutogradCollective() + communication = CUDAAGRSAttentionCPCommunication(collective=collective) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + + local_q = torch.zeros(1, 2, 1, 4, dtype=torch.bfloat16, requires_grad=True) + gathered_q = communication.all_gather_query(local_q, plan) + gathered_q.float().sum().backward() + assert torch.equal(local_q.grad, torch.full_like(local_q, 2)) + assert collective.reduce_scatter_calls == 1 + + full_out = torch.zeros(1, 2, 2, 4, dtype=torch.bfloat16, requires_grad=True) + full_lse = torch.zeros(1, 2, 2, dtype=torch.float32) + shard = communication.reduce_scatter_strict_result(full_out, full_lse, plan) + shard.out.float().sum().backward() + assert torch.equal(full_out.grad, torch.ones_like(full_out)) + assert collective.all_gather_calls == 2 + + +def test_cp_manifest_rejects_gap_wrong_owner_and_incomplete_gather(): + with pytest.raises(ValueError, match="gap-free"): + AttentionCPCommunicationPlan( + parallel=AttentionParallelSpec(tp_world_size=2, cp_world_size=2), + backend="p2p_nccl_reference", + status="implemented", + expected_blocks=( + AttentionCPBlockMetadata(0, 0, 2, 0, 0), + AttentionCPBlockMetadata(1, 3, 4, 1, 0), + ), + expected_kv_token_range=(0, 4), + query_token_ranges=((0, 1), (1, 2)), + ).validate() + + with pytest.raises(ValueError, match="owner_tp_rank"): + AttentionCPCommunicationPlan( + parallel=AttentionParallelSpec(tp_world_size=2, cp_world_size=2), + backend="p2p_nccl_reference", + status="implemented", + expected_blocks=( + AttentionCPBlockMetadata(0, 0, 2, 0, 1), + AttentionCPBlockMetadata(1, 2, 4, 1, 1), + ), + expected_kv_token_range=(0, 4), + query_token_ranges=((0, 1), (1, 2)), + ).validate() + + with pytest.raises(ValueError, match="expected KV token range|complete block manifest"): + sort_attention_cp_partial_states((_partial_state(0),), plan=_p2p_plan()) + + +def test_cp_manifest_rejects_wrong_local_cp_owner(): + communication = P2PNCCLAttentionCPCommunication( + dist_module=_FakeNCCLDistributed(rank=0), + validate_cuda_tensors=False, + ) + wrong_owner = AttentionCPPartialState( + out=torch.zeros(1, 2, 2, 4), + lse=torch.zeros(1, 2, 2, dtype=torch.float32), + block=AttentionCPBlockMetadata(0, 0, 2, 1, 0), + ) + + with pytest.raises(ValueError, match="wrong CP owner"): + communication.all_gather_partial_states((wrong_owner,), _p2p_plan()) + + +def test_p2p_nccl_reference_query_ag_preserves_cp_rank_order(): + local_q = torch.zeros(1, 2, 1, 4) + remote_q = torch.ones_like(local_q) + communication = P2PNCCLAttentionCPCommunication( + dist_module=_FakeNCCLDistributed(rank=0, receive_payloads=[remote_q]), + validate_cuda_tensors=False, + ) + + gathered = communication.all_gather_query(local_q, _p2p_plan()) + + assert gathered.shape == (1, 2, 2, 4) + assert torch.equal(gathered[:, :, :1, :], local_q) + assert torch.equal(gathered[:, :, 1:, :], remote_q) + + +def test_p2p_nccl_reference_gathers_kv_and_position_ids_in_owner_order(): + local_k = torch.zeros(1, 2, 2, 4) + local_v = torch.ones_like(local_k) + remote_k = torch.full_like(local_k, 7) + remote_v = torch.full_like(local_v, 9) + communication = P2PNCCLAttentionCPCommunication( + dist_module=_FakeNCCLDistributed( + rank=0, + receive_payloads=(remote_k, remote_v), + ), + validate_cuda_tensors=False, + ) + + global_k, global_v = communication.all_gather_kv(local_k, local_v, _p2p_plan()) + + assert torch.equal(global_k[:, :, :2], local_k) + assert torch.equal(global_k[:, :, 2:], remote_k) + assert torch.equal(global_v[:, :, :2], local_v) + assert torch.equal(global_v[:, :, 2:], remote_v) + + local_q_pos = torch.tensor([[2]], dtype=torch.long) + local_k_pos = torch.tensor([[0, 1]], dtype=torch.long) + remote_q_pos = torch.tensor([[3]], dtype=torch.long) + remote_k_pos = torch.tensor([[2, 3]], dtype=torch.long) + communication = P2PNCCLAttentionCPCommunication( + dist_module=_FakeNCCLDistributed( + rank=0, + receive_payloads=(remote_q_pos, remote_k_pos), + ), + validate_cuda_tensors=False, + ) + + global_q_pos, global_k_pos = communication.all_gather_position_ids( + local_q_pos, local_k_pos, _p2p_plan() + ) + + assert global_q_pos.tolist() == [[2, 3]] + assert global_k_pos.tolist() == [[0, 1, 2, 3]] + + +def test_cp_manifest_allows_sparse_global_block_indices_and_rejects_short_query_state(): + plan = AttentionCPCommunicationPlan( + parallel=AttentionParallelSpec(tp_world_size=2, cp_world_size=2), + backend="p2p_nccl_reference", + status="implemented", + expected_blocks=( + AttentionCPBlockMetadata(10, 0, 2, 0, 0), + AttentionCPBlockMetadata(20, 2, 4, 1, 0), + ), + expected_kv_token_range=(0, 4), + query_token_ranges=((0, 1), (1, 2)), + ) + plan.validate() + communication = P2PNCCLAttentionCPCommunication( + dist_module=_FakeNCCLDistributed(rank=0), + validate_cuda_tensors=False, + ) + short_query = AttentionCPPartialState( + out=torch.zeros(1, 2, 1, 4), + lse=torch.zeros(1, 2, 1, dtype=torch.float32), + block=plan.expected_blocks[0], + ) + + with pytest.raises(ValueError, match="complete query range"): + communication.all_gather_partial_states((short_query,), plan) + + +def test_p2p_nccl_reference_gathers_manifest_order_and_scatters_query_range(): + remote_out = torch.full((1, 2, 2, 4), 7.0) + remote_lse = torch.full((1, 2, 2), 3.0, dtype=torch.float32) + distributed = _FakeNCCLDistributed( + rank=0, + receive_payloads=(remote_out, remote_lse), + ) + communication = P2PNCCLAttentionCPCommunication( + dist_module=distributed, + validate_cuda_tensors=False, + ) + local = AttentionCPPartialState( + out=torch.zeros(1, 2, 2, 4), + lse=torch.zeros(1, 2, 2, dtype=torch.float32), + block=AttentionCPBlockMetadata(0, 0, 2, 0, 0), + ) + + gathered = communication.all_gather_partial_states((local,), _p2p_plan()) + + assert [state.block.global_block_index for state in gathered] == [0, 1] + torch.testing.assert_close(gathered[1].out, remote_out) + torch.testing.assert_close(gathered[1].lse, remote_lse) + + merged = AttentionCPMergedState( + out=torch.arange(16, dtype=torch.float32).reshape(1, 2, 2, 4), + lse=torch.arange(4, dtype=torch.float32).reshape(1, 2, 2), + ) + shard = communication.reduce_scatter_merged_state(merged, _p2p_plan()) + torch.testing.assert_close(shard.out, merged.out[:, :, 0:1, :]) + torch.testing.assert_close(shard.lse, merged.lse[:, :, 0:1]) + + +def test_p2p_nccl_reference_fails_closed_on_non_nccl_backend(): + class _FakeGlooDistributed(_FakeNCCLDistributed): + @staticmethod + def get_backend(group=None): + return "gloo" + + communication = P2PNCCLAttentionCPCommunication( + dist_module=_FakeGlooDistributed(rank=0), + validate_cuda_tensors=False, + ) + + with pytest.raises(AttentionCPCommunicationUnavailable, match="NCCL backend"): + communication.all_gather_partial_states((_partial_state(0),), _p2p_plan()) + + +def test_p2p_query_ranges_allow_empty_decode_shards(): + plan = AttentionCPCommunicationPlan( + parallel=AttentionParallelSpec(tp_world_size=2, cp_world_size=2), + backend="p2p_nccl_reference", + status="implemented", + expected_blocks=( + AttentionCPBlockMetadata(0, 0, 2, 0, 0), + AttentionCPBlockMetadata(1, 2, 4, 1, 0), + ), + expected_kv_token_range=(0, 4), + query_token_ranges=((0, 1), (1, 1)), + ) + plan.validate() + communication = P2PNCCLAttentionCPCommunication( + dist_module=_FakeNCCLDistributed(rank=0), + validate_cuda_tensors=False, + ) + merged = AttentionCPMergedState( + out=torch.zeros(1, 2, 1, 4), + lse=torch.zeros(1, 2, 1, dtype=torch.float32), + ) + + local = communication.reduce_scatter_merged_state(merged, plan) + + assert local.out.shape == (1, 2, 1, 4) + assert local.lse.shape == (1, 2, 1) + + +@pytest.mark.skipif( + not torch.cuda.is_available() + or not torch.distributed.is_available() + or not torch.distributed.is_initialized() + or "nccl" not in str(torch.distributed.get_backend()).lower() + or torch.distributed.get_world_size() != 2, + reason="requires an initialized two-rank NCCL process group", +) +def test_p2p_nccl_reference_real_process_group_smoke(): + rank = torch.distributed.get_rank() + local = AttentionCPPartialState( + out=torch.full((1, 2, 2, 4), float(rank), device="cuda"), + lse=torch.full((1, 2, 2), float(rank), dtype=torch.float32, device="cuda"), + block=AttentionCPBlockMetadata(rank, rank * 2, rank * 2 + 2, rank, 0), + ) + + gathered = P2PNCCLAttentionCPCommunication().all_gather_partial_states( + (local,), + _p2p_plan(cp_rank=rank), + ) + + assert [state.block.global_block_index for state in gathered] == [0, 1] + + +def test_flashinfer_pr7_real_backend_requires_cuda_before_importing_flashinfer(): + q, k, v = _qkv(query_len=1) + metadata = _metadata(query_len=1) + op = FlashInferQwen3PagedAttentionOp() + + with pytest.raises(FlashInferUnavailable, match="requires CUDA"): + op( + q, + k, + v, + metadata, + config=FlashInferPagedAttentionConfig(mode="decode", workspace_size_bytes=1024), + ) + + +def test_flashinfer_pr7_plan_and_cache_materialization_follow_logical_page_order(): + q, k, v = _qkv(batch=1, query_len=1) + positions = torch.full((1, 6), -1, dtype=torch.long) + positions[:, 4:6] = torch.tensor([0, 1], dtype=torch.long) + positions[:, 0:2] = torch.tensor([2, 3], dtype=torch.long) + positions[:, 2:4] = torch.tensor([4, 5], dtype=torch.long) + metadata = DecodeKVCacheMetadata( + cache_position=torch.tensor([[5]], dtype=torch.long), + kv_seq_lens=torch.tensor([6], dtype=torch.long), + block_table=torch.tensor([[2, 0, 1]], dtype=torch.long), + global_token_positions=positions, + query_position_ids=torch.tensor([[5]], dtype=torch.long), + key_position_ids=positions.clone(), + page_size=2, + q_rope_state="pre_rope", + k_cache_rope_state="pre_rope", + ) + + plan = build_flashinfer_paged_kv_plan( + metadata, + batch_size=1, + query_len=1, + cache_capacity=k.size(2), + device=q.device, + ) + k_pages, v_pages = materialize_flashinfer_paged_kv_cache(k, v, page_size=2) + + assert plan.paged_kv_indices.tolist() == [2, 0, 1] + torch.testing.assert_close(k_pages[2], k[0, :, 4:6, :].transpose(0, 1)) + torch.testing.assert_close(v_pages[0], v[0, :, 0:2, :].transpose(0, 1)) + + +def test_flashinfer_pr7_plan_rejects_position_metadata_mismatch(): + q, k, _ = _qkv(batch=1, query_len=1) + metadata = DecodeKVCacheMetadata( + cache_position=torch.tensor([[5]], dtype=torch.long), + kv_seq_lens=torch.tensor([6], dtype=torch.long), + block_table=torch.tensor([[2, 0, 1]], dtype=torch.long), + global_token_positions=torch.arange(6, dtype=torch.long).unsqueeze(0), + query_position_ids=torch.tensor([[5]], dtype=torch.long), + key_position_ids=torch.arange(6, dtype=torch.long).unsqueeze(0), + page_size=2, + q_rope_state="pre_rope", + k_cache_rope_state="pre_rope", + ) + + with pytest.raises(ValueError, match="reconstruct logical positions"): + build_flashinfer_paged_kv_plan( + metadata, + batch_size=1, + query_len=1, + cache_capacity=k.size(2), + device=q.device, + ) + + +def test_pr7_check_dry_run_writes_non_eligible_report(tmp_path): + output = tmp_path / "pr7-dry-run.json" + + assert check_script.main(["--dry-run", "--output", str(output)]) == 0 + + report = json.loads(output.read_text(encoding="utf-8")) + assert report["status"] == "dry_run" + assert report["passed"] is False + assert report["acceptance_eligible"] is False + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA for the real entrypoint") +def test_pr7_check_writes_not_available_report_for_missing_flashinfer(monkeypatch, tmp_path): + class MissingFlashInfer: + def __call__(self, *args, **kwargs): + raise FlashInferUnavailable("No module named 'flashinfer'") + + monkeypatch.setattr(check_script, "FlashInferQwen3PagedAttentionOp", MissingFlashInfer) + output = tmp_path / "pr7-not-available.json" + + assert ( + check_script.main(["--no-dry-run", "--device", "cuda", "--json", "--output", str(output)]) + == 1 + ) + + report = json.loads(output.read_text(encoding="utf-8")) + assert report["status"] == "not_available" + assert report["passed"] is False + assert report["acceptance_eligible"] is False + assert "Attention backend unavailable" in report["errors"][0] + + +def test_pr7_check_acceptance_errors_require_all_drift_and_invariance_fields(): + args = check_script._parse_args([]) + report = { + "device": "cuda:0", + "shape": {"q_heads": 16, "kv_heads": 4, "head_dim": 128}, + "candidate_provenance": { + "attention_mode": "decode", + "fallback": False, + "pos_encoding_mode": "ROPE_LLAMA", + "rope_theta": 1_000_000.0, + "rotary_dim": 128, + "arithmetic_semantics_verified": True, + "actual_split_kv_plans": [{"actual_split_boundaries": [[0, 4]]}], + "actual_split_kv_plan_set": { + "coverage": "complete_batch_tp_cp_owner_cartesian_product" + }, + }, + "drift": { + "out": {"max_abs": 0.0}, + "lse": {"max_abs": 0.0}, + "dlogp": {"max_abs": 0.0}, + }, + "batch_invariant_sweep": {"passed": False}, + "page_layout_invariant_sweep": {"passed": True}, + } + + assert check_script._acceptance_errors(report, args) == ["batch_invariant_sweep failed"] + + +def test_pr7_check_accepts_strict_fa4_production_core(): + args = check_script._parse_args(["--strict", "--device", "cuda"]) + report = { + "device": "cuda:0", + "shape": {"q_heads": 16, "kv_heads": 4, "head_dim": 128}, + "candidate_provenance": { + "attention_mode": "decode", + "fallback": False, + "strict_mode": True, + "strict_core_id": STRICT_ATTENTION_PRODUCTION_CORE_ID, + "strict_schedule": STRICT_ATTENTION_FA4_SCHEDULE_ID, + "actual_backend": "flash_attention_4.cute", + "native_attention_arithmetic": True, + "num_splits": 1, + "deterministic_backward": True, + "fa_api_source": "flash_attn.cute.interface", + "reference_only": False, + "strict_core_row_plans": [{"actual_split_kv_policy": "disabled"}], + "rope_backend": "rlkernel.cuda.rope_sm90", + "rope_theta": 1_000_000.0, + "rotary_dim": 128, + "arithmetic_semantics_verified": True, + }, + "drift": { + "out": {"max_abs": 0.0}, + "lse": {"max_abs": 0.0}, + "dlogp": {"max_abs": 0.0}, + }, + "batch_invariant_sweep": {"passed": True}, + "page_layout_invariant_sweep": {"passed": True}, + } + + assert check_script._acceptance_errors(report, args) == [] + + +def test_pr7_check_rejects_reference_core_as_production(): + args = check_script._parse_args(["--strict", "--device", "cuda"]) + report = { + "device": "cuda:0", + "shape": {"q_heads": 16, "kv_heads": 4, "head_dim": 128}, + "candidate_provenance": { + "attention_mode": "decode", + "fallback": False, + "strict_mode": True, + "strict_core_id": STRICT_ATTENTION_CORE_ID, + "strict_schedule": STRICT_ATTENTION_SCHEDULE_ID, + "actual_backend": "rlkernel.cuda.deterministic_attention", + "native_attention_arithmetic": False, + "reference_only": True, + "strict_core_row_plans": [{"actual_split_kv_policy": "disabled"}], + "rope_backend": "rlkernel.cuda.rope_sm90", + "rope_theta": 1_000_000.0, + "rotary_dim": 128, + "arithmetic_semantics_verified": True, + }, + "drift": { + "out": {"max_abs": 0.0}, + "lse": {"max_abs": 0.0}, + "dlogp": {"max_abs": 0.0}, + }, + "batch_invariant_sweep": {"passed": True}, + "page_layout_invariant_sweep": {"passed": True}, + } + + errors = check_script._acceptance_errors(report, args) + assert "strict runtime did not execute the FA4 production core" in errors + assert "strict runtime selected the reference core" in errors + + +def test_pr7_check_rejects_nonfinite_drift_and_wrong_tp_local_shape(): + args = check_script._parse_args([]) + report = { + "device": "cuda:0", + "shape": {"q_heads": 32, "kv_heads": 8, "head_dim": 128}, + "candidate_provenance": { + "attention_mode": "decode", + "fallback": False, + "pos_encoding_mode": "ROPE_LLAMA", + "rope_theta": 1_000_000.0, + "rotary_dim": 128, + "arithmetic_semantics_verified": True, + "actual_split_kv_plans": [{"actual_split_boundaries": [[0, 4]]}], + "actual_split_kv_plan_set": { + "coverage": "complete_batch_tp_cp_owner_cartesian_product" + }, + }, + "drift": { + "out": {"max_abs": float("nan")}, + "lse": {"max_abs": 0.0}, + "dlogp": {"max_abs": 0.0}, + }, + "batch_invariant_sweep": {"passed": True}, + "page_layout_invariant_sweep": {"passed": True}, + } + + errors = check_script._acceptance_errors(report, args) + assert any("finite and non-negative" in error for error in errors) + assert any("TP-local head shard" in error for error in errors) + + +def test_pr7_check_rejects_nonlocal_qwen3_head_arguments(): + with pytest.raises(SystemExit): + check_script._parse_args(["--q-heads", "32", "--kv-heads", "8"]) + + +def test_p2p_entrypoint_validates_qwen3_tp_local_shape_before_cuda_math(): + args = p2p_check_script.parse_args(["--q-heads", "32", "--kv-heads", "8"]) + + with pytest.raises(ValueError, match="TP=2 Qwen3-8B local heads"): + p2p_check_script.run_check( + args, + global_rank=0, + tp_rank=0, + cp_rank=0, + replica_index=0, + cp_group=None, + device=torch.device("cpu"), + ) + + +def test_strict_shared_core_entrypoint_requires_self_owned_ag_rs(): + with pytest.raises(SystemExit): + p2p_check_script.parse_args(["--strict-shared-core"]) + + args = p2p_check_script.parse_args(["--transport", "cuda_ag_rs", "--strict-shared-core"]) + assert args.strict_shared_core is True + + +@pytest.mark.parametrize( + ("argv", "message"), + [ + (["--batch", "0"], "batch must be positive"), + (["--atol", "inf"], "atol must be finite and non-negative"), + ( + ["--final-write-atol", "nan"], + "final_write_atol must be finite and non-negative", + ), + (["--repeats", "1"], "repeats must be at least 2"), + ], +) +def test_p2p_entrypoint_rejects_non_acceptance_arguments(argv, message): + args = p2p_check_script.parse_args(argv) + + with pytest.raises(ValueError, match=message): + p2p_check_script.run_check( + args, + global_rank=0, + tp_rank=0, + cp_rank=0, + replica_index=0, + cp_group=None, + device=torch.device("cpu"), + ) + + +def test_fa4_core_rejects_api_without_reduction_controls(): + def legacy_flash_attn(q, k, v, *, causal=False): + return q + + with pytest.raises(StrictFlashAttentionUnavailable, match="strict controls"): + StrictFlashAttention4Core(_op=legacy_flash_attn) + + +def test_fa4_core_fixes_reduction_controls_and_exports_fp32_lse(monkeypatch): + calls = [] + + def fake_fa4( + q, + k, + v, + *, + softmax_scale=None, + causal=False, + num_splits=0, + pack_gqa=None, + deterministic=False, + return_lse=False, + ): + calls.append( + { + "q_shape": tuple(q.shape), + "k_shape": tuple(k.shape), + "softmax_scale": softmax_scale, + "causal": causal, + "num_splits": num_splits, + "pack_gqa": pack_gqa, + "deterministic": deterministic, + "return_lse": return_lse, + } + ) + return q.clone(), torch.zeros( + q.size(0), q.size(2), q.size(1), dtype=torch.float32, device=q.device + ) + + core = StrictFlashAttention4Core(_op=fake_fa4, _package_version="4.test") + monkeypatch.setattr(core, "_validate_inputs", lambda *_args: None) + q = torch.randn(1, 4, 2, 8, dtype=torch.bfloat16) + k = torch.randn(1, 2, 3, 8, dtype=torch.bfloat16) + v = torch.randn(1, 2, 3, 8, dtype=torch.bfloat16) + result = core.forward_with_lse( + q, + k, + v, + causal=True, + scale=0.125, + query_position_ids=torch.tensor([[1, 2]]), + key_position_ids=torch.tensor([[0, 1, 2]]), + ) + + assert calls == [ + { + "q_shape": (1, 2, 4, 8), + "k_shape": (1, 3, 2, 8), + "softmax_scale": 0.125, + "causal": True, + "num_splits": 1, + "pack_gqa": True, + "deterministic": True, + "return_lse": True, + } + ] + assert result.out.shape == q.shape + assert result.lse.shape == q.shape[:3] + assert result.lse.dtype is torch.float32 + assert result.provenance["strict_core_id"] == STRICT_ATTENTION_PRODUCTION_CORE_ID + assert result.provenance["strict_schedule"] == STRICT_ATTENTION_FA4_SCHEDULE_ID + assert result.provenance["num_splits"] == 1 + assert result.provenance["dropout_p"] == 0.0 + assert result.provenance["native_attention_arithmetic"] is True + assert result.provenance["production_ready"] is True + + +def test_strict_paged_default_selects_fa4_production_core(monkeypatch): + core = _RecordingStrictCore() + core.core_id = STRICT_ATTENTION_PRODUCTION_CORE_ID + core.strict_schedule = STRICT_ATTENTION_FA4_SCHEDULE_ID + core.backend_id = "flash_attention_4.cute" + core.native_attention_arithmetic = True + core.num_splits = 1 + core.deterministic_backward = True + core.production_ready = True + core.reference_only = False + + original_forward = core.forward_with_lse + + def production_forward(*args, **kwargs): + result = original_forward(*args, **kwargs) + result.provenance.update( + { + "native_attention_arithmetic": True, + "production_ready": True, + "reference_only": False, + "num_splits": 1, + "deterministic_backward": True, + "fa_api_source": "flash_attn.cute.interface", + "fa_package_version": "4.test", + } + ) + return result + + core.forward_with_lse = production_forward + monkeypatch.setattr( + paged_attention_module, + "StrictFlashAttention4Core", + lambda *, split_kv: core, + ) + q, k, v = (tensor.to(torch.bfloat16) for tensor in _qkv(query_len=1)) + result = FlashInferQwen3PagedAttentionOp(flashinfer_module=_fake_flashinfer())( + q, + k, + v, + _metadata(query_len=1), + config=FlashInferPagedAttentionConfig( + mode="decode", + workspace_size_bytes=1024, + strict_mode=True, + strict_rope_op=_IdentityStrictRoPE(), + ), + ) + + assert len(core.calls) == q.size(0) + assert result.provenance["strict_core_id"] == STRICT_ATTENTION_PRODUCTION_CORE_ID + assert result.provenance["actual_backend"] == "flash_attention_4.cute" + assert result.provenance["num_splits"] == 1 + assert result.provenance["native_attention_arithmetic"] is True + assert result.provenance["reference_only"] is False