[TRTLLM-12467][feat] EPD improvements by venkywonka · Pull Request #13864 · NVIDIA/TensorRT-LLM

venkywonka · 2026-05-07T20:05:33Z

Summary

This PR restructures the multimodal disaggregated (Encoder/Prefill/Decode) serving pipeline to properly track encoder-output embedding lengths separately from prompt-side MM token counts.

Core problem: Previously multimodal_lengths (prompt-side counts that may include special/framing tokens) was used to split encoder outputs, causing shape mismatches (TRTLLM-12175).

Solution: Introduce multimodal_embedding_lengths as a first-class field throughout the pipeline:

# Before: tuple return, no embedding lengths
def get_prompt_token_ids(inputs, handles)
  -> (ids, lengths, positions)
# After: typed dataclass with embedding lengths
def build_disagg_prefill_multimodal_inputs(
    inputs, handles
) -> DisaggPrefillMultimodalInputs
  # .multimodal_lengths = prompt spans
  # .multimodal_embedding_lengths
  #   = encoder output rows

Key changes:

New DisaggPrefillMultimodalInputs dataclass replaces (ids, lengths, positions) tuple
MultimodalResult now carries sparse mm_embedding_request_indices for mixed text/MM batches
Validation extracted from MultimodalInput into reusable standalone functions
NemotronNano gains full EPD support with NanoV2VLMultimodalEncoder
Qwen3VL video+image placeholder expansion generalized
_is_disagg → _is_mm_disagg clarifies scope vs generic disagg serving
Dedicated EPD Accuracy e2e harness that runs video-mme benchmark for nano-omni-v3[fp8, nvfp4], qwen3vl-2b.

Out of Scope

The following are intentionally not addressed in this PR and are tracked as follow-up work:

EVS (Efficient Video Sampling) video token pruning + EPD — combining EVS-based video token pruning with the disaggregated Encoder/Prefill/Decode pipeline.
Video + audio + EPD — combining simultaneous video and audio modalities with the disaggregated Encoder/Prefill/Decode pipeline.
Full native E/PD multi-gpu disaggregation - E and PD currently only works if the workers are colocated and sharing same gpu. This is because, today only the cuda IPC handles are transported between E and PD workers - the actual gpu buffers aren't. Therefore the consumer uses the ipc handles to open gpu buffers owned by the producer. That read is non-garbage only if the consumer worker runs on the same device as the producer. If they are in different devices, the responsibility of moving the gpu tensors across devices is currently not done by TRTLLM natively.

Validation

gpu	model	benchmark	validation
8xH100	Nemotron Nano V3 Omni FP8	VideoMME	Threshold 62.785; AGG 69.667; E/PD 68.667
8xH200	Qwen3-VL-30B BF16	VideoMME	AGG baseline 62; E/PD 62
2xB100	Qwen3-VL-2B Instruct	VideoMME	Threshold 47.785; AGG 54.667; E/PD 54.667
2xB100	Nemotron Nano V3 Omni NVFP4	VideoMME	Threshold 61.285; AGG 68.334; E/P/D 69

Summary by CodeRabbit

Release Notes

Refactor
- Enhanced multimodal disaggregated serving infrastructure for improved embedding handling and shared tensor lifetime management across distributed execution.
- Standardized multimodal input processing APIs across model implementations.
Tests
- Added comprehensive integration tests for disaggregated multimodal serving scenarios including video processing.

venkywonka · 2026-05-13T01:55:01Z

/bot kill

venkywonka · 2026-05-13T01:55:08Z

/bot run

tensorrt-cicd · 2026-05-13T02:00:52Z

PR_Github #48062 [ kill ] triggered by Bot. Commit: 4223c2d Link to invocation

tensorrt-cicd · 2026-05-13T02:00:55Z

PR_Github #48062 [ kill ] completed with state SUCCESS. Commit: 4223c2d
Successfully killed previous jobs for commit 4223c2d

Link to invocation

tensorrt-cicd · 2026-05-13T02:01:21Z

PR_Github #48063 [ run ] triggered by Bot. Commit: 4223c2d Link to invocation

tensorrt-cicd · 2026-05-13T03:44:24Z

PR_Github #48063 [ run ] completed with state SUCCESS. Commit: 4223c2d
/LLM/main/L0_MergeRequest_PR pipeline #37895 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

