perf(#1464): fix MGTSD + RWKV7Block training-throughput timeouts - #1471
Merged
Merged
Conversation
…usion Issue #1464: MGTSDTests.Clone_ShouldProduceIdenticalOutput timed out at 120s. Throughput root cause: CreateDefaultMGTSDLayers built the denoising MLP hidden layers at outputSize = contextLength * hiddenDimension (168*128 = 21,504) - the flattened-sequence size mistaken for the per-position hidden width (same anti- pattern class as the PR #1455 fixes). Each hidden Dense became a 21,504 x 21,504 (~462M-param) matmul and the denoiser ~4B params, run DiffusionSteps x NumGranularities times per Predict. Fixed to hiddenDimension; class 120s -> ~16s. Fixing throughput surfaced a pre-existing (timeout-masked) correctness defect: the training path ran the layer stack as a 128-wide regressor while inference fed the denoising stack a 177-wide [x_t | cond | guidance | t] pack, so the shared BatchNorm could not cache both widths; per-granularity pack length also drifted. Redesigned to a paper-faithful x0-parameterized DDPM (Ho et al. 2020 / MG-TSD ICLR 2024): constant-width pack identical in train and inference; network predicts x_hat_0 with the x0 posterior at inference; DDPM x0-prediction training step with RevIN-normalized target; RevIN de-normalization of the forecast; deterministic seeded diffusion RNG for clone parity. All 21 MGTSDTests pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…er-step loop Issue #1464 (half 2). RWKV7Block.TimeMixingForward recomputed the token-shift lerp and the r/k/v/a/b projection matmuls once per timestep inside the sequential loop - seqLen separate small GEMMs per projection. Since none of these depend on the recurrent WKV state, they are now computed for the whole sequence in one batched GEMM each over [batch*seqLen, modelDim]; the loop only consumes per-step slices for the inherently sequential WKV state recurrence. All ops remain on the autodiff tape (the layer has no manual Backward; gradients flow through the tape), so projection-weight gradients are unchanged - RWKVForecaster Clone_ShouldProduceIdenticalOutput and Predict_ShouldBeDeterministic still pass. This is necessary but not sufficient for the memorization-task timeout: the dominant remaining cost is the per-step WKV recurrence itself (per-timestep tape micro-ops), which the follow-up fused Tensors WKV kernel addresses. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… kernel Issue #1464 (half 2), completing the RWKV7Block throughput work. The decomposed time-mixing loop recorded ~10 tape micro-ops per timestep for the WKV state recurrence (plus per-step groupnorm + output projection); on the RWKVForecaster memorization task (seqLen 512, 4 layers, 100 train iters) that per-step tape-dispatch overhead pushed the test past the 180s budget. TimeMixingForward now calls the fused, differentiable Engine.Rwkv7SequenceForward (AiDotNet.Tensors #477): the whole diagonal-decay + rank-1-injection + gated-readout recurrence runs as ONE tape node whose custom backward is the BPTT adjoint, so projection gradients are identical (clone-parity + Predict determinism + gradient-flow tests still pass). GroupNorm and the output projection are batched over all positions [batch*seqLen, modelDim], leaving NO per-timestep ops in time-mixing. LossStrictlyDecreasesOnMemorizationTask now passes within budget (was a 180s timeout). Requires AiDotNet.Tensors 0.86.7 (the #477 fused-recurrent-kernel release); the pin is bumped accordingly. AiDotNet.Native.* stay 0.86.6 (managed-only change). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ChannelMixingForward is purely position-wise (token-shift + SiLU-gated FFN, no recurrence), so it now runs as batched GEMMs over [batch*seqLen, modelDim] instead of ~8 Engine dispatches x seqLen x numLayers. Removes the last per-timestep dispatch hot path in RWKV7Block; the tape backs gradients automatically (no manual backward).
Issue #1464. RWKVForecaster stacked the older RWKVLayer, whose WKV recurrence is a scalar NumOps/Math.Exp loop over t x batch x heads x headDim x headDim (~4.2M element-ops/layer, detached from the tape). Profiling showed that scalar recurrence — NOT GEMM FLOPs (the projection/FFN GEMMs run at 30-100 GFLOP/s) — dominated the training step at ~8.3 s/iter, pushing LossStrictlyDecreasesOnMemorizationTask to a 180s timeout. Point the forecaster at RWKV7Block instead: the paper-faithful RWKV-7 block whose time-mixing runs through the fused, differentiable CpuEngine.Rwkv7SequenceForward kernel with batched token-shift/projections and batched channel-mixing. Per-iteration training drops from ~8.3 s to ~1.1 s (7-8x); the memorization test now passes in ~27 s, and all 27 RWKVForecasterTests pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Contributor
|
Warning Review limit reached
More reviews will be available in 31 minutes. Learn how PR review limits work. Your organization has run out of usage credits. Purchase more in the billing tab. ⌛ How to resolve this issue?After more reviews become available, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans include higher PR review limits than trial, open-source, and free plans. In all cases, reviews become available again over time. During sustained high-volume PR review activity, CodeRabbit may temporarily slow when the next review becomes available. Please see our Fair Usage Limits Policy for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (1)
WalkthroughVectorizes RWKV7Block sequence ops, integrates RWKV7Block into forecaster/layer construction, rewrites MGTSD training/inference to use deterministic sampling and pre-allocated denoising buffers, adjusts denoiser widths/activations, and bumps four AiDotNet package versions. ChangesRWKV-7 Sequence Vectorization
MGTSD Diffusion Training and Inference Optimization
Package Updates
Sequence Diagram(s)sequenceDiagram
participant Client
participant MGTSD
participant Denoiser
participant RNG
Client->>MGTSD: Predict(input)
MGTSD->>RNG: seeded sample t, ε (per-call counter)
MGTSD->>Denoiser: packed input [x_t | condHidden | guidance | t]
Denoiser-->>MGTSD: x0Pred
MGTSD->>MGTSD: update xt via DDPM posterior
MGTSD-->>Client: de-normalized forecast
sequenceDiagram
participant Forecaster
participant RWKV7Block
participant Engine
Forecaster->>RWKV7Block: forward(sequence)
RWKV7Block->>Engine: Rwkv7SequenceForward(batched seq)
Engine-->>RWKV7Block: recurrent outputs
RWKV7Block-->>Forecaster: batched layer output
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
The latest updates on your projects. Learn more about Vercel for GitHub. 2 Skipped Deployments
|
Owner
Author
|
@coderabbitai full review |
Contributor
✅ Actions performedFull review triggered. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
Directory.Packages.props (1)
8-48:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winBLOCKING: Update the version history comment to document 0.86.7.
The comment block ends at 0.86.6 but the package is now bumped to 0.86.7. The comment must document what 0.86.7 adds—specifically the fused RWKV-7 kernel (
Engine.Rwkv7SequenceForward) that gates the RWKV performance fix in this PR.Production-ready code requires complete documentation of dependency changes, especially when the version bump is a hard requirement for new functionality.
📝 Proposed documentation update
(corrupted GAMLSS and any model trained via AiModelBuilder's auto-detected GPU path). 0.86.6 materializes those immediately. --> + + Bumped 0.86.6 -> 0.86.7 for the fused RWKV-7 kernel + (Tensors PR `#XXX`): Engine.Rwkv7SequenceForward provides batched + token-shift, projections, and channel-mix for RWKV-7 blocks, + eliminating per-timestep scalar inner loops. Required for + RWKVForecaster performance fix (issue `#1464` / PR `#1471`). --> <PackageVersion Include="AiDotNet.Tensors" Version="0.86.7" />Replace
#XXXwith the actual Tensors PR number if known.🤖 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 `@Directory.Packages.props` around lines 8 - 48, Update the version-history comment block in Directory.Packages.props to include the 0.86.7 entry: add a short note that 0.86.7 introduces the fused RWKV-7 kernel (Engine.Rwkv7SequenceForward) which enables the RWKV performance fix gated by this PR, and include the Tensors PR number if available; ensure the new bullet is placed after the 0.86.6 explanation and mirrors the existing style/tense and any referenced issue/PR tags used in surrounding entries.src/Finance/Forecasting/Foundation/MGTSD.cs (1)
242-272:⚠️ Potential issue | 🟠 Major | ⚡ Quick winBLOCKING: Orphaned XML documentation creates dead code.
There are two consecutive XML documentation blocks before the
Train()method. The first block (lines 242-257, unchanged) is now orphaned because XML docs must immediately precede the code element they document. The second block (lines 258-272) is the actual documentation forTrain().The orphaned block should be removed as part of this PR since the PR's structural changes caused it to become dead documentation.
🧹 Proposed fix: Remove orphaned documentation
- /// <summary> - /// Tape-aware training forward. Runs the existing Layers stack as a - /// deterministic context → forecast regression head so - /// <c>NeuralNetworkBase.TrainWithTape</c> can record the tape, compute the - /// user-injected loss, and step the optimizer. - /// </summary> - /// <remarks> - /// Same bug family as CCDM / TimeDiff / TimeGrad / TSDiff: custom Train - /// hand-built DDPM noise perturbation and fed <c>_denoisingLayers</c> a - /// [noisyTarget | condHidden | t-embedding] tensor whose last-dim didn't - /// match the layers' baked-in hiddenDim input sizes, then called - /// <c>_optimizer.UpdateParameters(Layers)</c> without backward. - /// Multi-granularity DDPM reverse process with guidance weighting stays - /// in <see cref="ForwardNative"/> for probabilistic inference via - /// <see cref="Predict"/>/<see cref="Forecast"/>. - /// </remarks> /// <summary> /// Paper-faithful DDPM training step (x0-parameterization; MG-TSD §3, Ho et al. 2020). /// </summary>🤖 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 `@src/Finance/Forecasting/Foundation/MGTSD.cs` around lines 242 - 272, Remove the orphaned XML documentation block that precedes the real Train() doc comment so XML docs immediately precede the Train() method; specifically delete the first summary/remarks block referencing NeuralNetworkBase.TrainWithTape, _denoisingLayers and ForwardNative so the remaining DDPM/Train() summary (the paper-faithful DDPM training step) sits directly above the Train() method declaration.
🤖 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 `@Directory.Packages.props`:
- Line 49: Directory.Packages.props currently pins AiDotNet.Tensors to a
non-existent 0.86.7 which breaks restore; change the PackageVersion Include for
AiDotNet.Tensors to a published version that contains the required managed
symbol (or remove the hard pin and reference a matching NuGet range) and ensure
the managed/native packages are bumped in lockstep so
Engine.Rwkv7SequenceForward is supported by the managed DLL; specifically,
verify the chosen AiDotNet.Tensors package actually contains the
Rwkv7SequenceForward symbol in its managed DLL and update any Native.* package
pins to the exact same version, or add a TODO/comment to hold the RWKV-7 API
bump until a correct managed 0.86.7-equivalent is published.
In `@src/Finance/Forecasting/Foundation/MGTSD.cs`:
- Around line 531-543: The code can underflow when abarPrev ≈ 1; stabilize
oneMinusAbarPrev the same way as oneMinusAbarT by adding eps10 so divisions and
multiplications that use it are safe. Update the computation of oneMinusAbarPrev
(used by coefXt and sigmaT) to use NumOps.Add(NumOps.Subtract(NumOps.One,
abarPrev), eps10) and keep the existing t>0 conditional for sigmaT (so sigmaT
remains Zero at t==0), ensuring coefXt and sigmaT use the stabilized value.
In `@src/Helpers/LayerHelper.cs`:
- Around line 31801-31811: CreateDefaultMGTSDLayers is missing the required XML
doc elements—add a complete XML documentation block for the method that includes
<summary>, <param> entries for each parameter, a <returns> describing the
yielded layers, and a <remarks> section that includes a <para><b>For
Beginners:</b> ...</para> note following the same structure used by
CreateDefaultRWKVForecastingLayers; also verify the RWKV layer construction
inside CreateDefaultMGTSDLayers uses the correct RWKV7Block<T> constructor
signature (sequenceLength, modelDimension, numHeads and include optional
ffnMultiplier/activations parameters if applicable) so the docs and code are
consistent.
---
Outside diff comments:
In `@Directory.Packages.props`:
- Around line 8-48: Update the version-history comment block in
Directory.Packages.props to include the 0.86.7 entry: add a short note that
0.86.7 introduces the fused RWKV-7 kernel (Engine.Rwkv7SequenceForward) which
enables the RWKV performance fix gated by this PR, and include the Tensors PR
number if available; ensure the new bullet is placed after the 0.86.6
explanation and mirrors the existing style/tense and any referenced issue/PR
tags used in surrounding entries.
In `@src/Finance/Forecasting/Foundation/MGTSD.cs`:
- Around line 242-272: Remove the orphaned XML documentation block that precedes
the real Train() doc comment so XML docs immediately precede the Train() method;
specifically delete the first summary/remarks block referencing
NeuralNetworkBase.TrainWithTape, _denoisingLayers and ForwardNative so the
remaining DDPM/Train() summary (the paper-faithful DDPM training step) sits
directly above the Train() method declaration.
🪄 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: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: ba853e22-19d1-4067-a526-a9fa95af16b8
📒 Files selected for processing (5)
Directory.Packages.propssrc/Finance/Forecasting/Foundation/MGTSD.cssrc/Finance/Forecasting/StateSpace/RWKVForecaster.cssrc/Helpers/LayerHelper.cssrc/NeuralNetworks/Layers/SSM/RWKV7Block.cs
Owner
Author
|
@ Dependency unblocked + verified (Tensors 0.90.3)The fused differentiable RWKV-7 kernel this PR depends on ( Verification against 0.90.3 (net10.0, Release)
🤖 Generated with Claude Code |
…7 kernel The fused differentiable RWKV-7 sequence kernel (Engine.Rwkv7SequenceForward, ooples/AiDotNet.Tensors#514) this PR's RWKVForecaster wiring depends on shipped in AiDotNet.Tensors 0.90.0 (verified: absent in 0.87.0/0.88.0/0.89.0, present 0.90.0+). The prior 0.86.7 pin was a placeholder that never published. Bump managed + native packages to the latest published 0.90.3 (coreleased in lockstep). Verified against 0.90.3: src builds clean; MGTSDTests 21/21 (12s) and RWKVForecasterTests 27/27 (58s) pass at full scale. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
ooples
force-pushed
the
perf/1464-rwkv7-mgtsd-throughput
branch
from
June 1, 2026 10:29
0bda4b1 to
85ed73f
Compare
Apply the same 1e-10 stabilization to (1 - ᾱ_{t-1}) that (1 - ᾱ_t) already has, so
the posterior coefficients/σ_t can't underflow if a future noise schedule drives
ᾱ_{t-1} → 1 at early timesteps. No behavior change under the current linear beta
schedule (ᾱ_{t-1} < 1 for t > 0; t = 0 forces σ_t = 0). Addresses CodeRabbit review
on PR #1471. MGTSDTests 21/21 still pass.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #1464 —
MGTSDTests.Clone_ShouldProduceIdenticalOutputandRWKVForecasterTests.LossStrictlyDecreasesOnMemorizationTaskexceeded the test budget.MGTSD
Root cause:
CreateDefaultMGTSDLayerssized the denoiser hidden layers atcontextLength * hiddenDimension(168×128 = 21,504) — the flattened-sequence size mistaken for the per-position width — making each hidden Dense a ~462M-param matmul (~4B-param denoiser) runDiffusionSteps × NumGranularitiestimes perPredict. Fixed tohiddenDimension. This surfaced a pre-existing (timeout-masked) train/inference shape mismatch, fixed via a paper-faithful x0-parameterized DDPM redesign (constant-width pack identical in train/inference, x0-posterior sampling, RevIN de-normalization, deterministic seeded sampling). All 21 MGTSDTests pass (class 120s → ~16s).RWKV
The
RWKVForecasterstacked the older scalarRWKVLayer; profiling showed its scalarNumOps/Math.ExpWKV recurrence (~4.2M element-ops/layer) — not GEMM FLOPs — dominated the step (~8.3s/iter). Routed the forecaster through the fused, differentiableRWKV7Block(time-mix WKV viaEngine.Rwkv7SequenceForward, batched token-shift/projections/channel-mix). ~8.3s → ~1.1s/iter (7–8×); memorization test 180s timeout → ~27s; all 27 RWKVForecasterTests pass.Dependency
Requires AiDotNet.Tensors 0.86.7 (ooples/AiDotNet.Tensors#514, the fused RWKV-7 kernel). CI will not restore until that releases; pin is bumped accordingly (native packages stay 0.86.6 — managed-only change).
🤖 Generated with Claude Code
Summary by CodeRabbit
Chores
Refactor
New Features