[kunkunblueberry]

Summary

This PR has gone through several iterations around MammothModa2 AR -> DiT hidden-state transfer. The current final direction is intentionally narrower than the earlier request-end/full-payload experiments:

• keep the standard completed-AR -> DiT forwarding path;
• preserve AR hidden states in their source BF16/FP16 dtype across the existing EngineCore payload boundary;
• defer conversion until the DiT consumer moves the selected conditions to the model device/dtype;
• add focused regression coverage for the dtype/payload contract and DiT row-selection behavior.

The latest submitted code should be reviewed as a measured Layer-1 payload optimization, not as a new streaming transport runtime.

Related issues: #7075, #7087

Current code scope

The final precision optimization changes only MammothModa2 bridge code and tests:

Area	Change
vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py	Stop expanding full_hidden_states to FP32 in ar2dit; keep the AR source dtype and make the payload contiguous.
vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py	Keep selected text/image condition rows in the received dtype until the existing final .to(device=model_device, dtype=target_dtype) conversion.
tests/model_executor/stage_input_processors/test_mammoth_moda2.py	Cover FP16/BF16 EngineCore payload round trips, condition split dtype preservation, row alignment, and the standard completed-AR forwarding contract.

No benchmark scripts, raw logs, Nsight traces, or local evidence archives are committed in this update.

Why the earlier broader path was not retained

Earlier versions of this PR explored a request-end full-payload path. That direction attempted to retain AR hidden-state slices on the producing device, accumulate them across a request, and forward a request-end payload to DiT through a fuller cross-stage handoff.

That exploration was useful, but it crossed from a transfer optimization into request lifecycle design. In particular, it touched assumptions around:

• when a request is considered complete by AR;
• when DiT is allowed to consume the payload;
• how request identity is preserved across stage boundaries;
• multi-request scheduling and replay;
• cancellation and cleanup;
• buffer ownership and reuse.

The lifecycle risk is not theoretical. If a producer-side or streaming payload buffer is reused before all async copy/consumer dependencies are resolved, then a later request can observe stale readiness or consume another request's data. That class of problem needs an explicit Layer-2 contract, not a local patch hidden inside the MammothModa2 adapter.

The PR was therefore pulled back to the standard AR -> DiT forwarding path. The final retained optimization is only the precision/payload-volume point that has direct evidence and a much smaller compatibility surface.

Data-path change

Before this update, the effective handoff expanded the AR tensor before transport:

AR BF16/FP16 hidden states
  -> FP32 expansion in ar2dit / condition split
  -> EngineCore payload
  -> CPU decode
  -> DiT condition split
  -> model dtype

The current path is:

AR BF16/FP16 hidden states
  -> ar2dit keeps source dtype
  -> EngineCore raw-byte payload with dtype/shape metadata
  -> CPU decode preserving dtype
  -> DiT selects logical text/image rows
  -> final model-device/model-dtype conversion

This changes the transport representation only. It does not change token order, answer boundary, generated-token look-ahead handling, masks, or the arithmetic dtype used by the DiT model.

Correctness contract

The payload still represents a complete logical AR sequence:

Field	Contract
hidden states	2-D [logical_token_rows, hidden_size]
dtype	AR source dtype, BF16/FP16 in the tested paths
layout	contiguous before EngineCore serialization
token IDs	prompt IDs followed by generated IDs, excluding the final sampled look-ahead token
boundary	answer_start_index = len(prompt_token_ids)
text condition	question-side rows excluding visual/generative token IDs
image condition	answer-side rows with token_id >= gen_vocab_start_index

The EngineCore additional-information serializer already stores tensor payloads as raw bytes plus shape and dtype metadata, so BF16/FP16 can cross this boundary without introducing a new serializer protocol.

Prefix-cache and chunked-prefill reconstruction are deliberately not solved here. If AR execution skips cached prefix rows or splits prefill chunks, upstream must still reconstruct the complete logical hidden-state/token sequence before this MammothModa2 AR -> DiT boundary.