venkywonka · 2026-05-13T11:39:02Z

/bot run

tensorrt-cicd · 2026-05-13T11:46:38Z

PR_Github #48174 [ run ] triggered by Bot. Commit: 4223c2d Link to invocation

tensorrt-cicd · 2026-05-13T12:30:43Z

PR_Github #48174 [ run ] completed with state SUCCESS. Commit: 4223c2d
/LLM/main/L0_MergeRequest_PR pipeline #37994 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

venkywonka · 2026-05-19T23:46:23Z

/bot kill

venkywonka · 2026-05-19T23:46:42Z

/bot run

tensorrt-cicd · 2026-05-19T23:52:37Z

PR_Github #49276 [ kill ] triggered by Bot. Commit: 312a6d3 Link to invocation

tensorrt-cicd · 2026-05-19T23:52:41Z

PR_Github #49276 [ kill ] completed with state SUCCESS. Commit: 312a6d3
Successfully killed previous jobs for commit 312a6d3

Link to invocation

tensorrt-cicd · 2026-05-19T23:53:06Z

PR_Github #49277 [ run ] triggered by Bot. Commit: 312a6d3 Link to invocation

…t lists The EPD VideoMME tests run encoder + prefill/decode co-resident on one physical GPU, so they belong in single-GPU CI lists. Move them off the multi-GPU DGX lists (4x B200 / 2x H100) onto the matching single-GPU lists, whose hardware still satisfies the tests' skip_less_device_memory (80 GB) and Blackwell/Hopper arch markers: - qwen3vl_2b_instruct, nemotron_nano_v3_omni_nvfp4: l0_dgx_b200 -> l0_b200 - nemotron_nano_v3_omni_fp8: l0_dgx_h100 -> l0_h100 Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>

Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>

…t_multimodal Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>

The FakeProcessor.__call__ in test_tokenized_multimodal_overwrites_stale_embedding_lengths returned a single pre-expansion token ([999]), but the real call_with_token_ids fast path returns the EXPANDED token ids (placeholder tokens replaced with MM feature tokens). multimodal_hashing_process computes the MM masks on that __call__ output, so the test fed a 1-token sequence into mask math that expected num_mm_tokens=[4], tripping the assertion in _find_mm_embedding_lengths_from_masks (1 != 4). Route the fake through its own expand_prompt_token_ids_for_mm hook so it returns [10, 101, 102, 200, 103, 20], mirroring the real path; the stale multimodal_embedding_lengths=[999] is then correctly recomputed to [3]. Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>

…-only payload filter, alias identity - llm.py: reject multimodal E/P disagg early for the AutoDeploy backend - ad_executor.py: revert disagg-only MM helpers, keep multimodal_embedding_lengths reserved - modeling_utils.py: match vision-encoder registration by class identity, not name - model_engine.py: gate mm_encoder_only requests by real payload (multimodal_embedding_lengths) - test_engine.py: drop AD tests covering the reverted disagg handoff helpers Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>

…il-fast) LlavaNext declared MM E/P disaggregated support (@support_multimodal_disaggregated + build_disagg_prefill_multimodal_inputs) but LlavaNextModel.forward() raises NotImplementedError on the disagg branch, so a disagg request was accepted at the llm.py gate and then crashed in forward. Drop the decorator and the prefill-input builder so disagg requests are rejected cleanly at the API gate; LlavaNext stays aggregated-only (forward NotImplementedError remains a backstop for the env-var path). Completing or documenting LlavaNext disagg is tracked in TRTLLM-13140. Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>

…beddings Supersede the fail-fast revert: keep LlavaNext's disagg support surface (@support_multimodal_disaggregated + build_disagg_prefill_multimodal_inputs) and replace the forward() NotImplementedError gate with the attached-embeddings path. On the E/P prefill worker (DISAGG, mm_encoder=None) the encoder already ran, so forward() now consumes the handed-off embeddings via get_attached_multimodal_embeddings() before find_input_mm_embeds slicing -- mirroring the established qwen2vl/nemotron disagg idiom. Aggregated path unchanged. Advances TRTLLM-13140. Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>

