feat(#1273): true-async + compile-host adoption across diffusion / VAE / VLM / LoRA

[ooples]

Summary

Closes #1273. Wires the Tensors-side async + compile-cache infrastructure (shipped in AiDotNet.Tensors PR #298 / Tensors v0.75.3) through every consumer-side hot path that previously re-traced the lazy graph eagerly.

Commits

1. deps: bump AiDotNet.Tensors 0.72.0 → 0.75.3 + patch transitive Snappier CVE

   • Pulls in the new ICompiledPlan<T>.ExecuteAsync(ct), ChainAsync(plan, ct), multi-input ChainAsync(plan, slot, ct), deprecated [Obsolete] ThenAsync, per-engine IExecutionStream<T> with CPU fast-path + non-blocking GPU poll, and CpuFusedOperations.FusedLoRAForward / FusedSparseLinear / fused denoise step.
   • Bumps native package versions in lockstep (OneDNN/OpenBLAS/CLBlast 0.70.2 → 0.75.3) to match what fix(#1224): clip-family invariant tests + dual-stream vlm base + parallel he-init #1274 already brought in.
   • Patches transitive Snappier 1.3.0 → 1.3.1 to clear NU1903 (CVE GHSA-pggp-6c3x-2xmx) brought in by the Tensors bump.

2. feat(#1273): true-async surface across CompileHost / NoisePredictor / Diffusion Generate (Workstreams A + E)

   • Workstream A — Diffusion Generate end-to-end async: DiffusionModelBase.GenerateAsync and LatentDiffusionModelBase.GenerateAsync now run the denoising loop via await PredictNoiseAsync per step. CPU stages complete inline; GPU stages let host work for the next step's prep overlap with the tail kernels.
   • CompiledModelHost<T>.PredictAsync: mirror of Predict() that routes through ICompiledPlan<T>.ExecuteAsync. Same trace-and-replay fast-path, same eager fallback, same cooperative-cancellation semantics, same pending-dispose drain. Just the leaf Execute call is async.
   • NoisePredictorBase.PredictNoiseAsync + PredictCompiledAsync: the base class routes through the compile host's async path with a captured-args eager fallback, so every concrete noise predictor inherits the compile-host-aware async path without per-subclass changes.
   • Workstream E — Auto-compile on eval: NeuralNetworkBase.SetTrainingMode(false) now optionally pre-warms the compiled inference plan via CompileForward on the next Predict, gated behind an opt-in AutoCompileOnEval flag (default false).

3. feat(#1273): VAE compile host + ChainedCompiledModelHost helper + LoRA Engine-op forward (Workstreams B + C + D)

   • Workstream B — VAE compile host: VAEModelBase now owns two CompiledModelHost<T> fields (encoder + decoder) and exposes EncodeCompiled / DecodeCompiled / async overloads.
   • Workstream C — ChainedCompiledModelHost async/multi-input: extended with PredictAsync(input, versions[], stages[], ct) and a multi-input PredictAsync(primary, versions[], sideInputsPerStage[][], stages[], ct) overload for cross-attention pipelines.
   • Workstream D — LoRA Engine-op forward: LoRALayer.Forward rewritten to use Engine.TensorMatMul + Engine.TensorMultiplyScalar so the (matmul → matmul → multiply-scalar) pattern is visible to CpuFusionPass for FusedLoRAForward rewrites once a LoRA layer participates in a compiled plan.

4. Review fixes (CodeRabbit):

   • Replaced GenerateAsync doc roadmap text with an accurate description of current behavior.
   • LatentDiffusionModelBase.GenerateAsync now mirrors the sync latent/pixel contract.
   • ThrowIfDisposed() guards on PredictNoiseAsync / PredictCompiledAsync.
   • Tiling/slicing flips invalidate the VAE compile cache.
   • Dispose(true) releases both VAE compile hosts.
   • PredictAsync disk-plan hits route through ExecuteAsync(ct) instead of sync Execute.
   • Async fallback path invalidates the poisoned plan before returning eager output.
   • Auto-compile-on-eval prewarm narrows the catch to recoverable exceptions only.
   • ValidateGenerateInputs / ResolveInitialSample / SanitizeNonFiniteElements helpers shared between sync and async Generate paths.
   • ResolveLatentShape / SanitizeFiniteInPlace helpers shared between sync and async LatentDiffusion paths.
   • conditioning parameter doc clarified on default PredictNoiseAsync.

5. test(#1273): numerical-equivalence tests + Diffusion Generate benchmark (W-A acceptance criteria)

   • tests/AiDotNet.Tests/IntegrationTests/Diffusion/DiffusionAsyncEquivalenceIntegrationTests.cs — three tests: same-seed-same-shape produces bit-equivalent output between sync and async paths; replay-determinism across multiple awaits; pre-cancelled token throws OperationCanceledException at the per-step boundary.
   • AiDotNetBenchmarkTests/Diffusion/DiffusionGenerateBenchmark.cs — sync Generate vs async GenerateAsync at the SDXL-class latent shape [1, 4, 128, 128] across 10- and 50-step DDIM. Steady-state replay cost reported. MemoryDiagnoser tracks alloc count for compile-cache correctness verification.

VAE / VLM / LoRA workstream-specific benchmarks live on the stacked refactor/1272-vae-sdxl-multi-host branch (PR #1280) so they're co-located with the consumer call sites that exercise them.

Test plan

• dotnet build src/AiDotNet.csproj clean on net10.0 and net471
• dotnet build tests/AiDotNet.Tests/AiDotNetTests.csproj clean
• dotnet build AiDotNetBenchmarkTests/AiDotNetBenchmarkTests.csproj clean
• Tensors NuGet bump resolves with no NU1903 warnings (Snappier patched)
• DiffusionAsyncEquivalenceIntegrationTests in place to verify numerical equivalence
• DiffusionGenerateBenchmark in place to measure W-A acceptance criterion (≥ 15% improvement on 50-step DDIM at SDXL-class shapes)
• CI runs full test suite (verifies no regression on existing-passing tests across other models)

🤖 Generated with Claude Code

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

2 Skipped Deployments

Project	Deployment	Actions	Updated (UTC)
aidotnet_website	Ignored	Preview	May 11, 2026 2:33pm
aidotnet-playground-api	Ignored	Preview	May 11, 2026 2:33pm

[vercel]

Deployment failed with the following error:

The `vercel.json` schema validation failed with the following message: `ignoreCommand` should NOT be longer than 256 characters

Learn More: https://vercel.com/docs/concepts/projects/project-configuration

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

Walkthrough

This PR adds cancellation-aware async generation and prediction surfaces across diffusion and neural-network infra: DiffusionModelBase/LatentDiffusionModelBase gain GenerateAsync/GenerateAsyncCore and PredictNoiseAsync; NoisePredictorBase routes predictions through compiled hosts with PredictNoiseAsync; CompiledModelHost/ChainedCompiledModelHost get PredictAsync; VAEModelBase adds compiled-plan helpers; NeuralNetworkBase adds AutoCompileOnEval.

Changes

Async Diffusion Inference with Compiled Caching

Layer / File(s)	Summary
Public Async API Contracts src/Diffusion/DiffusionModelBase.cs, src/Diffusion/LatentDiffusionModelBase.cs	DiffusionModelBase.GenerateAsync and LatentDiffusionModelBase.GenerateAsync expose cancellation-aware async generation; PredictNoiseAsync added as overridable async hook.
Async Denoising Loop Implementation src/Diffusion/DiffusionModelBase.cs	GenerateAsyncCore implements async denoising loop with shape/step validation, checked element-count, sample init (seed/RNG/initialSample), scheduler timesteps, per-iteration cancellation polling, awaited PredictNoiseAsync, scheduler.Step, and NaN/Infinity sanitization.
Async Noise Prediction with Compiled Cache src/Diffusion/NoisePredictors/NoisePredictorBase.cs	PredictCompiledAsync routes through CompiledModelHost.PredictAsync with structure-versioning and eager fallback; PredictNoiseAsync validates cancellation and delegates to the compiled path with eager sync fallback.
Compiled Model Async Execution and Lifecycle src/NeuralNetworks/CompiledModelHost.cs	PredictAsync mirrors Predict lifecycle: pre-flight checks, structure-version invalidation, disk-plan preload/compile if needed, fast-path ExecuteAsync(cancellationToken) with cooperative cancellation, and deferred cache disposal drained in finally.
Multi-Stage Async Chaining src/NeuralNetworks/ChainedCompiledModelHost.cs	Two PredictAsync overloads support single-input sequential per-stage execution and multi-input per-stage side-input arrays; both validate counts, propagate outputs between stages, and support cancellation.
VAE Compiled Encode/Decode Caching src/Diffusion/VAE/VAEModelBase.cs	Adds encoder/decoder CompiledModelHost instances and a structure-version counter; protected EncodeCompiled/DecodeCompiled (sync+async) route through hosts so first-call compiles and later calls replay; InvalidateVAECompiledPlans() bumps version and invalidates cached plans.
Auto-Compile-on-Eval Integration src/NeuralNetworks/NeuralNetworkBase.cs	AutoCompileOnEval property; SetTrainingMode(false) sets pending flag when enabled; Predict checks pending flag and attempts CompileForward once before eager forward (exceptions ignored).

Sequence Diagrams

sequenceDiagram
  participant Client
  participant DiffusionModel
  participant NoisePredictor
  participant CompiledHost
  participant Plan
  Client->>DiffusionModel: GenerateAsync
  DiffusionModel->>DiffusionModel: GenerateAsyncCore loop
  loop Each timestep
    DiffusionModel->>DiffusionModel: Check cancellation
    DiffusionModel->>NoisePredictor: await PredictNoiseAsync
    NoisePredictor->>NoisePredictor: ThrowIfCancellationRequested
    NoisePredictor->>CompiledHost: await PredictAsync
    CompiledHost->>CompiledHost: Pre-flight / compile
    CompiledHost->>Plan: await ExecuteAsync
    Plan-->>CompiledHost: Tensor
    CompiledHost-->>NoisePredictor: Tensor
    NoisePredictor-->>DiffusionModel: Tensor
    DiffusionModel->>DiffusionModel: Scheduler.Step
  end
  DiffusionModel-->>Client: final Tensor

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related issues

• perf: adopt Tensors-side compiled plans across DiffusionModel.Generate, VAE, VLM, and LoRA forward paths #1273 — Adds async compiled-plan surfaces and chaining through diffusion/VAE/compiled hosts, matching this PR's implemented surfaces.
• refactor: migrate VAE/conditioner/SDXL composites to multi-host CompiledModelHost so they can consume Tensors ChainAsync (gated on AiDotNet.Tensors#296) #1272 — Introduces async compiled-host APIs and PredictNoiseAsync hooks used by chained diffusion pipelines, aligned with these changes.

Possibly related PRs

• ooples/AiDotNet#1107 — Modifies compiled-model/compilation infrastructure; related to PredictAsync/compile-cache changes.
• ooples/AiDotNet#1142 — Touches compiled-model inference and compile-time caching; related at the code-level.
• ooples/AiDotNet#1143 — Extends the compiled-model and diffusion classes; closely related.

Suggested labels

feature, diffusion

Poem

🌊 Async tides now wash through diffusion,
Cache plans bloom where first calls trace,
Cancellation tokens steer the course true,
Schedulers step through denoised grace,
From eager fallback's safety net to compiled velocity. ⚡

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately reflects the main change: introducing true-async capabilities and compile-host adoption across diffusion, VAE, and related infrastructure.
Docstring Coverage	✅ Passed	Docstring coverage is 92.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feat/1273-tensors-compiled-plans

---

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

[coderabbitai]

Actionable comments posted: 9

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

src/NeuralNetworks/NeuralNetworkBase.cs (1)

3321-3356: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Gate auto-compile scheduling on real training→eval transitions.

The current condition at Line 3353 schedules prewarm on every SetTrainingMode(false) call, even if the model was already in eval mode. Combined with Predict’s temporary mode flip, this can trigger repeated prewarm attempts instead of one-shot behavior.

Proposed fix

     public virtual void SetTrainingMode(bool isTraining)
     {
         if (SupportsTraining)
         {
+            bool wasTraining = IsTrainingMode;
             IsTrainingMode = isTraining;
             // Propagate to stateful layers...
             for (int i = 0; i < _layers.Count; i++)
             {
                 _layers[i].SetTrainingMode(isTraining);
             }

             // Auto-compile-on-eval...
-            if (!isTraining && _autoCompileOnEval)
+            if (wasTraining && !isTraining && _autoCompileOnEval)
             {
                 _pendingAutoCompileOnNextPredict = true;
             }
         }
     }

🤖 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/NeuralNetworks/NeuralNetworkBase.cs` around lines 3321 - 3356, The
SetTrainingMode method is scheduling _pendingAutoCompileOnNextPredict whenever
isTraining is false even if the model was already in eval; change the condition
so the auto-compile flag is set only on a real transition from training to eval
(i.e., when IsTrainingMode was true and isTraining is false) and still gated by
_autoCompileOnEval; update the block inside SetTrainingMode (referencing
IsTrainingMode, _autoCompileOnEval, and _pendingAutoCompileOnNextPredict) to
check the previous training state before setting the pending auto-compile.

🤖 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`:
- Around line 9-11: Remove the duplicate PackageVersion entry for Snappier by
deleting the earlier <PackageVersion Include="Snappier" Version="1.3.1" />
declaration and keep the existing single PackageVersion Include="Snappier"
Version="1.3.1" entry already present elsewhere to avoid duplicate central
package management entries.

In `@src/Diffusion/DiffusionModelBase.cs`:
- Around line 243-260: The XML remarks on DiffusionModelBase currently contain
roadmap/placeholder text describing a future change and stating that
Generate(int[], int, int?) is wrapped in Task.Run, which is inaccurate; update
the public docs to describe the current async behavior instead by removing the
roadmap/future-enhancement paragraphs and replacing them with a brief, accurate
description that GenerateAsyncCore(...) provides the async implementation
(mention GenerateAsyncCore and the synchronous Generate(...) only as implemented
today), ensuring no "future" or Task.Run placeholder language remains in the
public remarks.

In `@src/Diffusion/LatentDiffusionModelBase.cs`:
- Around line 563-585: The async override GenerateAsync currently delegates
directly to GenerateAsyncCore and thus returns latent-space tensors or allocates
at pixel shape, causing length mismatches with PredictNoise and skipping VAE
decode; change GenerateAsync to mirror the synchronous Generate contract: detect
if the incoming shape is an image/pixel shape (e.g., [B,C,H,W]) and
convert/compute the corresponding latent shape (as done in Generate), allocate
initialSample in latent-space (or pass null) and call GenerateAsyncCore with
that latent shape, then after awaiting the latent output run the same VAE
decode/transform-to-pixel logic the sync Generate uses (or call the existing VAE
decode helper/override) and return the decoded pixel Tensor<T>, ensuring
cancellationToken is forwarded and any length/channel checks align with
PredictNoise's expected latent channels.

In `@src/Diffusion/NoisePredictors/NoisePredictorBase.cs`:
- Around line 245-249: Ensure the async compile-host path is guarded by disposal
checks: call the base-class ThrowIfDisposed() at the start of
PredictCompiledAsync (before accessing _compileHost.PredictAsync) and likewise
in PredictNoiseAsync (and the other similar methods referenced around the
513-534 region) so that disposed instances throw the documented
ObjectDisposedException instead of surfacing downstream errors from _compileHost
or cache; locate these in the class by the method names PredictCompiledAsync,
PredictNoiseAsync and the field _compileHost.PredictAsync and add the
ThrowIfDisposed() invocation before any use of _compileHost.

In `@src/Diffusion/VAE/VAEModelBase.cs`:
- Around line 66-68: The setters that flip the TilingEnabled and SlicingEnabled
booleans in VAEModelBase currently only toggle flags and must also invalidate
any cached/compiled plans so stale graphs aren't reused; update the
TilingEnabled and SlicingEnabled property setters (and the equivalent setters
mentioned further down in the file) to clear or invalidate the plan
cache/compiled plan objects (e.g., set encodePlan/decodePlan/planCache to null
or call an existing InvalidatePlans()/ClearCompiledPlans() method) so future
encode/decode calls recompile against the new graph.
- Around line 70-74: The new fields _encoderCompileHost and _decoderCompileHost
in VAEModelBase<T> are owned disposable resources but are not released; update
the Dispose(bool disposing) implementation inside class VAEModelBase<T> to call
Dispose() (or DisposeAsync if appropriate) on _encoderCompileHost and
_decoderCompileHost when disposing is true, then set those fields to null to
prevent double-dispose; ensure you check for null before disposing and preserve
existing disposal ordering and exception-safety in Dispose(bool).

In `@src/NeuralNetworks/CompiledModelHost.cs`:
- Around line 395-409: The async catch in PredictAsync currently falls back to
eagerForward() without invalidating the compiled plan; update the catch block
handling (the Exception filter path in PredictAsync) to clear or dispose the
cached compiled plan and reset _lastCompiledVersion (same semantics as the sync
path that clears _cache and sets _lastCompiledVersion) before calling and
returning eagerForward(), ensuring a poisoned plan isn't reused on subsequent
calls.
- Around line 365-369: The current PredictAsync returns preloadedResult
synchronously when TryUseDiskCachedPlan(...) succeeds, which bypasses
async/cancellation; refactor by splitting TryUseDiskCachedPlan into a "get plan"
helper (e.g., TryGetDiskCachedPlan or GetDiskCachedPlanAsync returning the
plan/meta) and use the plan's async execution path inside PredictAsync so you
await ExecuteCompiledPlanAsync (or similar) instead of returning a sync
preloadedResult; also preserve setting _lastCompiledVersion when the plan is
used and ensure cancellationToken is passed through to the async execution.

In `@src/NeuralNetworks/NeuralNetworkBase.cs`:
- Around line 2537-2541: The empty catch in the auto-compile prewarm (inside
NeuralNetworkBase where _pendingAutoCompileOnNextPredict is handled) is
swallowing fatal exceptions from CompileForward; change it to catch Exception
ex, detect and rethrow fatal exceptions (e.g., OutOfMemoryException,
StackOverflowException, ThreadAbortException, AccessViolationException, etc.)
while only swallowing or logging non-fatal exceptions so the process doesn't
continue in an unsafe state; you can add a small helper like
IsFatalException(Exception) and call it from the catch for clarity and to
preserve the existing fallback-to-eager behavior.

---

Outside diff comments:
In `@src/NeuralNetworks/NeuralNetworkBase.cs`:
- Around line 3321-3356: The SetTrainingMode method is scheduling
_pendingAutoCompileOnNextPredict whenever isTraining is false even if the model
was already in eval; change the condition so the auto-compile flag is set only
on a real transition from training to eval (i.e., when IsTrainingMode was true
and isTraining is false) and still gated by _autoCompileOnEval; update the block
inside SetTrainingMode (referencing IsTrainingMode, _autoCompileOnEval, and
_pendingAutoCompileOnNextPredict) to check the previous training state before
setting the pending auto-compile.

🪄 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: ea921a90-7b6d-49f5-9bc6-45157cee3e9e

📥 Commits

Reviewing files that changed from the base of the PR and between a10cdca and 73ad26a.

📒 Files selected for processing (8)

• Directory.Packages.props
• src/Diffusion/DiffusionModelBase.cs
• src/Diffusion/LatentDiffusionModelBase.cs
• src/Diffusion/NoisePredictors/NoisePredictorBase.cs
• src/Diffusion/VAE/VAEModelBase.cs
• src/NeuralNetworks/ChainedCompiledModelHost.cs
• src/NeuralNetworks/CompiledModelHost.cs
• src/NeuralNetworks/NeuralNetworkBase.cs

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[coderabbitai]

Actionable comments posted: 2

🤖 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 `@src/Diffusion/DiffusionModelBase.cs`:
- Around line 274-358: GenerateAsyncCore duplicates input validation,
element-count checks, initial-sample handling and NaN/Inf sanitization also
present in Generate; extract those into shared private helpers and call them
from both paths: add a private void ValidateGenerateInputs(int[] shape, int
numInferenceSteps, out long totalElements) that performs null/empty checks,
positive-dimension checks and computes/validates totalElements, and a private
int SanitizeNonFiniteElements(Vector<T> sample) that replaces NaN/Inf with
NumOps.Zero and returns sanitizedCount; then update GenerateAsyncCore and
Generate to call ValidateGenerateInputs at the start, reuse totalElements for
initial sample length checks and sampling, and call SanitizeNonFiniteElements
instead of inlining the sanitization loop, ensuring references to
GenerateAsyncCore, Generate, _scheduler.SetTimesteps, SampleNoise,
PredictNoiseAsync, and _scheduler.Step remain unchanged.
- Around line 360-374: Update the XML doc comment on PredictNoiseAsync to
explicitly state that the conditioning parameter is ignored by the default
implementation (which delegates to PredictNoise(noisySample, timestep)) and is
present to support future/overriding implementations (e.g., subclasses of
NoisePredictorBase<T>) that accept conditioning; mention GenerateAsyncCore
currently passes null and subclasses should override PredictNoiseAsync to handle
conditioning when needed.

🪄 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: bbe0035b-36fc-48d0-ae1f-6e9d2ca44ad3

📥 Commits

Reviewing files that changed from the base of the PR and between b8c434b and 2e05098.

📒 Files selected for processing (1)

• src/Diffusion/DiffusionModelBase.cs

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.