A/B and profiling evidence

A100 end-to-end A/B

Environment:

• 2 x NVIDIA A100-PCIE-40GB;
• Torch 2.13.0+cu132, CUDA 13.2;
• model: /data/vllm-workspace/models/MammothModa2-Preview;
• resolution: 256 x 256;
• diffusion steps: 2;
• seed: 142;
• 2 warmups, 10 timed requests;
• baseline and candidate run sequentially with the same prompt and deployment.

Metric	Baseline	Candidate	Change
mean request time	6.367506 s	6.306694 s	-60.812 ms / -0.96%
standard deviation	80.246 ms	194.437 ms	higher variance
p50 request time	6.340186 s	6.254637 s	-1.35%
p95 request time	6.496411 s	6.654906 s	+2.44%
peak GPU0 memory	38195 MiB	38195 MiB	unchanged
peak GPU1 memory	7559 MiB	7559 MiB	unchanged

All corresponding timed output image hashes matched between baseline and candidate.

Interpretation: the end-to-end result is directionally positive, but because the sample variance is non-trivial, I do not want to overclaim this as a statistically significant E2E speedup. The stronger evidence is at the payload/transfer boundary.

100-run Layer-1 payload microbenchmark

Workload:

• rows: 1824;
• prompt rows: 256;
• generated rows: 1568;
• hidden size: 4096;
• source dtype: BF16;
• 100 timed samples.

Metric	Baseline: FP32 expansion	Candidate: dtype preserved	Change
tensor wire payload	29.88 MB	14.94 MB	-50.0%
serialize p50	4.896 ms	2.402 ms	-50.9%
deserialize p50	5.099 ms	2.372 ms	-53.5%
pinned H2D p50	2.492 ms	1.226 ms	-50.8%

The payload round trip preserved shape, dtype, contiguity, and values. The text/image condition shapes were unchanged.

Nsight Systems A/B

The Nsight comparison used the same A100 host, model, prompt, seed, resolution, and two diffusion steps. It captured stage-0 and stage-1 worker activity, CUDA kernels, memcpy records, and GPU allocation records.

Capture metric	Baseline	Candidate	Change
Stage-1 pageable H2D payload	4,358,152 B	2,179,080 B	-50.0%
Total H2D bytes	4,951,236 B	2,772,164 B	-44.0%
Total D2H bytes	2,588,298 B	2,588,298 B	unchanged
GPU H2D time	4.0349 ms	3.5950 ms	-10.9%
cudaMemcpyAsync API time	176.865 ms	159.443 ms	-9.85%
Stage-0 kernel time	3510.632 ms	3483.736 ms	-0.77%
Stage-1 kernel time	92.192 ms	92.097 ms	-0.10%
GPU allocation events	8,620	8,620	unchanged
cudaStreamSynchronize count	1,084	1,084	unchanged

This confirms the payload-volume effect. It also confirms the boundary of the change: D2H is not optimized by this patch, allocation behavior is unchanged, and the production path is not streaming.

V100 compatibility proxy

A separate FP16-only proxy ran on a shared Tesla V100S-PCIE-32GB. The V100 environment had an older vLLM/vLLM-Omni runtime and no matching Mammoth checkpoint/API, so this is supporting transport evidence rather than current-branch Mammoth end-to-end evidence.

Metric	Baseline	Candidate	Change
FP16 payload	227.875 MiB	113.938 MiB	-50.0%
peak allocation delta	366.422 MiB	234.056 MiB	-36.1%
adapter mean latency	113.49 ms	56.96 ms	-49.8%
serial staging -> event-chain consumer wait	74.90 ms	17.40 ms	-76.8%

This proxy supports the precision/payload direction and the future streaming hypothesis, but it is not used as an E2E production claim.

Why streaming is deferred

The current MammothModa2 condition path is still:

AR GPU hidden states
  -> existing AR runner D2H
  -> ar2dit
  -> EngineCore additional-information bytes
  -> pageable CPU tensor after decode
  -> DiT H2D
  -> DiT

