[Perf] Add reusable exact projection caching with MiniMax-H3 integration (#7987)

Signed-off-by: lishunyang12 <lishunyang12@users.noreply.github.com>
Signed-off-by: Armaan Amatya <armaanamatya2014@gmail.com>
Co-authored-by: lishunyang12 <lishunyang12@users.noreply.github.com>
Co-authored-by: Armaan Amatya <38102001+armaanamatya@users.noreply.github.com>
This commit is contained in:
authored and GitHub committed 2026-09-24 23:30:42 +08:00
1 parent 63bbd19a1d
commit 535bd35b42
11 files changed
+1886 -3

No files matched your search

@@ -57,3 +57,58 @@ cross-stage routing.
Test registry selection, checkpoint loading, minimal inference, input and output
contracts, and every declared optional capability.
## Exact conditioning projection reuse
Models with repeated conditioning projections can use `ExactProjectionCache`
from `vllm_omni.diffusion.cache`. It stores actual projection outputs for identical
conditioning tensors and unchanged weights, with a byte budget and TP-wide hit
coordination. It has no model names, sampling schedules, projection layouts, or
artifact formats built in. MiniMax-H3 is its first model integration.
Keep one cache per model instance and prepare each conditioning tensor once
before its projections. For a leaf linear with fixed SiLU preprocessing:
```python
from torch.nn.functional import silu
from vllm_omni.diffusion.cache import ExactProjectionCache
cache = ExactProjectionCache(max_bytes=256 * 1024**2)
# Inside an inference forward, with gradients disabled:
cache.prepare(conditioning)
modulation = cache.project(
"block.0.modulation",
projection,
conditioning,
lambda: projection(silu(conditioning)),
)
```
The original computation must depend only on the supplied conditioning tensor,
the leaf projection's weights/buffers, and fixed preprocessing. Include prompt,
guidance, modality, or other request-dependent conditioning in the actual input;
a timestep-only key is unsafe when those affect the result. Do not mutate the
input between `prepare` and `project`, or modify outputs returned on cache hits.
Use separate cache instances for concurrent forwards with different inputs.
Call `clear()` when loading or switching adapters, changing projection behavior,
or moving the model. Weight-version checks invalidate individual entries.
Compilation, gradient-enabled execution, projection hooks, and unversioned
inference weights use the original computation. Validate cold/warm parity,
invalidation, and rank-local misses when integrating another model.
Hook checks include global hooks and hooks on wrapped submodules. Numerical
settings are captured at each projection call, so entering an autocast context
after `prepare()` cannot reuse an output from a different precision.
Validated offline results can be supplied through `_lookup_precomputed()`;
this extension is consulted only at TP1 on runtime-cache misses. Artifact
validation, weight binding, and model-specific layouts remain in the adapter.
H3's adapter retains its main/Ref2VA sidecars and their eager-only admission.
The cache retains projection results and leaves all original weights available.
It does not reduce weight residency or cancel layerwise prefetch. A future
integration that skips weight transfers must coordinate with `OffloadPlan`
before prefetch and provide a safe miss path; moving weight storage currently
causes conservative cache invalidation.