[hongleng]

Summary

Extract the reusable GDN recurrent-step and causal-convolution PyTorch references from [#469](#469) into main, as agreed in [RFC #642](#642 (comment)).

This PR:

• Preserves the reference APIs and provider-specific FP32 arithmetic, reduction, state-rounding, and convolution dtype semantics.
• Fixes CPU batch-dependent GDN gating rounding by evaluating gates per active sequence, without relaxing bitwise invariance assertions.
• Adds independent CPU correctness, state-contract, and batch-invariance tests to CI.

The gating fix may change the last bits compared with the original batched reference. No CUDA backend, dispatch, or runtime paths are modified.

Validation

• CPU tests: 93 passed under AVX2 and AVX512 separately.
• CI: Full CPU reference suite and focused AVX2 regression (13 passed); CI-Pipeline green.
• Numerical accuracy: Independent FP64 oracles passed. Compared with the pre-fix implementation, maximum sampled FP32 differences were 1.19e-7 (output) and 2.38e-7 (state) on AVX2, with no difference on sampled AVX512 inputs.
• State and dtype contracts: Verified output/state batch invariance, BF16/FP32 convolution promotion, state isolation, and rounding.
• Checks: Pre-commit, Black, isort, Flake8, focused MyPy, and layout tests passed.

Historical extraction comparisons against #469 were bitwise on the tested CPU fixtures before the gating fix; that claim does not extend to the revised GDN implementation.

Scope and limitations

These references retain the vLLM 0.30.0 decode arithmetic profile, but GPU/provider fidelity and cross-ISA or cross-PyTorch-version bitwise equality are not claimed.

The shared numerical contract and miles-specific profiles belong to RFC #642 Task 20. Existing reference copies in [#507](#507) will need reconciliation during integration.

Merge Danger

Door: two-way

Blast Radius: references

This adds CPU references, tests, and CI coverage without modifying existing dispatch or runtime paths.

Summary by CodeRabbit

• New Features
  • Added PyTorch reference implementations for single-step causal convolution and gated delta recurrence, including paged state support and configurable output precision.
  • Inactive state entries produce zero outputs and remain unchanged.
• Tests
  • Added CPU coverage for numerical accuracy, data types, state handling, input validation, and activation behavior.
  • Added the reference tests to the continuous integration workflow, including an AVX2-specific batch-invariance regression check.

[coderabbitai]

Important

Review skipped

Review was skipped as selected files did not have any reviewable changes.

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: 1ec0c6bc-75b8-4079-b411-52bf315e690e

Commits

Reviewing files that changed from the base of the PR and between d38f69d and 8d5b8e2.

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

No actionable comments were generated in the recent review.

Walkthrough

This PR adds PyTorch reference implementations for paged-state gated-delta recurrence and single-token causal convolution. It exports both operators and their constants, adds CPU tests for numerical and state behavior, and runs the tests in CI.

Changes: Linear-attention references

Layer / File(s)	Summary
Gated-delta recurrence rl_engine/reference/linear_attn/gated_delta_rule.py, tests/reference/linear_attn/test_gdn_reference.py, tests/ops/linear_attn/test_gdn_state_contract.py	Adds a recurrent step with state-index validation, optional Q/K normalization, FP32 gating, and chunked reductions. Tests compare outputs with an FP64 oracle and check state updates, dtype behavior, gradients, and batch invariance.
Causal convolution rl_engine/reference/linear_attn/causal_conv1d.py, tests/reference/linear_attn/test_gdn_reference.py, tests/ops/linear_attn/test_gdn_state_contract.py	Adds a single-token paged-state convolution update with validation, cache-dtype rounding, sequential FP32 accumulation, optional bias, and SiLU/Swish. Tests cover layouts, widths, activations, and state updates.
Public API and CPU validation rl_engine/reference/linear_attn/__init__.py, tests/ops/linear_attn/test_gdn_state_contract.py, tests/reference/linear_attn/test_gdn_reference.py, .github/workflows/ci.yml	Exports both operators and constants. Adds shared contract and import-isolation tests, plus CPU CI steps for the linear-attention tests and the AVX2 batch-invariance regression.

Estimated code review effort: 4 (Complex) | ~45 minutes

Change: Feature

Possibly related issues

• NPU + verl + RL-Kernel + Qwen3.6 Series MoE #563: The roadmap includes a hybrid attention path with GDN, causal convolution, state, and backward work. This PR adds GDN and causal-convolution reference implementations.
• [RFC][Qwen3-Next-80B-A3B-Instruct][CUDA/ROCm] VIME/vLLM operator-level train-rollout consistency roadmap, ablation matrix, and integration plan #428: The roadmap covers GDN, causal convolution, and prompt-prefill to decode state transfer. This PR adds reference operators but not the state bridge.
• [RFC] miles + RL-Kernel + Qwen3.6 Series MoE on NVIDIA Hopper #642: The RFC includes Qwen3.6 GDN and causal-convolution work. This PR adds references with provider-specific arithmetic.
• [RFC][Nemotron-3-Nano-30B-A3B][CUDA/ROCm] ROADMAP #434: The roadmap requires a depthwise causal-convolution state update for its Mamba-2 path. This PR adds a single-token causal-convolution state reference.

Merge Risk: Minimal · previous review d38f6

No actionable merge-blocking issue is identified in this change.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 13.33% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta…	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly and concisely describes the main change: extracting reusable GDN and causal-convolution reference implementations into the main branch.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Docstring Coverage

Explanation

Docstring coverage is 13.33% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 15 functions across 5 files. (1 skipped: 1 unsupported.)

Resolution

Write docstrings for the functions missing them to satisfy the coverage threshold.

Finishing Touches

Generate unit tests (beta)

• Create a new PR

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

Tip

🎯 Customize your risk assessments

Guide how CodeRabbit assesses merge risk and security architecture risk. Highlight sensitive changes and explain your security boundaries using Merge Risk Instructions and Security Risk Instructions in your review settings. Security architecture risk assessments require the Advanced plan or above.

✅ Automatic approvals, with your rules

Let CodeRabbit approve pull requests once required threads are resolved, pre-merge checks pass, and your risk limits are met. Enable Automatic approvals in review settings, with overrides for teams or file paths.

Recent review info

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: 3ea9f8c6-1aff-468a-a7fc-3cef3042c012

Commits

Reviewing files that changed from the base of the PR and between 73ee54e and d38f69d.

Files selected for processing: 1

• .github/workflows/ci.yml

Included review availability: This review used your included allowance. Your plan provides up to 2 included reviews per hour; 0 remain after this review.

[github-actions]

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

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

⚠️ 5 warning(s) occurred during review.

---

📄 .github/workflows/ci.yml (L90-L90)

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

This makes AVX2 a hard prerequisite of the always-on unit-tests job: on any runner that does not expose AVX2 (self-hosted runners, fork runners, or a future ubuntu-latest host whose CPU flags differ), exit 1 fails the entire job and blocks unrelated PRs. The repo's existing pattern for hardware-dependent validation (gpu-ci.yml) isolates it into a dedicated, label-gated job rather than failing the main suite. Consider moving this step into a separate job gated on a needs-avx2-ci label/hardware gate (so it is explicitly reported as not-run rather than silently skipped, and does not gate the general unit-tests job).

---

📄 .github/workflows/ci.yml (L94-L94)

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

This assertion assumes torch.backends.cpu.get_cpu_capability() reflects the ATEN_CPU_CAPABILITY override. In PyTorch it reports the capability selected/detected for the process, and on AVX-512 capable GitHub ubuntu-latest runners it can report 'AVX512' even though ATEN_CPU_CAPABILITY=avx2 was set, which would fail this step for a correctly configured environment. Verify this behavior against the installed torch version; if it is hardware-based, assert the effective dispatch capability via the env you set (or drop this separate assertion since ATEN_CPU_CAPABILITY=avx2 already forces the backend and the cpuinfo guard already validates hardware), so the check does not create a flaky build failure.

---

📄 tests/ops/linear_attn/test_gdn_state_contract.py (L98-L98)

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

This guard does not actually validate the claimed "bias last" behaviour. (3.0 + 3.0) + 1e8 is evaluated in Python double precision and yields exactly 100000006.0, which trivially differs from 1e8 == 100000000.0, so the assertion passes regardless of whether the implementation accumulates the bias first or last. In fp32 the "bias last" result would round to 100000008.0 (spacing is 8 at 1e8), not 100000006.0. To actually catch a bias-last implementation, compute the negative case in fp32, e.g. torch.tensor(6.0) + torch.tensor(1e8) and compare against that value.

💡 Suggested Change

Before:

    assert out.item() != (3.0 + 3.0) + 1e8  # what adding the bias last would give

After:

    bias_last = (torch.tensor(3.0) + torch.tensor(3.0)) + torch.tensor(1e8)
    assert out.item() != bias_last.item()  # what adding the bias last would give

---

📄 tests/reference/linear_attn/test_gdn_reference.py (L204-L204)

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

rounded_state comes from op.forward(...) (updated cache, i.e. dtype, which is torch.bfloat16 when the test is parametrized on bf16) while state32 comes from forward_fp32, which for GDN returns an fp32 state (state32.dtype == torch.float32 in the other tests). Comparing a bf16 tensor to an fp32 tensor with torch.equal will fail on dtype, and the values also differ by bf16 rounding. This assertion should normalize both operands to the same dtype/precision (e.g. rounded_state vs state32.to(dtype)), otherwise the bf16 parametrization cannot pass.

💡 Suggested Change

Before:

    assert torch.equal(rounded_state, state32)

After:

    assert torch.equal(rounded_state, state32.to(dtype))

---

📄 tests/ops/linear_attn/test_gdn_state_contract.py (L77-L77)

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

True is a subclass of int in Python, so a validation of the form isinstance(heads, int) and heads > 0 will accept True as heads=1 and not raise. If the reference only rejects non-int/≤0 values, this parametrization documents a contract the implementation does not enforce and the test fails. Confirm the reference explicitly rejects booleans (e.g. type(heads) is int or isinstance(heads, bool) check); if it does not, either the test or the implementation should be aligned.

---

⚠️ Warnings:

• .github/workflows/ci.yml (review_limit_reached): reached the aggregate token budget before finishing; the grace round did not complete the task
• .github/workflows/ci.yml (token_budget_reached): stopped group ".github/workflows/ci.yml" mid-review: used 617633 tokens exceeds budget 500000
• rl_engine/reference/linear_attn/__init__.py,rl_engine/reference/linear_attn/causal_conv1d.py,rl_engine/reference/linear_attn/gated_delta_rule.py (review_limit_reached): reached the aggregate token budget before finishing; the grace round completed the task
• tests/ops/linear_attn/test_gdn_state_contract.py,tests/reference/linear_attn/test_gdn_reference.py (review_limit_reached): reached the aggregate token budget before finishing; the grace round did not complete the task
• tests/ops/linear_attn/test_gdn_state_contract.py (comment_refiled): comment filed against tests/reference/linear_attn/test_gdn_reference.py describes code in tests/ops/linear_attn/test_gdn_state_contract.py; re-filed

[github-actions]

@hongleng 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.

[Flink-ddd]

please resolve the CI error first

[hongleng]

@Flink-ddd Thanks for flagging this. The CI failure has been fixed in 73ee54e.

The issue was caused by batch-dependent FP32 rounding in GDN gating on AVX2. The fix evaluates gates per active sequence, preserving bitwise batch invariance without relaxing the assertions.

The latest CI-Pipeline is green, including the full CPU reference suite and the focused AVX2 regression (13 passed). I've also added regression tests and updated the PR description to clarify the numerical behavior and validation scope.

PTAL when convenient. Thanks!

[Flink-ddd]

The previous issues are resolved, and the changes follow the refactored layout. All 93 reference tests pass locally under both AVX2 and AVX512, and CI is green. LGTM. Please retain the gating fix when reconciling with test-qwennext and #507. also cc @maxiaosong1124 @KJLdefeated @inaniloquentee @ryankert01 PTAL

[maxiaosong1124]

LGTM!

[maxiaosong1124]

LGTM!