This path does not directly use SharedMemoryConnector for the Mammoth condition handoff. Replacing the default connector would therefore not fix this path without first changing the stage topology and lifecycle contract.

A production streaming design would need at least:

request_id + generation/epoch
staging-slot ownership
D2H-ready event
H2D-completion event
consumer acknowledgement/release
cancellation and failure cleanup
stale-event protection
bounded pool admission/backpressure
pageable fallback

Without those pieces, the optimization can corrupt multi-request behavior: request A can still have an async copy in flight while its staging slot is reused for request B; a delayed event from A can then mark B as ready, or B can consume A's payload. That is why I am not mixing streaming into this patch.

What this PR does not claim

This PR does not claim to implement:

• GPU-direct AR -> DiT transport;
• production streaming D2H/H2D;
• pinned staging pool ownership;
• cross-request transfer overlap;
• request cancellation/replay safety for async payload handles;
• prefix-cache hidden-state reconstruction;
• chunked-prefill hidden-state reassembly;
• scheduler or placement optimization.

Those are follow-up lifecycle/system problems. The current transfer-stage work found one low-risk optimization point with clear evidence; other candidates were either not on the current path, crossed into lifecycle architecture, or had poor expected return relative to risk.

Validation

Local/static:

python -m compileall on changed Python files
 git diff --check
pre-commit run --files on the changed runtime/test files

Relevant pre-commit hooks passed:

debug statements
end-of-file fixer
mixed line ending
trailing whitespace
ruff check
ruff format
typos
mypy-3.10
test marks
SPDX headers
forbidden imports
torch.cuda API guard
Buildkite no-files
suggestion

actionlint and markdownlint-cli2 were skipped for the focused local run because this final patch does not include workflow, docs, README, or benchmark files. The first full pre-commit attempt was blocked by local dependency/environment setup rather than by a code failure.

Functional/performance validation is summarized above. The benchmark scripts, raw logs, image outputs, Nsight traces, and local evidence archive were kept outside the commit and are available locally if deeper forensic review is needed.

AI assistance disclosure

I used AI assistance to inspect the code path, compare alternative designs, draft tests, summarize benchmark evidence, and prepare this self-review. The final patch is intentionally small and should still be reviewed normally by maintainers.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/input_output_modality_contracts.md.

Module owners: @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks

@kunkunblueberry, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit 8a72f7fa4efa produced:

• Priority: high. Prompt maintainer attention is suggested.

These are automated triage suggestions only — the final decision belongs to the maintainers.

[kunkunblueberry]

Self-review completed.

I reviewed the full MammothModa2 AR → DiT path, including the request-end hidden-state accumulation, full-payload handoff, DiT-side input handling, scheduler gating, and compatibility with the existing ar2dit() path.

For Phase 1, the main work is now complete and validated end to end: AR hidden-state slices remain device-resident during request execution, per-step host materialization is removed, and the final payload is materialized once at request end through the existing OmniPayload / transfer-manager / connector path. Focused regression tests pass (88 passed), and a 2-GPU A100 smoke test successfully completes MammothModa2 image generation.

The remaining limitation is the actual cross-stage transport. The existing generic vLLM-Omni connector path still uses host-mediated transfer; there is currently no directly reusable true GPU-to-GPU D2D transport interface that I can plug into this model path without changing the underlying communication infrastructure.

Therefore, I have intentionally kept this PR scoped to the model-side/request-end transfer optimization and left the generic connector backend unchanged. The remaining work is to profile the request-end D2H/H2D and serialization overhead, then determine whether a separate GPU-aware transport implementation (e.g. CUDA IPC / NCCL / UCX / asynchronous P2P transport) is justified.

So the current status is:

• Phase 1 AR → DiT integration: validated
• Per-step hidden-state D2H removal: implemented
• Existing connector path: preserved
• True D2D transport: not yet available in the current generic path
• Further transport optimization: pending profiling and follow-up infrastructure work

