Skip to content

fix: correct RWKV7 prefill state contract - #2297

Open
inureyes wants to merge 6 commits into
mainfrom
fix/issue-2286-rwkv7-chunk-contract
Open

inureyes wants to merge 6 commits into
mainfrom
fix/issue-2286-rwkv7-chunk-contract

Conversation

@inureyes

@inureyes inureyes commented Oct 13, 2026 •

Copy link
Copy Markdown
Member

Summary

  • Disable chunked prefill for RWKV7 because multi-token no-ID forwards intentionally reset model-owned recurrent state.
  • Add a loaded one-layer regression proving independent no-ID prompts reset and stable sequence-ID chunks retain equivalent recurrent state.
  • Restore the RWKV7 step-output shape contract by removing the singleton matrix-product axis before stacking timesteps.

Root cause found during validation

The loaded regression exposed an existing RWKV7 forward failure. Each recurrent step computed state @ r[..., None] as [B, H, D, 1] and stacked those tensors without removing the final singleton axis. This produced [B, T, H, D, 1], while the following per-head normalization expects [B, T, H, D] and weights shaped [H, D]. The pre-fix GPU run aborted with a native broadcast error between (2,4) and (1,1,2,4,1).

The fix squeezes only the matrix product's final axis, yielding the documented per-step [B, H, D] output and stacked [B, T, H, D] sequence tensor. Recurrent state remains [B, H, D, D].

Test plan

  • Pre-fix focused GPU run reproduced the RWKV7 output-shape failure.
  • GPU-locked rwkv7_ focused suite after the shape fix: 5 passed, including the loaded-model contract test in 0.64 seconds.
  • Root and surgery CUDA test-fast rebuild after the shape fix: passed in 2 minutes 22 seconds.
  • Contract gate: all pre-Python stages passed initially; after lowering competing build priority, both timing-sensitive guard tests passed, followed by all 123 Python tooling and 49 WebUI tests. No code or test was weakened for the retry.
  • git diff --check

Dependency refresh

The validation above was recorded before merging main commit 4b0c4519 from PR #2298. The current branch merged that MLX pin and overlay update at 9b05507b without changing src/models/rwkv7.rs (SHA-256 fa18e8e5e33f4af42fe7c7c8574e55db8578f90724d3a70d0cad4a7cf10adea7). Focused runtime revalidation on the new pin is pending; the checked results remain historical evidence for the reviewed model source.

Closes #2286

@inureyes inureyes added status:review Under review type:bug Bug fixes, error corrections, or issue resolutions priority:medium Medium priority area:models Model architectures, weights, loading, metadata area:inference Generation, sampling, decoding (incl. speculative, DRY) labels Oct 13, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:inference Generation, sampling, decoding (incl. speculative, DRY) area:models Model architectures, weights, loading, metadata priority:medium Medium priority status:review Under review type:bug Bug fixes, error corrections, or issue resolutions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fix(rwkv7): correct the no-ID chunked prefill contract

1 participant