[TRTLLM-14388][refactor] BREAKING: Remove 2-model code and tests#16380
[TRTLLM-14388][refactor] BREAKING: Remove 2-model code and tests#16380mikeiovine wants to merge 3 commits into
Conversation
|
/bot run --disable-fail-fast |
|
PR_Github #59235 [ run ] triggered by Bot. Commit: |
|
PR_Github #59235 [ run ] completed with state |
4764a86 to
97e26c8
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #59758 [ run ] triggered by Bot. Commit: |
97e26c8 to
dd496c8
Compare
4807177 to
cfe1571
Compare
|
/bot run --disable-fail-fast |
cfe1571 to
fb94cd2
Compare
|
PR_Github #59775 [ run ] completed with state
|
39ad75d to
e6bb6c9
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #59996 [ run ] triggered by Bot. Commit: |
|
PR_Github #59996 [ run ] completed with state
|
e6bb6c9 to
9bec289
Compare
|
/bot run |
|
PR_Github #61388 [ run ] triggered by Bot. Commit: |
|
PR_Github #61388 [ run ] completed with state
|
Signed-off-by: Mike Iovine <miovine@nvidia.com>
Signed-off-by: Mike Iovine <6158008+mikeiovine@users.noreply.github.com> Signed-off-by: Mike Iovine <miovine@nvidia.com>
f917b0b to
eedbdef
Compare
|
/bot run |
|
PR_Github #61624 [ run ] triggered by Bot. Commit: |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/attention_backend/sparse/dsa.py (1)
1531-1541: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winPer-step pinned host buffer allocation in a hot path.
host_block_tableis freshly allocated (withpin_memory=prefer_pinned()) on every call toprepare_for_indexer_k_cache, which runs every forward step viaprepare(). Every sibling staging buffer in this class (host_indexer_k_cache_block_offsets,host_slot_mapping_fp8, etc.) is pre-allocated once increate_buffers_for_indexer/__post_init__and reused. Repeated pinned-memory allocation/free is a comparatively costly page-locking operation on a per-forward-step path.♻️ Suggested fix: pre-allocate a persistent staging buffer
- # pool_indices already has correct values; set padding to -1. - # Stage through a fresh pinned buffer: an async H2D from pageable - # memory would block the host behind the busy execution stream. - host_block_table = torch.empty((pool_indices.shape[0], max_blocks_used), - dtype=pool_indices.dtype, - pin_memory=prefer_pinned()) - host_block_table.copy_(pool_indices[:, :max_blocks_used]) + # pool_indices already has correct values; set padding to -1. + # Stage through a persistent pinned buffer (allocated once in + # create_buffers_for_indexer) to avoid a per-step pinned alloc. + host_block_table = self.host_block_table_staging[:pool_indices.shape[0], :max_blocks_used] + host_block_table.copy_(pool_indices[:, :max_blocks_used])And add a
self.host_block_table_stagingbuffer of shape[max_num_sequences, max_blocks_per_seq]alongside the other persistent host buffers increate_buffers_for_indexer.🤖 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 `@tensorrt_llm/_torch/attention_backend/sparse/dsa.py` around lines 1531 - 1541, Pre-allocate a persistent host_block_table_staging buffer in create_buffers_for_indexer (or the corresponding __post_init__ setup) with shape [max_num_sequences, max_blocks_per_seq], matching pool_indices’ dtype and pinned-memory policy. Update prepare_for_indexer_k_cache to reuse this buffer for copying and padding instead of allocating host_block_table on every call, while preserving the existing per-step contents and masking behavior.
🤖 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.
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/dsa.py`:
- Around line 1531-1541: Pre-allocate a persistent host_block_table_staging
buffer in create_buffers_for_indexer (or the corresponding __post_init__ setup)
with shape [max_num_sequences, max_blocks_per_seq], matching pool_indices’ dtype
and pinned-memory policy. Update prepare_for_indexer_k_cache to reuse this
buffer for copying and padding instead of allocating host_block_table on every
call, while preserving the existing per-step contents and masking behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ba4defc3-cdbf-42a2-905c-473d215f6d0f
📒 Files selected for processing (83)
.pre-commit-config.yamldocs/source/blogs/tech_blog/blog06_Llama4_maverick_eagle_guide.mddocs/source/features/torch_compile_and_piecewise_cuda_graph.mddocs/source/legacy/advanced/speculative-decoding.mdexamples/auto_deploy/model_registry/configs/super_v3_mtp.yamlexamples/llm-api/llm_speculative_decoding.pyexamples/llm-api/quickstart_advanced.pyexamples/models/core/deepseek_v4/README.mdexamples/models/core/exaone/README.mdexamples/models/core/mistral_large_3/README.mdexamples/models/core/nemotron/README_nemotron_super_v3.mdlegacy-files.txtpyproject.tomlruff-legacy.tomltensorrt_llm/_torch/attention_backend/interface.pytensorrt_llm/_torch/attention_backend/sparse/dsa.pytensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/auto_deploy/llm_args.pytensorrt_llm/_torch/auto_deploy/shim/ad_executor.pytensorrt_llm/_torch/models/modeling_auto.pytensorrt_llm/_torch/models/modeling_deepseekv3.pytensorrt_llm/_torch/models/modeling_glm.pytensorrt_llm/_torch/models/modeling_nemotron_h.pytensorrt_llm/_torch/models/modeling_qwen3vl.pytensorrt_llm/_torch/models/modeling_speculative.pytensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/cuda_graph_runner.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/model_loader.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.pytensorrt_llm/_torch/pyexecutor/sampler/sampler.pytensorrt_llm/_torch/speculative/__init__.pytensorrt_llm/_torch/speculative/drafter.pytensorrt_llm/_torch/speculative/drafting_loops.pytensorrt_llm/_torch/speculative/eagle3.pytensorrt_llm/_torch/speculative/interface.pytensorrt_llm/_torch/speculative/model_drafter.pytensorrt_llm/_torch/speculative/mtp.pytensorrt_llm/_torch/speculative/utils.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/integration/defs/.test_durationstests/integration/defs/.test_durations_aws_dfwtests/integration/defs/accuracy/test_disaggregated_serving.pytests/integration/defs/accuracy/test_llm_api_autodeploy.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/defs/disaggregated/test_ad_disagg.pytests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_gptoss_eagle_triton.yamltests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_gptoss_eagle_trtllm.yamltests/integration/defs/disaggregated/test_disaggregated_single_gpu.pytests/integration/defs/examples/serve/test_spec_decoding_metrics.pytests/integration/defs/examples/test_ad_speculative_decoding.pytests/integration/defs/perf/pytorch_model_config.pytests/integration/defs/perf/test_perf_sanity.pytests/integration/defs/test_e2e.pytests/integration/test_lists/qa/llm_function_core.txttests/integration/test_lists/qa/llm_function_rtx6k.txttests/integration/test_lists/qa/llm_spark_func.ymltests/integration/test_lists/test-db/l0_b200.ymltests/integration/test_lists/test-db/l0_dgx_b200.ymltests/integration/test_lists/test-db/l0_dgx_h100.ymltests/integration/test_lists/test-db/l0_gb200_multi_gpus.ymltests/integration/test_lists/test-db/l0_h100.ymltests/integration/test_lists/test-db/l0_rtx_pro_6000.ymltests/integration/test_lists/waives.txttests/scripts/perf-sanity/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yamltests/scripts/perf-sanity/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yamltests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yamltests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yamltests/scripts/perf/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yamltests/scripts/perf/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yamltests/scripts/perf/disaggregated/gb200_stress-gpt-oss-120b-fp4_8k1k_ctx1_tp1_gen1_tp4_eplb0_eagle3_ccb-NIXL.yamltests/scripts/perf/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yamltests/scripts/perf/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yamltests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.pytests/unittest/_torch/executor/test_dual_pool_kv_cache.pytests/unittest/_torch/executor/test_kv_cache_estimation.pytests/unittest/_torch/executor/test_model_loader_mx.pytests/unittest/_torch/executor/test_py_executor.pytests/unittest/_torch/executor/test_seq_slot_sizing.pytests/unittest/_torch/helpers.pytests/unittest/_torch/modeling/test_modeling_llama.py
💤 Files with no reviewable changes (34)
- tests/unittest/_torch/executor/test_dual_pool_kv_cache.py
- examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml
- tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml
- tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_gptoss_eagle_triton.yaml
- tensorrt_llm/_torch/speculative/drafting_loops.py
- legacy-files.txt
- tensorrt_llm/_torch/models/modeling_auto.py
- pyproject.toml
- tests/integration/defs/perf/pytorch_model_config.py
- examples/models/core/nemotron/README_nemotron_super_v3.md
- tensorrt_llm/usage/llm_args_golden_manifest.json
- tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_gptoss_eagle_trtllm.yaml
- tests/integration/defs/disaggregated/test_ad_disagg.py
- examples/llm-api/quickstart_advanced.py
- examples/models/core/mistral_large_3/README.md
- tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml
- tests/unittest/_torch/executor/test_kv_cache_estimation.py
- tests/scripts/perf/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml
- tests/integration/test_lists/test-db/l0_b200.yml
- tests/scripts/perf/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml
- tests/unittest/_torch/helpers.py
- tensorrt_llm/_torch/attention_backend/interface.py
- .pre-commit-config.yaml
- tests/unittest/_torch/modeling/test_modeling_llama.py
- tensorrt_llm/_torch/speculative/model_drafter.py
- tests/scripts/perf/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml
- tests/integration/defs/perf/test_perf_sanity.py
- ruff-legacy.toml
- tests/scripts/perf-sanity/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml
- tests/scripts/perf-sanity/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml
- tensorrt_llm/_torch/models/modeling_nemotron_h.py
- tests/scripts/perf/disaggregated/gb200_stress-gpt-oss-120b-fp4_8k1k_ctx1_tp1_gen1_tp4_eplb0_eagle3_ccb-NIXL.yaml
- tests/integration/defs/test_e2e.py
- tests/scripts/perf/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml
🚧 Files skipped from review as they are similar to previous changes (30)
- docs/source/blogs/tech_blog/blog06_Llama4_maverick_eagle_guide.md
- docs/source/features/torch_compile_and_piecewise_cuda_graph.md
- tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
- tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py
- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py
- tests/unittest/_torch/executor/test_seq_slot_sizing.py
- tests/integration/test_lists/qa/llm_spark_func.yml
- tensorrt_llm/_torch/auto_deploy/llm_args.py
- tensorrt_llm/_torch/speculative/drafter.py
- tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml
- tests/integration/test_lists/test-db/l0_dgx_h100.yml
- tests/integration/test_lists/qa/llm_function_rtx6k.txt
- tests/integration/defs/.test_durations_aws_dfw
- tests/unittest/_torch/executor/test_py_executor.py
- tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
- tensorrt_llm/_torch/speculative/mtp.py
- tests/integration/defs/accuracy/test_llm_api_autodeploy.py
- tests/integration/defs/examples/test_ad_speculative_decoding.py
- tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
- tensorrt_llm/_torch/models/modeling_deepseekv3.py
- tests/integration/test_lists/test-db/l0_dgx_b200.yml
- tensorrt_llm/_torch/speculative/interface.py
- tests/integration/test_lists/test-db/l0_h100.yml
- tensorrt_llm/_torch/models/modeling_speculative.py
- tests/integration/test_lists/qa/llm_function_core.txt
- tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
- tensorrt_llm/llmapi/llm_args.py
- tensorrt_llm/_torch/speculative/utils.py
- tests/integration/defs/accuracy/test_llm_api_pytorch.py
- tensorrt_llm/_torch/pyexecutor/model_engine.py
|
PR_Github #61624 [ run ] completed with state
|
Description
Remove 2-model code and related tests. We are past the deprecation period and all feature gaps have been closed.
Some 2-model code is left behind for a few niche things that can't use 1-model (ngram, user-provided).
Test Coverage
Existing tests.
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.Dev Engineer Review
ModelDrafter,MTPDraftModel, separate draft-engine orchestration, and related CUDA-graph/KV-cache handling.eagle3_one_model,mtp_eagle_one_model, draft-engine memory types, and one-model mode variants. This is API-breaking and requires downstream configuration updates.fb94cd2,e6bb6c9, andf917b0bfailed in the L0 merge-request pipeline; investigate, fix, and rerun with NVIDIA-team ownership.QA Engineer Review
Modified test code covers EAGLE3, MTP, speculative metadata, KV-cache reuse, CUDA graphs, draft-length scheduling, disaggregated serving, AutoDeploy, model loading, metrics, and executor behavior. Obsolete two-model MTP and draft-token tree tests were removed; remaining tests were updated to use unified configuration names.
Modified test-list and QA files include:
tests/integration/test_lists/test-db/l0_b200.ymltests/integration/test_lists/test-db/l0_dgx_b200.ymltests/integration/test_lists/test-db/l0_dgx_h100.ymltests/integration/test_lists/test-db/l0_gb200_multi_gpus.ymltests/integration/test_lists/test-db/l0_h100.ymltests/integration/test_lists/test-db/l0_rtx_pro_6000.ymltests/integration/test_lists/qa/llm_function_core.txttests/integration/test_lists/qa/llm_function_rtx6k.txttests/integration/test_lists/qa/llm_spark_func.ymltests/integration/test_lists/waives.txtEntries were renamed or replaced to remove
one_model/two_modelvariants, and obsolete two-model coverage was removed. Test paths and parameter formats should be validated against collected pytest node IDs, with no duplicate or unintentionally lost coverage.Verdict: needs follow-up pending investigation of the failed L0 CI runs and confirmation of CBTS coverage.