[hsliuustc0106]

Reviewed head 226af4cc5b2c01aee60908b09f4c24a4301a4205 for #7075 / v0.30.0.

The three available targeted suites passed 89 tests (stage-input processor, GPU AR runner, and scheduling coordinator). No new correctness regression was confirmed, but the lifecycle and memory tradeoffs need further validation.

[P2, inherited lifecycle follow-up] Reconcile accumulated states on preemption/resume. The GPU snapshot path feeds an append-only accumulator that survives preemption. Recomputed positions can be appended again; a method-level replay probe produced 9 hidden rows where 5 were expected. The adapter's length check rejects the payload, and the connector skips delivery. The legacy host-side path also has this alignment limitation, so this is an inherited follow-up rather than a newly introduced regression. Please add scheduler-driven preemption/resume and cancellation coverage, checking downstream completion and release of retained tensors.

A synthetic 4,160 × 3,584 BF16 trajectory using the snapshot/accumulator methods retained 28.44 MiB on GPU and peaked at 56.88 MiB during concatenation. Releasing the probe references returned memory to baseline; this does not establish engine-level completion/cancellation cleanup. Please account for active and paused requests and concurrent completion—the AR profile permits 100 sequences.

Requested A/B evaluation and detailed profiling

Please compare a pinned baseline with the final PR head on the same hardware, checkpoint revision, runtime versions, device placement, memory budgets, prompts, seeds, resolution, steps, and CFG. Keep unrelated features fixed. Include representative short/long prompts and image sizes, single-request latency, and a concurrency sweep within supported limits.

Please provide:

• Correctness comparisons of token alignment and AR→DiT conditioning, plus quantitative image comparisons; exercise completion, cancellation, and forced preemption/resume.
• Unprofiled repeated-run results with sample counts, warmup policy, variability, p50/p95 latency, throughput, and AR/transfer/DiT/E2E breakdowns.
• Per-device peak allocated/reserved memory and the retained-state timeline through generation, request-end concatenation, payload transfer, and cleanup. Explain the capacity impact separately from latency.
• Representative baseline/head Nsight Systems or PyTorch Profiler traces, commands, and a written analysis of D2H/H2D copy counts and bytes, copy/synchronization time, GPU snapshot kernels, concatenation, CPU FP32 conversion, serialization, connector queueing, and any actual compute/transfer overlap. Explain which costs were removed, deferred, or increased, and connect them to the measured E2E result.

The PR's reported command references tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py, which is absent at this head. Please correct the reproducible command or include the intended coverage.

Please also coordinate with #7119: both PRs add the same stage-input test file, and enabling prefix caching bypasses deferred D2H. Preserve both test suites and evaluate that combination separately.

[kunkunblueberry]

Performance Status

Phase 1 has now been validated end to end on 2 × A100 40GB.

The request-end implementation successfully removes the per-step hidden-state host materialization from the MammothModa2 AR path. In a worker-local PyTorch profiler capture, stage-0 aten::to calls decreased from 15,857 to 14,802, a reduction of 1,055 calls, which is consistent with the 1,057 generated-token AR workload. This provides evidence that the repeated per-step tensor materialization has been removed.

However, this result should not yet be interpreted as an end-to-end latency improvement or as elimination of the cross-stage host transfer. The existing connector path is unchanged and still performs host-mediated transport after the request-end payload is materialized.

A first Nsight comparison was collected, but the baseline and optimized captures covered unequal sets of spawned CUDA workers (cudaLaunchKernel: 79,319 vs 410,185). Therefore, the resulting D2H/H2D memcpy counts are not considered comparable performance evidence and are intentionally excluded from any speedup claim.

The current status is therefore:

• AR → DiT correctness: validated
• Request-end full-payload integration: validated
• Per-step hidden-state host materialization: removed
• End-to-end performance improvement: not yet quantified
• Existing connector host-mediated transfer: unchanged

The remaining performance work is to obtain a worker-complete, apples-to-apples Nsight trace and attribute:

