[inaniloquentee]

Roadmap item 24 ??joint compressed/window attention with sink

Target: dsv4-rocm. Task commit: 159b0592be933b65aac76ff8b9fe411fa3f27456.
Based directly on latest main 8588158650aff58c9c1f527674de94bad06c3d10; no dependency on pending item-13/23 implementation commits. The target predates main, so review task-only changes against that main SHA.

Reuse and arithmetic

• All train-forward/recompute/prefill/decode records use the SAME unmodified, SHA-pinned aiter sparse training-forward kernel (commit 326cffa).
• Fixed H32/K32 launch recipe, one KV partition; compressed KV + full recent window + sink share the denominator. Sink has no V; BF16 P/O and FP32 natural-log LSE are explicit.
• Reuses the actual vLLM mixed FP8/BF16 cache decoder, with explicit OCP/FNUZ format and 64/128-token pages.
• No new numerical GPU kernel or silent mathematical fallback. The pinned source is provisioned explicitly by a hash-verifying helper, never downloaded at runtime.

Contracts and lifetime

• Logical request/position/kind keys, complete windows, C4/C128 group boundaries, canonical KV ordering, capacity and version checks.
• Owned fixed-capacity Graph buffers; real kernel capture/replay; output/saved-state snapshots, stream-event reuse barriers and host-submission lock.
• Failed mutable workspaces are poisoned, not retried. Shared explicit vLLM/Megatron forward boundary.

MI300X validation

• 236 passed, zero skipped, twice in fresh processes (28.00s with coverage; 12.57s repeat).
• 94 new tests: 54 CPU, 40 actual ROCm; 142 contract regressions.
• FP64 O/LSE oracle separately from exact-byte modes/chunks/permutations; 640 selected KV entries and 128 heads; C128 boundaries; real Graph, two streams, concurrent host threads, record lifetime; native packed-cache tests.
• Runtime statement coverage 294/300 (98%); branches 27/30 (90%). Native instructions excluded.
• Ruff passes. Source fingerprints, benchmarks and detailed report committed under reports/joint-attention-mi300x-20261029.{md,json}.

Explicit scope boundary

forward_record is detached saved forward state, NOT an implicit autograd implementation. Backward/dQ/dKV/d_sink are the separate item 25 PR, which follows next. No full engine/model, CP/DCP, multi-rank transport, larger/split-K profile, or whole-cache-pipeline Graph signoff. Packed/contiguous equality uses the same dequantized BF16 values, not lossless equivalence to unquantized training KV. Checked API overhead is reported separately from native/kernel Graph core timings; no unsupported speedup claim.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

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

• Source mode: Auto
• Repository: Unavailable
• Branch: Unavailable
• File: No configuration file from this repository applied
• Commit: Unavailable
• Layers: defaults
• Review profile: CHILL
• Plan: Advanced
• Run ID: ee98f439-460f-44cc-a7e2-562276f2094f

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

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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

[github-actions]

🔍 OpenCodeReview found 7 issue(s) in this PR.

• ✅ Successfully posted inline: 0 comment(s)
• 📋 Routed to summary by policy: 7 comment(s)

⚠️ 5 warning(s) occurred during review.

---

📄 tests/models/deepseek_v4/test_joint_attention_rocm.py (L181-L182)

⚠️ GitHub could not post this as an inline comment: Routed to summary (severity medium · category test)

The size fault does not trigger the guard it intends to test. With block_size=64, the validity check requires cache.shape[1] >= block_size * 584 == 37376; the sliced cache is [:, :128] (128 columns), so it is rejected by the same generic "cache must be contiguous uint8" check used by the other faults rather than by anything specific to size. All six parametrizations therefore assert the same coarse JointAttentionError from overlapping early checks (e.g. range with 128 also fails only the slot bound, stride fails contiguity) — the guard matrix gives little discriminating coverage. Consider asserting on distinct messages/conditions per fault, or at minimum set faults so each targets a distinct validated clause (e.g. size keeps >= block_size*584 but keeps a non-576 multiple stride/padding).

---

📄 rl_engine/contracts/operators/joint_attention.py (L82-L84)

⚠️ GitHub could not post this as an inline comment: Routed to summary (severity medium · category bug)

workspace_bytes does not match the actual allocations in JointAttentionPlan.__init__. In the plan, _q and _out are both bf16 [max_queries, num_heads, head_dim] (2 bytes/elt each), but this formula counts only one 4-byte buffer (4 * max_queries*num_heads*head_dim), omitting the second bf16 O buffer entirely. _kv (bf16) is counted as 2 * max_kv * head_dim — consistent with bytes, but the Q/O pair is off by 2x-bytes and one buffer. Concrete impact: with the documented example JointAttentionContract(64, 512, 32, 1024, 512, 128) the reported workspace is ~4.4 MB while the real resident Q+O is ~4.2 MB and KV adds ~1 MB, so the workspace_bytes <= max_workspace_bytes gate in __post_init__ under-reports and can admit configs whose real resident allocation exceeds max_workspace_bytes. Either enumerate every buffer (Q, O, KV, sink, indices, active, LSE) with its true element size, or document that O is excluded/derived and keep the accounting consistent.

---

📄 rl_engine/models/deepseek_v4/joint_attention.py (L386-L389)

⚠️ GitHub could not post this as an inline comment: Routed to summary (severity medium · category bug)

