feat(attention): add decode-stage KV-cache CP replay - #260
Conversation
Signed-off-by: JLiu4Coding <lzwgre@126.com>
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughChangesThe testing package adds decode-attention metadata, paged-KV replay, logical full-prefill reference execution, drift comparison, cache validation, RoPE handling, context-parallel ownership tracking, and comprehensive replay tests. The design documentation defines the corresponding contracts and validation rules. Decode attention contracts and validation
Reference and paged replay execution
Replay fixtures and behavioral coverage
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant compare_decode_kv_replay
participant run_decode_full_prefill_reference
participant run_decode_kv_replay
participant MergeBackend
compare_decode_kv_replay->>run_decode_full_prefill_reference: reconstruct and execute logical KV reference
compare_decode_kv_replay->>run_decode_kv_replay: execute visible paged KV blocks
run_decode_kv_replay->>MergeBackend: merge FP32 partial attention states
compare_decode_kv_replay-->>compare_decode_kv_replay: compare output, LSE, and selected log-prob drift
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
|
@coderabbitai review |
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (3)
docs/design/ws2-attention-decode-replay.md (1)
23-25: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the RoPE restrictions enforced by validation.
_validate_decode_inputsrejectsrope_cast_at != "after_rope"and anyrope_rotary_dimthat differs fromhead_dim. Partial rotary is therefore unsupported today. Add these two constraints to the validation list so callers learn the limit from the design document.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/design/ws2-attention-decode-replay.md` around lines 23 - 25, Update the metadata validation list in the design document to state that rope_cast_at must be "after_rope" and rope_rotary_dim must equal head_dim. Explicitly indicate that partial rotary is unsupported, alongside the existing validation constraints.rl_engine/testing/attention_comparison.py (2)
690-709: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueVectorize the slot reconstruction.
The nested loop calls
.item()once per cached token. On CUDA each call is a host sync. The mapping is a pure function ofblock_table,page_size, andkv_seq_lens, so it can be built with tensor ops. Validation already guarantees the logical positions equalarange(sequence_length).♻️ Proposed vectorized slot index
- physical_slots: list[int] = [] - logical_positions: list[int] = [] - logical_block_count = math.ceil(sequence_length / metadata.page_size) - for logical_block in range(logical_block_count): - physical_page = int(metadata.block_table[batch_index, logical_block].item()) - tokens_in_block = min( - metadata.page_size, - sequence_length - logical_block * metadata.page_size, - ) - for page_offset in range(tokens_in_block): - slot = physical_page * metadata.page_size + page_offset - physical_slots.append(slot) - logical_positions.append(int(metadata.global_token_positions[batch_index, slot].item())) - - slot_index = torch.tensor(physical_slots, device=inputs.k_cache.device, dtype=torch.long) - logical_position_tensor = torch.tensor( - logical_positions, - device=inputs.k_cache.device, - dtype=torch.long, - ) + logical_block_count = math.ceil(sequence_length / metadata.page_size) + device = inputs.k_cache.device + logical_index = torch.arange(sequence_length, device=device, dtype=torch.long) + pages = metadata.block_table[batch_index, :logical_block_count].to(device=device).long() + slot_index = ( + pages[logical_index // metadata.page_size] * metadata.page_size + + logical_index % metadata.page_size + ) + logical_position_tensor = metadata.global_token_positions[batch_index, slot_index].long()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rl_engine/testing/attention_comparison.py` around lines 690 - 709, Replace the nested Python loops in the slot reconstruction with device-side tensor operations derived from block_table, page_size, and sequence_length/kv_seq_lens, eliminating per-token .item() calls and CUDA synchronization. Construct slot_index in logical token order and derive logical_position_tensor as the corresponding arange(sequence_length) while preserving the existing device and dtype behavior.
264-266: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueValidation runs three times per comparison.
compare_decode_kv_replayvalidates, thenrun_decode_full_prefill_referenceandrun_decode_kv_replayvalidate again._validate_decode_inputsuses per-batch Python loops and scalar.item()reads, so the cost is paid three times. Consider an internal_validatedflag or a private unvalidated entry point if the harness is used on larger batches.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rl_engine/testing/attention_comparison.py` around lines 264 - 266, Eliminate the repeated input validation in compare_decode_kv_replay by adding an internal validated path or _validated flag to run_decode_full_prefill_reference and run_decode_kv_replay. Validate inputs once in compare_decode_kv_replay, then have both execution functions skip _validate_decode_inputs while preserving normal validation for standalone callers.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rl_engine/testing/attention_comparison.py`:
- Line 339: Update the merge-order collection around merge_orders so it is
nested per batch, matching the batch structure of cp_block_owners. Initialize
one inner list per batch, append each query’s order to that batch’s list, and
expose the resulting B × Sq structure through logical_merge_orders at all
affected locations.
- Around line 714-724: The pre-RoPE K path currently casts using the query dtype
from _decode_rope_output_dtype. Update the dtype handling around
_decode_rope_output_dtype and _validate_decode_inputs so K uses the K-cache
dtype, either by adding a K-cache dtype default or validating q.dtype ==
k_cache.dtype; ensure the pre-RoPE Q path continues using the query dtype and
the pre-RoPE K path preserves the cache dtype.
In `@tests/test_attention_comparison.py`:
- Around line 383-385: Update the output drift assertion in
compare_decode_kv_replay to use a bfloat16-scale tolerance rather than 1.0e-6,
since report.drifts[0].out compares bfloat16 results that may round differently.
Keep the existing tight 1.0e-6 tolerance for report.drifts[0].lse.
---
Nitpick comments:
In `@docs/design/ws2-attention-decode-replay.md`:
- Around line 23-25: Update the metadata validation list in the design document
to state that rope_cast_at must be "after_rope" and rope_rotary_dim must equal
head_dim. Explicitly indicate that partial rotary is unsupported, alongside the
existing validation constraints.
In `@rl_engine/testing/attention_comparison.py`:
- Around line 690-709: Replace the nested Python loops in the slot
reconstruction with device-side tensor operations derived from block_table,
page_size, and sequence_length/kv_seq_lens, eliminating per-token .item() calls
and CUDA synchronization. Construct slot_index in logical token order and derive
logical_position_tensor as the corresponding arange(sequence_length) while
preserving the existing device and dtype behavior.
- Around line 264-266: Eliminate the repeated input validation in
compare_decode_kv_replay by adding an internal validated path or _validated flag
to run_decode_full_prefill_reference and run_decode_kv_replay. Validate inputs
once in compare_decode_kv_replay, then have both execution functions skip
_validate_decode_inputs while preserving normal validation for standalone
callers.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: cf3bf3d9-a32f-420f-a386-a20ea4905588
📒 Files selected for processing (4)
docs/design/ws2-attention-decode-replay.mdrl_engine/testing/__init__.pyrl_engine/testing/attention_comparison.pytests/test_attention_comparison.py
| _validate_decode_inputs(inputs) | ||
| outs: list[torch.Tensor] = [] | ||
| lses: list[torch.Tensor] = [] | ||
| merge_orders: list[list[int]] = [] |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
logical_merge_orders loses the batch dimension.
merge_orders.append(order) runs inside the query loop, so the list is flat with B * Sq entries. cp_block_owners next to it is nested per batch. A provenance consumer cannot map a merge order back to its batch when Sq > 1, and the two fields disagree in shape. Nest the merge orders per batch.
🐛 Proposed fix for per-batch merge-order provenance
- merge_orders: list[list[int]] = []
+ merge_orders: list[list[list[int]]] = []
cp_block_owners: list[list[int]] = []
for batch_index in range(inputs.q.size(0)):
q, k, v, logical_positions = _decode_logical_qkv(inputs, batch_index)
owners = _logical_block_owners(inputs, batch_index)
cp_block_owners.append(owners)
batch_out: list[torch.Tensor] = []
batch_lse: list[torch.Tensor] = []
+ batch_orders: list[list[int]] = [] out, lse = _merge_partial_states(states, backend=merge_backend)
batch_out.append(out.to(inputs.output_dtype))
batch_lse.append(lse)
- merge_orders.append(order)
+ batch_orders.append(order)
+ merge_orders.append(batch_orders)
outs.append(torch.cat(batch_out, dim=2))
lses.append(torch.cat(batch_lse, dim=2))Also applies to: 386-386, 409-411
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rl_engine/testing/attention_comparison.py` at line 339, Update the
merge-order collection around merge_orders so it is nested per batch, matching
the batch structure of cp_block_owners. Initialize one inner list per batch,
append each query’s order to that batch’s list, and expose the resulting B × Sq
structure through logical_merge_orders at all affected locations.
| rope = NativeRoPEOp() | ||
| output_dtype = _decode_rope_output_dtype(inputs) | ||
| if metadata.q_rope_state == "pre_rope": | ||
| q = rope.forward_fp32( | ||
| q, | ||
| metadata.query_position_ids[batch_index : batch_index + 1], | ||
| theta=inputs.rope_theta, | ||
| ).to(output_dtype) | ||
| if metadata.k_cache_rope_state == "pre_rope": | ||
| key_positions = metadata.key_position_ids[batch_index, slot_index].unsqueeze(0) | ||
| k = rope.forward_fp32(k, key_positions, theta=inputs.rope_theta).to(output_dtype) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Check whether any caller builds decode inputs with mixed q / k_cache dtypes.
rg -n -C 6 'DecodeAttentionInputs\(' --type=pyRepository: RL-Align/RL-Kernel
Length of output: 156
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate and inspect the target file around the relevant lines, plus related helpers/tests.
printf 'Files with attention_comparison.py:\n'
fd -a 'attention_comparison\.py$' . || true
printf '\nTarget outline:\n'
ast-grep outline rl_engine/testing/attention_comparison.py --match "_decode_rope_output_dtype" --view expanded || true
ast-grep outline rl_engine/testing/attention_comparison.py --match "DecodeAttentionInputs" --view expanded || true
printf '\nRelevant snippets:\n'
sed -n '680,750p' rl_engine/testing/attention_comparison.py
printf '\n--- helper region ---\n'
sed -n '1120,1290p' rl_engine/testing/attention_comparison.py
printf '\nDecodeAttentionInputs references (all files):\n'
rg -n -C 8 'DecodeAttentionInputs|_decode_rope_output_dtype|rope_state' . --glob '*.py' || true
printf '\nInput dtype construction references:\n'
rg -n -C 5 'q_dtype|k_cache_dtype|dtype|q_cache|rope_theta|batch_index|slot_index' rl_engine/testing/attention_comparison.py --glob '*.py' | sed -n '1,260p'Repository: RL-Align/RL-Kernel
Length of output: 50374
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf 'Imports at top of core file:\n'
sed -n '1,45p' rl_engine/testing/attention_comparison.py
printf '\nDecode inputs builder in tests and dtype helpers:\n'
sed -n '35,95p' tests/test_attention_comparison.py
printf '\nAll tensor dtype constructors/replacements in compare files/tests:\n'
rg -n -C 3 '(float32|float64|bfloat16|half|float16)\)|dtype=.*torch|torch\.empty_like|torch\.randn_like|torch\.empty|torch\.randn' rl_engine/testing/attention_comparison.py tests/test_attention_comparison.py \
| sed -n '1,260p'
printf '\nAttentionComparisonInputs fields and constructors:\n'
rg -n -C 8 'class AttentionComparisonInputs|rope_output_dtype|rope_cast_at|q: torch.Tensor|k: torch.Tensor' rl_engine/testing/attention_comparison.py tests/test_attention_comparison.py | sed -n '1,240p'Repository: RL-Align/RL-Kernel
Length of output: 37149
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from dataclasses import dataclass
from typing import Any, Literal
RoPEState = Literal["pre_rope", "post_rope"]
`@dataclass`(frozen=True)
class TensorLike:
dtype: Any
device: Any
`@dataclass`(frozen=True)
class DecodeAttentionInputs:
q: TensorLike
k_cache: TensorLike
v_cache: TensorLike
rope_output_dtype: Any = None
def _decode_rope_output_dtype(inputs):
return inputs.q.dtype if inputs.rope_output_dtype is None else inputs.rope_output_dtype
for case in (
("matching bfloat16", "bfloat16", "bfloat16"),
("mixed k_cache bfloat16", "float32", "bfloat16"),
("mixed k_cache float16", "float32", "float16"),
):
name, q_dtype, k_dtype = case
inputs = DecodeAttentionInputs(
q=TensorLike(q_dtype, "cpu"),
k_cache=TensorLike(k_dtype, "cpu"),
v_cache=TensorLike(k_dtype, "cpu"),
)
print(name, "q_dtype=", q_dtype, "k_cache_dtype=", k_dtype, "_decode_rope_output_dtype=", _decode_rope_output_dtype(inputs))
PYRepository: RL-Align/RL-Kernel
Length of output: 448
Use the K cache dtype for pre-RoPE K casting.
_decode_rope_output_dtype currently returns inputs.q.dtype, so a pre-RoPE K cache is cast to the query dtype before rotation. Add a separate K cache dtype default for _decode_rope_output_dtype, validate q.dtype == k_cache.dtype in _validate_decode_inputs, or split query/K output dtype semantics so the K path does not mix dtypes.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rl_engine/testing/attention_comparison.py` around lines 714 - 724, The
pre-RoPE K path currently casts using the query dtype from
_decode_rope_output_dtype. Update the dtype handling around
_decode_rope_output_dtype and _validate_decode_inputs so K uses the K-cache
dtype, either by adding a K-cache dtype default or validating q.dtype ==
k_cache.dtype; ensure the pre-RoPE Q path continues using the query dtype and
the pre-RoPE K path preserves the cache dtype.
| report = compare_decode_kv_replay(inputs) | ||
| assert report.drifts[0].out.max_abs <= 1.0e-6 | ||
| assert report.drifts[0].lse.max_abs <= 1.0e-6 |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
The bfloat16 output tolerance is below the bfloat16 resolution.
output_dtype is torch.bfloat16, so report.drifts[0].out compares two bfloat16 tensors. bfloat16 keeps 8 mantissa bits, so neighbouring values near 1.0 differ by about 7.8e-3. The reference downcasts each per-query result, while the replay accumulates in fp32 and downcasts once. The two paths can round to different bfloat16 values. The 1.0e-6 bound then requires bit-identical output and passes only for this seed. Use a bfloat16-scale bound for out. The lse bound can stay tight because both paths keep LSE in fp32.
💚 Proposed tolerance fix
report = compare_decode_kv_replay(inputs)
- assert report.drifts[0].out.max_abs <= 1.0e-6
+ assert report.drifts[0].out.max_abs <= 1.0e-2
assert report.drifts[0].lse.max_abs <= 1.0e-6📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| report = compare_decode_kv_replay(inputs) | |
| assert report.drifts[0].out.max_abs <= 1.0e-6 | |
| assert report.drifts[0].lse.max_abs <= 1.0e-6 | |
| report = compare_decode_kv_replay(inputs) | |
| assert report.drifts[0].out.max_abs <= 1.0e-2 | |
| assert report.drifts[0].lse.max_abs <= 1.0e-6 |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/test_attention_comparison.py` around lines 383 - 385, Update the output
drift assertion in compare_decode_kv_replay to use a bfloat16-scale tolerance
rather than 1.0e-6, since report.drifts[0].out compares bfloat16 results that
may round differently. Keep the existing tight 1.0e-6 tolerance for
report.drifts[0].lse.
| ).to(output_dtype) | ||
| if metadata.k_cache_rope_state == "pre_rope": | ||
| key_positions = metadata.key_position_ids[batch_index, slot_index].unsqueeze(0) | ||
| k = rope.forward_fp32(k, key_positions, theta=inputs.rope_theta).to(output_dtype) |
There was a problem hiding this comment.
This casts pre-RoPE k_cache with the Q RoPE output dtype. If rollout stores KV in a different dtype from Q, the replay no longer checks the real cache boundary. Can we keep Q/K RoPE output dtypes separate or validate they are intentionally the same?
There was a problem hiding this comment.
Thanks a lot. I have fixed in 2eb4362. I split the decode RoPE dtype contract into q_rope_output_dtype and k_cache_rope_output_dtype. They default independently to q.dtype and k_cache.dtype, respectively, so a pre-RoPE cached K no longer inherits the query dtype. I also added mixed FP32-Q/BF16-KV pre/post-RoPE equivalence coverage and provenance checks.
| raise ValueError("q_rope_state must be 'pre_rope' or 'post_rope'") | ||
| if metadata.k_cache_rope_state not in {"pre_rope", "post_rope"}: | ||
| raise ValueError("k_cache_rope_state must be 'pre_rope' or 'post_rope'") | ||
| if metadata.prefix_cache_enabled and not metadata.prefix_cache_key: |
There was a problem hiding this comment.
Prefix cache validation only checks that a key exists. It does not verify that the reused prefix actually has the same token positions/pages/content identity, so a bad shared prefix can pass the harness. Can we add a real prefix identity check or narrow the doc/contract?
There was a problem hiding this comment.
Also fixed in 2eb4362. Prefix-cache validation now requires prefix_length and prefix_cache_fingerprint, in addition to the cache key. The harness recomputes a physical-layout-invariant SHA-256 fingerprint over the logical prefix positions and cached K/V content, including the cached-K RoPE state and tensor dtypes. A stale or mismatched prefix now fails before replay. Tests cover equivalent physical page layouts and modified prefix content.
|
|
||
| _validate_decode_inputs(inputs) | ||
| reference = run_decode_full_prefill_reference(inputs) | ||
| candidate = run_decode_kv_replay(inputs, merge_backend=merge_backend) |
There was a problem hiding this comment.
The decode TE path is not covered like the PR2 TE path. Please add a decode fake-TE test that proves the same sorted partial states go through the TE merge helpers and that unavailable TE is reported/handled consistently.
There was a problem hiding this comment.
Fixed in 2eb4362 as well. compare_decode_kv_replay now always runs the RL-Kernel reference candidate and optionally runs the capability-probed TE merge oracle through include_transformer_engine=True. I added fake-TE coverage verifying that sorted logical partial states pass through the TE correction helpers and match RL-Kernel, plus an unavailable-TE test confirming that the core result is preserved and the reason is recorded in report.unavailable.
Signed-off-by: JLiu4Coding <lzwgre@126.com>
|
Hi @inaniloquentee, thank you for the detailed review. I’ve addressed all three requested changes in 2eb4362: |
2 participants
Implements PR6: Decode-stage KV-cache CP attention replay from #235. This is currently stacked on #253 and will be retargeted to main after #253 is merged.
Summary
Sq=1and few-query attention.(out, lse)partial states in fixed logical block order and downcast only at final write.out, attention-domainlse, and optional active-token selected-logprob drift.Scope and design
DecodeKVCacheMetadataseparates physical page layout from logical sequence identity. The replay reconstructs logical blocks from the page table, computes per-block partial attention state, and merges byglobal_block_index; physical page order and CP ownership do not select the numerical reduction order.RoPE identity is explicit for decode queries and cached K. Pre-RoPE inputs are materialized with the declared query/key positions, while post-RoPE inputs are consumed directly. Inconsistent cache/query positions, cached-key positions, duplicate pages, missing logical positions, and invalid prefix-cache provenance fail loudly.
This PR reuses Transformer Engine only as an optional validation oracle for the FP32
(out, lse)merge, through the capability-probed context-parallel correction helpers introduced in #253. Transformer Engine does not construct or interpret the KV cache, determine the logical block order, or replace RL-Kernel’s deterministic reference implementation.This PR models CP block ownership on one device for attributable validation. A later integration with PR3's transport can gather the same partial states without changing the deterministic merge contract.
Validation
pre-commit run --files rl_engine/testing/attention_comparison.py rl_engine/testing/__init__.py tests/test_attention_comparison.py docs/design/ws2-attention-decode-replay.mdpytest tests/test_attention_comparison.py -q— 17 passedpytest tests/test_attention_comparison.py rl_engine/tests/test_dispatch.py tests/test_kv_cache_attention.py -q -k "not large and not gpu"— 34 passed, 4 deselectedmypy --ignore-missing-imports rl_engine/— passed for 90 source filespytest tests/test_attention.py -q -k "not large and not gpu"— 24 passed, 2 deselectedpytest tests/test_attention_correctness.py -q -rs— 127 GPU-only cases skipped because CUDA/ROCm is unavailable locallymkdocs build --strict -f mkdocs.yaml— passedgit diff --check— passedDCO
Signed-off-by.Summary by CodeRabbit
New Features
Documentation
Tests