1. AR hidden-state D2H at request end
2. SharedMemoryConnector transport overhead
3. DiT-side H2D
4. Total contribution of these transfers to end-to-end latency

This profiling work will determine whether the remaining connector transport is significant enough to justify a separate Phase 2 investigation of GPU-aware asynchronous/P2P transport.

[hsliuustc0106]

fix DCO

[hsliuustc0106]

is there any A/B test result?

[vllm-omni-review-bot]

Omni ReviewBot: three questions on the performance claim

@kunkunblueberry this PR reads as a performance or value claim:

• title: [Performance] Optimize MammothModa2 AR → DiT hidden-state transfer

Before the full evidence checklist, three short questions:

1. Bottleneck — what is the current bottleneck, and which profile, trace or per-stage measurement shows it?
2. Value — what does the change buy the user or the system (latency, throughput, memory, cost), and at which workload?
3. A/B or ablation — is there a same-workload, same-head/config comparison that isolates each main claim on its own? For stacked optimizations, one number per item rather than a blended delta.

When you answer, the evidence that settles it is: base and head SHA, hardware, model, workload, warm-up and repeat count, mean or percentiles with their spread, and a correctness/quality-equivalence signal; an end-to-end claim also needs stage attribution.

[kunkunblueberry]

Thanks for the detailed review. I went back through the implementation and re-evaluated the design from the architecture, lifecycle, memory, and performance perspectives rather than continuing to optimize the current path incrementally.

For this round, I decided to step back from the request-end GPU accumulation approach and restore MammothModa2 to the standard stage-forwarding behavior:

AR completes
  -> ar2dit
  -> construct the complete DiT input
  -> DiT executes once

The main reason is that Mammoth's DiT is not an incremental consumer. It needs the complete full_hidden_states, the complete token sequence, and answer_start_index together to construct the text/image conditioning and masks. So although the intermediate hidden states can theoretically be transferred incrementally, the downstream computation still has a hard dependency on the complete AR result.

I also looked carefully at the memory trade-off. The full AR result has to be retained somewhere until AR completion in either design. The previous optimization mainly changed the residency:

baseline:
GPU hidden -> CPU materialization -> CPU retained state

request-end optimization:
GPU hidden -> GPU retained state -> one final D2H

So the optimization does not eliminate the retention requirement; it moves the retention from CPU to GPU and changes when the host transfer/materialization happens. In return, it introduces additional GPU memory pressure, request lifecycle complexity, replay/preemption considerations, positional reconciliation, and more complicated failure handling.

I did measure a real local effect from the optimization: the AR-side aten::to count dropped from 15,857 to 14,802, and cumulative CPU time for those operations decreased by roughly 60–76 ms. However, the end-to-end A/B results were not stable enough to attribute a meaningful speedup to that change. Some runs showed a small improvement, some showed several seconds of regression, and the available Nsight traces did not fully cover the relevant subprocesses, so I don't think it would be technically justified to present those results as a reliable E2E performance gain.

Based on that, I decided that continuing to add lifecycle- or generic-runtime-level machinery for the request-end optimization would not have a good complexity/performance trade-off for this PR.

I therefore removed the optimization-specific pieces, including the retained GPU accumulator, positional reconciliation, the positional payload contract, the request-end BF16 payload path, and the associated tests. I also restored the scheduler/runner/connector changes so that the generic runtime no longer contains Mammoth-specific behavior.

One architectural principle I tried to follow throughout this revision is:

Mammoth-specific logic
    -> existing Omni abstraction
    -> existing generic runtime

rather than making the generic scheduler, runner, coordinator, or connector aware of Mammoth-specific semantics.

I explicitly revisited the places where it would have been possible to “fix” the problem by modifying the generic layer. In particular, the request-end implementation exposed questions around preemption/resume, retained state reconciliation, and producer-side error propagation. My conclusion was that these are not good reasons to introduce Mammoth-specific behavior into generic infrastructure when the existing framework already provides the standard completed-stage forwarding path.