…presence Replace the module-level DISAGG env-flag checks in LlavaNextModel with a runtime decision: the constructor uses _is_mm_disagg() and the forward branches on whether self.mm_encoder is set. This drops the stale module-level constant and lets the prefill path follow the actual encoder state instead of a snapshot of the env var taken at import time. Restructure test_single_request_chat_multiple_images so the disagg prefill LLM is built with env_overrides TLLM_MULTIMODAL_DISAGGREGATED=1 (mm_encoder=None), exercising the real prefill-worker path that consumes handed-off embeddings, with a separate encoder-loaded reference for the output comparison. Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>

Parametrize the standalone encode/prefill-decode test with NVIDIA-Nemotron-Nano-12B-v2-VL-FP8 (no_pd_disagg and pd_disagg), covering the nemotron-nano-VL E/P path enabled in this PR. The FP8 12B variant fits encoder + prefill + decode on a single H100, and the case is gated to Hopper+ (SM90) so non-Hopper runs skip cleanly. Run it in post-merge on l0_h100 and deselect it from the pre-merge multimodal glob, since the larger model is not needed on every pre-merge. Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>

The rebase onto upstream/main merged upstream's new test_serialize_item_* cases (which call serialize_item) into this file, but the import-block conflict was resolved to the PR's pre-rebase form that didn't import serialize_item. Restore the import so the merged file is consistent (fixes ruff F821 undefined-name). Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>

The llms_and_encoder fixture built the MultimodalEncoder without trust_remote_code=True, so the nemotron_nano_v2_vl_fp8 cases failed at executor init with ValueError (NVIDIA-Nemotron-Nano-12B-v2-VL-FP8 ships custom HF code). Every other model/encoder creation in this file already passes trust_remote_code=True; align this fixture too. Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>

The pd_disagg variant of test_single_request_chat_multiple_images holds ref_llm + llm_decode + encoder live from the llms_and_encoder fixture and then builds a prefill_llm worker in the test body. Three simultaneous 12B FP8 instances exceed a single 80GB H100 during KV-cache estimation, so the pd_disagg case OOMs at fixture setup. Skip the pd_disagg variant for NVIDIA-Nemotron-Nano-12B-v2-VL-FP8 on single-GPU CI; the no_pd_disagg encode path (which loads two instances and fits) continues to provide nemotron-nano-VL E/P coverage. Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>

modeling_step3p7vl.py was missed when the helper was centralised and renamed from _is_disagg() to _is_mm_disagg() in modeling_multimodal_utils. The stale import caused an ImportError at package init time, breaking all pytest collection. Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>

…_payload_keys Two fixes for test_mm_encoder_standalone timeouts (RC-1): 1. model_engine.py: Replace `get_multimodal_embedding_lengths() is None` guard with `not _has_mm_payload_keys()` in `_forward_step_mm_encoder_only`. The payload-keys check is semantically correct -- it distinguishes mrope-only warmup requests from real vision-content requests directly, rather than relying on the proxy `multimodal_embedding_lengths` key which is set for all real image inputs by registry.py regardless. 2. test_mm_encoder_standalone.py: Move `prefill_llm` creation from the test body into the `llms_and_encoder` module-scoped fixture's `_instantiate_models` call. Previously, each of the 10 (model_dir x pd_disagg) module contexts created `prefill_llm` serially inside the test (~90s each), totalling ~15 minutes against an 8.7-minute CI stage limit. Building it in the fixture together with ref_llm / llm_decode / encoder collapses the cumulative init time to a single parallel pass. Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>

… encoder-only filter Two regressions from af6b39e: 1. test_mm_encoder_standalone.py: prefill_llm creation was moved into _instantiate_models (ThreadPoolExecutor), which deadlocks when TLLM_MULTIMODAL_DISAGGREGATED=1 triggers subprocess spawning that depends on main-thread process state. Fix: create prefill_llm sequentially after the parallel _instantiate_models call. The module-scoped fixture still amortizes init cost across tests. 2. model_engine.py: _forward_step_mm_encoder_only dropped the `multimodal_embedding_lengths is None` guard when switching to the _has_mm_payload_keys filter. A request with vision payload keys but no pre-computed lengths would append None to mm_embedding_lengths, causing sum(None) TypeError in split_lengths. Restore the None guard as a second filter after _has_mm_payload_keys. Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>

