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

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

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

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

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

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