So the current revision is intentionally much smaller and more conservative. The production path is back to the model's standard:

completed AR output
    -> `ar2dit`
    -> complete `OmniTokensPrompt`
    -> DiT

and the tests are focused on validating that actual topology rather than introducing a synthetic per-step DiT streaming test that does not match Mammoth's execution semantics.

For validation, I have completed the local/static checks and replaced the optimization-specific tests with focused CPU adapter/topology/forwarding coverage. I have not yet run the Linux pytest suite or the 2xA100 E2E test because the available GPU resources are currently occupied. I don't want to claim those results before actually running them.

I will run the Linux test suite first and then the fixed 2xA100 E2E validation with the same checkpoint, device placement, resolution, diffusion steps, and seed, and submit the resulting test evidence with the next update as soon as the resources are available.

For the remaining performance question, I think there is still a potentially interesting direction, but I would prefer to discuss it separately rather than expanding this PR again. Instead of retaining the whole payload on GPU, a possible future design would be to pipeline:

AR GPU
  -> async D2H
  -> CPU staging / queue
  -> immediately forward through the transport path
  -> H2D
  -> DiT-side staging

The potential benefit there would come from overlapping AR computation with data movement rather than simply reducing the number of host materializations. It may also have better concurrency and memory characteristics because the complete hidden state would not need to remain resident on the AR GPU. However, I have not assumed that the current Omni architecture supports this cleanly; it may require changes to the stage-transfer abstraction and generic runtime lifecycle.

I would therefore prefer to keep that as a separate design/performance discussion after this PR is validated, rather than introducing architectural changes without first establishing that the existing abstraction can support them.

At this point, I believe the current implementation is substantially cleaner and closer to the existing Omni architecture, and I have tried to keep the scope limited to the minimum model-specific changes required for MammothModa2. I think the main functionality is now addressed, but I would still appreciate reviewers' assessment of whether the current design and scope are appropriate before considering any further architectural changes.

[hsliuustc0106]

Re-review for #7075 / v0.30.0. Tests below ran on ae160e398b0d4dba28483f9c4c2577a038e3c762 against base 6fb7b0a05e3758e41e3871804598cea18a0b42fb. Before posting I checked the new head 68f1d9fcdf5303b0a00bf8dfbff53e70513cc739: both findings remain in the source. The new low-precision payload changes have not been rerun in this review; the earlier test counts must not be read as results for the new head.

[P1] Restore the existing Talker streaming-window handler. The scheduler deletion is still present, with the current path falling through at lines 818–819. Removing update_streaming_prompt_for_condition bypasses sequence validation, bounded-window rollover, and KV/encoder-cache release/reset for sender-only MiniCPM Talker continuations. In the overflow regression test, num_computed_tokens stays 74 instead of resetting to zero. All 27 scheduler-streaming tests pass on the tested base; nine fail on the tested PR head. Please preserve this generic runtime logic when reverting the Mammoth-specific optimization.

[P2] Fix the forwarding test's client fixture. The new head's fixture is still a SimpleNamespace without _call_custom_process_input, which the real client calls. The test raises AttributeError before reaching ar2dit. Use an appropriately initialized client or bind its real helper in the fixture.

Validation on ae160e398b0d: 36 passed, 10 failed, using Python 3.12, torch 2.13.0, vLLM 0.28.0:

PYTHONDONTWRITEBYTECODE=1 PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest \
  -p pytest_mock -p no:cacheprovider -o addopts='' -q \
  tests/core/sched/test_omni_ar_scheduler_streaming.py \
  tests/model_executor/stage_input_processors/test_mammoth_moda2.py \
  tests/distributed/omni_connectors/test_basic_connectors.py