gather_packed_joint_kv hardcodes head width 512 and the packed bytes-per-token 584 (block_size * 584, torch.empty((pages.numel(), block_size, 512), ...), needed = ...*512*2, and the empty-result (0, 512)), while JointAttentionContract explicitly allows head_dim in (128, 512). There is no check that the contract/backend expects 512 on this path, and no head_dim argument, so a caller using a D128 joint contract gets a silently wrong-width tensor (or a stride check that still passes, since only block_size * 584 is validated). Concretely, a D128 profile would produce KV rows of 512 columns that fail or silently mis-shape in forward_record's kv.shape[1] == c.head_dim check with a confusing error far from the cause. Pass the expected head_dim (and derive 584 from it) or accept it as a parameter and reject mismatches explicitly at the top of this function.

---

📄 rl_engine/models/deepseek_v4/joint_attention.py (L389-L390)

⚠️ GitHub could not post this as an inline comment: Routed to summary (severity low · category maintainability)

needed double-counts the intermediate page buffer: pages.numel() * block_size * 512 * 2 is the full decoded out tensor (allocated as (pages.numel(), block_size, 512) bf16), and slots.numel() * 512 * 2 is added as if it were a second workspace, but the actual final index_select result is also slots.numel() * 512 * 2 while out is still alive. So the real peak is pages*block_size*512*2 + slots*512*2 (both buffers live) — which the current formula happens to match numerically — but the two terms are unnamed and the semantics are easy to break on edit; more importantly the budget does not account for the flatten intermediate being a view (fine) while the index_select output is separate. Please name these terms (decoded_pages, row_snapshot) so the budget reflects the actual two live buffers and an editor cannot silently drop one.

---

📄 rl_engine/models/deepseek_v4/joint_attention.py (L236-L237)

⚠️ GitHub could not post this as an inline comment: Routed to summary (severity low · category maintainability)

eager(device) and the plan's event plumbing degrade silently for non-CUDA devices: eager only rejects capture on cuda, _wait/_record no-op when self._event is None, and _transaction uses nullcontext() off-CUDA. The backend contract (AiterJointAttentionBackend) requires ROCm/gfx942, but nothing in JointAttentionPlan/forward_record enforces self.device.type == "cuda", so a plan constructed with any alternative backend (including the CPU ReferenceBackend used in tests) gets no cross-thread/stream ordering guarantee at all while still being accepted by forward_record. This makes the "an event after all returned clones orders the next stream's reuse" guarantee conditional and undocumented at the boundary. Require backend.device_type == "cuda" (or an explicit supports_events capability) in __init__, or document that stream ordering is only provided for CUDA devices.

---

📄 benchmarks/benchmark_joint_attention.py (L26-L31)

⚠️ GitHub could not post this as an inline comment: Routed to summary (severity low · category documentation)

event_us and wall_us come from two different loops. The event-based loop records start/end and calls end.synchronize() before the wall loop begins, then re-runs fn() repeats more times and syncs at the end of the wall loop. This means the committed report's upstream_forward/captured_core/eager_forward_record comparisons mix two measurement methodologies: event_us is GPU-only elapsed time, wall_us includes CPU dispatch overhead. This is by design for a CPU-vs-GPU split, but the report JSON presents them side-by-side without labeling the boundary. Minor; consider labeling which field is the GPU-only measure vs. end-to-end wall time so readers don't conflate them.

---

📄 benchmarks/benchmark_joint_attention.py (L48-L54)

⚠️ GitHub could not post this as an inline comment: Routed to summary (severity low · category maintainability)

The benchmarked sparse-density is topk per request (e.g. 6 for the (32, 64, 512, 2, 4) shape). But the report row reports c.selected_capacity in the selected column, which is 6 only because it happens to equal topk for the first three shapes, and 640 for the last where topk=512. When window=128, topk=512 (256 selected? no — topk*1=512 requests × 512 keys... actually selected_capacity derives from topk and count). Whatever the exact derivation, the selected column in the committed report describes the contract capacity, not the actual per-request number of selected indices built. This makes the report's cross-shape selected=6 column ambiguous — it is the same for m=1/32/128 even though the built inputs differ in per-request key count only via pos = topk*4-1 (constant per shape). Recommend emitting the actual per-request len(selected[i]) alongside selected_capacity so readers can distinguish workload from capacity.

---

⚠️ Warnings:

• tests/models/deepseek_v4/joint_attention_helpers.py,tests/models/deepseek_v4/test_joint_attention.py,tests/models/deepseek_v4/test_joint_attention_rocm.py (review_limit_reached): reached the aggregate token budget before finishing; the grace round did not complete the task
• tests/models/deepseek_v4/joint_attention_helpers.py (token_budget_reached): stopped group "tests/models/deepseek_v4/joint_attention_helpers.py,tests/models/deepseek_v4/test_joint_attention.py,tests/models/deepseek_v4/test_joint_attention_rocm.py" mid-review: used 580954 tokens exceeds budget 500000
• rl_engine/integrations/orchestrators/vime/patches/rocm_attention/sources.json (token_budget_reached): stopped dispatch: used 580954 tokens + group estimate 49168 = projected 630122 exceeds budget 500000
• rl_engine/backends/rocm/attention/joint_attention.py,rl_engine/contracts/operators/joint_attention.py,rl_engine/integrations/engines/rollout/vllm/joint_attention.py,rl_engine/integrations/engines/train/megatron/joint_attention.py,rl_engine/models/deepseek_v4/joint_attention.py (review_limit_reached): reached the aggregate token budget before finishing; the grace round did not complete the task
• benchmarks/benchmark_joint_attention.py,reports/joint-attention-mi300x-20261029.json (review_limit_reached): reached the aggregate token budget before finishing; the grace round did not complete the task

[github-actions]

@inaniloquentee The automated OCR review has completed.

Please review and address all automated review bot findings, then push your updates. Maintainers will begin manual review after the bot findings have been addressed.