The EPD prefill-worker additions to test_mm_encoder_standalone.py (a dedicated prefill_llm built with env_overrides={TLLM_MULTIMODAL_DISAGGREGATED:1}) exposed a pre-existing parent-process os.environ leak in LLM._process_env_overrides, which deterministically failed test_single_request_chat_multiple_images[qwen3_30b_a3b_fp8-no_pd_disagg] on H100 and B200 (ref_llm built with mm_encoder=None). Revert test_mm_encoder_standalone.py to its pre-EPD (merge-base) state and drop the now-dangling l0_h100 nemotron_nano_v2_vl_fp8 entry. The pre-existing qwen3_30b_a3b_fp8 coverage in l0_b200 is retained. EPD product code is unchanged. Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>

venkywonka · 2026-06-08T17:27:35Z

/bot run

tensorrt-cicd · 2026-06-08T17:33:14Z

PR_Github #52829 [ run ] triggered by Bot. Commit: 77b0714 Link to invocation

tensorrt-cicd · 2026-06-09T01:06:36Z

PR_Github #52829 [ run ] completed with state SUCCESS. Commit: 77b0714
/LLM/main/L0_MergeRequest_PR pipeline #42082 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

venkywonka · 2026-06-09T01:32:36Z

/bot run

tensorrt-cicd · 2026-06-09T01:44:24Z

PR_Github #52905 [ run ] triggered by Bot. Commit: 77b0714 Link to invocation

tensorrt-cicd · 2026-06-09T07:06:27Z

PR_Github #52905 [ run ] completed with state FAILURE. Commit: 77b0714
/LLM/main/L0_MergeRequest_PR pipeline #42154 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

venkywonka · 2026-06-09T13:32:38Z

/bot run

tensorrt-cicd · 2026-06-09T13:39:12Z

PR_Github #53071 [ run ] triggered by Bot. Commit: 77b0714 Link to invocation

tensorrt-cicd · 2026-06-09T16:13:28Z

PR_Github #53071 [ run ] completed with state SUCCESS. Commit: 77b0714
/LLM/main/L0_MergeRequest_PR pipeline #42283 completed with status: 'SUCCESS'

CI Report

Link to invocation

The rc18->rc19 rebase kept our rc18-based modeling_kimi_k25.py wholesale, which missed rc19's own evolution of the (shared-lineage) file: - NVIDIA#13864 (EPD): get_prompt_token_ids() -> build_disagg_prefill_multimodal_inputs() returning DisaggPrefillMultimodalInputs. rc19's inputs/registry.py now REQUIRES this method for multimodal disagg; our old file only had the removed get_prompt_token_ids -> disaggregated multimodal Kimi would fail registry validation. - NVIDIA#14741: removed DI's lazy slow-tokenizer shim as redundant. - NVIDIA#15219: native vocab_size_padded delegating property. Re-base our file on rc19's modeling_kimi_k25.py and re-apply only the DI deltas that rc19 lacks: - tokenizer denylist discard ("kimi_k25") for correct gapped special-token ids (kept even though the slow-tokenizer shim was dropped; harmless no-op if upstream resolves the denylist entry); - root-logger filter silencing the per-encode "Calling super().encode" spam; - vision-tower MLP row chunking (TLLM_KIMI_VISION_MLP_CHUNK) to bound peak activation on large images; - keep MoonViT3d LayerNorm params in float32 (TRT-LLM LayerNorm.forward upcasts activations -> bf16 weights crash F.layer_norm); - data.image placeholder deficit top-up (drop the now-removed _ensure_k25_slow_tokenizer call); - spec-decode delegations rc19 lacks (spec_worker, spec_config, set_guided_decoder; rc19 already forwards the rest); - preserve original_input_ids across fuse_input_embeds for the Eagle3 drafter. Net +110 lines over rc19 base. Aggregated and disaggregated multimodal paths both supported now. Needs a vision smoke before prod. Signed-off-by: Bo <bo@deepinfra.com>

github-actions Bot assigned venkywonka May 7, 2026

venkywonka force-pushed the venky/epd-for-video branch from d76c722 to e0b5f70 Compare May 7, 2026 20:11

venkywonka mentioned this pull request May 11, 2026

[TRTLLM-12417][feat] Exact multimodal cache hashing + EPD unlock for videos #13724

Closed

1 task

venkywonka force-pushed the venky/epd-for-video branch from e0b5f70 to 0732119 Compare May 12, 2026 23:39

venkywonka force-pushed the venky/epd-for-video branch from 4223c2d to bd4556e Compare May 13, 2026 17:09