A/B and profiling follow-up. GPU accumulation is removed, but the new head now preserves FP16/BF16 across ar2dit and condition splitting. Please evaluate this final, narrower change independently: fixed baseline/head/checkpoint, hardware/placement, runtime, prompts, seeds, resolution, steps and CFG; warmup followed by repeated unprofiled requests with variability. Include output/conditioning equivalence, CPU staging and per-device peak memory, transfer/DiT/E2E timing, and matched worker-complete traces explaining serialization, dtype conversions, D2H/H2D bytes and synchronization. Earlier accumulator measurements do not establish this new path's benefit. Please rerun the targeted suites and real-model forwarding after fixing the two regressions.

AR-runtime review from @tzhouam or @yinpeiqi would help assess the scheduler regression. Preserve both PRs' stage-input tests when integrating with #7119.

[kunkunblueberry]

Re-review update for #7075 / #7102 at head 43d1e54.

Reviewer findings

• P1 is restored, not specialized. 43d1e54 restores the generic Talker sender-only streaming-window handler that had been deleted. The scheduler diff from PR base bad5098 to 43d1e54 is empty. This preserves the existing sequence validation, bounded-window rollover, and cache release/reset path instead of adding Mammoth-specific scheduler behavior.
• P2 now exercises real client dispatch. The Mammoth forwarding test creates an uninitialized StageEngineCoreClient, sets only the two fields needed by the production method, and calls client.process_engine_inputs(...). The exercised path is process_engine_inputs -> _call_custom_process_input -> ar2dit, rather than calling the adapter through a SimpleNamespace.

Current-head validation

On the A100 test environment (Python 3.12.13, torch 2.13.0+cu132, vLLM 0.28.0), the exact reviewer-targeted command passed on this head:

PYTHONDONTWRITEBYTECODE=1 PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest \
  -p pytest_mock -p no:cacheprovider -o addopts='' -q \
  tests/core/sched/test_omni_ar_scheduler_streaming.py \
  tests/model_executor/stage_input_processors/test_mammoth_moda2.py \
  tests/distributed/omni_connectors/test_basic_connectors.py

50 passed, 16 warnings in 1.50s

The warnings are the known vLLM/vLLM-Omni version metadata warning and pytest configuration warnings, not test failures. compileall over the touched Python code and git diff --check against the PR base also passed. The exact GitHub PR head has passing DCO, pre-commit, Python 3.11 build, Python 3.12 build, and ReadTheDocs checks.

Performance evidence and its boundary

The executable contents of the two Mammoth production runtime files in this head are identical to the saved A100 profiling candidate; their only difference is two SPDX header lines per file. Reusing those artifacts is therefore source-semantic, not a claim that a fresh A/B was run.

The fixed 2x A100 measurements establish the payload effect: stage-1 pageable H2D payload -50.0%, total H2D bytes -44.0%, and GPU H2D time -10.9%. The 10-run E2E result was 6.367506 s -> 6.306694 s (-0.96%), but with substantial variance; it is directional evidence only, not a statistically significant end-to-end latency claim. Timed output hashes and peak GPU memory were unchanged.

I also attempted one current-head fixed-seed AR-to-DiT smoke on the replacement server. Its image has vLLM 0.28.0 but lacks vllm.entrypoints.generate.base.protocol, which the current Omni API imports before model initialization. The attempt stopped before CUDA/model execution. I did not alter the shared environment or claim that fresh forwarding passed; the saved candidate output and the source-equivalence check are recorded separately. A full fresh integration smoke should be rerun only in the compatible vLLM environment.

Scope

This PR preserves the existing completed-AR stage handoff and only removes the unnecessary FP32 expansion in Mammoth-specific payload handling. It does not claim streaming D2H/H2D, prefix-cache/chunked-prefill reconstruction, cross-request overlap, staging-pool ownership, cancellation handling, or async lifecycle safety. Those require a separate Layer-2 request-lifecycle design and tests.

[kunkunblueberry]

Decision-log update for head 43d1e54

This updates the earlier smoke caveat: the current head has now been exercised through a real two-stage Mammoth request.

Reviewer regression validation

The requested suites were rerun on the detached current-head worktree:

