Repository navigation
[Model] Cosmos3: skip text-length synchronization on full cache hits - #5669
EchoHayate wants to merge 3 commits into
Conversation
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
hsliuustc0106
added
diffusion
|
Updated validation with the actual dependency and a real Cosmos3-Nano run.
The PR description now replaces the stale “cache_dit is not installed” limitation with the full environment and real-model evidence. |
230da28 to
2970d4e
Compare
Signed-off-by: Allen Wu <allenwu2795@gmail.com> Co-authored-by: TRAE CLI <traecli@bytedance.com>
2970d4e to
bea9437
Compare
|
Refreshed on the vLLM 0.28 mainline and completed a fresh real-checkpoint A100 A/B.
I am treating the E2E latency result as directional rather than a stable throughput claim because this is two process-level rounds on one GPU. The PR body now contains the full workload and evidence boundary. @MaciejBalaNV @linzhenpl07 PTAL when convenient. |
Omni ReviewBot: no human activity for 12 days@EchoHayate this pull request has had no human commit, comment or review since 2026-08-27. Per repository policy it may be closed if it stays inactive. To keep it moving, any one of these is enough: push an update, reply to the open blocker, or post the current plan and timeline. |
Signed-off-by: Allen Wu <allenwu2795@gmail.com> Co-authored-by: TRAE CLI <traecli@bytedance.com>
Author self-review
Current head reviewed: @MaciejBalaNV @linzhenpl07 PTAL when convenient. A current-head A100 rerun can |
Omni ReviewBot: no human activity for 14 days@EchoHayate this pull request has had no human commit, comment or review since 2026-09-09. Please consider marking this PR as draft until work can resume. The author or a maintainer decides whether to change the PR state. To keep it moving, any one of these is enough: push an update, reply to the open blocker, or post the current plan and timeline. |
September 30 lifecycle follow-upI intend to continue this optimization, with its integration point updated The read-only audit of main Current PR head: The next bounded step is to adapt and revalidate the guard in shared No new tests, benchmarks, or CI runs were performed. The previous A100 result |
Omni ReviewBot routing recordAssigned Strict on zcode (GLM-5.3-Flash) under experiment |
vllm-omni-review-bot
left a comment
vllm-omni-review-bot
left a comment
There was a problem hiding this comment.
Omni ReviewBot review
3 actionable finding(s).
CI at
40815a48a6bc(2026-10-09T08:47:42.736076+00:00): required check(s) blocking:buildkite/vllm-omni(missing).
Full review analysis
Scan:
| Category | Result |
|---|---|
| Tests / verification | 3 finding(s) below |
| Security | no finding reported |
| Docs / comments | no finding reported |
| Behavior / compatibility | no finding reported |
| Correctness | no finding reported |
Validated:
- [resolved] SPDX flip on both :1 headers is hook-mandated at this head: tools/pre_commit/check_spdx_header.py:18-19 'Treat it as stale and rewrite it to the Omni line', :20 LEGACY_COPYRIGHT_TEXT = 'SPDX-FileCopyrightText: Copyright contributors to the vLLM project', :42 COPYRIGHT_TEXT = '...vLLM-Omni project'; residual: 9 untouched cosmos3 siblings keep the old line until next touched — no ask (round-1 no_issue stands).
- [claim-verified] 'Cosmos3 cache ownership, reset behavior, and branch keys remain unchanged' — reset_cache (transformer_cosmos3.py:1606-1608) untouched, branch-cache keys in pipeline_cosmos3.py:1301/2866-2896 untouched by the diff.
- [claim-verified] PR's own numbers internally consistent: 804.42-770.50=33.92 ms; 33.92/804.42=4.22%; per-process (801.51-767.49)/801.51=4.24%, (807.70-788.72)/807.70=2.35% — all match the description table; directional-only framing is honestly scoped.
- [claim-verified] Focused-regression claim ('2 passed') corresponds to exactly the 2 added tests; no other test file touched.
- [claim-verified] PR's own numbers check out: 804.42-770.50 = -33.92 ms = -4.22%; r1 801.51->767.49 = -4.24%; r2 807.70->788.72 = -2.35%; evidence is explicitly labeled historical (bea9437) and directional, not claimed as a current-head rerun — honest framing.
- [claim-verified] PR claim 'Focused regressions: 2 passed' maps to exactly the two added tests; consistent with their structure (autouse single-rank TP fixture :18-34, _get_ulysses_state monkeypatched to (1,0,None)); not rerun in this read-only review — static reachability analysis only.
3 actionable finding(s).
Verdict: REQUEST CHANGES
See inline comments below.
🤖 This review was generated by InferMatrix Copilot, an open-source repo-maintenance agent for PR review, CI debugging and issue triage. Try it on your own repo, and ⭐ star it if it helped!
| # cached alongside UND K/V by both the bespoke and session-state paths. | ||
| need_kv = self.cached_kv is None | ||
| if need_kv or self.cached_freqs_gen is None: | ||
| if need_kv: |
There was a problem hiding this comment.
[P0] This diff moves the identical-text-length validation out of the top of `Cosmos3…
Evidence and suggested fix
This diff moves the identical-text-length validation out of the top of Cosmos3VFMTransformer.forward (removed hunk at :1685) into the new if need_kv: branch at :1779-1787, but it cannot merge as-is: the merge gate reports DIRTY/CONFLICTING against base, and main's #6922 has since moved this exact logic into _gen_preprocess() (repo-wide grep for _gen_preprocess at 40815a4 returns 0 matches), with both text-length scalar reads still unconditional there. Rebase onto current main and re-home the guard in the shared _gen_preprocess() preprocessing, preserving (a) validation on every cache-miss/build path, (b) freqs-only reconstruction without re-validation, and coverage of both the normal and SeaCache/session entry paths; then re-run tests/diffusion/models/cosmos3/test_cosmos3_transformer.py against the moved site. The guard's behavior itself checked out on this head — full cache hits never consume text_mask (:1780/:1789 are the only forward-body uses, both in the miss/freqs branch) and every generation resets the cache (pipeline_cosmos3.py:2680-2681, :3112-3113) — so the optimization is sound once it lands at main's current location.
Evidence: Unmet requirement: the PR must apply onto its base, but main has restructured the exact code region this PR edits. PR-time tree (head 40815a4, changed by this diff): vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py:1777-1779 need_kv = self.cached_kv is None / if need_kv or self.cached_freqs_gen is None: / if need_kv: with the guard raise at :1783-1787 now miss-branch-only, and repo-wide grep for _gen_preprocess at this head: 0 matches (verified). Decisive main-side code lies OUTSIDE this PR tree, present in the PR-time checkout of refs/heads/main (verified by grep, unchanged by this diff): vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py:1856 prep = self._gen_preprocess(, :1901 def _gen_preprocess(, :1955 f"Cosmos3 requires identical real text lengths within a batch " — i.e. main #6922 moved this exact validation into _gen_preprocess(), which is why the gate step reports MERGE STATE: DIRTY/CONFLICTING and the PR cannot merge as-is. Supporting in-tree soundness facts (unchanged-context checks): transformer_cosmos3.py:1780 text_lengths = text_mask.sum(dim=1) and :1789 text_mask, are the only forward-body text_mask uses, both inside the miss/freqs branch (full cache hits never consume it); pipeline_cosmos3.py:2681 self._kv_reset_und(kv_state) and :3112-3113 self.transformer.reset_cache() / self._cosmos3_branch_caches = {} reset the cache every generation.
There was a problem hiding this comment.
Refreshed in 25e0b25c545d06ce8ed228b395a79f40cd059f37 against main
198c2105930bd1bac0b0966ffd0997cdd84d4790. The guard is now in shared
_gen_preprocess(), with regression coverage through normal forward and the
production SeaCache extractor plus actual session-adapter restoration.
One detail differs from the suggested fix: frequency-only reconstruction
must still validate. The session adapter can restore K/V without
frequencies, and RoPE then consumes the current mask. The companion finding
in thread 4228321394 identified this correctly; the new tests reproduce the
missing rejection with the inner need_kv guard. The final guard is in the
outer reconstruction branch, so full hits still avoid scalar reads.
The focused/session suite passes 29 cases. Full two-file runs retain the same
14 platform-construction failures and one accelerator skip on main and head;
the description includes the baseline comparison and static-check limitations.
AI assistance: Codex/TraeCode assisted with verification and this reply.
| # cached alongside UND K/V by both the bespoke and session-state paths. | ||
| need_kv = self.cached_kv is None | ||
| if need_kv or self.cached_freqs_gen is None: | ||
| if need_kv: |
There was a problem hiding this comment.
[P2] The moved guard runs only under if need_kv: (1779), but the enclosing branch…
Evidence and suggested fix
The moved guard runs only under if need_kv: (1779), but the enclosing branch (1778) also fires for freq-only reconstruction — cached_kv present with cached_freqs_gen None — a state this repo's own adapter restores and pins as expected behavior, not just hand-assembly: state_cosmos3_adapter.py:150-151 assigns kv and freqs_gen independently, get_branch_freqs_gen returns None when absent (:142-146), and tests/diffusion/models/cosmos3/test_session_memory_equivalence.py:107-114 asserts that exact end state on the transformer. In that cell _compute_rope_freqs still consumes the current text_mask per item (transformer_cosmos3.py:1529) while the cached K/V were trimmed to a previous batch's max_real_len (:1810), so a nonuniform mask now silently builds text positions against a differently-trimmed cache where the base tree raised. Hoist the guard to the outer branch: it then runs exactly when the mask is consumed (miss or freq-only, once per reconstruction) while the full cache-hit path — the actual optimization target, pinned by the .item() tripwire test — never enters this branch and stays sync-free. Add a test pinning the freq-only state with a nonuniform mask.
Evidence: Trigger: transformer in kv-present/freqs-None state, forward with a nonuniform text_mask. PR head: transformer_cosmos3.py:1778 if need_kv or self.cached_freqs_gen is None: with :1779 if need_kv: — guard skipped in the freq-only cell, then :1788 freqs_und, freqs_gen = self._compute_rope_freqs( consumes the current mask per item at :1529 real_len = int(text_lengths[b].item()) while :1810 self.cached_kv = [(k[:, :max_real_len], v[:, :max_real_len]) for k, v in cached_kv_full] shows the cache was trimmed to a previous batch's length; the base tree raised here before any cache logic (removed hunk: raise ValueError(f"Cosmos3 requires identical real text lengths within a batch ...")). The state is repo-supported, unchanged by this diff: state_cosmos3_adapter.py:150-151 transformer.cached_kv = self.get_branch_kv(is_negative)/transformer.cached_freqs_gen = self.get_branch_freqs_gen(is_negative)with :142-146 returning None when freqs are absent, and tests/diffusion/models/cosmos3/test_session_memory_equivalence.py:107-114test_load_into_transformer_without_freqs_sets_kvassertingassert tr.cached_freqs_gen is None` with kv set — the exact cell where base protection is lost. Full-hit path unaffected: with cached_freqs_gen present the outer branch is not entered, so the tripwire test remains green.
Suggestion: text_lengths = text_mask.sum(dim=1)
min_real_len = int(text_lengths.min().item())
max_real_len = int(text_lengths.max().item())
if min_real_len != max_real_len:
raise ValueError(
f"Cosmos3 requires identical real text lengths within a batch "
f"(got min={min_real_len}, max={max_real_len})."
)
freqs_und, freqs_gen = self._compute_rope_freqs(
There was a problem hiding this comment.
Fixed in 25e0b25c545d06ce8ed228b395a79f40cd059f37: identical-length
validation now runs under if need_kv or self.cached_freqs_gen is None,
before RoPE reconstruction in shared _gen_preprocess().
The new tests restore K/V without frequencies via Cosmos3StateAdapter,
then verify that a nonuniform mask raises the original ValueError, does
not call RoPE or publish frequencies, and preserves K/V identity. Uniform
frequency-only reconstruction preserves K/V and matches the reference
output/frequencies exactly; full-hit tripwires reject any .item() or cache
rebuild. Both normal forward and the SeaCache extractor are exercised.
Before the guard correction, the new cases had four genuine failures:
two missing exceptions and two missing pre-RoPE validation-read pairs.
The final focused/session suite passes 29 cases. The earlier statement
that frequency-only reconstruction could skip validation was incorrect and
has been corrected in the PR description.
AI assistance: Codex/TraeCode assisted with the fix, tests and this reply.
| # cached alongside UND K/V by both the bespoke and session-state paths. | ||
| need_kv = self.cached_kv is None | ||
| if need_kv or self.cached_freqs_gen is None: | ||
| if need_kv: |
There was a problem hiding this comment.
[P3] The headline perf number has no current-head evidence: the A100 medians in the…
Evidence and suggested fix
The headline perf number has no current-head evidence: the A100 medians in the description (-33.92 ms / -4.22%, and the per-process -4.24% / -2.35%) come from historical commit bea9437 and the author explicitly labels them directional, not a rerun at this head — the review notes 'current-head evidence stays open.' The kept blocker's rerun ask covers only the focused transformer regressions, not the benchmark, so this gap is not closed by the blocker as filed. Since the required rebase onto main (re-homing the guard into _gen_preprocess()) changes the code the measurement was taken against, rerun the same A100 benchmark at the rebased head and update the description's numbers; until that rerun lands, keep the table explicitly labeled historical/directional so it is not read as a measurement of the merged code.
Evidence: [claim] Historical A100 medians are self-consistent (-33.92ms / 804.42ms = -4.22%); the author explicitly labels them directional and not a current-head rerun — current-head evidence stays open and is folded into the blocker's rerun ask. Corroborated by [claim-verified]: 'evidence is explicitly labeled historical (bea9437) and directional, not claimed as a current-head rerun — honest framing.'
There was a problem hiding this comment.
The description now explicitly binds the A100 table to historical head
bea9437d / base 94b1546c, before both the shared-preprocessing integration
and the frequency-only guard correction. It is not evidence for current head
25e0b25c545d06ce8ed228b395a79f40cd059f37.
This refresh adds CPU regression evidence, not new performance measurements.
The current-head A100 benchmark and official hardware CI remain outstanding;
the historical -4.22% stays directional, with no fresh E2E speedup claim.
AI assistance: Codex/TraeCode assisted with evidence review and this reply.
Omni ReviewBot: finding feedback**[p0] This diff moves the identical-text-length validation out of the top of See the review for details. If you are the PR author and disagree, react 👎 here; the maintainer will see your disagreement. |
Merge main at 198c210. Validate text lengths for both K/V misses and frequency-only reconstruction, while full cache hits avoid scalar reads. Preserve both parents' regression coverage and exercise real forward, SeaCache extractor, and session restoration paths. Signed-off-by: Allen Wu <allenwu2795@gmail.com> Co-authored-by: TRAE CLI <traecli@bytedance.com>
EchoHayate
changed the title
Author self-review — October 12 refreshCurrent head reviewed:
The three inline review items are answered in their original threads. AI |
Omni ReviewBot triage noteAutomated triage of commit
These are automated triage suggestions only — the final decision belongs to the maintainers. |
3 participants
Purpose
Avoid redundant text-length reductions and two Python scalar reads on Cosmos3
full text-cache hits. The optimization now lives in
_gen_preprocess(), theshared preprocessing used by normal forward and the SeaCache extractor after
the upstream refactor.
Fix and boundaries
Run the existing identical-text-length validation inside
if need_kv or self.cached_freqs_gen is None, before RoPE reconstruction:restore this state, and RoPE reconstruction consumes the current mask.
and its two
.item()reads.unchanged. This is not a new cache-identity or invalidation mechanism.
This corrects the previous description's claim that frequency-only restoration
should skip validation. Preserving that behavior would silently accept
nonuniform text lengths where the base implementation raised
ValueError.The new tests cover that error through both entry paths, actual session-adapter
restoration, and unchanged K/V identity.
The scope is the existing Cosmos3 text-cache contract, not all synchronization
in preprocessing, full SeaCache skip-policy qualification, new modality support,
or distributed execution qualification.
October 12 refresh and validation
Base:
198c2105930bd1bac0b0966ffd0997cdd84d4790Head:
25e0b25c545d06ce8ed228b395a79f40cd059f37The merge preserves both parents' existing test functions, including upstream
GEN residual/sequence-parallel tests. Only the transformer and its test file
differ from the pinned main.
Focused regressions plus the entire session-adapter test file:
Result: 29 passed, 106 deselected, 17 warnings, no skips.
The new tests use deterministic synthetic one-layer UND/GEN weights and the
existing SDPA PyTorch implementation on CPU, with real preprocessing, RoPE,
extractor block/postprocessing closures and session restoration. These are not
checkpoint, GPU synchronization, or quality measurements.
The pre-fix reconstruction guard produced 4 failed / 8 passed in the new
cases: two missing
ValueErrors and two missing pre-RoPE validation-read pairs.Earlier attention/linear fixture setup failures were not counted as regression
RED.
Running the same command without
-kon both complete source trees gives:The failure identifiers match exactly. All 14 fail during
UnspecifiedOmniPlatform.get_diffusion_attn_backend_clsconstruction, beforethe changed preprocessing. The identical skip needs an accelerator. This is
a baseline comparison, not a passing full-suite result.
Changed-file pre-commit passes applicable checks except
mypy-3.10, whichreports missing annotations for
full_refsandfull_pointersin the unchangedupstream SP test. Running mypy 1.11.1 against the pristine pinned-main test file
reproduces exactly those two diagnostics. No suppressions or unrelated fixes
were added. The repository wrapper excludes production model trees from mypy;
this is not production-model type-check evidence. The scoped diff passes
git diff 198c2105930bd1bac0b0966ffd0997cdd84d4790 HEAD --check.Local environment: macOS CPU, Python 3.12, PyTorch 2.11.0, vLLM 0.31.0 source
overlay, transformers 5.14.1, diffusers 0.39.0, cache-dit 1.5.0 and pytest 9.1.1.
Pinned main requires diffusers 0.40.0; no dependencies were installed for this
refresh. These are scoped CPU results, not fully provisioned current-main CI.
Historical A100 E2E A/B — not this head
The following real-checkpoint evidence belongs to historical head
bea9437don base
94b1546c. It predates the shared-preprocessing integration and thefrequency-only guard correction, and is not a measurement of
25e0b25c.Environment:
public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:94b1546c...nvidia/Cosmos3-Nanorevision411f42a8...Fixed workload:
A red cube on a white table, studio lighting.blurry, distortedThe historical combined median moved by -33.92 ms (-4.22%). Both
process-level comparisons favored the patch (-4.24% and -2.35%). All four
generated output PNGs were bitwise identical, and reported peak GPU memory
was identical at 32,302 MiB.
This remains directional evidence from two process-level rounds on one A100,
not a stable throughput claim. A current-head A100 rerun and official hardware
CI are still outstanding; no new E2E speedup or quality claim is made here.
AI assistance
Codex/TraeCode assisted with upstream inspection, conflict resolution,
regression tests, diff review, baseline comparisons and this write-up.
The evidence above distinguishes executed CPU checks from historical hardware
results and validation still outstanding.