[WS2][PR7][Attention] Add FlashInfer RoPE-fused paged attention scaffold

[inaniloquentee]

Summary

Implements the #235 PR7 scaffold for fused attention backend alignment, with the distributed TP=2, CP=2 attention boundary made explicit.

This is a draft PR stacked on #260 / PR6. The branch is based on the latest fetched #260 PR head (2eb4362) and currently targets feat/ws2-attention-single-gpu-harness-pr2 because #260's head is available as a PR ref rather than an origin branch. After #260 lands, this PR should be retargeted so the diff only contains PR7.

What This Adds

• Adds an opt-in FlashInfer paged-attention candidate for Qwen3 rollout prefill/decode:
  • FlashInferQwen3PagedAttentionOp
  • FlashInferPagedAttentionConfig
  • FlashInferRoPEFusionConfig
  • FlashInferSplitKVPolicy
• Locks the RoPE-fused path to Qwen3 semantics:
  • pos_encoding_mode="ROPE_LLAMA"
  • rope_theta=1_000_000.0
  • rope_scale=1.0
  • rotary_dim=head_dim
  • pre-RoPE Q and pre-RoPE K cache only
• Surfaces split-KV as explicit provenance:
  • disabled
  • fixed:<N>
  • auto, rejected when batch invariance is required
• Adds page-table and position-identity guards before FlashInfer execution:
  • physical page bounds
  • logical block_table/global_token_positions reconstruction
  • key_position_ids identity
• Adds a dry-run and real-CUDA validation script:
  • plan/provenance mode without CUDA or FlashInfer
  • real out/lse drift path when CUDA + FlashInfer are available
  • batch-invariant sweep for real runs

CP/TP Communication Interface

Issue #235 targets Qwen3-8B, TP=2, CP=2, BF16, so PR7 cannot look like a single-rank-only adapter. This PR exposes the communication-operator interface now, but the real self-owned CUDA communication kernels are not implemented in this scaffold.

This is not a p2p NCCL interface. The intended communication layer is custom CUDA AG/RS, with compute and communication decoupled.

Added interface surface:

rl_engine/kernels/ops/cuda/attention/cp_comm.py
  AttentionParallelSpec(tp_world_size=2, cp_world_size=2)
  AttentionCPBlockMetadata(global_block_index, kv_block_start, kv_block_end, owner_cp_rank, owner_tp_rank)
  AttentionCPPartialState(out, lse, block)
  AttentionCPMergedState(out, lse)
  AttentionCPCommunicationPlan(backend="cuda_ag_rs", status="interface_only")
  CUDAAGRSAttentionCPCommunication.all_gather_partial_states(...)
  CUDAAGRSAttentionCPCommunication.reduce_scatter_merged_state(...)
  sort_attention_cp_partial_states(...)

The intended future flow is:

local attention over rank-owned KV blocks
  -> AttentionCPPartialState(
       out: [B, Hq, Sq, D],
       lse: [B, Hq, Sq] fp32,
       global_block_index,
       kv_block_start / kv_block_end,
       owner_cp_rank / owner_tp_rank,
     )
  -> custom CUDA AG communication operator
  -> sort by global_block_index
  -> PR3 FP32 online-softmax merge
  -> custom CUDA RS communication operator

Important details:

• FlashInferPagedAttentionConfig defaults to tp_world_size=2, cp_world_size=2, backend="cuda_ag_rs", status="interface_only".
• require_cp_comm=True fails before execution. This prevents this scaffold from pretending real communication exists.
• CUDAAGRSAttentionCPCommunication.all_gather_partial_states(...) validates local partial states and then raises AttentionCPCommunicationUnavailable.
• CUDAAGRSAttentionCPCommunication.reduce_scatter_merged_state(...) validates the merged state and then raises AttentionCPCommunicationUnavailable.
• The merge order is part of the contract: gathered partial states must be sorted by AttentionCPBlockMetadata.global_block_index.
• Duplicate global_block_index is rejected.
• Merge accumulation is required to be FP32 and the AG-carried state must include attention-domain lse.
• FlashInfer split-KV is backend-local reduction inside one rank. It is not the CP cross-rank merge and does not define CP ordering.

Key Files

• rl_engine/kernels/ops/cuda/attention/cp_comm.py
• rl_engine/kernels/ops/cuda/attention/flashinfer_paged_attention.py
• rl_engine/kernels/ops/cuda/attention/__init__.py
• scripts/ws2_pr7_flashinfer_attention_check.py
• tests/test_flashinfer_pr7_attention.py
• docs/design/ws2-attention-pr7-flashinfer-rope-splitk.md

Non-Claims

This PR does not implement custom CUDA AG/RS communication kernels, does not enable FlashInfer by default, does not implement the TE training lane, does not replace PR3 CP merge semantics, does not prove real H-card batch invariance locally, and does not implement training backward. Real CUDA/H-card validation is still required before promotion beyond candidate status.

Validation

Pre-submit checks run locally from a clean worktree after the communication-interface correction:

pre-commit run --all-files
  passed

python -m ruff check 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
  passed

python -m mypy --ignore-missing-imports rl_engine/
  passed; no issues found in 92 source files

python -m pytest tests/test_flashinfer_pr7_attention.py -q
  12 passed

python -m pytest tests/test_flashinfer_pr7_attention.py tests/test_attention_comparison.py -q
  33 passed

python -m pytest rl_engine/tests/test_dispatch.py -v
  6 passed

PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest tests/test_attention_correctness.py -q -rs
  127 skipped locally because CUDA/ROCm is unavailable

python -m pytest tests/test_attention.py -v -k "not large and not gpu"
  24 passed, 2 deselected

python -m pytest tests/test_kv_cache_attention.py -v -k "not large and not gpu"
  15 passed

python scripts/ws2_pr7_flashinfer_attention_check.py --dry-run --json
  passed; reports TP=2, CP=2, backend=cuda_ag_rs, pattern=ag_rs, compute_communication=decoupled, status=interface_only

python scripts/ws2_pr7_flashinfer_attention_check.py --dry-run --mode prefill --query-len 16 --split-kv-policy fixed --fixed-split-size 4 --json
  passed; reports fixed split-KV and CP/TP custom CUDA AG/RS communication interface provenance

python -m mkdocs build --strict -f mkdocs.yaml
  passed

git diff --check
  passed

DCO

• Commit includes Signed-off-by: inaniloquentee <3051000145@qq.com>.

Part of #235
Depends on #260

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 3a3a37c4-eb95-47e1-87a4-cfbd0bd9d732

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}