PYTHONDONTWRITEBYTECODE=1 PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest \
  -p pytest_mock -p no:cacheprovider -o addopts='' -q \
  tests/core/sched/test_omni_ar_scheduler_streaming.py \
  tests/model_executor/stage_input_processors/test_mammoth_moda2.py \
  tests/distributed/omni_connectors/test_basic_connectors.py

50 passed, 16 warnings in 1.36s

P1 remains a restoration: the scheduler has no net diff from PR base bad5098. P2 uses a real StageEngineCoreClient instance and reaches production client dispatch.

Current-head Mammoth forwarding / E2E smoke

The fixed MammothModa2-Preview request ran on 2x NVIDIA A100-PCIE-40GB with Python 3.12.13 and torch 2.13.0+cu132:

• 256x256, seed 142, two diffusion steps;
• current detached PR source used by both spawned stage workers;
• stage 0 on GPU 0 and stage 1 on GPU 1;
• both engines initialized, the request generated, and both workers shut down cleanly;
• elapsed time including cold initialization: 140.054 s;
• new output: 256x256, SHA256 eb0fb9831ad93db3968c47b7c2a16f772acd29dc70eed17fcb54f5d6aca2f570;
• the SHA256 exactly matches the saved fixed-seed candidate output.

The replacement server still carries vLLM 0.28 while this rebased Omni head targets vLLM 0.29. To avoid claiming an untested compatibility shim, the smoke used an ephemeral physical runner that changed only the saved harness import of clone_sampling_params to the equivalent copy.deepcopy function. This bypasses the new OpenAI package import unavailable in vLLM 0.28; it does not alter product source, deployment YAML, model weights, request data, stage topology, or the AR-to-DiT path. The temporary runner was deleted after the run. This is valid current-source forwarding evidence, but it is not a substitute for a future vLLM-0.29 compatibility-matrix run.

Static gates and hygiene

• compileall over the touched production and test files: passed.
• git diff --check against the PR base: passed.
• GitHub checks on this exact head: DCO, pre-commit, Python 3.11 build, Python 3.12 build, and ReadTheDocs all pass.
• I also invoked local pre-commit for the PR diff. It did not reach a hook because first-time actionlint bootstrap could not download Go modules from proxy.golang.org on this Windows host. No repository file was changed by that failed bootstrap; the exact-head GitHub pre-commit result is the passing authoritative gate.

No A/B or Nsight workload was rerun. The prior transfer measurements remain applicable because the two Mammoth runtime files have unchanged executable content relative to the saved candidate, apart from SPDX headers. No streaming, prefix-cache/chunked-prefill lifecycle, or cross-request async-safety claim is added.

[linyueqian]

Approving. The change is small and the justification is the part that matters, so I checked it rather than taking it on trust. The comment being deleted claimed the float32 widening existed because "the serializer uses numpy, which has no bf16". That is no longer true: _SAFE_TORCH_DTYPES in kv_transfer_manager.py lists bfloat16 explicitly, and this PR adds test_bfloat16_tensor_serialization proving OmniSerializer round-trips a bf16 tensor with both dtype and values intact. Pinning the assumption with a test instead of asserting it in a comment is the right way to retire a stale one.

The other half also holds. Dropping .to(dtype=torch.float32) from _split_ar_conditions is safe because forward already casts both conditions at lines 266-267 with .to(device=model_device, dtype=target_dtype), where target_dtype comes from the refiner's own parameters, so the selected rows reach the model in exactly the dtype they did before. What changes is only that they stop being widened while sitting on the host and crossing the IPC boundary, which is where the doubled staging footprint and H2D traffic came from.

CI: buildkite/vllm-omni green at this head. Intel and AMD carry the standing inherited reds, covered by ci_approval_override_head.

Validation: static review of the four-file author diff taken from the GitHub file list rather than a merge-base diff, since this branch has main merged into it and a merge-base comparison shows 85 unrelated files. Verified the serializer dtype table, the new round-trip test, and the downstream cast in forward. Fork head, so the branch was not executed.