From d95f66541d67a0ed7cf931847ec30112fc5e41fe Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+Tabrizian@users.noreply.github.com> Date: Thu, 23 Jul 2026 16:38:19 -0700 Subject: [PATCH 01/13] [None][fix] Stage DSA indexer block_table H2D through pinned memory (#16757) Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> --- .../_torch/attention_backend/sparse/dsa.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index f422b1da1ec3..8466bbf59900 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -1530,11 +1530,17 @@ def prepare_for_indexer_k_cache(self): 1) // tokens_per_block max_blocks_used = num_blocks_per_seq.max().item( ) if self.num_seqs > 0 else 1 - # pool_indices already has correct values; set padding to -1 - host_block_table = pool_indices[:, :max_blocks_used].clone() - for i in range(self.num_seqs): - if num_blocks_per_seq[i] < max_blocks_used: - host_block_table[i, num_blocks_per_seq[i]:] = -1 + # 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]) + pad_cols = torch.arange(max_blocks_used, dtype=num_blocks_per_seq.dtype) + host_block_table.masked_fill_( + pad_cols.unsqueeze(0) + >= num_blocks_per_seq[:self.num_seqs].unsqueeze(1), -1) # Copy to GPU self.block_table[:self.num_seqs, :max_blocks_used].copy_( host_block_table, non_blocking=True) From 35134c2793454e15e58a91da85a8ec656cc94412 Mon Sep 17 00:00:00 2001 From: dpitman-nvda Date: Thu, 23 Jul 2026 21:09:39 -0400 Subject: [PATCH 02/13] [None][infra] Unwaive test_proxy_fast_death tests (#16808) Signed-off-by: Derek Pitman --- tests/integration/test_lists/waives.txt | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 79384bbde5ce..80e697faf3be 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -366,10 +366,6 @@ unittest/auto_deploy/multigpu/custom_ops SKIP (https://nvbugs/6403920) unittest/disaggregated/test_kv_transfer.py SKIP (https://nvbugs/6403793) unittest/disaggregated/test_kv_transfer.py::test_transfer_worker_v2[tp1_pp4_to_tp2_pp2] SKIP (https://nvbugs/6445316) unittest/disaggregated/test_kv_transfer.py::test_transfer_worker_v2[tp4_pp1_to_tp2_pp2] SKIP (https://nvbugs/6426834) -unittest/executor/test_proxy_fast_death.py SKIP (https://nvbugs/6499882) -unittest/executor/test_proxy_fast_death.py::test_shutdown_does_not_block_on_dead_engine SKIP (https://nvbugs/6499882) -unittest/executor/test_proxy_fast_death.py::test_shutdown_does_not_shut_down_external_session SKIP (https://nvbugs/6499882) -unittest/executor/test_proxy_fast_death.py::test_shutdown_keeps_blocking_semantics_when_engine_alive SKIP (https://nvbugs/6499882) unittest/executor/test_rpc.py::TestRpcCorrectness::test_incremental_task_async SKIP (https://nvbugs/5741476) unittest/executor/test_rpc_proxy.py SKIP (https://nvbugs/5605741) unittest/executor/test_rpc_worker.py SKIP (https://nvbugs/5605741) From e5287d6080af834aa1be1b62e9107832a1cb6391 Mon Sep 17 00:00:00 2001 From: TensorRT LLM AI Agent <296075020+trtllm-agent@users.noreply.github.com> Date: Fri, 24 Jul 2026 09:48:39 +0800 Subject: [PATCH 03/13] [None][infra] Waive 2 failed cases for main in pre-merge 49489 (#16819) Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 80e697faf3be..75133bae9fcb 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -176,6 +176,8 @@ full:B300/disaggregated/test_auto_scaling.py::test_service_discovery[http-kv_cac full:B300/disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_genpp2[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6322073) full:B300/test_e2e.py::test_qwen_e2e_cpprunner_large_new_tokens[DeepSeek-R1-Distill-Qwen-1.5B-DeepSeek-R1-Distill-Qwen-1.5B] SKIP (https://nvbugs/6414760) full:DGX_B200/accuracy/test_disaggregated_serving.py::TestQwen3NextInstruct::test_auto_dtype[use_py_transceiver=True] SKIP (https://nvbugs/6501837) +full:DGX_B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV4ProDSpark::test_gsm8k_dep8_megamoe_deepgemm SKIP (https://nvbugs/6506920) +full:DGX_H100/unittest/_torch/multi_gpu/test_linear.py::test_column_linear[2-unbalanced] SKIP (https://nvbugs/6506918) full:DGX_H200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_fp8_blockscale[disable_skip_indexer] SKIP (https://nvbugs/6476233) full:DGX_H200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_fp8_blockscale[latency_default] SKIP (https://nvbugs/6476233) full:GB200/accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy SKIP (https://nvbugs/6276923) From 3c07ada3c5af0fd31c6c91628392330ba38cc99a Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Fri, 24 Jul 2026 09:52:41 +0800 Subject: [PATCH 04/13] [TRTLLM-14474][chore] Remove legacy python relics and refresh docs after the backend removal (#16612) Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .pre-commit-config.yaml | 1074 ------------- AGENTS.md | 12 +- benchmarks/README.md | 34 - benchmarks/prepare_dataset.py | 118 -- benchmarks/utils/__init__.py | 0 benchmarks/utils/convert_nemo_dataset.py | 47 - benchmarks/utils/generate_rand_loras.py | 48 - benchmarks/utils/prepare_real_data.py | 434 ------ benchmarks/utils/prepare_synthetic_data.py | 164 -- benchmarks/utils/utils.py | 168 --- docs/source/_ext/llmapi_config_telemetry.py | 23 +- docs/source/commands/trtllm-bench.rst | 32 +- docs/source/developer-guide/perf-overview.md | 8 +- docs/source/developer-guide/telemetry.md | 304 +--- docs/source/examples/customization.md | 55 +- .../paragraf/create_standalone_package.py | 8 - .../sample_performance_alignment.sh | 5 +- examples/quantization/README.md | 258 +--- examples/quantization/quantize.py | 208 --- legacy-files.txt | 537 ------- pyproject.toml | 537 ------- ruff-legacy-baseline.json | 200 +-- ruff-legacy.toml | 537 ------- setup.py | 2 - tensorrt_llm/_deprecation.py | 64 - tensorrt_llm/bench/benchmark/__init__.py | 2 +- tensorrt_llm/bench/benchmark/utils/general.py | 8 +- tensorrt_llm/bench/build/__init__.py | 0 tensorrt_llm/bench/build/utils.py | 34 - tensorrt_llm/bench/dataset/prepare_dataset.py | 17 +- tensorrt_llm/bench/dataset/utils.py | 4 + .../bench/tuning}/__init__.py | 0 .../bench/{build => tuning}/dataclasses.py | 299 ++-- .../{build/tuning.py => tuning/heuristics.py} | 97 +- .../{build/build.py => tuning/settings.py} | 35 +- tensorrt_llm/bench/tuning/utils.py | 34 + tensorrt_llm/llmapi/__init__.py | 1 - tensorrt_llm/models/convert_utils.py | 34 - tensorrt_llm/models/unet/pp/__init__.py | 0 tensorrt_llm/quantization/__init__.py | 4 +- tensorrt_llm/quantization/image_processing.py | 97 -- .../quantization/quantize_by_modelopt.py | 1332 ----------------- tensorrt_llm/serialization.py | 1 - .../smoke/test_ad_allreduce_strategies.py | 14 +- .../singlegpu/smoke/test_ad_trtllm_bench.py | 5 +- .../others/test_quantize_calib_dataset.py | 83 - tests/unittest/tools/test_prepare_dataset.py | 11 +- 47 files changed, 397 insertions(+), 6592 deletions(-) delete mode 100644 benchmarks/README.md delete mode 100644 benchmarks/prepare_dataset.py delete mode 100644 benchmarks/utils/__init__.py delete mode 100644 benchmarks/utils/convert_nemo_dataset.py delete mode 100644 benchmarks/utils/generate_rand_loras.py delete mode 100644 benchmarks/utils/prepare_real_data.py delete mode 100644 benchmarks/utils/prepare_synthetic_data.py delete mode 100644 benchmarks/utils/utils.py delete mode 100644 examples/quantization/quantize.py delete mode 100644 tensorrt_llm/_deprecation.py delete mode 100644 tensorrt_llm/bench/build/__init__.py delete mode 100644 tensorrt_llm/bench/build/utils.py rename {benchmarks => tensorrt_llm/bench/tuning}/__init__.py (100%) rename tensorrt_llm/bench/{build => tuning}/dataclasses.py (65%) rename tensorrt_llm/bench/{build/tuning.py => tuning/heuristics.py} (79%) rename tensorrt_llm/bench/{build/build.py => tuning/settings.py} (77%) create mode 100644 tensorrt_llm/bench/tuning/utils.py delete mode 100755 tensorrt_llm/models/unet/pp/__init__.py delete mode 100644 tensorrt_llm/quantization/image_processing.py delete mode 100755 tensorrt_llm/quantization/quantize_by_modelopt.py delete mode 100644 tests/unittest/others/test_quantize_calib_dataset.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7dccd521319c..37c51fbbb377 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -9,14 +9,6 @@ common-files: &common_files | .devcontainer/make_env.py | .github/scripts/label_community_user.py | .github/scripts/pr_checklist_check.py | - benchmarks/__init__.py | - benchmarks/prepare_dataset.py | - benchmarks/utils/__init__.py | - benchmarks/utils/convert_nemo_dataset.py | - benchmarks/utils/generate_rand_loras.py | - benchmarks/utils/prepare_real_data.py | - benchmarks/utils/prepare_synthetic_data.py | - benchmarks/utils/utils.py | cpp/conanfile.py | cpp/kernels/fmha_v2/conftest.py | cpp/kernels/fmha_v2/fmha_test.py | @@ -48,24 +40,10 @@ common-files: &common_files | examples/apps/fastapi_server.py | examples/disaggregated/clients/disagg_client.py | examples/disaggregated/slurm/benchmark/submit.py | - examples/dora/normalize_weights.py | - examples/eagle/convert_checkpoint.py | - examples/eval_long_context.py | - examples/generate_checkpoint_config.py | - examples/generate_xgrammar_tokenizer_info.py | - examples/hf_lora_convert.py | examples/infinitebench/args.py | examples/infinitebench/compute_scores.py | examples/infinitebench/construct_synthetic_dataset.py | examples/infinitebench/eval_utils.py | - examples/llm-api/_tensorrt_engine/llm_eagle2_decoding.py | - examples/llm-api/_tensorrt_engine/llm_eagle_decoding.py | - examples/llm-api/_tensorrt_engine/llm_inference_customize.py | - examples/llm-api/_tensorrt_engine/llm_inference_kv_events.py | - examples/llm-api/_tensorrt_engine/llm_lookahead_decoding.py | - examples/llm-api/_tensorrt_engine/llm_medusa_decoding.py | - examples/llm-api/_tensorrt_engine/llm_quantization.py | - examples/llm-api/_tensorrt_engine/quickstart_example.py | examples/llm-api/llm_guided_decoding.py | examples/llm-api/llm_inference.py | examples/llm-api/llm_inference_async.py | @@ -85,122 +63,17 @@ common-files: &common_files | examples/llm-api/quickstart_example.py | examples/llm-api/quickstart_multimodal.py | examples/llm-api/star_attention.py | - examples/llm-eval/lm-eval-harness/lm_eval_tensorrt_llm.py | examples/longbench/eval_longbench_v1.py | - examples/medusa/convert_checkpoint.py | - examples/mmlu.py | - examples/models/contrib/baichuan/convert_checkpoint.py | - examples/models/contrib/bloom/convert_checkpoint.py | - examples/models/contrib/chatglm-6b/tokenization_chatglm.py | - examples/models/contrib/chatglm2-6b/tokenization_chatglm.py | - examples/models/contrib/chatglm3-6b-32k/tokenization_chatglm.py | - examples/models/contrib/cogvlm/convert_checkpoint.py | - examples/models/contrib/dbrx/convert_checkpoint.py | - examples/models/contrib/deepseek_v1/__init__.py | - examples/models/contrib/deepseek_v1/convert_checkpoint.py | - examples/models/contrib/deepseek_v2/convert_checkpoint.py | - examples/models/contrib/dit/convert_checkpoint.py | - examples/models/contrib/dit/diffusion.py | - examples/models/contrib/dit/sample.py | - examples/models/contrib/dit/utils_modelopt.py | - examples/models/contrib/dit/vae_decoder_trt.py | - examples/models/contrib/falcon/convert_checkpoint.py | - examples/models/contrib/gptj/convert_checkpoint.py | - examples/models/contrib/gptneox/convert_checkpoint.py | - examples/models/contrib/grok/convert_checkpoint.py | - examples/models/contrib/mmdit/convert_checkpoint.py | - examples/models/contrib/mmdit/sample.py | - examples/models/contrib/mpt/convert_checkpoint.py | - examples/models/contrib/opt/convert_checkpoint.py | - examples/models/contrib/sdxl/build_sdxl_unet.py | - examples/models/contrib/sdxl/pipeline_stable_diffusion_xl.py | - examples/models/contrib/sdxl/run_sdxl.py | - examples/models/contrib/stdit/aspect.py | - examples/models/contrib/stdit/convert_checkpoint.py | - examples/models/contrib/stdit/pipeline_tllm.py | - examples/models/contrib/stdit/sample.py | - examples/models/contrib/stdit/scheduler.py | - examples/models/contrib/stdit/text_encoder.py | - examples/models/contrib/stdit/utils.py | - examples/models/contrib/stdit/vae.py | - examples/models/contrib/stdit/video_transforms.py | - examples/models/core/bert/__init__.py | - examples/models/core/bert/convert_checkpoint.py | - examples/models/core/bert/run.py | - examples/models/core/bert/utils.py | - examples/models/core/commandr/convert_checkpoint.py | - examples/models/core/enc_dec/__init__.py | - examples/models/core/enc_dec/convert_checkpoint.py | - examples/models/core/enc_dec/helper.py | - examples/models/core/enc_dec/run.py | - examples/models/core/gemma/convert_checkpoint.py | - examples/models/core/glm-4-9b/convert_checkpoint.py | - examples/models/core/glm-4-9b/tokenization_chatglm.py | - examples/models/core/gpt/convert_checkpoint.py | - examples/models/core/gpt/merge_ptuning_tables.py | - examples/models/core/gpt/nemo_lora_convert.py | - examples/models/core/gpt/nemo_prompt_convert.py | - examples/models/core/gpt/run_hf.py | examples/models/core/gpt_oss/openai_chat_client_function_calling.py | - examples/models/core/internlm2/convert_checkpoint.py | examples/models/core/kimi_k2/kimi_k2_tool_calling_example.py | - examples/models/core/llama/convert_checkpoint.py | - examples/models/core/llama/summarize_long.py | - examples/models/core/mamba/convert_checkpoint.py | - examples/models/core/mllama/convert_checkpoint.py | - examples/models/core/multimodal/__init__.py | - examples/models/core/multimodal/build_multimodal_engine.py | - examples/models/core/multimodal/eval.py | - examples/models/core/multimodal/run.py | - examples/models/core/multimodal/utils.py | - examples/models/core/nemotron_nas/calibration_utils.py | - examples/models/core/nemotron_nas/convert_checkpoint.py | - examples/models/core/phi/convert_checkpoint.py | - examples/models/core/qwen/convert_checkpoint.py | - examples/models/core/qwen2audio/run.py | - examples/models/core/qwen2audio/run_chat.py | - examples/models/core/qwen2audio/utils.py | - examples/models/core/qwenvl/run.py | - examples/models/core/qwenvl/run_chat.py | - examples/models/core/qwenvl/show_pic.py | - examples/models/core/qwenvl/vit_onnx_trt.py | - examples/models/core/recurrentgemma/convert_checkpoint.py | - examples/models/core/vit/convert_checkpoint.py | - examples/models/core/whisper/convert_checkpoint.py | - examples/models/core/whisper/distil_whisper/convert_from_distil_whisper.py | - examples/models/core/whisper/run.py | - examples/models/core/whisper/tokenizer.py | - examples/models/core/whisper/whisper_utils.py | - examples/ngram/run_dtm_ngram.py | - examples/openai_triton/manual_plugin/build.py | - examples/openai_triton/manual_plugin/fmha_triton.py | - examples/openai_triton/manual_plugin/plugin.py | - examples/openai_triton/manual_plugin/run.py | - examples/openai_triton/plugin_autogen/build_engine.py | - examples/openai_triton/plugin_autogen/kernel_config.py | - examples/openai_triton/plugin_autogen/run_engine.py | - examples/python_plugin/build_lookup.py | - examples/python_plugin/plugin_lib/__init__.py | - examples/python_plugin/plugin_lib/lookup_kernel.py | - examples/python_plugin/plugin_lib/lookup_plugin.py | - examples/python_plugin/run_lookup.py | - examples/quantization/quantize.py | examples/quantization/quantize_mixed_precision_moe.py | examples/ray_orchestrator/llm_inference_async_ray.py | examples/ray_orchestrator/llm_inference_distributed_ray.py | - examples/redrafter/convert_checkpoint.py | - examples/run.py | examples/scaffolding/contrib/AsyncGeneration/stream_generation_controller.py | examples/scaffolding/contrib/DeepConf/run_generation.py | examples/scaffolding/contrib/Dynasor/scaffolding_dynasor_run.py | examples/scaffolding/contrib/TreeInference/run_mcts_example.py | examples/scaffolding/contrib/TreeInference/run_tot_example.py | - examples/scaffolding/contrib/mcp/e2b/e2bserver.py | - examples/scaffolding/contrib/mcp/e2b/main.py | - examples/scaffolding/contrib/mcp/mcptest.py | - examples/scaffolding/contrib/mcp/weather/weather.py | - examples/scaffolding/contrib/mcp/websearch/main.py | - examples/scaffolding/contrib/mcp/websearch/websearch.py | examples/scaffolding/run_basic_generation.py | examples/scaffolding/run_best_of_n_with_reward.py | examples/scaffolding/run_majority_vote_aime24.py | @@ -210,8 +83,6 @@ common-files: &common_files | examples/serve/openai_completion_client.py | examples/serve/openai_completion_client_for_lora.py | examples/serve/openai_completion_client_json_schema.py | - examples/summarize.py | - examples/utils.py | examples/wide_ep/ep_load_balancer/generate_eplb_config.py | examples/wide_ep/ep_load_balancer/report_load_statistics.py | examples/wide_ep/ep_load_balancer/utils.py | @@ -223,7 +94,6 @@ common-files: &common_files | scripts/check_test_list.py | scripts/dco_check.py | scripts/format_test_list.py | - scripts/generate_duration.py | scripts/generate_lock_file.py | scripts/get_wheel_from_package.py | scripts/git_replace.py | @@ -234,7 +104,6 @@ common-files: &common_files | setup.py | tensorrt_llm/__init__.py | tensorrt_llm/_ray_utils.py | - tensorrt_llm/_tensorrt_engine/__init__.py | tensorrt_llm/_torch/__init__.py | tensorrt_llm/_torch/attention_backend/__init__.py | tensorrt_llm/_torch/attention_backend/flashinfer.py | @@ -438,7 +307,6 @@ common-files: &common_files | tensorrt_llm/_torch/pyexecutor/guided_decoder.py | tensorrt_llm/_torch/pyexecutor/handle_additional_outputs.py | tensorrt_llm/_torch/pyexecutor/handle_logits.py | - tensorrt_llm/_torch/pyexecutor/kv_cache_connector.py | tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py | tensorrt_llm/_torch/pyexecutor/layerwise_nvtx_marker.py | tensorrt_llm/_torch/pyexecutor/llm_request.py | @@ -476,11 +344,6 @@ common-files: &common_files | tensorrt_llm/bench/benchmark/utils/asynchronous.py | tensorrt_llm/bench/benchmark/utils/general.py | tensorrt_llm/bench/benchmark/utils/processes.py | - tensorrt_llm/bench/build/__init__.py | - tensorrt_llm/bench/build/build.py | - tensorrt_llm/bench/build/dataclasses.py | - tensorrt_llm/bench/build/tuning.py | - tensorrt_llm/bench/build/utils.py | tensorrt_llm/bench/dataclasses/__init__.py | tensorrt_llm/bench/dataclasses/configuration.py | tensorrt_llm/bench/dataclasses/engine.py | @@ -490,13 +353,9 @@ common-files: &common_files | tensorrt_llm/bench/dataclasses/statistics.py | tensorrt_llm/bench/utils/__init__.py | tensorrt_llm/bench/utils/data.py | - tensorrt_llm/builder.py | tensorrt_llm/commands/__init__.py | tensorrt_llm/commands/bench.py | - tensorrt_llm/commands/build.py | tensorrt_llm/commands/eval.py | - tensorrt_llm/commands/prune.py | - tensorrt_llm/commands/refit.py | tensorrt_llm/commands/serve.py | tensorrt_llm/evaluate/__init__.py | tensorrt_llm/evaluate/cnn_dailymail.py | @@ -532,23 +391,7 @@ common-files: &common_files | tensorrt_llm/inputs/multimodal.py | tensorrt_llm/inputs/registry.py | tensorrt_llm/inputs/utils.py | - tensorrt_llm/layers/__init__.py | - tensorrt_llm/layers/activation.py | - tensorrt_llm/layers/attention.py | - tensorrt_llm/layers/cast.py | - tensorrt_llm/layers/conv.py | - tensorrt_llm/layers/embedding.py | - tensorrt_llm/layers/language_adapter.py | - tensorrt_llm/layers/linear.py | - tensorrt_llm/layers/lora.py | - tensorrt_llm/layers/mlp.py | - tensorrt_llm/layers/moe.py | - tensorrt_llm/layers/normalization.py | - tensorrt_llm/layers/pooling.py | - tensorrt_llm/layers/recurrent.py | - tensorrt_llm/layers/ssm.py | tensorrt_llm/llmapi/__init__.py | - tensorrt_llm/llmapi/build_cache.py | tensorrt_llm/llmapi/disagg_utils.py | tensorrt_llm/llmapi/kv_cache_type.py | tensorrt_llm/llmapi/llm.py | @@ -571,179 +414,17 @@ common-files: &common_files | tensorrt_llm/metrics/enums.py | tensorrt_llm/models/__init__.py | tensorrt_llm/models/automodel.py | - tensorrt_llm/models/baichuan/__init__.py | - tensorrt_llm/models/baichuan/config.py | - tensorrt_llm/models/baichuan/convert.py | - tensorrt_llm/models/baichuan/model.py | - tensorrt_llm/models/bert/__init__.py | - tensorrt_llm/models/bert/config.py | - tensorrt_llm/models/bert/convert.py | - tensorrt_llm/models/bert/model.py | - tensorrt_llm/models/bloom/__init__.py | - tensorrt_llm/models/bloom/model.py | - tensorrt_llm/models/chatglm/__init__.py | - tensorrt_llm/models/chatglm/config.py | - tensorrt_llm/models/chatglm/convert.py | - tensorrt_llm/models/chatglm/model.py | - tensorrt_llm/models/clip/__init__.py | - tensorrt_llm/models/clip/model.py | - tensorrt_llm/models/cogvlm/__init__.py | - tensorrt_llm/models/cogvlm/config.py | - tensorrt_llm/models/cogvlm/convert.py | - tensorrt_llm/models/cogvlm/model.py | - tensorrt_llm/models/commandr/__init__.py | - tensorrt_llm/models/commandr/config.py | - tensorrt_llm/models/commandr/model.py | tensorrt_llm/models/convert_utils.py | - tensorrt_llm/models/dbrx/__init__.py | - tensorrt_llm/models/dbrx/config.py | - tensorrt_llm/models/dbrx/model.py | - tensorrt_llm/models/deepseek_v1/__init__.py | - tensorrt_llm/models/deepseek_v1/config.py | - tensorrt_llm/models/deepseek_v1/convert.py | - tensorrt_llm/models/deepseek_v1/model.py | - tensorrt_llm/models/deepseek_v2/__init__.py | - tensorrt_llm/models/deepseek_v2/config.py | - tensorrt_llm/models/deepseek_v2/convert.py | - tensorrt_llm/models/deepseek_v2/model.py | - tensorrt_llm/models/dit/__init__.py | - tensorrt_llm/models/dit/model.py | - tensorrt_llm/models/eagle/__init__.py | - tensorrt_llm/models/eagle/config.py | - tensorrt_llm/models/eagle/model.py | - tensorrt_llm/models/enc_dec/__init__.py | - tensorrt_llm/models/enc_dec/model.py | - tensorrt_llm/models/falcon/__init__.py | - tensorrt_llm/models/falcon/config.py | - tensorrt_llm/models/falcon/convert.py | - tensorrt_llm/models/falcon/model.py | - tensorrt_llm/models/gemma/__init__.py | - tensorrt_llm/models/gemma/config.py | - tensorrt_llm/models/gemma/convert.py | - tensorrt_llm/models/gemma/model.py | - tensorrt_llm/models/gemma/smoothquant.py | - tensorrt_llm/models/gemma/utils/__init__.py | - tensorrt_llm/models/gemma/utils/layers.py | - tensorrt_llm/models/gemma/utils/modules.py | - tensorrt_llm/models/gemma/utils/params.py | - tensorrt_llm/models/gemma/utils/positional_embeddings.py | - tensorrt_llm/models/gemma/utils/sampler.py | - tensorrt_llm/models/gemma/utils/transformer.py | - tensorrt_llm/models/gemma/weight.py | - tensorrt_llm/models/generation_mixin.py | - tensorrt_llm/models/gpt/__init__.py | - tensorrt_llm/models/gpt/config.py | - tensorrt_llm/models/gpt/convert.py | - tensorrt_llm/models/gpt/model.py | - tensorrt_llm/models/gptj/__init__.py | - tensorrt_llm/models/gptj/config.py | - tensorrt_llm/models/gptj/convert.py | - tensorrt_llm/models/gptj/model.py | - tensorrt_llm/models/gptneox/__init__.py | - tensorrt_llm/models/gptneox/model.py | - tensorrt_llm/models/grok/__init__.py | - tensorrt_llm/models/grok/convert.py | - tensorrt_llm/models/grok/model.py | - tensorrt_llm/models/grok/weight.py | - tensorrt_llm/models/llama/__init__.py | - tensorrt_llm/models/llama/config.py | - tensorrt_llm/models/llama/convert.py | - tensorrt_llm/models/llama/model.py | - tensorrt_llm/models/mamba/__init__.py | - tensorrt_llm/models/mamba/config.py | - tensorrt_llm/models/mamba/convert.py | - tensorrt_llm/models/mamba/model.py | - tensorrt_llm/models/medusa/__init__.py | - tensorrt_llm/models/medusa/config.py | - tensorrt_llm/models/medusa/model.py | - tensorrt_llm/models/medusa/weight.py | - tensorrt_llm/models/mllama/__init__.py | - tensorrt_llm/models/mllama/config.py | - tensorrt_llm/models/mllama/model.py | - tensorrt_llm/models/mmdit_sd3/__init__.py | - tensorrt_llm/models/mmdit_sd3/config.py | - tensorrt_llm/models/mmdit_sd3/model.py | - tensorrt_llm/models/model_weights_loader.py | tensorrt_llm/models/modeling_utils.py | - tensorrt_llm/models/mpt/__init__.py | - tensorrt_llm/models/mpt/model.py | - tensorrt_llm/models/multimodal_encoders/__init__.py | - tensorrt_llm/models/multimodal_encoders/config.py | - tensorrt_llm/models/multimodal_encoders/model.py | - tensorrt_llm/models/nemotron_nas/__init__.py | - tensorrt_llm/models/nemotron_nas/config.py | - tensorrt_llm/models/nemotron_nas/convert.py | - tensorrt_llm/models/nemotron_nas/layer_config.py | - tensorrt_llm/models/nemotron_nas/model.py | - tensorrt_llm/models/opt/__init__.py | - tensorrt_llm/models/opt/model.py | - tensorrt_llm/models/phi/__init__.py | - tensorrt_llm/models/phi/config.py | - tensorrt_llm/models/phi/convert.py | - tensorrt_llm/models/phi/model.py | - tensorrt_llm/models/phi3/__init__.py | - tensorrt_llm/models/phi3/config.py | - tensorrt_llm/models/phi3/convert.py | - tensorrt_llm/models/phi3/model.py | - tensorrt_llm/models/phi3/split_weights.py | - tensorrt_llm/models/qwen/__init__.py | - tensorrt_llm/models/qwen/config.py | - tensorrt_llm/models/qwen/convert.py | - tensorrt_llm/models/qwen/model.py | - tensorrt_llm/models/qwen/utils.py | - tensorrt_llm/models/recurrentgemma/__init__.py | - tensorrt_llm/models/recurrentgemma/model.py | - tensorrt_llm/models/redrafter/__init__.py | - tensorrt_llm/models/redrafter/drafter.py | - tensorrt_llm/models/redrafter/model.py | - tensorrt_llm/models/redrafter/redrafter_helper.py | - tensorrt_llm/models/stdit/__init__.py | - tensorrt_llm/models/stdit/config.py | - tensorrt_llm/models/stdit/model.py | - tensorrt_llm/models/unet/__init__.py | - tensorrt_llm/models/unet/attention.py | - tensorrt_llm/models/unet/embeddings.py | - tensorrt_llm/models/unet/pp/__init__.py | - tensorrt_llm/models/unet/pp/attention.py | - tensorrt_llm/models/unet/pp/conv2d.py | - tensorrt_llm/models/unet/pp/groupnorm.py | - tensorrt_llm/models/unet/pp/unet_pp.py | - tensorrt_llm/models/unet/resnet.py | - tensorrt_llm/models/unet/unet_2d_blocks.py | - tensorrt_llm/models/unet/unet_2d_condition.py | - tensorrt_llm/models/unet/weights.py | - tensorrt_llm/network.py | - tensorrt_llm/parameter.py | - tensorrt_llm/plugin/__init__.py | - tensorrt_llm/plugin/plugin.py | tensorrt_llm/quantization/__init__.py | tensorrt_llm/quantization/functional.py | - tensorrt_llm/quantization/image_processing.py | - tensorrt_llm/quantization/layers.py | tensorrt_llm/quantization/mode.py | - tensorrt_llm/quantization/quantize.py | - tensorrt_llm/quantization/quantize_by_modelopt.py | tensorrt_llm/quantization/utils/__init__.py | tensorrt_llm/quantization/utils/fp4_utils.py | tensorrt_llm/quantization/utils/fp8_utils.py | tensorrt_llm/ray_stub.py | tensorrt_llm/runtime/__init__.py | - tensorrt_llm/runtime/enc_dec_model_runner.py | - tensorrt_llm/runtime/generation.py | - tensorrt_llm/runtime/kv_cache_manager.py | - tensorrt_llm/runtime/medusa_utils.py | tensorrt_llm/runtime/memory_pools/__init__.py | - tensorrt_llm/runtime/memory_pools/memory_pools_allocator.py | - tensorrt_llm/runtime/memory_pools/pool.py | - tensorrt_llm/runtime/memory_pools/pools_kv_cache_manager.py | - tensorrt_llm/runtime/model_runner.py | - tensorrt_llm/runtime/model_runner_cpp.py | - tensorrt_llm/runtime/multimodal_model_runner.py | - tensorrt_llm/runtime/processor_wrapper/__init__.py | - tensorrt_llm/runtime/processor_wrapper/mllama_processor_wrapper.py | - tensorrt_llm/runtime/processor_wrapper/processor_wrapper.py | - tensorrt_llm/runtime/redrafter_utils.py | - tensorrt_llm/runtime/session.py | tensorrt_llm/scaffolding/__init__.py | tensorrt_llm/scaffolding/benchmark.py | tensorrt_llm/scaffolding/contrib/AsyncGeneration/stream_generation.py | @@ -798,12 +479,7 @@ common-files: &common_files | tensorrt_llm/tokenizer/tokenizer.py | tensorrt_llm/tools/__init__.py | tensorrt_llm/tools/importlib_utils.py | - tensorrt_llm/tools/multimodal_builder.py | - tensorrt_llm/tools/onnx_utils.py | tensorrt_llm/tools/plugin_gen/__init__.py | - tensorrt_llm/tools/plugin_gen/core.py | - tensorrt_llm/tools/plugin_gen/plugin_gen.py | - tensorrt_llm/tools/plugin_gen/shape_infer.py | tensorrt_llm/tools/ppl.py | tensorrt_llm/tools/profiler/nsys_profile_tools/gputrc2graph.py | tensorrt_llm/version.py | @@ -812,9 +488,7 @@ common-files: &common_files | tests/integration/defs/accuracy/accuracy_core.py | tests/integration/defs/accuracy/scripts/collect_evaluated_accuracies.py | tests/integration/defs/accuracy/scripts/compute_theta_and_thresholds.py | - tests/integration/defs/accuracy/test_cli_flow.py | tests/integration/defs/accuracy/test_disaggregated_serving.py | - tests/integration/defs/accuracy/test_llm_api.py | tests/integration/defs/accuracy/test_llm_api_autodeploy.py | tests/integration/defs/accuracy/test_llm_api_pytorch.py | tests/integration/defs/accuracy/test_llm_api_pytorch_ray.py | @@ -830,53 +504,22 @@ common-files: &common_files | tests/integration/defs/disaggregated/test_disaggregated_etcd.py | tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py | tests/integration/defs/disaggregated/test_workers.py | - tests/integration/defs/examples/run_llm_fp8_quant_llama_70b.py | tests/integration/defs/examples/run_llm_quickstart_atexit.py | tests/integration/defs/examples/serve/test_serve.py | tests/integration/defs/examples/serve/test_serve_negative.py | tests/integration/defs/examples/test_ad_guided_decoding.py | - tests/integration/defs/examples/test_bert.py | - tests/integration/defs/examples/test_bindings.py | - tests/integration/defs/examples/test_chatglm.py | - tests/integration/defs/examples/test_commandr.py | - tests/integration/defs/examples/test_draft_target_model.py | - tests/integration/defs/examples/test_eagle.py | - tests/integration/defs/examples/test_enc_dec.py | - tests/integration/defs/examples/test_exaone.py | - tests/integration/defs/examples/test_gemma.py | tests/integration/defs/examples/test_gpt.py | - tests/integration/defs/examples/test_gptj.py | - tests/integration/defs/examples/test_granite.py | - tests/integration/defs/examples/test_internlm.py | - tests/integration/defs/examples/test_llama.py | tests/integration/defs/examples/test_llm_api_with_mpi.py | - tests/integration/defs/examples/test_mamba.py | - tests/integration/defs/examples/test_medusa.py | - tests/integration/defs/examples/test_mistral.py | - tests/integration/defs/examples/test_mixtral.py | - tests/integration/defs/examples/test_multimodal.py | - tests/integration/defs/examples/test_nemotron.py | - tests/integration/defs/examples/test_nemotron_nas.py | - tests/integration/defs/examples/test_ngram.py | - tests/integration/defs/examples/test_openai.py | tests/integration/defs/examples/test_phi.py | - tests/integration/defs/examples/test_qwen.py | - tests/integration/defs/examples/test_qwen2audio.py | - tests/integration/defs/examples/test_qwenvl.py | tests/integration/defs/examples/test_ray.py | - tests/integration/defs/examples/test_recurrentgemma.py | - tests/integration/defs/examples/test_redrafter.py | - tests/integration/defs/examples/test_whisper.py | tests/integration/defs/llmapi/__init__.py | tests/integration/defs/llmapi/_run_llmapi_llm.py | tests/integration/defs/llmapi/test_llm_api_connector.py | tests/integration/defs/llmapi/test_llm_api_qa.py | - tests/integration/defs/llmapi/test_llm_e2e.py | tests/integration/defs/llmapi/test_llm_examples.py | tests/integration/defs/local_venv.py | tests/integration/defs/perf/__init__.py | tests/integration/defs/perf/allowed_configs.py | - tests/integration/defs/perf/build.py | tests/integration/defs/perf/create_perf_comparison_report.py | tests/integration/defs/perf/data.py | tests/integration/defs/perf/data_export.py | @@ -900,34 +543,18 @@ common-files: &common_files | tests/integration/defs/test_sanity.py | tests/integration/defs/test_unittests.py | tests/integration/defs/triton_server/__init__.py | - tests/integration/defs/triton_server/build_engines.py | tests/integration/defs/triton_server/common.py | tests/integration/defs/triton_server/conftest.py | - tests/integration/defs/triton_server/local_venv.py | - tests/integration/defs/triton_server/rcca/bug_4323566/inflight_batcher_llm_client_with_end_id.py | - tests/integration/defs/triton_server/runner_interface.py | tests/integration/defs/triton_server/test_list_parser.py | - tests/integration/defs/triton_server/test_triton.py | - tests/integration/defs/triton_server/test_triton_llm.py | - tests/integration/defs/triton_server/test_triton_memleak.py | - tests/integration/defs/triton_server/test_triton_multi_node.py | - tests/integration/defs/triton_server/test_triton_rcca.py | tests/integration/defs/triton_server/trt_test_alternative.py | tests/integration/defs/trt_test_alternative.py | tests/integration/defs/utils/__init__.py | tests/integration/defs/utils/periodic_junit.py | tests/integration/defs/utils/timeout_manager.py | tests/microbenchmarks/all_reduce.py | - tests/microbenchmarks/build_time_benchmark.py | - tests/microbenchmarks/build_time_dashboard.py | tests/scripts/allreduce_perf/allreduce_heuristic_code_gen.py | tests/scripts/allreduce_perf/allreduce_perf_viz.py | tests/scripts/iteration_log_parser.py | - tests/scripts/perf-sanity/parse_benchmark_results.py | - tests/scripts/perf-sanity/run_benchmark_serve.py | - tests/unittest/_torch/attention/sparse/test_dsa_indexer.py | - tests/unittest/_torch/attention/sparse/test_flash_mla.py | - tests/unittest/_torch/attention/sparse/test_rocketkv.py | tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py | tests/unittest/_torch/attention/test_attention.py | tests/unittest/_torch/attention/test_attention_mla.py | @@ -948,7 +575,6 @@ common-files: &common_files | tests/unittest/_torch/misc/test_virtual_memory.py | tests/unittest/_torch/modeling/test_modeling_bert.py | tests/unittest/_torch/modeling/test_modeling_clip.py | - tests/unittest/_torch/modeling/test_modeling_exaone4.py | tests/unittest/_torch/modeling/test_modeling_gemma3.py | tests/unittest/_torch/modeling/test_modeling_gpt_oss.py | tests/unittest/_torch/modeling/test_modeling_llama.py | @@ -972,8 +598,6 @@ common-files: &common_files | tests/unittest/_torch/modules/test_moe_routing.py | tests/unittest/_torch/modules/test_rotary_embedding.py | tests/unittest/_torch/modules/test_triton_linear.py | - tests/unittest/_torch/modules/tests_lora_modules/test_lora_attention_pytorch_flow_vs_trt.py | - tests/unittest/_torch/modules/tests_lora_modules/test_lora_plugin_vs_lora_op.py | tests/unittest/_torch/multi_gpu/test_allreduce.py | tests/unittest/_torch/multi_gpu/test_alltoall.py | tests/unittest/_torch/multi_gpu/test_ar_residual_norm.py | @@ -1000,24 +624,10 @@ common-files: &common_files | tests/unittest/_torch/sampler/test_beam_search.py | tests/unittest/_torch/sampler/test_best_of_n.py | tests/unittest/_torch/sampler/test_trtllm_sampler.py | - tests/unittest/_torch/speculative/test_draft_target.py | - tests/unittest/_torch/speculative/test_draft_token_tree_sampling.py | - tests/unittest/_torch/speculative/test_draft_token_tree_verification.py | - tests/unittest/_torch/speculative/test_dynamic_spec_decode.py | tests/unittest/_torch/speculative/test_eagle3.py | - tests/unittest/_torch/speculative/test_kv_cache_reuse.py | - tests/unittest/_torch/speculative/test_mtp.py | - tests/unittest/_torch/speculative/test_ngram.py | - tests/unittest/_torch/speculative/test_save_state.py | - tests/unittest/_torch/speculative/test_spec_gate.py | - tests/unittest/_torch/speculative/test_torch_rejection_sampling.py | - tests/unittest/_torch/speculative/test_user_provided.py | tests/unittest/_torch/test_connector.py | tests/unittest/_torch/test_torch_multi_arange.py | tests/unittest/_torch/thop/parallel/deep_gemm_tests.py | - tests/unittest/_torch/thop/parallel/test_causal_conv1d_op.py | - tests/unittest/_torch/thop/parallel/test_cublas_mm.py | - tests/unittest/_torch/thop/parallel/test_custom_ops.py | tests/unittest/_torch/thop/parallel/test_dsv3_fused_a_gemm.py | tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py | tests/unittest/_torch/thop/parallel/test_finegrained_mixed_dtype_gemm.py | @@ -1031,11 +641,6 @@ common-files: &common_files | tests/unittest/_torch/thop/parallel/test_fp8_per_tensor_scale_tllmg_gemm.py | tests/unittest/_torch/thop/parallel/test_fp8_quantize.py | tests/unittest/_torch/thop/parallel/test_fp8_rowwise_linear.py | - tests/unittest/_torch/thop/parallel/test_fused_qk_norm_rope.py | - tests/unittest/_torch/thop/parallel/test_logits_bitmask_op.py | - tests/unittest/_torch/thop/parallel/test_mamba2_chunk_ss_update.py | - tests/unittest/_torch/thop/parallel/test_mamba_conv1d_op.py | - tests/unittest/_torch/thop/parallel/test_noaux_tc.py | tests/unittest/_torch/thop/parallel/test_scaled_mm.py | tests/unittest/_torch/thop/parallel/test_selective_scan_op.py | tests/unittest/_torch/thop/parallel/test_tinygemm2.py | @@ -1079,12 +684,10 @@ common-files: &common_files | tests/unittest/llmapi/apps/_test_openai_chat_harmony.py | tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py | tests/unittest/llmapi/apps/_test_openai_completions.py | - tests/unittest/llmapi/apps/_test_openai_consistent_chat.py | tests/unittest/llmapi/apps/_test_openai_lora.py | tests/unittest/llmapi/apps/_test_openai_metrics.py | tests/unittest/llmapi/apps/_test_openai_misc.py | tests/unittest/llmapi/apps/_test_openai_mmencoder.py | - tests/unittest/llmapi/apps/_test_openai_multi_chat.py | tests/unittest/llmapi/apps/_test_openai_multi_gpu.py | tests/unittest/llmapi/apps/_test_openai_multi_nodes.py | tests/unittest/llmapi/apps/_test_openai_perf_metrics.py | @@ -1107,15 +710,12 @@ common-files: &common_files | tests/unittest/llmapi/run_llm_exit.py | tests/unittest/llmapi/run_llm_with_postproc.py | tests/unittest/llmapi/test_additional_model_outputs.py | - tests/unittest/llmapi/test_build_cache.py | tests/unittest/llmapi/test_executor.py | tests/unittest/llmapi/test_gc_utils.py | tests/unittest/llmapi/test_llm.py | tests/unittest/llmapi/test_llm_args.py | tests/unittest/llmapi/test_llm_download.py | tests/unittest/llmapi/test_llm_kv_cache_events.py | - tests/unittest/llmapi/test_llm_models.py | - tests/unittest/llmapi/test_llm_multi_gpu.py | tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py | tests/unittest/llmapi/test_llm_pytorch.py | tests/unittest/llmapi/test_llm_quant.py | @@ -1126,24 +726,14 @@ common-files: &common_files | tests/unittest/llmapi/test_serialization.py | tests/unittest/llmapi/test_utils.py | tests/unittest/others/__init__.py | - tests/unittest/others/test_builder.py | tests/unittest/others/test_convert_spec_decoding_mask_to_packed_mask.py | - tests/unittest/others/test_debugging_api.py | tests/unittest/others/test_exception.py | tests/unittest/others/test_export.py | - tests/unittest/others/test_graph_rewriter.py | - tests/unittest/others/test_kv_cache_manager.py | tests/unittest/others/test_kv_cache_transceiver.py | tests/unittest/others/test_kv_cache_update.py | - tests/unittest/others/test_layer.py | tests/unittest/others/test_mapping.py | - tests/unittest/others/test_model_dtype.py | - tests/unittest/others/test_module.py | tests/unittest/others/test_multimodal_registry.py | - tests/unittest/others/test_plugins.py | - tests/unittest/others/test_precision_control.py | tests/unittest/others/test_pretrained_config.py | - tests/unittest/others/test_session.py | tests/unittest/others/test_time_breakdown.py | tests/unittest/profile_utils.py | tests/unittest/scaffolding/__init__.py | @@ -1152,141 +742,14 @@ common-files: &common_files | tests/unittest/scaffolding/test_scaffolding.py | tests/unittest/scaffolding/test_task_collection.py | tests/unittest/scaffolding/test_worker.py | - tests/unittest/test_model_runner_cpp.py | tests/unittest/test_pip_install.py | tests/unittest/tools/__init__.py | - tests/unittest/tools/plugin_gen/__init__.py | - tests/unittest/tools/plugin_gen/kernel_config.py | - tests/unittest/tools/plugin_gen/test_core.py | - tests/unittest/tools/plugin_gen/test_plugin_gen.py | - tests/unittest/tools/plugin_gen/test_shape_infer.py | tests/unittest/tools/test_prepare_dataset.py | tests/unittest/tools/test_test_to_stage_mapping.py | - tests/unittest/trt/__init__.py | - tests/unittest/trt/attention/test_bert_attention.py | - tests/unittest/trt/attention/test_gpt_attention.py | - tests/unittest/trt/attention/test_gpt_attention_IFB.py | - tests/unittest/trt/attention/test_gpt_attention_no_cache.py | - tests/unittest/trt/attention/test_sage_attention.py | - tests/unittest/trt/functional/__init__.py | - tests/unittest/trt/functional/test_alibi.py | - tests/unittest/trt/functional/test_allreduce_norm.py | - tests/unittest/trt/functional/test_allreduce_prepost_residual_norm.py | - tests/unittest/trt/functional/test_arange.py | - tests/unittest/trt/functional/test_argmax.py | - tests/unittest/trt/functional/test_assertion.py | - tests/unittest/trt/functional/test_avg_pool2d.py | - tests/unittest/trt/functional/test_cast.py | - tests/unittest/trt/functional/test_conv2d.py | - tests/unittest/trt/functional/test_conv3d.py | - tests/unittest/trt/functional/test_cos.py | - tests/unittest/trt/functional/test_cumsum.py | - tests/unittest/trt/functional/test_dora.py | - tests/unittest/trt/functional/test_einsum.py | - tests/unittest/trt/functional/test_embedding_single_gpu.py | - tests/unittest/trt/functional/test_exp.py | - tests/unittest/trt/functional/test_expand.py | - tests/unittest/trt/functional/test_flatten.py | - tests/unittest/trt/functional/test_flip.py | - tests/unittest/trt/functional/test_fp4_gemm.py | - tests/unittest/trt/functional/test_fp4_gemm_ootb.py | - tests/unittest/trt/functional/test_gather.py | - tests/unittest/trt/functional/test_gather_nd.py | - tests/unittest/trt/functional/test_geglu.py | - tests/unittest/trt/functional/test_gelu.py | - tests/unittest/trt/functional/test_gemm_swiglu.py | - tests/unittest/trt/functional/test_group_norm.py | - tests/unittest/trt/functional/test_identity.py | - tests/unittest/trt/functional/test_index_select.py | - tests/unittest/trt/functional/test_interpolate.py | - tests/unittest/trt/functional/test_logsoftmax.py | - tests/unittest/trt/functional/test_lora.py | - tests/unittest/trt/functional/test_low_latency_gemm.py | - tests/unittest/trt/functional/test_mamba_conv1d.py | - tests/unittest/trt/functional/test_masked_scatter.py | - tests/unittest/trt/functional/test_masked_select.py | - tests/unittest/trt/functional/test_matmul.py | - tests/unittest/trt/functional/test_meshgrid2d.py | - tests/unittest/trt/functional/test_moe.py | - tests/unittest/trt/functional/test_nccl.py | - tests/unittest/trt/functional/test_nonzero.py | - tests/unittest/trt/functional/test_outer.py | - tests/unittest/trt/functional/test_pad.py | - tests/unittest/trt/functional/test_permute.py | - tests/unittest/trt/functional/test_pp_reduce_scatter.py | - tests/unittest/trt/functional/test_quant.py | - tests/unittest/trt/functional/test_rearrange.py | - tests/unittest/trt/functional/test_repeat.py | - tests/unittest/trt/functional/test_repeat_interleave.py | - tests/unittest/trt/functional/test_rg_lru.py | - tests/unittest/trt/functional/test_sample.py | - tests/unittest/trt/functional/test_scatter.py | - tests/unittest/trt/functional/test_scatter_nd.py | - tests/unittest/trt/functional/test_select.py | - tests/unittest/trt/functional/test_selective_scan.py | - tests/unittest/trt/functional/test_sigmoid.py | - tests/unittest/trt/functional/test_silu.py | - tests/unittest/trt/functional/test_sin.py | - tests/unittest/trt/functional/test_slice.py | - tests/unittest/trt/functional/test_softplus.py | - tests/unittest/trt/functional/test_split.py | - tests/unittest/trt/functional/test_squeeze.py | - tests/unittest/trt/functional/test_swiglu.py | - tests/unittest/trt/functional/test_topk.py | - tests/unittest/trt/functional/test_transpose.py | - tests/unittest/trt/functional/test_unbind.py | - tests/unittest/trt/functional/test_unsqueeze.py | - tests/unittest/trt/functional/test_view.py | - tests/unittest/trt/functional/test_where.py | - tests/unittest/trt/model/__init__.py | - tests/unittest/trt/model/eagle/test_decode_draft_tokens_plugin.py | - tests/unittest/trt/model/eagle/test_prepare_drafter_inputs_plugin.py | - tests/unittest/trt/model/eagle/test_sample_accept_draft_tokens_plugin.py | - tests/unittest/trt/model/redrafter/test_beams2tree.py | - tests/unittest/trt/model/redrafter/test_draft_token.py | - tests/unittest/trt/model/redrafter/test_draft_token_indices.py | - tests/unittest/trt/model/redrafter/test_gather_beams.py | - tests/unittest/trt/model/redrafter/test_mask.py | - tests/unittest/trt/model/redrafter/test_packed_position_ids.py | - tests/unittest/trt/model/redrafter/test_prefix_match_indices.py | - tests/unittest/trt/model/redrafter/test_prepare_input.py | - tests/unittest/trt/model/redrafter/test_process_logits.py | - tests/unittest/trt/model/redrafter/test_top1.py | - tests/unittest/trt/model/redrafter/test_unpack_gen_data.py | - tests/unittest/trt/model/redrafter/test_validate.py | - tests/unittest/trt/model/test_gpt.py | - tests/unittest/trt/model/test_gpt_e2e.py | - tests/unittest/trt/model/test_llama.py | - tests/unittest/trt/model/test_mamba.py | - tests/unittest/trt/model/test_mistral.py | - tests/unittest/trt/model/test_nemotron_nas.py | - tests/unittest/trt/model/test_phi.py | - tests/unittest/trt/model/test_unet.py | - tests/unittest/trt/model_api/test_model_api_multi_gpu.py | - tests/unittest/trt/model_api/test_model_level_api.py | - tests/unittest/trt/model_api/test_model_quantization.py | - tests/unittest/trt/python_plugin/plugin_wrapper_utils.py | - tests/unittest/trt/python_plugin/test_plugin_wrapper.py | - tests/unittest/trt/quantization/__init__.py | - tests/unittest/trt/quantization/_utils.py | - tests/unittest/trt/quantization/test_fp8_quantization.py | - tests/unittest/trt/quantization/test_fp8_rowwise_gemm.py | - tests/unittest/trt/quantization/test_functional.py | - tests/unittest/trt/quantization/test_mode.py | - tests/unittest/trt/quantization/test_moe_weight_only_quant_matmul.py | - tests/unittest/trt/quantization/test_qserve_gemm.py | - tests/unittest/trt/quantization/test_quant.py | - tests/unittest/trt/quantization/test_quant_layer.py | - tests/unittest/trt/quantization/test_smooth_quant_gemm.py | - tests/unittest/trt/quantization/test_smooth_quant_layer_norm.py | - tests/unittest/trt/quantization/test_smooth_quant_rms_norm.py | - tests/unittest/trt/quantization/test_weight_only_groupwise_quant_matmul.py | - tests/unittest/trt/quantization/test_weight_only_quant_matmul.py | tests/unittest/utils/__init__.py | tests/unittest/utils/cpp_paths.py | tests/unittest/utils/llm_data.py | tests/unittest/utils/runtime_defaults.py | - tests/unittest/utils/test_medusa_utils.py | tests/unittest/utils/test_prebuilt_whl_cpp_extensions.py | tests/unittest/utils/test_util.py | tests/unittest/utils/torch_ref.py | @@ -1316,14 +779,6 @@ legacy-files: &legacy_files | .devcontainer/make_env.py | .github/scripts/label_community_user.py | .github/scripts/pr_checklist_check.py | - benchmarks/__init__.py | - benchmarks/prepare_dataset.py | - benchmarks/utils/__init__.py | - benchmarks/utils/convert_nemo_dataset.py | - benchmarks/utils/generate_rand_loras.py | - benchmarks/utils/prepare_real_data.py | - benchmarks/utils/prepare_synthetic_data.py | - benchmarks/utils/utils.py | cpp/conanfile.py | cpp/kernels/fmha_v2/conftest.py | cpp/kernels/fmha_v2/fmha_test.py | @@ -1355,24 +810,10 @@ legacy-files: &legacy_files | examples/apps/fastapi_server.py | examples/disaggregated/clients/disagg_client.py | examples/disaggregated/slurm/benchmark/submit.py | - examples/dora/normalize_weights.py | - examples/eagle/convert_checkpoint.py | - examples/eval_long_context.py | - examples/generate_checkpoint_config.py | - examples/generate_xgrammar_tokenizer_info.py | - examples/hf_lora_convert.py | examples/infinitebench/args.py | examples/infinitebench/compute_scores.py | examples/infinitebench/construct_synthetic_dataset.py | examples/infinitebench/eval_utils.py | - examples/llm-api/_tensorrt_engine/llm_eagle2_decoding.py | - examples/llm-api/_tensorrt_engine/llm_eagle_decoding.py | - examples/llm-api/_tensorrt_engine/llm_inference_customize.py | - examples/llm-api/_tensorrt_engine/llm_inference_kv_events.py | - examples/llm-api/_tensorrt_engine/llm_lookahead_decoding.py | - examples/llm-api/_tensorrt_engine/llm_medusa_decoding.py | - examples/llm-api/_tensorrt_engine/llm_quantization.py | - examples/llm-api/_tensorrt_engine/quickstart_example.py | examples/llm-api/llm_guided_decoding.py | examples/llm-api/llm_inference.py | examples/llm-api/llm_inference_async.py | @@ -1392,122 +833,17 @@ legacy-files: &legacy_files | examples/llm-api/quickstart_example.py | examples/llm-api/quickstart_multimodal.py | examples/llm-api/star_attention.py | - examples/llm-eval/lm-eval-harness/lm_eval_tensorrt_llm.py | examples/longbench/eval_longbench_v1.py | - examples/medusa/convert_checkpoint.py | - examples/mmlu.py | - examples/models/contrib/baichuan/convert_checkpoint.py | - examples/models/contrib/bloom/convert_checkpoint.py | - examples/models/contrib/chatglm-6b/tokenization_chatglm.py | - examples/models/contrib/chatglm2-6b/tokenization_chatglm.py | - examples/models/contrib/chatglm3-6b-32k/tokenization_chatglm.py | - examples/models/contrib/cogvlm/convert_checkpoint.py | - examples/models/contrib/dbrx/convert_checkpoint.py | - examples/models/contrib/deepseek_v1/__init__.py | - examples/models/contrib/deepseek_v1/convert_checkpoint.py | - examples/models/contrib/deepseek_v2/convert_checkpoint.py | - examples/models/contrib/dit/convert_checkpoint.py | - examples/models/contrib/dit/diffusion.py | - examples/models/contrib/dit/sample.py | - examples/models/contrib/dit/utils_modelopt.py | - examples/models/contrib/dit/vae_decoder_trt.py | - examples/models/contrib/falcon/convert_checkpoint.py | - examples/models/contrib/gptj/convert_checkpoint.py | - examples/models/contrib/gptneox/convert_checkpoint.py | - examples/models/contrib/grok/convert_checkpoint.py | - examples/models/contrib/mmdit/convert_checkpoint.py | - examples/models/contrib/mmdit/sample.py | - examples/models/contrib/mpt/convert_checkpoint.py | - examples/models/contrib/opt/convert_checkpoint.py | - examples/models/contrib/sdxl/build_sdxl_unet.py | - examples/models/contrib/sdxl/pipeline_stable_diffusion_xl.py | - examples/models/contrib/sdxl/run_sdxl.py | - examples/models/contrib/stdit/aspect.py | - examples/models/contrib/stdit/convert_checkpoint.py | - examples/models/contrib/stdit/pipeline_tllm.py | - examples/models/contrib/stdit/sample.py | - examples/models/contrib/stdit/scheduler.py | - examples/models/contrib/stdit/text_encoder.py | - examples/models/contrib/stdit/utils.py | - examples/models/contrib/stdit/vae.py | - examples/models/contrib/stdit/video_transforms.py | - examples/models/core/bert/__init__.py | - examples/models/core/bert/convert_checkpoint.py | - examples/models/core/bert/run.py | - examples/models/core/bert/utils.py | - examples/models/core/commandr/convert_checkpoint.py | - examples/models/core/enc_dec/__init__.py | - examples/models/core/enc_dec/convert_checkpoint.py | - examples/models/core/enc_dec/helper.py | - examples/models/core/enc_dec/run.py | - examples/models/core/gemma/convert_checkpoint.py | - examples/models/core/glm-4-9b/convert_checkpoint.py | - examples/models/core/glm-4-9b/tokenization_chatglm.py | - examples/models/core/gpt/convert_checkpoint.py | - examples/models/core/gpt/merge_ptuning_tables.py | - examples/models/core/gpt/nemo_lora_convert.py | - examples/models/core/gpt/nemo_prompt_convert.py | - examples/models/core/gpt/run_hf.py | examples/models/core/gpt_oss/openai_chat_client_function_calling.py | - examples/models/core/internlm2/convert_checkpoint.py | examples/models/core/kimi_k2/kimi_k2_tool_calling_example.py | - examples/models/core/llama/convert_checkpoint.py | - examples/models/core/llama/summarize_long.py | - examples/models/core/mamba/convert_checkpoint.py | - examples/models/core/mllama/convert_checkpoint.py | - examples/models/core/multimodal/__init__.py | - examples/models/core/multimodal/build_multimodal_engine.py | - examples/models/core/multimodal/eval.py | - examples/models/core/multimodal/run.py | - examples/models/core/multimodal/utils.py | - examples/models/core/nemotron_nas/calibration_utils.py | - examples/models/core/nemotron_nas/convert_checkpoint.py | - examples/models/core/phi/convert_checkpoint.py | - examples/models/core/qwen/convert_checkpoint.py | - examples/models/core/qwen2audio/run.py | - examples/models/core/qwen2audio/run_chat.py | - examples/models/core/qwen2audio/utils.py | - examples/models/core/qwenvl/run.py | - examples/models/core/qwenvl/run_chat.py | - examples/models/core/qwenvl/show_pic.py | - examples/models/core/qwenvl/vit_onnx_trt.py | - examples/models/core/recurrentgemma/convert_checkpoint.py | - examples/models/core/vit/convert_checkpoint.py | - examples/models/core/whisper/convert_checkpoint.py | - examples/models/core/whisper/distil_whisper/convert_from_distil_whisper.py | - examples/models/core/whisper/run.py | - examples/models/core/whisper/tokenizer.py | - examples/models/core/whisper/whisper_utils.py | - examples/ngram/run_dtm_ngram.py | - examples/openai_triton/manual_plugin/build.py | - examples/openai_triton/manual_plugin/fmha_triton.py | - examples/openai_triton/manual_plugin/plugin.py | - examples/openai_triton/manual_plugin/run.py | - examples/openai_triton/plugin_autogen/build_engine.py | - examples/openai_triton/plugin_autogen/kernel_config.py | - examples/openai_triton/plugin_autogen/run_engine.py | - examples/python_plugin/build_lookup.py | - examples/python_plugin/plugin_lib/__init__.py | - examples/python_plugin/plugin_lib/lookup_kernel.py | - examples/python_plugin/plugin_lib/lookup_plugin.py | - examples/python_plugin/run_lookup.py | - examples/quantization/quantize.py | examples/quantization/quantize_mixed_precision_moe.py | examples/ray_orchestrator/llm_inference_async_ray.py | examples/ray_orchestrator/llm_inference_distributed_ray.py | - examples/redrafter/convert_checkpoint.py | - examples/run.py | examples/scaffolding/contrib/AsyncGeneration/stream_generation_controller.py | examples/scaffolding/contrib/DeepConf/run_generation.py | examples/scaffolding/contrib/Dynasor/scaffolding_dynasor_run.py | examples/scaffolding/contrib/TreeInference/run_mcts_example.py | examples/scaffolding/contrib/TreeInference/run_tot_example.py | - examples/scaffolding/contrib/mcp/e2b/e2bserver.py | - examples/scaffolding/contrib/mcp/e2b/main.py | - examples/scaffolding/contrib/mcp/mcptest.py | - examples/scaffolding/contrib/mcp/weather/weather.py | - examples/scaffolding/contrib/mcp/websearch/main.py | - examples/scaffolding/contrib/mcp/websearch/websearch.py | examples/scaffolding/run_basic_generation.py | examples/scaffolding/run_best_of_n_with_reward.py | examples/scaffolding/run_majority_vote_aime24.py | @@ -1517,8 +853,6 @@ legacy-files: &legacy_files | examples/serve/openai_completion_client.py | examples/serve/openai_completion_client_for_lora.py | examples/serve/openai_completion_client_json_schema.py | - examples/summarize.py | - examples/utils.py | examples/wide_ep/ep_load_balancer/generate_eplb_config.py | examples/wide_ep/ep_load_balancer/report_load_statistics.py | examples/wide_ep/ep_load_balancer/utils.py | @@ -1530,7 +864,6 @@ legacy-files: &legacy_files | scripts/check_test_list.py | scripts/dco_check.py | scripts/format_test_list.py | - scripts/generate_duration.py | scripts/generate_lock_file.py | scripts/get_wheel_from_package.py | scripts/git_replace.py | @@ -1541,7 +874,6 @@ legacy-files: &legacy_files | setup.py | tensorrt_llm/__init__.py | tensorrt_llm/_ray_utils.py | - tensorrt_llm/_tensorrt_engine/__init__.py | tensorrt_llm/_torch/__init__.py | tensorrt_llm/_torch/attention_backend/__init__.py | tensorrt_llm/_torch/attention_backend/flashinfer.py | @@ -1745,7 +1077,6 @@ legacy-files: &legacy_files | tensorrt_llm/_torch/pyexecutor/guided_decoder.py | tensorrt_llm/_torch/pyexecutor/handle_additional_outputs.py | tensorrt_llm/_torch/pyexecutor/handle_logits.py | - tensorrt_llm/_torch/pyexecutor/kv_cache_connector.py | tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py | tensorrt_llm/_torch/pyexecutor/layerwise_nvtx_marker.py | tensorrt_llm/_torch/pyexecutor/llm_request.py | @@ -1783,11 +1114,6 @@ legacy-files: &legacy_files | tensorrt_llm/bench/benchmark/utils/asynchronous.py | tensorrt_llm/bench/benchmark/utils/general.py | tensorrt_llm/bench/benchmark/utils/processes.py | - tensorrt_llm/bench/build/__init__.py | - tensorrt_llm/bench/build/build.py | - tensorrt_llm/bench/build/dataclasses.py | - tensorrt_llm/bench/build/tuning.py | - tensorrt_llm/bench/build/utils.py | tensorrt_llm/bench/dataclasses/__init__.py | tensorrt_llm/bench/dataclasses/configuration.py | tensorrt_llm/bench/dataclasses/engine.py | @@ -1797,13 +1123,9 @@ legacy-files: &legacy_files | tensorrt_llm/bench/dataclasses/statistics.py | tensorrt_llm/bench/utils/__init__.py | tensorrt_llm/bench/utils/data.py | - tensorrt_llm/builder.py | tensorrt_llm/commands/__init__.py | tensorrt_llm/commands/bench.py | - tensorrt_llm/commands/build.py | tensorrt_llm/commands/eval.py | - tensorrt_llm/commands/prune.py | - tensorrt_llm/commands/refit.py | tensorrt_llm/commands/serve.py | tensorrt_llm/evaluate/__init__.py | tensorrt_llm/evaluate/cnn_dailymail.py | @@ -1839,23 +1161,7 @@ legacy-files: &legacy_files | tensorrt_llm/inputs/multimodal.py | tensorrt_llm/inputs/registry.py | tensorrt_llm/inputs/utils.py | - tensorrt_llm/layers/__init__.py | - tensorrt_llm/layers/activation.py | - tensorrt_llm/layers/attention.py | - tensorrt_llm/layers/cast.py | - tensorrt_llm/layers/conv.py | - tensorrt_llm/layers/embedding.py | - tensorrt_llm/layers/language_adapter.py | - tensorrt_llm/layers/linear.py | - tensorrt_llm/layers/lora.py | - tensorrt_llm/layers/mlp.py | - tensorrt_llm/layers/moe.py | - tensorrt_llm/layers/normalization.py | - tensorrt_llm/layers/pooling.py | - tensorrt_llm/layers/recurrent.py | - tensorrt_llm/layers/ssm.py | tensorrt_llm/llmapi/__init__.py | - tensorrt_llm/llmapi/build_cache.py | tensorrt_llm/llmapi/disagg_utils.py | tensorrt_llm/llmapi/kv_cache_type.py | tensorrt_llm/llmapi/llm.py | @@ -1878,179 +1184,17 @@ legacy-files: &legacy_files | tensorrt_llm/metrics/enums.py | tensorrt_llm/models/__init__.py | tensorrt_llm/models/automodel.py | - tensorrt_llm/models/baichuan/__init__.py | - tensorrt_llm/models/baichuan/config.py | - tensorrt_llm/models/baichuan/convert.py | - tensorrt_llm/models/baichuan/model.py | - tensorrt_llm/models/bert/__init__.py | - tensorrt_llm/models/bert/config.py | - tensorrt_llm/models/bert/convert.py | - tensorrt_llm/models/bert/model.py | - tensorrt_llm/models/bloom/__init__.py | - tensorrt_llm/models/bloom/model.py | - tensorrt_llm/models/chatglm/__init__.py | - tensorrt_llm/models/chatglm/config.py | - tensorrt_llm/models/chatglm/convert.py | - tensorrt_llm/models/chatglm/model.py | - tensorrt_llm/models/clip/__init__.py | - tensorrt_llm/models/clip/model.py | - tensorrt_llm/models/cogvlm/__init__.py | - tensorrt_llm/models/cogvlm/config.py | - tensorrt_llm/models/cogvlm/convert.py | - tensorrt_llm/models/cogvlm/model.py | - tensorrt_llm/models/commandr/__init__.py | - tensorrt_llm/models/commandr/config.py | - tensorrt_llm/models/commandr/model.py | tensorrt_llm/models/convert_utils.py | - tensorrt_llm/models/dbrx/__init__.py | - tensorrt_llm/models/dbrx/config.py | - tensorrt_llm/models/dbrx/model.py | - tensorrt_llm/models/deepseek_v1/__init__.py | - tensorrt_llm/models/deepseek_v1/config.py | - tensorrt_llm/models/deepseek_v1/convert.py | - tensorrt_llm/models/deepseek_v1/model.py | - tensorrt_llm/models/deepseek_v2/__init__.py | - tensorrt_llm/models/deepseek_v2/config.py | - tensorrt_llm/models/deepseek_v2/convert.py | - tensorrt_llm/models/deepseek_v2/model.py | - tensorrt_llm/models/dit/__init__.py | - tensorrt_llm/models/dit/model.py | - tensorrt_llm/models/eagle/__init__.py | - tensorrt_llm/models/eagle/config.py | - tensorrt_llm/models/eagle/model.py | - tensorrt_llm/models/enc_dec/__init__.py | - tensorrt_llm/models/enc_dec/model.py | - tensorrt_llm/models/falcon/__init__.py | - tensorrt_llm/models/falcon/config.py | - tensorrt_llm/models/falcon/convert.py | - tensorrt_llm/models/falcon/model.py | - tensorrt_llm/models/gemma/__init__.py | - tensorrt_llm/models/gemma/config.py | - tensorrt_llm/models/gemma/convert.py | - tensorrt_llm/models/gemma/model.py | - tensorrt_llm/models/gemma/smoothquant.py | - tensorrt_llm/models/gemma/utils/__init__.py | - tensorrt_llm/models/gemma/utils/layers.py | - tensorrt_llm/models/gemma/utils/modules.py | - tensorrt_llm/models/gemma/utils/params.py | - tensorrt_llm/models/gemma/utils/positional_embeddings.py | - tensorrt_llm/models/gemma/utils/sampler.py | - tensorrt_llm/models/gemma/utils/transformer.py | - tensorrt_llm/models/gemma/weight.py | - tensorrt_llm/models/generation_mixin.py | - tensorrt_llm/models/gpt/__init__.py | - tensorrt_llm/models/gpt/config.py | - tensorrt_llm/models/gpt/convert.py | - tensorrt_llm/models/gpt/model.py | - tensorrt_llm/models/gptj/__init__.py | - tensorrt_llm/models/gptj/config.py | - tensorrt_llm/models/gptj/convert.py | - tensorrt_llm/models/gptj/model.py | - tensorrt_llm/models/gptneox/__init__.py | - tensorrt_llm/models/gptneox/model.py | - tensorrt_llm/models/grok/__init__.py | - tensorrt_llm/models/grok/convert.py | - tensorrt_llm/models/grok/model.py | - tensorrt_llm/models/grok/weight.py | - tensorrt_llm/models/llama/__init__.py | - tensorrt_llm/models/llama/config.py | - tensorrt_llm/models/llama/convert.py | - tensorrt_llm/models/llama/model.py | - tensorrt_llm/models/mamba/__init__.py | - tensorrt_llm/models/mamba/config.py | - tensorrt_llm/models/mamba/convert.py | - tensorrt_llm/models/mamba/model.py | - tensorrt_llm/models/medusa/__init__.py | - tensorrt_llm/models/medusa/config.py | - tensorrt_llm/models/medusa/model.py | - tensorrt_llm/models/medusa/weight.py | - tensorrt_llm/models/mllama/__init__.py | - tensorrt_llm/models/mllama/config.py | - tensorrt_llm/models/mllama/model.py | - tensorrt_llm/models/mmdit_sd3/__init__.py | - tensorrt_llm/models/mmdit_sd3/config.py | - tensorrt_llm/models/mmdit_sd3/model.py | - tensorrt_llm/models/model_weights_loader.py | tensorrt_llm/models/modeling_utils.py | - tensorrt_llm/models/mpt/__init__.py | - tensorrt_llm/models/mpt/model.py | - tensorrt_llm/models/multimodal_encoders/__init__.py | - tensorrt_llm/models/multimodal_encoders/config.py | - tensorrt_llm/models/multimodal_encoders/model.py | - tensorrt_llm/models/nemotron_nas/__init__.py | - tensorrt_llm/models/nemotron_nas/config.py | - tensorrt_llm/models/nemotron_nas/convert.py | - tensorrt_llm/models/nemotron_nas/layer_config.py | - tensorrt_llm/models/nemotron_nas/model.py | - tensorrt_llm/models/opt/__init__.py | - tensorrt_llm/models/opt/model.py | - tensorrt_llm/models/phi/__init__.py | - tensorrt_llm/models/phi/config.py | - tensorrt_llm/models/phi/convert.py | - tensorrt_llm/models/phi/model.py | - tensorrt_llm/models/phi3/__init__.py | - tensorrt_llm/models/phi3/config.py | - tensorrt_llm/models/phi3/convert.py | - tensorrt_llm/models/phi3/model.py | - tensorrt_llm/models/phi3/split_weights.py | - tensorrt_llm/models/qwen/__init__.py | - tensorrt_llm/models/qwen/config.py | - tensorrt_llm/models/qwen/convert.py | - tensorrt_llm/models/qwen/model.py | - tensorrt_llm/models/qwen/utils.py | - tensorrt_llm/models/recurrentgemma/__init__.py | - tensorrt_llm/models/recurrentgemma/model.py | - tensorrt_llm/models/redrafter/__init__.py | - tensorrt_llm/models/redrafter/drafter.py | - tensorrt_llm/models/redrafter/model.py | - tensorrt_llm/models/redrafter/redrafter_helper.py | - tensorrt_llm/models/stdit/__init__.py | - tensorrt_llm/models/stdit/config.py | - tensorrt_llm/models/stdit/model.py | - tensorrt_llm/models/unet/__init__.py | - tensorrt_llm/models/unet/attention.py | - tensorrt_llm/models/unet/embeddings.py | - tensorrt_llm/models/unet/pp/__init__.py | - tensorrt_llm/models/unet/pp/attention.py | - tensorrt_llm/models/unet/pp/conv2d.py | - tensorrt_llm/models/unet/pp/groupnorm.py | - tensorrt_llm/models/unet/pp/unet_pp.py | - tensorrt_llm/models/unet/resnet.py | - tensorrt_llm/models/unet/unet_2d_blocks.py | - tensorrt_llm/models/unet/unet_2d_condition.py | - tensorrt_llm/models/unet/weights.py | - tensorrt_llm/network.py | - tensorrt_llm/parameter.py | - tensorrt_llm/plugin/__init__.py | - tensorrt_llm/plugin/plugin.py | tensorrt_llm/quantization/__init__.py | tensorrt_llm/quantization/functional.py | - tensorrt_llm/quantization/image_processing.py | - tensorrt_llm/quantization/layers.py | tensorrt_llm/quantization/mode.py | - tensorrt_llm/quantization/quantize.py | - tensorrt_llm/quantization/quantize_by_modelopt.py | tensorrt_llm/quantization/utils/__init__.py | tensorrt_llm/quantization/utils/fp4_utils.py | tensorrt_llm/quantization/utils/fp8_utils.py | tensorrt_llm/ray_stub.py | tensorrt_llm/runtime/__init__.py | - tensorrt_llm/runtime/enc_dec_model_runner.py | - tensorrt_llm/runtime/generation.py | - tensorrt_llm/runtime/kv_cache_manager.py | - tensorrt_llm/runtime/medusa_utils.py | tensorrt_llm/runtime/memory_pools/__init__.py | - tensorrt_llm/runtime/memory_pools/memory_pools_allocator.py | - tensorrt_llm/runtime/memory_pools/pool.py | - tensorrt_llm/runtime/memory_pools/pools_kv_cache_manager.py | - tensorrt_llm/runtime/model_runner.py | - tensorrt_llm/runtime/model_runner_cpp.py | - tensorrt_llm/runtime/multimodal_model_runner.py | - tensorrt_llm/runtime/processor_wrapper/__init__.py | - tensorrt_llm/runtime/processor_wrapper/mllama_processor_wrapper.py | - tensorrt_llm/runtime/processor_wrapper/processor_wrapper.py | - tensorrt_llm/runtime/redrafter_utils.py | - tensorrt_llm/runtime/session.py | tensorrt_llm/scaffolding/__init__.py | tensorrt_llm/scaffolding/benchmark.py | tensorrt_llm/scaffolding/contrib/AsyncGeneration/stream_generation.py | @@ -2105,12 +1249,7 @@ legacy-files: &legacy_files | tensorrt_llm/tokenizer/tokenizer.py | tensorrt_llm/tools/__init__.py | tensorrt_llm/tools/importlib_utils.py | - tensorrt_llm/tools/multimodal_builder.py | - tensorrt_llm/tools/onnx_utils.py | tensorrt_llm/tools/plugin_gen/__init__.py | - tensorrt_llm/tools/plugin_gen/core.py | - tensorrt_llm/tools/plugin_gen/plugin_gen.py | - tensorrt_llm/tools/plugin_gen/shape_infer.py | tensorrt_llm/tools/ppl.py | tensorrt_llm/tools/profiler/nsys_profile_tools/gputrc2graph.py | tensorrt_llm/version.py | @@ -2119,9 +1258,7 @@ legacy-files: &legacy_files | tests/integration/defs/accuracy/accuracy_core.py | tests/integration/defs/accuracy/scripts/collect_evaluated_accuracies.py | tests/integration/defs/accuracy/scripts/compute_theta_and_thresholds.py | - tests/integration/defs/accuracy/test_cli_flow.py | tests/integration/defs/accuracy/test_disaggregated_serving.py | - tests/integration/defs/accuracy/test_llm_api.py | tests/integration/defs/accuracy/test_llm_api_autodeploy.py | tests/integration/defs/accuracy/test_llm_api_pytorch.py | tests/integration/defs/accuracy/test_llm_api_pytorch_ray.py | @@ -2137,53 +1274,22 @@ legacy-files: &legacy_files | tests/integration/defs/disaggregated/test_disaggregated_etcd.py | tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py | tests/integration/defs/disaggregated/test_workers.py | - tests/integration/defs/examples/run_llm_fp8_quant_llama_70b.py | tests/integration/defs/examples/run_llm_quickstart_atexit.py | tests/integration/defs/examples/serve/test_serve.py | tests/integration/defs/examples/serve/test_serve_negative.py | tests/integration/defs/examples/test_ad_guided_decoding.py | - tests/integration/defs/examples/test_bert.py | - tests/integration/defs/examples/test_bindings.py | - tests/integration/defs/examples/test_chatglm.py | - tests/integration/defs/examples/test_commandr.py | - tests/integration/defs/examples/test_draft_target_model.py | - tests/integration/defs/examples/test_eagle.py | - tests/integration/defs/examples/test_enc_dec.py | - tests/integration/defs/examples/test_exaone.py | - tests/integration/defs/examples/test_gemma.py | tests/integration/defs/examples/test_gpt.py | - tests/integration/defs/examples/test_gptj.py | - tests/integration/defs/examples/test_granite.py | - tests/integration/defs/examples/test_internlm.py | - tests/integration/defs/examples/test_llama.py | tests/integration/defs/examples/test_llm_api_with_mpi.py | - tests/integration/defs/examples/test_mamba.py | - tests/integration/defs/examples/test_medusa.py | - tests/integration/defs/examples/test_mistral.py | - tests/integration/defs/examples/test_mixtral.py | - tests/integration/defs/examples/test_multimodal.py | - tests/integration/defs/examples/test_nemotron.py | - tests/integration/defs/examples/test_nemotron_nas.py | - tests/integration/defs/examples/test_ngram.py | - tests/integration/defs/examples/test_openai.py | tests/integration/defs/examples/test_phi.py | - tests/integration/defs/examples/test_qwen.py | - tests/integration/defs/examples/test_qwen2audio.py | - tests/integration/defs/examples/test_qwenvl.py | tests/integration/defs/examples/test_ray.py | - tests/integration/defs/examples/test_recurrentgemma.py | - tests/integration/defs/examples/test_redrafter.py | - tests/integration/defs/examples/test_whisper.py | tests/integration/defs/llmapi/__init__.py | tests/integration/defs/llmapi/_run_llmapi_llm.py | tests/integration/defs/llmapi/test_llm_api_connector.py | tests/integration/defs/llmapi/test_llm_api_qa.py | - tests/integration/defs/llmapi/test_llm_e2e.py | tests/integration/defs/llmapi/test_llm_examples.py | tests/integration/defs/local_venv.py | tests/integration/defs/perf/__init__.py | tests/integration/defs/perf/allowed_configs.py | - tests/integration/defs/perf/build.py | tests/integration/defs/perf/create_perf_comparison_report.py | tests/integration/defs/perf/data.py | tests/integration/defs/perf/data_export.py | @@ -2207,34 +1313,18 @@ legacy-files: &legacy_files | tests/integration/defs/test_sanity.py | tests/integration/defs/test_unittests.py | tests/integration/defs/triton_server/__init__.py | - tests/integration/defs/triton_server/build_engines.py | tests/integration/defs/triton_server/common.py | tests/integration/defs/triton_server/conftest.py | - tests/integration/defs/triton_server/local_venv.py | - tests/integration/defs/triton_server/rcca/bug_4323566/inflight_batcher_llm_client_with_end_id.py | - tests/integration/defs/triton_server/runner_interface.py | tests/integration/defs/triton_server/test_list_parser.py | - tests/integration/defs/triton_server/test_triton.py | - tests/integration/defs/triton_server/test_triton_llm.py | - tests/integration/defs/triton_server/test_triton_memleak.py | - tests/integration/defs/triton_server/test_triton_multi_node.py | - tests/integration/defs/triton_server/test_triton_rcca.py | tests/integration/defs/triton_server/trt_test_alternative.py | tests/integration/defs/trt_test_alternative.py | tests/integration/defs/utils/__init__.py | tests/integration/defs/utils/periodic_junit.py | tests/integration/defs/utils/timeout_manager.py | tests/microbenchmarks/all_reduce.py | - tests/microbenchmarks/build_time_benchmark.py | - tests/microbenchmarks/build_time_dashboard.py | tests/scripts/allreduce_perf/allreduce_heuristic_code_gen.py | tests/scripts/allreduce_perf/allreduce_perf_viz.py | tests/scripts/iteration_log_parser.py | - tests/scripts/perf-sanity/parse_benchmark_results.py | - tests/scripts/perf-sanity/run_benchmark_serve.py | - tests/unittest/_torch/attention/sparse/test_dsa_indexer.py | - tests/unittest/_torch/attention/sparse/test_flash_mla.py | - tests/unittest/_torch/attention/sparse/test_rocketkv.py | tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py | tests/unittest/_torch/attention/test_attention.py | tests/unittest/_torch/attention/test_attention_mla.py | @@ -2255,7 +1345,6 @@ legacy-files: &legacy_files | tests/unittest/_torch/misc/test_virtual_memory.py | tests/unittest/_torch/modeling/test_modeling_bert.py | tests/unittest/_torch/modeling/test_modeling_clip.py | - tests/unittest/_torch/modeling/test_modeling_exaone4.py | tests/unittest/_torch/modeling/test_modeling_gemma3.py | tests/unittest/_torch/modeling/test_modeling_gpt_oss.py | tests/unittest/_torch/modeling/test_modeling_llama.py | @@ -2279,8 +1368,6 @@ legacy-files: &legacy_files | tests/unittest/_torch/modules/test_moe_routing.py | tests/unittest/_torch/modules/test_rotary_embedding.py | tests/unittest/_torch/modules/test_triton_linear.py | - tests/unittest/_torch/modules/tests_lora_modules/test_lora_attention_pytorch_flow_vs_trt.py | - tests/unittest/_torch/modules/tests_lora_modules/test_lora_plugin_vs_lora_op.py | tests/unittest/_torch/multi_gpu/test_allreduce.py | tests/unittest/_torch/multi_gpu/test_alltoall.py | tests/unittest/_torch/multi_gpu/test_ar_residual_norm.py | @@ -2307,24 +1394,10 @@ legacy-files: &legacy_files | tests/unittest/_torch/sampler/test_beam_search.py | tests/unittest/_torch/sampler/test_best_of_n.py | tests/unittest/_torch/sampler/test_trtllm_sampler.py | - tests/unittest/_torch/speculative/test_draft_target.py | - tests/unittest/_torch/speculative/test_draft_token_tree_sampling.py | - tests/unittest/_torch/speculative/test_draft_token_tree_verification.py | - tests/unittest/_torch/speculative/test_dynamic_spec_decode.py | tests/unittest/_torch/speculative/test_eagle3.py | - tests/unittest/_torch/speculative/test_kv_cache_reuse.py | - tests/unittest/_torch/speculative/test_mtp.py | - tests/unittest/_torch/speculative/test_ngram.py | - tests/unittest/_torch/speculative/test_save_state.py | - tests/unittest/_torch/speculative/test_spec_gate.py | - tests/unittest/_torch/speculative/test_torch_rejection_sampling.py | - tests/unittest/_torch/speculative/test_user_provided.py | tests/unittest/_torch/test_connector.py | tests/unittest/_torch/test_torch_multi_arange.py | tests/unittest/_torch/thop/parallel/deep_gemm_tests.py | - tests/unittest/_torch/thop/parallel/test_causal_conv1d_op.py | - tests/unittest/_torch/thop/parallel/test_cublas_mm.py | - tests/unittest/_torch/thop/parallel/test_custom_ops.py | tests/unittest/_torch/thop/parallel/test_dsv3_fused_a_gemm.py | tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py | tests/unittest/_torch/thop/parallel/test_finegrained_mixed_dtype_gemm.py | @@ -2338,11 +1411,6 @@ legacy-files: &legacy_files | tests/unittest/_torch/thop/parallel/test_fp8_per_tensor_scale_tllmg_gemm.py | tests/unittest/_torch/thop/parallel/test_fp8_quantize.py | tests/unittest/_torch/thop/parallel/test_fp8_rowwise_linear.py | - tests/unittest/_torch/thop/parallel/test_fused_qk_norm_rope.py | - tests/unittest/_torch/thop/parallel/test_logits_bitmask_op.py | - tests/unittest/_torch/thop/parallel/test_mamba2_chunk_ss_update.py | - tests/unittest/_torch/thop/parallel/test_mamba_conv1d_op.py | - tests/unittest/_torch/thop/parallel/test_noaux_tc.py | tests/unittest/_torch/thop/parallel/test_scaled_mm.py | tests/unittest/_torch/thop/parallel/test_selective_scan_op.py | tests/unittest/_torch/thop/parallel/test_tinygemm2.py | @@ -2386,12 +1454,10 @@ legacy-files: &legacy_files | tests/unittest/llmapi/apps/_test_openai_chat_harmony.py | tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py | tests/unittest/llmapi/apps/_test_openai_completions.py | - tests/unittest/llmapi/apps/_test_openai_consistent_chat.py | tests/unittest/llmapi/apps/_test_openai_lora.py | tests/unittest/llmapi/apps/_test_openai_metrics.py | tests/unittest/llmapi/apps/_test_openai_misc.py | tests/unittest/llmapi/apps/_test_openai_mmencoder.py | - tests/unittest/llmapi/apps/_test_openai_multi_chat.py | tests/unittest/llmapi/apps/_test_openai_multi_gpu.py | tests/unittest/llmapi/apps/_test_openai_multi_nodes.py | tests/unittest/llmapi/apps/_test_openai_perf_metrics.py | @@ -2414,15 +1480,12 @@ legacy-files: &legacy_files | tests/unittest/llmapi/run_llm_exit.py | tests/unittest/llmapi/run_llm_with_postproc.py | tests/unittest/llmapi/test_additional_model_outputs.py | - tests/unittest/llmapi/test_build_cache.py | tests/unittest/llmapi/test_executor.py | tests/unittest/llmapi/test_gc_utils.py | tests/unittest/llmapi/test_llm.py | tests/unittest/llmapi/test_llm_args.py | tests/unittest/llmapi/test_llm_download.py | tests/unittest/llmapi/test_llm_kv_cache_events.py | - tests/unittest/llmapi/test_llm_models.py | - tests/unittest/llmapi/test_llm_multi_gpu.py | tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py | tests/unittest/llmapi/test_llm_pytorch.py | tests/unittest/llmapi/test_llm_quant.py | @@ -2433,24 +1496,14 @@ legacy-files: &legacy_files | tests/unittest/llmapi/test_serialization.py | tests/unittest/llmapi/test_utils.py | tests/unittest/others/__init__.py | - tests/unittest/others/test_builder.py | tests/unittest/others/test_convert_spec_decoding_mask_to_packed_mask.py | - tests/unittest/others/test_debugging_api.py | tests/unittest/others/test_exception.py | tests/unittest/others/test_export.py | - tests/unittest/others/test_graph_rewriter.py | - tests/unittest/others/test_kv_cache_manager.py | tests/unittest/others/test_kv_cache_transceiver.py | tests/unittest/others/test_kv_cache_update.py | - tests/unittest/others/test_layer.py | tests/unittest/others/test_mapping.py | - tests/unittest/others/test_model_dtype.py | - tests/unittest/others/test_module.py | tests/unittest/others/test_multimodal_registry.py | - tests/unittest/others/test_plugins.py | - tests/unittest/others/test_precision_control.py | tests/unittest/others/test_pretrained_config.py | - tests/unittest/others/test_session.py | tests/unittest/others/test_time_breakdown.py | tests/unittest/profile_utils.py | tests/unittest/scaffolding/__init__.py | @@ -2459,141 +1512,14 @@ legacy-files: &legacy_files | tests/unittest/scaffolding/test_scaffolding.py | tests/unittest/scaffolding/test_task_collection.py | tests/unittest/scaffolding/test_worker.py | - tests/unittest/test_model_runner_cpp.py | tests/unittest/test_pip_install.py | tests/unittest/tools/__init__.py | - tests/unittest/tools/plugin_gen/__init__.py | - tests/unittest/tools/plugin_gen/kernel_config.py | - tests/unittest/tools/plugin_gen/test_core.py | - tests/unittest/tools/plugin_gen/test_plugin_gen.py | - tests/unittest/tools/plugin_gen/test_shape_infer.py | tests/unittest/tools/test_prepare_dataset.py | tests/unittest/tools/test_test_to_stage_mapping.py | - tests/unittest/trt/__init__.py | - tests/unittest/trt/attention/test_bert_attention.py | - tests/unittest/trt/attention/test_gpt_attention.py | - tests/unittest/trt/attention/test_gpt_attention_IFB.py | - tests/unittest/trt/attention/test_gpt_attention_no_cache.py | - tests/unittest/trt/attention/test_sage_attention.py | - tests/unittest/trt/functional/__init__.py | - tests/unittest/trt/functional/test_alibi.py | - tests/unittest/trt/functional/test_allreduce_norm.py | - tests/unittest/trt/functional/test_allreduce_prepost_residual_norm.py | - tests/unittest/trt/functional/test_arange.py | - tests/unittest/trt/functional/test_argmax.py | - tests/unittest/trt/functional/test_assertion.py | - tests/unittest/trt/functional/test_avg_pool2d.py | - tests/unittest/trt/functional/test_cast.py | - tests/unittest/trt/functional/test_conv2d.py | - tests/unittest/trt/functional/test_conv3d.py | - tests/unittest/trt/functional/test_cos.py | - tests/unittest/trt/functional/test_cumsum.py | - tests/unittest/trt/functional/test_dora.py | - tests/unittest/trt/functional/test_einsum.py | - tests/unittest/trt/functional/test_embedding_single_gpu.py | - tests/unittest/trt/functional/test_exp.py | - tests/unittest/trt/functional/test_expand.py | - tests/unittest/trt/functional/test_flatten.py | - tests/unittest/trt/functional/test_flip.py | - tests/unittest/trt/functional/test_fp4_gemm.py | - tests/unittest/trt/functional/test_fp4_gemm_ootb.py | - tests/unittest/trt/functional/test_gather.py | - tests/unittest/trt/functional/test_gather_nd.py | - tests/unittest/trt/functional/test_geglu.py | - tests/unittest/trt/functional/test_gelu.py | - tests/unittest/trt/functional/test_gemm_swiglu.py | - tests/unittest/trt/functional/test_group_norm.py | - tests/unittest/trt/functional/test_identity.py | - tests/unittest/trt/functional/test_index_select.py | - tests/unittest/trt/functional/test_interpolate.py | - tests/unittest/trt/functional/test_logsoftmax.py | - tests/unittest/trt/functional/test_lora.py | - tests/unittest/trt/functional/test_low_latency_gemm.py | - tests/unittest/trt/functional/test_mamba_conv1d.py | - tests/unittest/trt/functional/test_masked_scatter.py | - tests/unittest/trt/functional/test_masked_select.py | - tests/unittest/trt/functional/test_matmul.py | - tests/unittest/trt/functional/test_meshgrid2d.py | - tests/unittest/trt/functional/test_moe.py | - tests/unittest/trt/functional/test_nccl.py | - tests/unittest/trt/functional/test_nonzero.py | - tests/unittest/trt/functional/test_outer.py | - tests/unittest/trt/functional/test_pad.py | - tests/unittest/trt/functional/test_permute.py | - tests/unittest/trt/functional/test_pp_reduce_scatter.py | - tests/unittest/trt/functional/test_quant.py | - tests/unittest/trt/functional/test_rearrange.py | - tests/unittest/trt/functional/test_repeat.py | - tests/unittest/trt/functional/test_repeat_interleave.py | - tests/unittest/trt/functional/test_rg_lru.py | - tests/unittest/trt/functional/test_sample.py | - tests/unittest/trt/functional/test_scatter.py | - tests/unittest/trt/functional/test_scatter_nd.py | - tests/unittest/trt/functional/test_select.py | - tests/unittest/trt/functional/test_selective_scan.py | - tests/unittest/trt/functional/test_sigmoid.py | - tests/unittest/trt/functional/test_silu.py | - tests/unittest/trt/functional/test_sin.py | - tests/unittest/trt/functional/test_slice.py | - tests/unittest/trt/functional/test_softplus.py | - tests/unittest/trt/functional/test_split.py | - tests/unittest/trt/functional/test_squeeze.py | - tests/unittest/trt/functional/test_swiglu.py | - tests/unittest/trt/functional/test_topk.py | - tests/unittest/trt/functional/test_transpose.py | - tests/unittest/trt/functional/test_unbind.py | - tests/unittest/trt/functional/test_unsqueeze.py | - tests/unittest/trt/functional/test_view.py | - tests/unittest/trt/functional/test_where.py | - tests/unittest/trt/model/__init__.py | - tests/unittest/trt/model/eagle/test_decode_draft_tokens_plugin.py | - tests/unittest/trt/model/eagle/test_prepare_drafter_inputs_plugin.py | - tests/unittest/trt/model/eagle/test_sample_accept_draft_tokens_plugin.py | - tests/unittest/trt/model/redrafter/test_beams2tree.py | - tests/unittest/trt/model/redrafter/test_draft_token.py | - tests/unittest/trt/model/redrafter/test_draft_token_indices.py | - tests/unittest/trt/model/redrafter/test_gather_beams.py | - tests/unittest/trt/model/redrafter/test_mask.py | - tests/unittest/trt/model/redrafter/test_packed_position_ids.py | - tests/unittest/trt/model/redrafter/test_prefix_match_indices.py | - tests/unittest/trt/model/redrafter/test_prepare_input.py | - tests/unittest/trt/model/redrafter/test_process_logits.py | - tests/unittest/trt/model/redrafter/test_top1.py | - tests/unittest/trt/model/redrafter/test_unpack_gen_data.py | - tests/unittest/trt/model/redrafter/test_validate.py | - tests/unittest/trt/model/test_gpt.py | - tests/unittest/trt/model/test_gpt_e2e.py | - tests/unittest/trt/model/test_llama.py | - tests/unittest/trt/model/test_mamba.py | - tests/unittest/trt/model/test_mistral.py | - tests/unittest/trt/model/test_nemotron_nas.py | - tests/unittest/trt/model/test_phi.py | - tests/unittest/trt/model/test_unet.py | - tests/unittest/trt/model_api/test_model_api_multi_gpu.py | - tests/unittest/trt/model_api/test_model_level_api.py | - tests/unittest/trt/model_api/test_model_quantization.py | - tests/unittest/trt/python_plugin/plugin_wrapper_utils.py | - tests/unittest/trt/python_plugin/test_plugin_wrapper.py | - tests/unittest/trt/quantization/__init__.py | - tests/unittest/trt/quantization/_utils.py | - tests/unittest/trt/quantization/test_fp8_quantization.py | - tests/unittest/trt/quantization/test_fp8_rowwise_gemm.py | - tests/unittest/trt/quantization/test_functional.py | - tests/unittest/trt/quantization/test_mode.py | - tests/unittest/trt/quantization/test_moe_weight_only_quant_matmul.py | - tests/unittest/trt/quantization/test_qserve_gemm.py | - tests/unittest/trt/quantization/test_quant.py | - tests/unittest/trt/quantization/test_quant_layer.py | - tests/unittest/trt/quantization/test_smooth_quant_gemm.py | - tests/unittest/trt/quantization/test_smooth_quant_layer_norm.py | - tests/unittest/trt/quantization/test_smooth_quant_rms_norm.py | - tests/unittest/trt/quantization/test_weight_only_groupwise_quant_matmul.py | - tests/unittest/trt/quantization/test_weight_only_quant_matmul.py | tests/unittest/utils/__init__.py | tests/unittest/utils/cpp_paths.py | tests/unittest/utils/llm_data.py | tests/unittest/utils/runtime_defaults.py | - tests/unittest/utils/test_medusa_utils.py | tests/unittest/utils/test_prebuilt_whl_cpp_extensions.py | tests/unittest/utils/test_util.py | tests/unittest/utils/torch_ref.py | diff --git a/AGENTS.md b/AGENTS.md index 20fa71f70f7f..cf6c717d0928 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,7 +1,7 @@ # AGENTS.md TensorRT-LLM: open-source library for optimized LLM inference on NVIDIA GPUs. -Python and C++ codebase supporting TensorRT engine-based and PyTorch-based execution paths. +Python and C++ codebase with PyTorch and AutoDeploy execution paths. > If a `CLAUDE.local.md` file exists alongside this file, read and respect it — it contains developer-specific overrides that supplement this shared guidance. @@ -55,17 +55,16 @@ See [architecture diagram](.github/tava_architecture_diagram.md) for the full Me |---------|--------|-------------|----------| | **PyTorch** | Default | `TorchLlmArgs` | `_torch/pyexecutor/` → `PyExecutor` → PyTorch Engine | | **AutoDeploy** | Beta | `_torch/auto_deploy/` shim | `_torch/auto_deploy/shim/ad_executor.py` → adapts `PyExecutor` → graph transforms + torch.export | -| **TensorRT** | Legacy | `TrtLlmArgs` | `builder.py` → `trtllm.Executor` → TensorRT Engine | ### Shared C++ Core (via Nanobind) -Both PyTorch and TensorRT backends share these C++ components: +Both backends share these C++ components: - **Scheduling pipeline**: Scheduler → BatchManager (in-flight batching) → KV Cache Manager - **Decoding pipeline**: Decoder (token generation orchestration) → Sampling ### Request Flow ```text -HuggingFace Model → LLM API → Executor (PyTorch/AutoDeploy/TensorRT) +HuggingFace Model → LLM API → Executor (PyTorch/AutoDeploy) → Scheduler → Model Forward → Decoder → Sampling → Generated Tokens ``` @@ -84,7 +83,7 @@ HuggingFace Model → LLM API → Executor (PyTorch/AutoDeploy/TensorRT) | `tensorrt_llm/models/modeling_utils.py` | Base classes for all models (`PretrainedConfig`, `PretrainedModel`) | | `tensorrt_llm/executor/executor.py` | Execution abstraction (`GenerationExecutor`) | | `tensorrt_llm/models/automodel.py` | Auto-discovery and model registry | -| `tensorrt_llm/_torch/models/` | PyTorch backend model implementations (distinct from `models/` used by TensorRT backend) | +| `tensorrt_llm/_torch/models/` | PyTorch backend model implementations (distinct from the top-level `models/` package) | | `tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md` | Attention, MLA, backend families, sparse backends, metadata contracts, and KV-cache behavior - **read before modifying `tensorrt_llm/_torch/modules/attention.py`, `tensorrt_llm/_torch/modules/mla.py`, or `tensorrt_llm/_torch/attention_backend/`** | | `tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md` | MoE architecture, backends, communication, development patterns — **read before modifying MoE code** | | `CODING_GUIDELINES.md` | C++ and Python coding standards (referenced throughout, must read before contributing) | @@ -93,7 +92,7 @@ HuggingFace Model → LLM API → Executor (PyTorch/AutoDeploy/TensorRT) | Pattern | Key Points | |---------|------------| -| **Config hierarchy** | `BaseLlmArgs` → `TrtLlmArgs` / `TorchLlmArgs`, model-specific defaults override generics, Pydantic validation | +| **Config hierarchy** | `BaseLlmArgs` → `TorchLlmArgs`, model-specific defaults override generics, Pydantic validation | | **Model architecture** | Each model: `Config` (inherits `PretrainedConfig`) + `ForCausalLM` (inherits `PretrainedModel`) | | **Model defaults** | Architecture-specific overrides in `llm_utils.py` (attention kernels, quant, spec decoding, cache) | | **Attention backends** | `TorchLlmArgs.attn_backend` selects kernel: `TRTLLM` (default), `FlashInfer`, `FlashAttention` | @@ -125,7 +124,6 @@ Key files: - **Avoid broad exception handling** — catch specific exceptions, not bare `except:` (see `CODING_GUIDELINES.md`). - **One concern per PR** — avoid scope creep. If a PR touches unrelated areas, split it. - **User-facing configuration classes** - when editing or defining any user-facing configuration classes (particularly `BaseLlmArgs` or any class used in its fields), you **MUST** follow the Pydantic guidelines in `CODING_GUIDELINES.md`. -- **TensorRT backend is legacy** — `TrtLlmArgs` / `backend="tensorrt"` and all exclusive tooling (`trtllm-build`, `trtllm-refit`, `convert_checkpoint.py`, `ModelRunner*`) are legacy. Bug fixes OK; new features target PyTorch or AutoDeploy. ## Development Workflow diff --git a/benchmarks/README.md b/benchmarks/README.md deleted file mode 100644 index 87e2c06e6d5d..000000000000 --- a/benchmarks/README.md +++ /dev/null @@ -1,34 +0,0 @@ -# TensorRT-LLM Benchmarking - -For benchmarking TensorRT-LLM, use -[`trtllm-bench`](../docs/source/developer-guide/perf-benchmarking.md) and -`trtllm-serve`. - -This directory keeps the dataset preparation tools consumed by `trtllm-bench`: - -- `prepare_dataset.py` — generate benchmark datasets from real data or with - synthetic normal/uniform token-length distributions: - - ```bash - python3 prepare_dataset.py \ - --tokenizer \ - --output preprocessed_dataset.json \ - dataset \ - --dataset-name \ - --dataset-split \ - --dataset-input-key \ - --dataset-prompt-key \ - --dataset-output-key \ - [--num-requests 100] \ - [--max-input-len 1000] \ - [--output-len-dist 100,10] - ``` - - Synthetic variants: `python3 prepare_dataset.py ... token-norm-dist ...` and - `... token-unif-dist ...`. Run with `--help` for the full option list. - -- `utils/prepare_real_data.py`, `utils/prepare_synthetic_data.py` — the - subcommand implementations. -- `utils/generate_rand_loras.py` — generate random LoRA adapters for - LoRA benchmarking. -- `utils/convert_nemo_dataset.py` — convert NeMo chat datasets. diff --git a/benchmarks/prepare_dataset.py b/benchmarks/prepare_dataset.py deleted file mode 100644 index a1f1358f11ff..000000000000 --- a/benchmarks/prepare_dataset.py +++ /dev/null @@ -1,118 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import logging -from typing import Optional, Tuple - -import click -from pydantic import BaseModel, model_validator -from transformers import AutoTokenizer -from utils.prepare_real_data import dataset -from utils.prepare_synthetic_data import token_norm_dist, token_unif_dist - - -class RootArgs(BaseModel): - tokenizer: str - output: str - random_seed: int - task_id: int - std_out: bool - trust_remote_code: bool = False - rand_task_id: Optional[Tuple[int, int]] - lora_dir: Optional[str] = None - - @model_validator(mode='after') - def validate_tokenizer(self): - try: - tokenizer = AutoTokenizer.from_pretrained( - self.tokenizer, - padding_side='left', - trust_remote_code=self.trust_remote_code) - except EnvironmentError as e: - raise ValueError( - f"Cannot find a tokenizer from the given string because of {e}\nPlease set tokenizer to the directory that contains the tokenizer, or set to a model name in HuggingFace." - ) - tokenizer.pad_token = tokenizer.eos_token - self.tokenizer = tokenizer - - return self - - -@click.group(deprecated=True) -@click.option( - "--tokenizer", - required=True, - type=str, - help= - "Tokenizer dir for the benchmarked model, or the model name from HuggingFace." -) -@click.option("--output", - type=str, - help="Output json filename.", - default="preprocessed_dataset.json") -@click.option( - "--stdout", - is_flag=True, - help="Print output to stdout with a JSON dataset entry on each line.", - default=False) -@click.option("--random-seed", - required=False, - type=int, - help="random seed for token_ids", - default=420) -@click.option("--task-id", type=int, default=-1, help="LoRA task id") -@click.option("--rand-task-id", - type=int, - default=None, - nargs=2, - help="Random LoRA Tasks") -@click.option("--lora-dir", - type=str, - default=None, - help="Directory containing LoRA adapters") -@click.option("--log-level", - default="info", - type=click.Choice(['info', 'debug']), - help="Logging level.") -@click.option("--trust-remote-code", - is_flag=True, - default=False, - envvar="TRUST_REMOTE_CODE", - help="Trust remote code.") -@click.pass_context -def cli(ctx, **kwargs): - """This script generates benchmark dataset input (e.g. for trtllm-bench).""" - if kwargs['log_level'] == 'info': - logging.basicConfig(level=logging.INFO) - elif kwargs['log_level'] == 'debug': - logging.basicConfig(level=logging.DEBUG) - else: - raise ValueError(f"Unsupported logging level {kwargs['log_level']}") - - ctx.obj = RootArgs(tokenizer=kwargs['tokenizer'], - output=kwargs['output'], - std_out=kwargs['stdout'], - random_seed=kwargs['random_seed'], - task_id=kwargs['task_id'], - rand_task_id=kwargs['rand_task_id'], - lora_dir=kwargs['lora_dir'], - trust_remote_code=kwargs['trust_remote_code']) - - -cli.add_command(dataset) -cli.add_command(token_norm_dist) -cli.add_command(token_unif_dist) - -if __name__ == "__main__": - cli() diff --git a/benchmarks/utils/__init__.py b/benchmarks/utils/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/benchmarks/utils/convert_nemo_dataset.py b/benchmarks/utils/convert_nemo_dataset.py deleted file mode 100644 index 6f4884347677..000000000000 --- a/benchmarks/utils/convert_nemo_dataset.py +++ /dev/null @@ -1,47 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -#!/usr/bin/env python3 - -import argparse -import json - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("input") - parser.add_argument("output") - - args = parser.parse_args() - - output_o = [] - - with open(args.input, 'r') as infile: - for _l in infile: - l = _l.strip() - if len(l) == 0: - continue - o = json.loads(l) - output_o.append({ - "input": o["prompt"], - "instruction": "", - "output": o["completion"] - }) - - with open(args.output, 'w') as outfile: - json.dump(output_o, outfile) - - -if __name__ == "__main__": - main() diff --git a/benchmarks/utils/generate_rand_loras.py b/benchmarks/utils/generate_rand_loras.py deleted file mode 100644 index 12eb1fdc3648..000000000000 --- a/benchmarks/utils/generate_rand_loras.py +++ /dev/null @@ -1,48 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -#!/usr/bin/env python3 - -import argparse -import os -from pathlib import Path - -import numpy as np - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("input_lora") - parser.add_argument("output") - parser.add_argument("num_loras", type=int) - - args = parser.parse_args() - - lora_path = Path(args.input_lora) - weights_path = lora_path / "model.lora_weights.npy" - config_path = lora_path / "model.lora_config.npy" - - weights = np.load(weights_path) - config = np.load(config_path) - - for i in range(args.num_loras): - out_path = Path(args.output) / str(i) - os.makedirs(out_path, exist_ok=True) - w = np.random.normal(0, 2, weights.shape).astype(weights.dtype) - np.save(out_path / "model.lora_weights.npy", w) - np.save(out_path / "model.lora_config.npy", config) - - -if __name__ == "__main__": - main() diff --git a/benchmarks/utils/prepare_real_data.py b/benchmarks/utils/prepare_real_data.py deleted file mode 100644 index 4441c57ee4bc..000000000000 --- a/benchmarks/utils/prepare_real_data.py +++ /dev/null @@ -1,434 +0,0 @@ -import logging -import random -import re -import tempfile -from pathlib import Path -from typing import Optional - -import click -import datasets -from PIL import Image -from pydantic import BaseModel, model_validator -from utils.utils import (get_norm_dist_lengths, multimodal_dataset_dump, - print_multimodal_dataset, print_text_dataset, - text_dataset_dump) - - -def validate_output_len_dist(ctx, param, value): - """Validate the --output-len-dist option.""" - if value is None: - return value - m = re.match(r"(\d+),(\d+)", value) - if m: - return int(m.group(1)), int(m.group(2)) - else: - raise AssertionError( - "Incorrect specification for --output-len-dist. Correct format: --output-len-dist ," - ) - - -class DatasetConfig(BaseModel): - """Dataset configurations.""" - """Name of the dataset on HuggingFace.""" - name: Optional[str] = None - """Config name of the dataset if existing.""" - config_name: Optional[str] = None - """Split of the dataset. Typical values: train, validation, test. Setting to None will include all splits.""" - split: Optional[str] - """The dataset dictionary used for the input sentence.""" - input_key: Optional[str] = None - """The dataset dictionary key used for the prompt of the input sentence. Must not be set when prompt is set.""" - image_key: Optional[str] = None - """The dataset dictionary key used for the images.""" - prompt_key: Optional[str] = None - """The prompt sentence to be added to the input sentence. Must not be set when prompt_key is set.""" - prompt: Optional[str] = None - """The dataset dictionary key used to derive the output sequence length. Set to None if the dataset does not have a key for output.""" - output_key: Optional[str] - """The local path to the dataset to be loaded when using a local cache.""" - local_path: Optional[str] = None - - @model_validator(mode='after') - def check_prompt(self) -> 'DatasetConfig': - if self.prompt_key and self.prompt: - raise AssertionError( - "--prompt-key and --prompt cannot be set at the same time.") - if (not self.prompt_key) and (not self.prompt): - raise AssertionError("Either --prompt-key or --prompt must be set.") - return self - - @model_validator(mode='after') - def check_name_and_local_path(self) -> 'DatasetConfig': - if self.name and self.local_path: - raise AssertionError( - "--dataset-name and --dataset-local-path cannot be set at the same time." - ) - if (not self.name) and (not self.local_path): - raise AssertionError( - "Either --dataset-name or --dataset-local-path must be set.") - return self - - @property - def query(self): - """Generate the query for HuggingFace `datasets.load_dataset()`""" - first_arg = self.local_path if self.local_path else self.name - - if self.config_name: - return [first_arg, self.config_name] - else: - return [first_arg] - - @property - def display_name(self) -> str: - """Returns a human-readable identifier for error messages.""" - # model_validator ensures exactly one of name or local_path is set - if self.name is not None: - return self.name - return self.local_path - - def get_prompt(self, req): - """Get the prompt sentence from the given request.""" - if self.prompt_key: - assert self.prompt_key in req, ( - f"Dataset {self.display_name} does not have key '{self.prompt_key}'. " - "Please set --prompt-key to one of the available keys: " - f"{req.keys()}") - return req[self.prompt_key] - else: - return self.prompt - - def get_input(self, req): - """Get the input sentence from the given request.""" - assert self.input_key in req, ( - f"Dataset {self.display_name} does not have key '{self.input_key}'. " - "Please set --input-key to one of the available keys: " - f"{req.keys()}") - return req[self.input_key] - - def get_images(self, req): - """Get the images from the given request.""" - image_keys = [self.image_key - ] + [f"{self.image_key}_{i}" for i in range(1, 8)] - assert any(key in req for key in image_keys), ( - f"Dataset {self.display_name} does not have key '{self.image_key}'. " - "Please set --dataset-image-key to one of the available keys: " - f"{req.keys()}") - images = [] - for key in image_keys: - if key in req and req[key] is not None: - images.append(req[key]) - return images - - def get_output(self, req): - """Get the output sentence from the given request.""" - if self.output_key is None: - raise RuntimeError( - "--output-key is not set. Please either:\n" - "1. Define output length through --output-len-dist.\n" - f"2. If the dataset {self.display_name} has key for golden output and " - "you wish to set output length to the length of the golden " - "output, set --output-key.") - assert self.output_key in req, ( - f"Dataset {self.display_name} does not have key '{self.output_key}'. " - "Please set --output-key to one of the available keys: " - f"{req.keys()}") - return req[self.output_key] - - -def _create_dataset_load_error(e: ValueError) -> ValueError: - """Create a more informative ValueError from a dataset loading error. - - Args: - e: The original ValueError from datasets.load_dataset(). - Returns: - A new ValueError with additional context. - """ - error_msg = str(e) - if "Config" in error_msg: - error_msg += "\n Please add the config name to the dataset config yaml." - elif "split" in error_msg: - error_msg += "\n Please specify supported split in the dataset config yaml." - return ValueError(error_msg) - - -def load_dataset(dataset_config: DatasetConfig): - """Load dataset from local path or HuggingFace. - Args: - dataset_config: A `DatasetConfig` object that defines the dataset to load. - Returns: - Dataset iterator. - Raises: - ValueError: When dataset loading fails due to incorrect dataset config setting. - """ - if dataset_config.local_path: - return load_dataset_from_local(dataset_config) - else: - return load_dataset_from_hf(dataset_config) - - -def load_dataset_from_hf(dataset_config: DatasetConfig): - """Load dataset from HuggingFace. - - Args: - dataset_config: A `DatasetConfig` object that defines the dataset to load. - Returns: - Dataset iterator. - Raises: - ValueError: When dataset loading fails due to incorrect dataset config setting. - """ - logging.debug( - f"Loading dataset from HF: query={dataset_config.query}, split={dataset_config.split}" - ) - - try: - dataset = iter( - datasets.load_dataset(*dataset_config.query, - split=dataset_config.split, - streaming=True, - trust_remote_code=True)) - except ValueError as e: - raise _create_dataset_load_error(e) - - logging.debug("Finished loading HF dataset") - - return dataset - - -def load_dataset_from_local(dataset_config: DatasetConfig): - """Load dataset from local path. - - Args: - dataset_config: A `DatasetConfig` object that defines the dataset to load. - Returns: - Dataset iterator. - Raises: - FileNotFoundError: When local dataset path does not exist. - ValueError: When dataset loading fails due to incorrect dataset config setting. - """ - - local_path = Path(dataset_config.local_path) - - if not local_path.exists(): - raise FileNotFoundError( - f"Local dataset path {local_path} does not exist.") - - logging.debug( - f"Loading dataset from local path: path={local_path}, query={dataset_config.query}, split={dataset_config.split}" - ) - - # If it's a directory we can use the normal loader, otherwise custom loader - # depends on the file extension - if local_path.is_dir(): - try: - dataset = datasets.load_dataset(*dataset_config.query, - split=dataset_config.split, - trust_remote_code=True) - except ValueError as e: - raise _create_dataset_load_error(e) - else: - format_map = { - ".json": "json", - ".jsonl": "json", - ".csv": "csv", - ".parquet": "parquet", - } - - file_extension = local_path.suffix - dataset_type = format_map.get(file_extension) - - if dataset_type is None: - raise ValueError(f"Unsupported file extension: {file_extension}") - - try: - dataset = datasets.load_dataset(dataset_type, - data_files=str(local_path), - split=dataset_config.split) - except ValueError as e: - raise _create_dataset_load_error(e) - - logging.debug("Finished loading local dataset") - - return iter(dataset) - - -@click.command() -@click.option("--dataset-name", type=str, help="Dataset name in HuggingFace.") -@click.option("--dataset-config-name", - type=str, - default=None, - help="Dataset config name in HuggingFace (if exists).") -@click.option("--dataset-split", - type=str, - required=True, - help="Split of the dataset to use.") -@click.option("--dataset-input-key", - type=str, - help="The dataset dictionary key for input.") -@click.option("--dataset-image-key", - type=str, - default="image", - help="The dataset dictionary key for images.") -@click.option("--dataset-prompt-key", - type=str, - default=None, - help="The dataset dictionary key for prompt (if exists).") -@click.option( - "--dataset-local-path", - type=str, - default=None, - help= - "The local path to the dataset to be loaded when using an offline cache.") -@click.option( - "--dataset-prompt", - type=str, - default=None, - help="The prompt string when there is no prompt key for the dataset.") -@click.option("--dataset-output-key", - type=str, - default=None, - help="The dataset dictionary key for output (if exists).") -@click.option( - "--num-requests", - type=int, - default=None, - help= - "Number of requests to be generated. Will be capped to min(dataset.num_rows, num_requests)." -) -@click.option( - "--max-input-len", - type=int, - default=None, - help= - "Maximum input sequence length for a given request. This will be used to filter out the requests with long input sequence length. Default will include all the requests." -) -@click.option( - "--output-len-dist", - type=str, - default=None, - callback=validate_output_len_dist, - help= - "Output length distribution. Default will be the length of the golden output from the dataset. Format: ,. E.g. 100,10 will randomize the output length with mean=100 and variance=10." -) -@click.pass_obj -def dataset(root_args, **kwargs): - """Prepare dataset from real dataset.""" - dataset_config = DatasetConfig(**{ - k[8:]: v - for k, v in kwargs.items() if k.startswith('dataset_') - }) - - input_ids = [] - input_lens = [] - output_lens = [] - task_ids = [] - req_cnt = 0 - modality = None - multimodal_texts = [] - multimodal_image_paths = [] - for req in load_dataset(dataset_config): - if any(key in req for key in ['image', 'image_1', 'video']): - # multimodal input - if 'video' in req and req['video'] is not None: - assert "Not supported yet" - assert kwargs['output_len_dist'] is not None, ( - "Output length distribution must be set for multimodal requests." - ) - modality = 'image' - text = dataset_config.get_prompt(req) - images = dataset_config.get_images(req) - image_paths = [] - for image in images: - if image is not None: - if isinstance(image, str): - image_paths.append(image) - elif isinstance(image, Image.Image): - with tempfile.NamedTemporaryFile( - suffix=".jpg", delete=False) as tmp_file: - logging.debug(f"Saving image to {tmp_file.name}") - image = image.convert("RGB") - image.save(tmp_file, "JPEG") - filepath = tmp_file.name - image_paths.append(filepath) - else: - raise ValueError(f"Invalid image path: {image}") - multimodal_texts.append(text) - multimodal_image_paths.append(image_paths) - else: - # text input - prompt = dataset_config.get_prompt( - req) + ' ' + dataset_config.get_input(req) - logging.debug(f"Input sequence: {prompt}") - line = root_args.tokenizer.encode(prompt) - if kwargs['max_input_len'] and len(line) > kwargs['max_input_len']: - continue - input_ids.append(line) - input_lens.append(len(line)) - - # output if fetch from golden - if kwargs['output_len_dist'] is None: - output_lens.append( - len( - root_args.tokenizer.encode( - dataset_config.get_output(req)))) - - # lora task id - task_id = root_args.task_id - if root_args.rand_task_id is not None: - min_id, max_id = root_args.rand_task_id - task_id = random.randint(min_id, max_id) - task_ids.append(task_id) - - req_cnt += 1 - if kwargs['num_requests'] and req_cnt >= kwargs['num_requests']: - break - - if kwargs['num_requests'] and (len(input_ids) if modality is None else len( - multimodal_texts)) < kwargs['num_requests']: - logging.warning( - f"Number of requests={len(input_ids) if modality is None else len(multimodal_texts)} is" - f" smaller than the num-requests user set={kwargs['num_requests']}." - ) - - # output if randomized - if kwargs['output_len_dist'] is not None: - osl_mean, osl_stdev = kwargs['output_len_dist'] - output_lens = get_norm_dist_lengths( - osl_mean, osl_stdev, - len(input_ids) if modality is None else len(multimodal_texts), - root_args.random_seed) - logging.debug(f"Input lengths: {[len(i) for i in input_ids]}") - logging.debug(f"Output lengths: {output_lens}") - if modality is not None: - logging.debug(f"Modality: {modality}") - - if modality is not None: - if not root_args.std_out: - multimodal_dataset_dump( - multimodal_texts, multimodal_image_paths, output_lens, task_ids, - { - "workload_type": "dataset", - "tokenizer": root_args.tokenizer.__class__.__name__, - "num_requests": len(task_ids), - "max_output_len": max(output_lens) - }, root_args.output) - else: - print_multimodal_dataset( - multimodal_texts, - multimodal_image_paths, - output_lens, - ) - else: - if not root_args.std_out: - text_dataset_dump( - input_lens, input_ids, output_lens, task_ids, { - "workload_type": "dataset", - "tokenizer": root_args.tokenizer.__class__.__name__, - "num_requests": len(input_ids), - "max_input_len": max(input_lens), - "max_output_len": max(output_lens) - }, root_args.output) - else: - print_text_dataset( - input_ids, - output_lens, - ) diff --git a/benchmarks/utils/prepare_synthetic_data.py b/benchmarks/utils/prepare_synthetic_data.py deleted file mode 100644 index b072b712d085..000000000000 --- a/benchmarks/utils/prepare_synthetic_data.py +++ /dev/null @@ -1,164 +0,0 @@ -import random -import warnings - -import click -from utils.utils import (gen_random_tokens, get_norm_dist_lengths, - get_unif_dist_lengths, print_text_dataset, - text_dataset_dump) - - -def _generate_task_ids_and_lora_config(root_args, num_reqs): - """Generate task IDs and determine LoRA configuration based on root_args.""" - if root_args.rand_task_id is None: - task_ids = [root_args.task_id for _ in range(num_reqs)] - else: - min_id, max_id = root_args.rand_task_id - task_ids = [random.randint(min_id, max_id) for _ in range(num_reqs)] - - use_task_ids = root_args.task_id != -1 or root_args.rand_task_id is not None - - # Determine if LoRA should be used (requires both task IDs and lora_dir) - use_lora = use_task_ids and root_args.lora_dir is not None - - # Warn if task IDs are specified but no LoRA directory is provided - if use_task_ids and not use_lora: - warnings.warn( - "Task IDs require LoRA directory. Use --lora-dir or omit task IDs.", - UserWarning) - - return (task_ids, task_ids if use_task_ids else None, { - "lora_dir": root_args.lora_dir - } if use_lora else None) - - -@click.command() -@click.option("--num-requests", - required=True, - type=int, - help='Number of requests to be generated') -@click.option('--input-mean', - required=True, - type=int, - help='normal dist mean for input tokens') -@click.option('--input-stdev', - required=True, - type=int, - help='normal dist stdev for input tokens') -@click.option('--output-mean', - required=True, - type=int, - help='normal dist mean for output tokens') -@click.option('--output-stdev', - required=True, - type=int, - help='normal dist stdev for output tokens') -@click.pass_obj -def token_norm_dist(root_args, **kwargs): - """Prepare synthetic dataset by generating random tokens with normal dist lengths.""" - input_ids = [] - input_lens = [] - output_lens = [] - - input_lens = get_norm_dist_lengths(kwargs['input_mean'], - kwargs['input_stdev'], - kwargs['num_requests'], - root_args.random_seed) - - num_reqs = len(input_lens) - output_lens = get_norm_dist_lengths(kwargs['output_mean'], - kwargs['output_stdev'], num_reqs, - root_args.random_seed) - - max_input_len = max(input_lens) - max_output_len = max(output_lens) - - input_ids = gen_random_tokens(input_lens, root_args.tokenizer, - root_args.random_seed) - - task_ids, print_task_ids, lora_config = _generate_task_ids_and_lora_config( - root_args, num_reqs) - - if not root_args.std_out: - text_dataset_dump( - input_lens, input_ids, output_lens, task_ids, { - "workload_type": "token-norm-dist", - "input_mean": kwargs['input_mean'], - "input_stdev": kwargs['input_stdev'], - "output_mean": kwargs['output_mean'], - "output_stdev": kwargs['output_stdev'], - "num_requests": kwargs['num_requests'], - "tokenize_vocabsize": root_args.tokenizer.vocab_size, - "max_input_len": max_input_len, - "max_output_len": max_output_len - }, root_args.output) - else: - print_text_dataset(input_ids, - output_lens, - task_ids=print_task_ids, - lora_config=lora_config) - - -@click.command() -@click.option("--num-requests", - required=True, - type=int, - help='Number of requests to be generated') -@click.option('--input-min', - required=True, - type=int, - help='uniform dist (inclusive) min for input tokens') -@click.option('--input-max', - required=True, - type=int, - help='normal dist (inclusive) max for input tokens') -@click.option('--output-min', - required=True, - type=int, - help='normal dist (inclusive) min for output tokens') -@click.option('--output-max', - required=True, - type=int, - help='normal dist (inclusive) max for output tokens') -@click.pass_obj -def token_unif_dist(root_args, **kwargs): - """Prepare synthetic dataset by generating random tokens with normal uniformly lengths.""" - input_ids = [] - input_lens = [] - output_lens = [] - - input_lens = get_unif_dist_lengths(kwargs['input_min'], kwargs['input_max'], - kwargs['num_requests'], - root_args.random_seed) - - num_reqs = len(input_lens) - output_lens = get_unif_dist_lengths(kwargs['output_min'], - kwargs['output_max'], num_reqs, - root_args.random_seed) - - max_input_len = max(input_lens) - max_output_len = max(output_lens) - - input_ids = gen_random_tokens(input_lens, root_args.tokenizer, - root_args.random_seed) - - task_ids, print_task_ids, lora_config = _generate_task_ids_and_lora_config( - root_args, num_reqs) - - if not root_args.std_out: - text_dataset_dump( - input_lens, input_ids, output_lens, task_ids, { - "workload_type": "token-unif-dist", - "input_min": kwargs['input_min'], - "input_max": kwargs['input_max'], - "output_min": kwargs['output_min'], - "output_max": kwargs['output_max'], - "num_requests": kwargs['num_requests'], - "tokenize_vocabsize": root_args.tokenizer.vocab_size, - "max_input_len": max_input_len, - "max_output_len": max_output_len - }, root_args.output) - else: - print_text_dataset(input_ids, - output_lens, - task_ids=print_task_ids, - lora_config=lora_config) diff --git a/benchmarks/utils/utils.py b/benchmarks/utils/utils.py deleted file mode 100644 index c395cf6c9449..000000000000 --- a/benchmarks/utils/utils.py +++ /dev/null @@ -1,168 +0,0 @@ -import json -import math -import os -import random -from typing import List, Union - -import numpy as np -from pydantic import BaseModel - - -class TextSample(BaseModel): - input_len: int - input_ids: List[int] - output_len: int - task_id: int - - -class MultimodalSample(BaseModel): - task_id: int - prompt: str - media_paths: List[str] - output_len: int - - -class Workload(BaseModel): - metadata: dict - samples: List[Union[TextSample, MultimodalSample]] = [] - - def __init__(self, **kwargs) -> None: - super().__init__(**kwargs) - self.setup_workload_name() - - def setup_workload_name(self): - # Keys to ignore - ignore_keys = ['tokenizer'] - # Create a string by concatenating keys and values with "__" - workload_name = '__'.join(f'{key}:{value}' - for key, value in self.metadata.items() - if key not in ignore_keys) - self.metadata.setdefault('workload_name', workload_name) - - -def text_dataset_dump(input_lens, input_ids, output_lens, task_ids, metadata, - output_file): - samples = [] - for i in range(len(input_ids)): - samples.append( - TextSample(input_len=input_lens[i], - input_ids=input_ids[i], - output_len=output_lens[i], - task_id=task_ids[i])) - workload = Workload(metadata=metadata, samples=samples) - os.makedirs(os.path.dirname(output_file), exist_ok=True) - with open(output_file, 'w') as f: - json.dump(workload.model_dump(), f) - - -def multimodal_dataset_dump(multimodal_texts, multimodal_image_paths, - output_lens, task_ids, metadata, output_file): - samples = [] - for i in range(len(multimodal_texts)): - samples.append( - MultimodalSample(task_id=task_ids[i], - prompt=multimodal_texts[i], - media_paths=multimodal_image_paths[i], - output_len=output_lens[i])) - workload = Workload(metadata=metadata, samples=samples) - os.makedirs(os.path.dirname(output_file), exist_ok=True) - with open(output_file, 'w') as f: - json.dump(workload.model_dump(), f) - - -def print_text_dataset(input_ids, output_lens, task_ids=None, lora_config=None): - for i, input_tokens in enumerate(input_ids): - d = { - "task_id": i, - "input_ids": input_tokens, - "output_tokens": output_lens[i] - } - - # Add LoRA request if task_ids indicate LoRA usage - if task_ids is not None and lora_config is not None: - task_id = task_ids[i] - if task_id != -1: # -1 means no LoRA - d["lora_request"] = { - "lora_name": - f"lora_{task_id}", - "lora_int_id": - task_id, - "lora_path": - os.path.join(lora_config.get("lora_dir", "loras"), - str(task_id)) - } - - print(json.dumps(d, separators=(',', ':'), ensure_ascii=False)) - - -def print_multimodal_dataset(multimodal_texts, multimodal_image_paths, - output_lens): - for i, (text, image_paths) in enumerate( - zip(multimodal_texts, multimodal_image_paths)): - d = { - "task_id": i, - "prompt": text, - "media_paths": image_paths, - "output_tokens": output_lens[i] - } - print(json.dumps(d, separators=(',', ':'), ensure_ascii=False)) - - -def get_list_of_delays(delay_dist, mean_time_bet_reqs, num_reqs, random_seed): - if delay_dist == "constant": - delays = [mean_time_bet_reqs] * num_reqs - elif delay_dist == "exponential_dist": - delays = get_exponential_dist_delays(mean_time_bet_reqs, num_reqs, - random_seed) - - return delays - - -def get_exponential_dist_delays(mean_time_bet_reqs, num_reqs, random_seed): - # set seed for determinism - np.random.seed(random_seed) - return np.random.exponential(mean_time_bet_reqs, num_reqs).tolist() - - -def get_norm_dist_lengths(mean, stdev, num_reqs, random_seed): - # set seed for determinism - np.random.seed(random_seed) - numbers_list = np.random.normal(loc=mean, scale=stdev, - size=num_reqs).tolist() - return [max(1, math.ceil(x)) for x in numbers_list] - - -def get_unif_dist_lengths(min_len, max_len, num_reqs, random_seed): - # set seed for determinism - rng = np.random.default_rng(random_seed) - numbers = rng.integers(low=min_len, high=max_len + 1, size=num_reqs) - return numbers.tolist() - - -def gen_random_tokens(ip_lens, tokenizer, random_seed): - - def get_sample_from_population(population_range, sample_size): - # random.sample can not sample a value more than once. hence the check - if sample_size < len(population_range): - sample = random.sample(population_range, sample_size) - else: - sample = random.choices(population_range, k=sample_size) - - return sample - - input_ids = [] - random.seed(random_seed) - for ip_len in ip_lens: - start_ids = get_sample_from_population(range(0, tokenizer.vocab_size), - ip_len) - # Make sure it does not contain EOS token - eos_id = tokenizer.encode(tokenizer.eos_token, add_special_tokens=False) - while set(eos_id).issubset(start_ids): - tmp_id = (eos_id[0] + 1) % tokenizer.vocab_size - start_ids = [ - tmp_id if element == eos_id[0] else element - for element in start_ids - ] - input_ids.append(start_ids) - - return input_ids diff --git a/docs/source/_ext/llmapi_config_telemetry.py b/docs/source/_ext/llmapi_config_telemetry.py index 01509c74a628..0b0bf9e63427 100644 --- a/docs/source/_ext/llmapi_config_telemetry.py +++ b/docs/source/_ext/llmapi_config_telemetry.py @@ -76,19 +76,16 @@ def _table(rows: list[dict]) -> str: def generate_telemetry_reference(repo_root: Path | str, output_path: Path | str) -> None: repo_root = Path(repo_root) golden = json.loads((repo_root / _GOLDEN_REL).read_text()) - content = [_REFERENCE_PREAMBLE] - for args_class in ("TorchLlmArgs", "TrtLlmArgs"): - rows = golden.get(args_class, []) - content.extend( - [ - f"### `{args_class}`", - "", - f"{len(rows)} captured fields.", - "", - _table(rows), - "", - ] - ) + rows = golden.get("TorchLlmArgs", []) + content = [ + _REFERENCE_PREAMBLE, + "### `TorchLlmArgs`", + "", + f"{len(rows)} captured fields.", + "", + _table(rows), + "", + ] output = Path(output_path) output.parent.mkdir(parents=True, exist_ok=True) output.write_text("\n".join(content)) diff --git a/docs/source/commands/trtllm-bench.rst b/docs/source/commands/trtllm-bench.rst index a55899a961d1..309422df3fba 100644 --- a/docs/source/commands/trtllm-bench.rst +++ b/docs/source/commands/trtllm-bench.rst @@ -20,10 +20,10 @@ Syntax Dataset preparation ------------------ -prepare_dataset.py -^^^^^^^^^^^^^^^^^^ +prepare-dataset +^^^^^^^^^^^^^^^ -trtllm-bench is designed to work with the `prepare_dataset.py `_ script, which generates benchmark datasets in the required format. The prepare_dataset script supports: +trtllm-bench ships a ``prepare-dataset`` subcommand which generates benchmark datasets in the required format. It supports: **Dataset Types:** @@ -38,17 +38,17 @@ trtllm-bench is designed to work with the `prepare_dataset.py prepare-dataset [OPTIONS] **Options** @@ -60,12 +60,10 @@ prepare_dataset * - Option - Description - * - ``--tokenizer`` - - Tokenizer directory or HuggingFace model name (required) * - ``--output`` - Output JSON filename (default: preprocessed_dataset.json) * - ``--stdout`` - - Print output to stdout with JSON dataset entry on each line (**required for trtllm-bench**) + - Print output to stdout with a JSON dataset entry on each line instead of writing a file * - ``--random-seed`` - Random seed for token generation (default: 420) * - ``--task-id`` @@ -77,14 +75,14 @@ prepare_dataset * - ``--log-level`` - Logging level: info or debug (default: info) -dataset -""""""" +real-dataset +"""""""""""" Process real datasets from various sources. .. code-block:: bash - python prepare_dataset.py dataset [OPTIONS] + trtllm-bench --model prepare-dataset real-dataset [OPTIONS] **Options** @@ -108,14 +106,14 @@ Process real datasets from various sources. - Input format: json, jsonl, csv, or txt (default: auto-detect) -token_norm_dist +token-norm-dist """"""""""""""" Generate synthetic datasets with normal token distribution. .. code-block:: bash - python prepare_dataset.py token_norm_dist [OPTIONS] + trtllm-bench --model prepare-dataset token-norm-dist [OPTIONS] **Options** @@ -139,14 +137,14 @@ Generate synthetic datasets with normal token distribution. - Normal distribution standard deviation for output tokens (required) -token_unif_dist +token-unif-dist """"""""""""""" Generate synthetic datasets with uniform token distribution .. code-block:: bash - python prepare_dataset.py token_unif_dist [OPTIONS] + trtllm-bench --model prepare-dataset token-unif-dist [OPTIONS] **Options** diff --git a/docs/source/developer-guide/perf-overview.md b/docs/source/developer-guide/perf-overview.md index dbf999d4ca03..58de0107b3d8 100644 --- a/docs/source/developer-guide/perf-overview.md +++ b/docs/source/developer-guide/perf-overview.md @@ -268,7 +268,7 @@ Testing was performed using the PyTorch backend - this workflow does not require | Stage | Description | Command | | :- | - | - | -| [Dataset](#preparing-a-dataset) | Create a synthetic dataset | `python benchmarks/prepare_dataset.py --tokenizer=$model_name --stdout token-norm-dist --num-requests=$num_requests --input-mean=$isl --output-mean=$osl --input-stdev=0 --output-stdev=0 > $dataset_file` | +| [Dataset](#preparing-a-dataset) | Create a synthetic dataset | `trtllm-bench --model $model_name prepare-dataset --output $dataset_file token-norm-dist --num-requests=$num_requests --input-mean=$isl --output-mean=$osl --input-stdev=0 --output-stdev=0` | | [Run](#running-the-benchmark) | Run a benchmark with a dataset | `trtllm-bench --model $model_name throughput --dataset $dataset_file --backend pytorch --config $llm_options` | ### Variables @@ -281,18 +281,18 @@ Testing was performed using the PyTorch backend - this workflow does not require | `$pp_size` | Pipeline parallel mapping degree to run the benchmark with | | `$ep_size` | Expert parallel mapping degree to run the benchmark with | | `$model_name` | HuggingFace model name eg. meta-llama/Llama-2-7b-hf or use the path to a local weights directory | -| `$dataset_file` | Location of the dataset file generated by `prepare_dataset.py` | +| `$dataset_file` | Location of the dataset file generated by `trtllm-bench prepare-dataset` | | `$num_requests` | The number of requests to generate for dataset generation | | `$seq_len` | A sequence length of ISL + OSL | | `$llm_options` | (optional) A yaml file containing additional options for the LLM API | ### Preparing a Dataset -In order to prepare a dataset, you can use the provided [script](source:benchmarks/prepare_dataset.py). +In order to prepare a dataset, use the `trtllm-bench prepare-dataset` subcommand. To generate a synthetic dataset, run the following command: ```shell -python benchmarks/prepare_dataset.py --tokenizer=$model_name --stdout token-norm-dist --num-requests=$num_requests --input-mean=$isl --output-mean=$osl --input-stdev=0 --output-stdev=0 > $dataset_file +trtllm-bench --model $model_name prepare-dataset --output $dataset_file token-norm-dist --num-requests=$num_requests --input-mean=$isl --output-mean=$osl --input-stdev=0 --output-stdev=0 ``` The command will generate a text file located at the path specified `$dataset_file` where all requests are of the same diff --git a/docs/source/developer-guide/telemetry.md b/docs/source/developer-guide/telemetry.md index fc0e1fc8bbe1..9ca7b7aaa49f 100644 --- a/docs/source/developer-guide/telemetry.md +++ b/docs/source/developer-guide/telemetry.md @@ -28,7 +28,7 @@ unset or when the safety sanitizer rejects the runtime value. ### `TorchLlmArgs` -261 captured fields. +269 captured fields. | Captured key | Annotation | Kind | Converter | Allowed values | |--------------|------------|------|-----------|----------------| @@ -71,6 +71,7 @@ unset or when the safety sanitizer rejects the runtime value. | `cuda_graph_config.mode` | `Literal['decode']` | `categorical` | | `decode`, `encode` | | `cuda_graph_config.num_tokens` | `Optional[List[Annotated[int, Gt(gt=0)]]]` | `value` | | | | `cuda_graph_config.seq_lens` | `Optional[List[Annotated[int, Gt(gt=0)]]]` | `value` | | | +| `disable_mm_encoder` | `` | `value` | | | | `disable_overlap_scheduler` | `` | `value` | | | | `dtype` | `` | `categorical` | allowlist | `auto`, `float16`, `bfloat16`, `float32` | | `dwdp_config.contention_opt` | `` | `value` | | | @@ -104,7 +105,7 @@ unset or when the safety sanitizer rejects the runtime value. | `iter_stats_max_iterations` | `Optional[int]` | `value` | | | | `kv_cache_config.attention_dp_events_gather_period_ms` | `` | `value` | | | | `kv_cache_config.avg_seq_len` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` | | | -| `kv_cache_config.block_reuse_policy` | `Literal['all_reusable', 'per_request']` | `categorical` | | `all_reusable`, `per_request` | +| `kv_cache_config.block_reuse_policy` | `Literal['all_reusable', 'per_request', 'per_conversation']` | `categorical` | | `all_reusable`, `per_request`, `per_conversation` | | `kv_cache_config.copy_on_partial_reuse` | `` | `value` | | | | `kv_cache_config.cross_kv_cache_fraction` | `Optional[float]` | `value` | | | | `kv_cache_config.disk_cache_size` | `Optional[Annotated[int, Ge(ge=0)]]` | `value` | | | @@ -131,7 +132,7 @@ unset or when the safety sanitizer rejects the runtime value. | `kv_cache_config.secondary_offload_min_priority` | `Optional[int]` | `value` | | | | `kv_cache_config.sink_token_length` | `Optional[int]` | `value` | | | | `kv_cache_config.tokens_per_block` | `` | `value` | | | -| `kv_cache_config.use_kv_cache_manager_v2` | `` | `value` | | | +| `kv_cache_config.use_kv_cache_manager_v2` | `Union[bool, Literal['auto']]` | `value` | | `auto` | | `kv_cache_config.use_uvm` | `` | `value` | | | | `kv_connector_config.connector` | `Optional[str]` | `categorical` | allowlist | `lmcache`, `lmcache-mp`, `kvbm` | | `layer_wise_benchmarks_config.calibration_layer_indices` | `Optional[List[int]]` | `value` | | | @@ -156,11 +157,13 @@ unset or when the safety sanitizer rejects the runtime value. | `moe_config.use_low_precision_moe_combine` | `` | `value` | | | | `moe_expert_parallel_size` | `Optional[int]` | `value` | | | | `moe_tensor_parallel_size` | `Optional[int]` | `value` | | | +| `multimodal_config.encoder_cache_max_bytes` | `` | `value` | | | | `multimodal_config.encoder_side_stream_max_ahead` | `` | `value` | | | | `multimodal_config.video_pruning_rate` | `Optional[float]` | `value` | | | | `mx_config.preshard_strategy` | `` | `categorical` | allowlist | `per_module` | | `mx_config.server_query_timeout_s` | `Optional[Annotated[int, Ge(ge=0)]]` | `value` | | | | `num_postprocess_workers` | `` | `value` | | | +| `num_serve_frontends` | `` | `value` | | | | `nvfp4_gemm_config.allowed_backends` | `List[Literal['cutlass', 'cublaslt', 'cutedsl', 'cuda_core', 'marlin']]` | `value` | | `cutlass`, `cublaslt`, `cutedsl`, `cuda_core`, `marlin` | | `orchestrator_type` | `Optional[Literal['rpc', 'ray']]` | `categorical` | | `rpc`, `ray` | | `peft_cache_config.device_cache_percent` | `` | `value` | | | @@ -208,6 +211,7 @@ unset or when the safety sanitizer rejects the runtime value. | `sparse_attention_config.algorithm` | `Literal['dsa']` | `categorical` | | `dsa`, `deepseek_v4`, `minimax_m3`, `rocket`, `skip_softmax` | | `sparse_attention_config.compress_ratios` | `List[int]` | `value` | | | | `sparse_attention_config.enable_heuristic_topk` | `` | `value` | | | +| `sparse_attention_config.implementation` | `Literal['triton', 'msa']` | `categorical` | | `triton`, `msa` | | `sparse_attention_config.index_head_dim` | `Optional[int]` | `value` | | | | `sparse_attention_config.index_n_heads` | `Optional[int]` | `value` | | | | `sparse_attention_config.index_topk` | `Optional[int]` | `value` | | | @@ -216,6 +220,8 @@ unset or when the safety sanitizer rejects the runtime value. | `sparse_attention_config.indexer_rope_interleave` | `` | `value` | | | | `sparse_attention_config.kernel_size` | `Optional[int]` | `value` | | | | `sparse_attention_config.kt_cache_dtype` | `Optional[str]` | `categorical` | allowlist | `bfloat16`, `float8_e5m2` | +| `sparse_attention_config.num_attention_heads` | `Optional[int]` | `value` | | | +| `sparse_attention_config.num_key_value_heads` | `Optional[int]` | `value` | | | | `sparse_attention_config.page_size` | `Optional[int]` | `value` | | | | `sparse_attention_config.prompt_budget` | `Optional[int]` | `value` | | | | `sparse_attention_config.q_split_threshold` | `` | `value` | | | @@ -238,7 +244,8 @@ unset or when the safety sanitizer rejects the runtime value. | `speculative_config.acceptance_rate_window_size` | `Optional[Annotated[int, Ge(ge=0)]]` | `value` | | | | `speculative_config.allow_advanced_sampling` | `` | `value` | | | | `speculative_config.begin_thinking_phase_token` | `` | `value` | | | -| `speculative_config.decoding_type` | `Literal['AUTO']` | `categorical` | | `AUTO`, `DFlash`, `Draft_Target`, `Eagle3`, `Eagle`, `Lookahead`, `MTP`, `Medusa`, `NGram`, `PARD`, `SA`, `SaveState`, `User_Provided` | +| `speculative_config.block_size` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` | | | +| `speculative_config.decoding_type` | `Literal['AUTO']` | `categorical` | | `AUTO`, `DFlash`, `DSpark`, `Draft_Target`, `Eagle3`, `Eagle`, `Lookahead`, `MTP`, `Medusa`, `NGram`, `PARD`, `SA`, `SaveState`, `User_Provided` | | `speculative_config.dynamic_tree_max_topK` | `Optional[int]` | `value` | | | | `speculative_config.eagle3_layers_to_capture` | `Optional[Set[int]]` | `value` | | | | `speculative_config.eagle3_model_arch` | `Literal['llama3', 'mistral_large3']` | `categorical` | | `llama3`, `mistral_large3` | @@ -251,6 +258,8 @@ unset or when the safety sanitizer rejects the runtime value. | `speculative_config.is_keep_all` | `` | `value` | | | | `speculative_config.is_public_pool` | `` | `value` | | | | `speculative_config.is_use_oldest` | `` | `value` | | | +| `speculative_config.markov_head_type` | `Optional[Literal['vanilla', 'gated', 'rnn']]` | `categorical` | | `vanilla`, `gated`, `rnn` | +| `speculative_config.markov_rank` | `Optional[int]` | `value` | | | | `speculative_config.mask_token_id` | `Optional[int]` | `value` | | | | `speculative_config.max_concurrency` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` | | | | `speculative_config.max_draft_len` | `Optional[Annotated[int, Ge(ge=0)]]` | `value` | | | @@ -292,290 +301,3 @@ unset or when the safety sanitizer rejects the runtime value. | `use_cute_dsl_bf16_gemm` | `` | `value` | | | | `use_cute_dsl_blockscaling_bmm` | `` | `value` | | | | `use_cute_dsl_blockscaling_mm` | `` | `value` | | | - -### `TrtLlmArgs` - -280 captured fields. - -| Captured key | Annotation | Kind | Converter | Allowed values | -|--------------|------------|------|-----------|----------------| -| `backend` | `Optional[str]` | `categorical` | allowlist | `pytorch`, `tensorrt`, `_autodeploy` | -| `batching_type` | `Optional[tensorrt_llm.llmapi.llm_args.BatchingType]` | `categorical` | | `STATIC`, `INFLIGHT` | -| `build_config.dry_run` | `` | `value` | | | -| `build_config.enable_debug_output` | `` | `value` | | | -| `build_config.force_num_profiles` | `Optional[int]` | `value` | | | -| `build_config.gather_context_logits` | `` | `value` | | | -| `build_config.gather_generation_logits` | `` | `value` | | | -| `build_config.kv_cache_type` | `Optional[tensorrt_llm.llmapi.kv_cache_type.KVCacheType]` | `categorical` | | `continuous`, `paged`, `disabled` | -| `build_config.lora_config.lora_ckpt_source` | `Literal['hf', 'nemo']` | `categorical` | | `hf`, `nemo` | -| `build_config.lora_config.max_cpu_loras` | `Optional[int]` | `value` | | | -| `build_config.lora_config.max_lora_rank` | `` | `value` | | | -| `build_config.lora_config.max_loras` | `Optional[int]` | `value` | | | -| `build_config.lora_config.swap_gate_up_proj_lora_b_weight` | `` | `value` | | | -| `build_config.max_batch_size` | `` | `value` | | | -| `build_config.max_beam_width` | `` | `value` | | | -| `build_config.max_draft_len` | `` | `value` | | | -| `build_config.max_encoder_input_len` | `` | `value` | | | -| `build_config.max_input_len` | `` | `value` | | | -| `build_config.max_num_tokens` | `` | `value` | | | -| `build_config.max_prompt_embedding_table_size` | `` | `value` | | | -| `build_config.max_seq_len` | `Optional[int]` | `value` | | | -| `build_config.monitor_memory` | `` | `value` | | | -| `build_config.opt_batch_size` | `` | `value` | | | -| `build_config.opt_num_tokens` | `Optional[int]` | `value` | | | -| `build_config.plugin_config.bert_attention_plugin` | `Optional[Literal['auto', 'float16', 'float32', 'bfloat16', 'int32', None]]` | `categorical` | | `auto`, `float16`, `float32`, `bfloat16`, `int32`, `None` | -| `build_config.plugin_config.bert_context_fmha_fp32_acc` | `` | `value` | | | -| `build_config.plugin_config.context_fmha` | `` | `value` | | | -| `build_config.plugin_config.dora_plugin` | `` | `value` | | | -| `build_config.plugin_config.fp8_rowwise_gemm_plugin` | `Optional[Literal['auto', 'float16', 'float32', 'bfloat16', 'int32', None]]` | `categorical` | | `auto`, `float16`, `float32`, `bfloat16`, `int32`, `None` | -| `build_config.plugin_config.fuse_fp4_quant` | `` | `value` | | | -| `build_config.plugin_config.gemm_allreduce_plugin` | `Optional[Literal['float16', 'bfloat16', None]]` | `categorical` | | `float16`, `bfloat16`, `None` | -| `build_config.plugin_config.gemm_plugin` | `Optional[Literal['auto', 'float16', 'float32', 'bfloat16', 'int32', 'fp8', 'nvfp4', None]]` | `categorical` | | `auto`, `float16`, `float32`, `bfloat16`, `int32`, `fp8`, `nvfp4`, `None` | -| `build_config.plugin_config.gemm_swiglu_plugin` | `Optional[Literal['fp8', None]]` | `categorical` | | `fp8`, `None` | -| `build_config.plugin_config.gpt_attention_plugin` | `Optional[Literal['auto', 'float16', 'float32', 'bfloat16', 'int32', None]]` | `categorical` | | `auto`, `float16`, `float32`, `bfloat16`, `int32`, `None` | -| `build_config.plugin_config.identity_plugin` | `Optional[Literal['auto', 'float16', 'float32', 'bfloat16', 'int32', None]]` | `categorical` | | `auto`, `float16`, `float32`, `bfloat16`, `int32`, `None` | -| `build_config.plugin_config.layernorm_quantization_plugin` | `Optional[Literal['auto', 'float16', 'float32', 'bfloat16', 'int32', None]]` | `categorical` | | `auto`, `float16`, `float32`, `bfloat16`, `int32`, `None` | -| `build_config.plugin_config.lora_plugin` | `Optional[Literal['auto', 'float16', 'float32', 'bfloat16', 'int32', None]]` | `categorical` | | `auto`, `float16`, `float32`, `bfloat16`, `int32`, `None` | -| `build_config.plugin_config.low_latency_gemm_plugin` | `Optional[Literal['fp8', None]]` | `categorical` | | `fp8`, `None` | -| `build_config.plugin_config.low_latency_gemm_swiglu_plugin` | `Optional[Literal['fp8', None]]` | `categorical` | | `fp8`, `None` | -| `build_config.plugin_config.mamba_conv1d_plugin` | `Optional[Literal['auto', 'float16', 'float32', 'bfloat16', 'int32', None]]` | `categorical` | | `auto`, `float16`, `float32`, `bfloat16`, `int32`, `None` | -| `build_config.plugin_config.manage_weights` | `` | `value` | | | -| `build_config.plugin_config.moe_plugin` | `Optional[Literal['auto', 'float16', 'float32', 'bfloat16', 'int32', None]]` | `categorical` | | `auto`, `float16`, `float32`, `bfloat16`, `int32`, `None` | -| `build_config.plugin_config.multiple_profiles` | `` | `value` | | | -| `build_config.plugin_config.nccl_plugin` | `Optional[Literal['auto', 'float16', 'float32', 'bfloat16', 'int32', None]]` | `categorical` | | `auto`, `float16`, `float32`, `bfloat16`, `int32`, `None` | -| `build_config.plugin_config.norm_quant_fusion` | `` | `value` | | | -| `build_config.plugin_config.paged_kv_cache` | `Optional[bool]` | `value` | | | -| `build_config.plugin_config.paged_state` | `` | `value` | | | -| `build_config.plugin_config.pp_reduce_scatter` | `` | `value` | | | -| `build_config.plugin_config.qserve_gemm_plugin` | `Optional[Literal['auto', 'float16', 'float32', 'bfloat16', 'int32', None]]` | `categorical` | | `auto`, `float16`, `float32`, `bfloat16`, `int32`, `None` | -| `build_config.plugin_config.quantize_per_token_plugin` | `` | `value` | | | -| `build_config.plugin_config.quantize_tensor_plugin` | `` | `value` | | | -| `build_config.plugin_config.reduce_fusion` | `` | `value` | | | -| `build_config.plugin_config.remove_input_padding` | `` | `value` | | | -| `build_config.plugin_config.rmsnorm_quantization_plugin` | `Optional[Literal['auto', 'float16', 'float32', 'bfloat16', 'int32', None]]` | `categorical` | | `auto`, `float16`, `float32`, `bfloat16`, `int32`, `None` | -| `build_config.plugin_config.smooth_quant_gemm_plugin` | `Optional[Literal['auto', 'float16', 'float32', 'bfloat16', 'int32', None]]` | `categorical` | | `auto`, `float16`, `float32`, `bfloat16`, `int32`, `None` | -| `build_config.plugin_config.smooth_quant_plugins` | `` | `value` | | | -| `build_config.plugin_config.streamingllm` | `` | `value` | | | -| `build_config.plugin_config.tokens_per_block` | `` | `value` | | | -| `build_config.plugin_config.use_fp8_context_fmha` | `` | `value` | | | -| `build_config.plugin_config.use_fused_mlp` | `` | `value` | | | -| `build_config.plugin_config.use_paged_context_fmha` | `` | `value` | | | -| `build_config.plugin_config.user_buffer` | `` | `value` | | | -| `build_config.plugin_config.weight_only_groupwise_quant_matmul_plugin` | `Optional[Literal['auto', 'float16', 'float32', 'bfloat16', 'int32', None]]` | `categorical` | | `auto`, `float16`, `float32`, `bfloat16`, `int32`, `None` | -| `build_config.plugin_config.weight_only_quant_matmul_plugin` | `Optional[Literal['auto', 'float16', 'float32', 'bfloat16', 'int32', None]]` | `categorical` | | `auto`, `float16`, `float32`, `bfloat16`, `int32`, `None` | -| `build_config.speculative_decoding_mode` | `` | `categorical` | | `NONE`, `DRAFT_TOKENS_EXTERNAL`, `MEDUSA`, `LOOKAHEAD_DECODING`, `EXPLICIT_DRAFT_TOKENS`, `EAGLE`, `NGRAM`, `USER_PROVIDED`, `SAVE_HIDDEN_STATES`, `AUTO` | -| `build_config.strongly_typed` | `` | `value` | | | -| `build_config.use_mrope` | `` | `value` | | | -| `build_config.use_refit` | `` | `value` | | | -| `build_config.use_strip_plan` | `` | `value` | | | -| `build_config.weight_sparsity` | `` | `value` | | | -| `build_config.weight_streaming` | `` | `value` | | | -| `cache_transceiver_config.backend` | `Optional[Literal['DEFAULT', 'UCX', 'NIXL', 'MOONCAKE', 'MPI']]` | `categorical` | | `DEFAULT`, `UCX`, `NIXL`, `MOONCAKE`, `MPI` | -| `cache_transceiver_config.kv_cache_bounce_size_mb` | `` | `value` | | | -| `cache_transceiver_config.kv_transfer_poll_interval_ms` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` | | | -| `cache_transceiver_config.kv_transfer_sender_future_timeout_ms` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` | | | -| `cache_transceiver_config.kv_transfer_timeout_ms` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` | | | -| `cache_transceiver_config.max_tokens_in_buffer` | `Optional[int]` | `value` | | | -| `cache_transceiver_config.transceiver_runtime` | `Optional[Literal['CPP', 'PYTHON', 'auto']]` | `categorical` | | `CPP`, `PYTHON`, `auto` | -| `calib_config.calib_batch_size` | `` | `value` | | | -| `calib_config.calib_batches` | `` | `value` | | | -| `calib_config.calib_max_seq_length` | `` | `value` | | | -| `calib_config.device` | `Literal['cuda', 'cpu']` | `categorical` | | `cuda`, `cpu` | -| `calib_config.random_seed` | `` | `value` | | | -| `calib_config.tokenizer_max_seq_length` | `` | `value` | | | -| `context_parallel_size` | `` | `value` | | | -| `cp_config.block_size` | `Optional[int]` | `value` | | | -| `cp_config.cp_anchor_size` | `Optional[int]` | `value` | | | -| `cp_config.cp_type` | `` | `categorical` | | `ULYSSES`, `STAR`, `RING`, `HELIX` | -| `cp_config.fifo_version` | `Optional[int]` | `value` | | | -| `cp_config.tokens_per_block` | `Optional[int]` | `value` | | | -| `cp_config.use_nccl_for_alltoall` | `Optional[bool]` | `value` | | | -| `dtype` | `` | `categorical` | allowlist | `auto`, `float16`, `bfloat16`, `float32` | -| `embedding_parallel_mode` | `Literal['NONE', 'SHARDING_ALONG_VOCAB', 'SHARDING_ALONG_HIDDEN']` | `categorical` | | `NONE`, `SHARDING_ALONG_VOCAB`, `SHARDING_ALONG_HIDDEN` | -| `enable_attention_dp` | `` | `value` | | | -| `enable_build_cache.max_cache_storage_gb` | `` | `value` | | | -| `enable_build_cache.max_records` | `` | `value` | | | -| `enable_chunked_prefill` | `` | `value` | | | -| `enable_energy_metrics` | `` | `value` | | | -| `enable_lm_head_tp_in_adp` | `` | `value` | | | -| `enable_lora` | `` | `value` | | | -| `enable_prompt_adapter` | `` | `value` | | | -| `enable_tqdm` | `` | `value` | | | -| `extended_runtime_perf_knob_config.cuda_graph_cache_size` | `` | `value` | | | -| `extended_runtime_perf_knob_config.cuda_graph_mode` | `` | `value` | | | -| `extended_runtime_perf_knob_config.enable_context_fmha_fp32_acc` | `` | `value` | | | -| `extended_runtime_perf_knob_config.multi_block_mode` | `` | `value` | | | -| `fail_fast_on_attention_window_too_large` | `` | `value` | | | -| `fast_build` | `` | `value` | | | -| `gather_generation_logits` | `` | `value` | | | -| `gpus_per_node` | `Optional[int]` | `value` | | | -| `guided_decoding_backend` | `Optional[Literal['xgrammar', 'llguidance']]` | `categorical` | | `xgrammar`, `llguidance` | -| `iter_stats_max_iterations` | `Optional[int]` | `value` | | | -| `kv_cache_config.attention_dp_events_gather_period_ms` | `` | `value` | | | -| `kv_cache_config.avg_seq_len` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` | | | -| `kv_cache_config.block_reuse_policy` | `Literal['all_reusable', 'per_request']` | `categorical` | | `all_reusable`, `per_request` | -| `kv_cache_config.copy_on_partial_reuse` | `` | `value` | | | -| `kv_cache_config.cross_kv_cache_fraction` | `Optional[float]` | `value` | | | -| `kv_cache_config.disk_cache_size` | `Optional[Annotated[int, Ge(ge=0)]]` | `value` | | | -| `kv_cache_config.disk_prefetch_num_reqs` | `` | `value` | | | -| `kv_cache_config.dtype` | `` | `categorical` | allowlist | `auto`, `float16`, `bfloat16`, `float32`, `fp8`, `nvfp4` | -| `kv_cache_config.enable_block_reuse` | `` | `value` | | | -| `kv_cache_config.enable_kv_pool_rebalance` | `` | `value` | | | -| `kv_cache_config.enable_partial_reuse` | `` | `value` | | | -| `kv_cache_config.enable_swa_scratch_reuse` | `` | `value` | | | -| `kv_cache_config.event_buffer_max_size` | `` | `value` | | | -| `kv_cache_config.free_gpu_memory_fraction` | `Optional[float]` | `value` | | | -| `kv_cache_config.host_cache_size` | `Optional[int]` | `value` | | | -| `kv_cache_config.iteration_stats_interval` | `` | `value` | | | -| `kv_cache_config.kv_cache_event_hash_algo` | `Literal['auto', 'v1_block_key', 'v2_sha256', 'v2_sha256_64']` | `categorical` | | `auto`, `v1_block_key`, `v2_sha256`, `v2_sha256_64` | -| `kv_cache_config.mamba_ssm_cache_dtype` | `Literal['auto', 'float16', 'bfloat16', 'float32']` | `categorical` | | `auto`, `float16`, `bfloat16`, `float32` | -| `kv_cache_config.mamba_ssm_philox_rounds` | `` | `value` | | | -| `kv_cache_config.mamba_ssm_stochastic_rounding` | `` | `value` | | | -| `kv_cache_config.mamba_state_cache_interval` | `` | `value` | | | -| `kv_cache_config.max_attention_window` | `Optional[List[int]]` | `value` | | | -| `kv_cache_config.max_gpu_total_bytes` | `` | `value` | | | -| `kv_cache_config.max_tokens` | `Optional[int]` | `value` | | | -| `kv_cache_config.max_util_for_resume` | `` | `value` | | | -| `kv_cache_config.pool_ratio` | `Optional[List[float]]` | `value` | | | -| `kv_cache_config.secondary_offload_min_priority` | `Optional[int]` | `value` | | | -| `kv_cache_config.sink_token_length` | `Optional[int]` | `value` | | | -| `kv_cache_config.tokens_per_block` | `` | `value` | | | -| `kv_cache_config.use_kv_cache_manager_v2` | `` | `value` | | | -| `kv_cache_config.use_uvm` | `` | `value` | | | -| `load_format` | `Literal['auto', 'dummy']` | `categorical` | | `auto`, `dummy` | -| `lora_config.lora_ckpt_source` | `Literal['hf', 'nemo']` | `categorical` | | `hf`, `nemo` | -| `lora_config.max_cpu_loras` | `Optional[int]` | `value` | | | -| `lora_config.max_lora_rank` | `` | `value` | | | -| `lora_config.max_loras` | `Optional[int]` | `value` | | | -| `lora_config.swap_gate_up_proj_lora_b_weight` | `` | `value` | | | -| `max_batch_size` | `Optional[int]` | `value` | | | -| `max_beam_width` | `Optional[int]` | `value` | | | -| `max_input_len` | `Optional[int]` | `value` | | | -| `max_num_tokens` | `Optional[int]` | `value` | | | -| `max_prompt_adapter_token` | `` | `value` | | | -| `max_seq_len` | `Optional[int]` | `value` | | | -| `moe_cluster_parallel_size` | `Optional[int]` | `value` | | | -| `moe_expert_parallel_size` | `Optional[int]` | `value` | | | -| `moe_tensor_parallel_size` | `Optional[int]` | `value` | | | -| `normalize_log_probs` | `` | `value` | | | -| `num_postprocess_workers` | `` | `value` | | | -| `orchestrator_type` | `Optional[Literal['rpc', 'ray']]` | `categorical` | | `rpc`, `ray` | -| `peft_cache_config.device_cache_percent` | `` | `value` | | | -| `peft_cache_config.host_cache_size` | `` | `value` | | | -| `peft_cache_config.max_adapter_size` | `` | `value` | | | -| `peft_cache_config.max_pages_per_block_device` | `` | `value` | | | -| `peft_cache_config.max_pages_per_block_host` | `` | `value` | | | -| `peft_cache_config.num_copy_streams` | `` | `value` | | | -| `peft_cache_config.num_device_module_layer` | `` | `value` | | | -| `peft_cache_config.num_ensure_workers` | `` | `value` | | | -| `peft_cache_config.num_host_module_layer` | `` | `value` | | | -| `peft_cache_config.num_put_workers` | `` | `value` | | | -| `peft_cache_config.optimal_adapter_size` | `` | `value` | | | -| `perf_metrics_max_requests` | `` | `value` | | | -| `pipeline_parallel_size` | `` | `value` | | | -| `pp_partition` | `Optional[List[int]]` | `value` | | | -| `prometheus_metrics_config.e2e_request_latency_buckets` | `Optional[List[float]]` | `value` | | | -| `prometheus_metrics_config.request_decode_time_buckets` | `Optional[List[float]]` | `value` | | | -| `prometheus_metrics_config.request_inference_time_buckets` | `Optional[List[float]]` | `value` | | | -| `prometheus_metrics_config.request_prefill_time_buckets` | `Optional[List[float]]` | `value` | | | -| `prometheus_metrics_config.request_queue_time_buckets` | `Optional[List[float]]` | `value` | | | -| `prometheus_metrics_config.time_per_output_token_buckets` | `Optional[List[float]]` | `value` | | | -| `prometheus_metrics_config.time_to_first_token_buckets` | `Optional[List[float]]` | `value` | | | -| `quant_config.clamp_val` | `Optional[List[float]]` | `value` | | | -| `quant_config.group_size` | `Optional[int]` | `value` | | | -| `quant_config.has_zero_point` | `` | `value` | | | -| `quant_config.kv_cache_quant_algo` | `Optional[tensorrt_llm.quantization.mode.QuantAlgo]` | `categorical` | | `W8A16`, `W4A16`, `W4A16_AWQ`, `W4A8_AWQ`, `W8A16_GPTQ`, `W4A16_GPTQ`, `W8A8_SQ_PER_CHANNEL`, `W8A8_SQ_PER_TENSOR_PLUGIN`, `W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN`, `W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN`, `W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN`, `W4A8_QSERVE_PER_GROUP`, `W4A8_QSERVE_PER_CHANNEL`, `FP8`, `FP8_PER_CHANNEL_PER_TOKEN`, `FP8_BLOCK_SCALES`, `INT8`, `MIXED_PRECISION`, `NVFP4`, `W4A8_NVFP4_FP8`, `W4A8_MXFP4_FP8`, `W4A8_MXFP4_MXFP8`, `W4A16_MXFP4`, `MXFP8`, `W4A16_NVFP4`, `NVFP4_AWQ`, `NVFP4_ARC`, `NO_QUANT` | -| `quant_config.mamba_ssm_philox_rounds` | `` | `value` | | | -| `quant_config.mamba_ssm_stochastic_rounding` | `` | `value` | | | -| `quant_config.pre_quant_scale` | `` | `value` | | | -| `quant_config.quant_algo` | `Optional[tensorrt_llm.quantization.mode.QuantAlgo]` | `categorical` | | `W8A16`, `W4A16`, `W4A16_AWQ`, `W4A8_AWQ`, `W8A16_GPTQ`, `W4A16_GPTQ`, `W8A8_SQ_PER_CHANNEL`, `W8A8_SQ_PER_TENSOR_PLUGIN`, `W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN`, `W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN`, `W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN`, `W4A8_QSERVE_PER_GROUP`, `W4A8_QSERVE_PER_CHANNEL`, `FP8`, `FP8_PER_CHANNEL_PER_TOKEN`, `FP8_BLOCK_SCALES`, `INT8`, `MIXED_PRECISION`, `NVFP4`, `W4A8_NVFP4_FP8`, `W4A8_MXFP4_FP8`, `W4A8_MXFP4_MXFP8`, `W4A16_MXFP4`, `MXFP8`, `W4A16_NVFP4`, `NVFP4_AWQ`, `NVFP4_ARC`, `NO_QUANT` | -| `quant_config.smoothquant_val` | `` | `value` | | | -| `quant_config.use_meta_recipe` | `` | `value` | | | -| `reasoning_parser` | `Optional[str]` | `categorical` | allowlist | `auto`, `deepseek-r1`, `laguna`, `qwen3`, `qwen3_5`, `minimax_m2`, `minimax_m2_append_think`, `nano-v3`, `gemma4`, `kimi_k2`, `kimi_k25` | -| `request_stats_max_iterations` | `Optional[int]` | `value` | | | -| `return_perf_metrics` | `` | `value` | | | -| `scheduler_config.capacity_scheduler_policy` | `` | `categorical` | | `MAX_UTILIZATION`, `GUARANTEED_NO_EVICT`, `STATIC_BATCH` | -| `scheduler_config.context_chunking_policy` | `Optional[tensorrt_llm.llmapi.llm_args.ContextChunkingPolicy]` | `categorical` | | `FIRST_COME_FIRST_SERVED`, `EQUAL_PROGRESS`, `FORCE_CHUNK` | -| `scheduler_config.dynamic_batch_config.dynamic_batch_moving_average_window` | `` | `value` | | | -| `scheduler_config.dynamic_batch_config.enable_batch_size_tuning` | `` | `value` | | | -| `scheduler_config.dynamic_batch_config.enable_max_num_tokens_tuning` | `` | `value` | | | -| `scheduler_config.enable_prefix_aware_scheduling` | `` | `value` | | | -| `scheduler_config.use_python_scheduler` | `` | `value` | | | -| `scheduler_config.waiting_queue_policy` | `` | `categorical` | | `fcfs`, `priority` | -| `skip_tokenizer_init` | `` | `value` | | | -| `sparse_attention_config.algorithm` | `Literal['dsa']` | `categorical` | | `dsa`, `deepseek_v4`, `minimax_m3`, `rocket`, `skip_softmax` | -| `sparse_attention_config.compress_ratios` | `List[int]` | `value` | | | -| `sparse_attention_config.enable_heuristic_topk` | `` | `value` | | | -| `sparse_attention_config.index_head_dim` | `Optional[int]` | `value` | | | -| `sparse_attention_config.index_n_heads` | `Optional[int]` | `value` | | | -| `sparse_attention_config.index_topk` | `Optional[int]` | `value` | | | -| `sparse_attention_config.indexer_k_dtype` | `Literal['fp8', 'fp4']` | `categorical` | | `fp8`, `fp4` | -| `sparse_attention_config.indexer_max_chunk_size` | `Optional[int]` | `value` | | | -| `sparse_attention_config.indexer_rope_interleave` | `` | `value` | | | -| `sparse_attention_config.kernel_size` | `Optional[int]` | `value` | | | -| `sparse_attention_config.kt_cache_dtype` | `Optional[str]` | `categorical` | allowlist | `bfloat16`, `float8_e5m2` | -| `sparse_attention_config.page_size` | `Optional[int]` | `value` | | | -| `sparse_attention_config.prompt_budget` | `Optional[int]` | `value` | | | -| `sparse_attention_config.q_split_threshold` | `` | `value` | | | -| `sparse_attention_config.seq_len_threshold` | `Optional[int]` | `value` | | | -| `sparse_attention_config.skip_indexer_for_short_seqs` | `` | `value` | | | -| `sparse_attention_config.sparse_block_size` | `` | `value` | | | -| `sparse_attention_config.sparse_disable_index_value` | `` | `value` | | | -| `sparse_attention_config.sparse_index_dim` | `` | `value` | | | -| `sparse_attention_config.sparse_init_blocks` | `` | `value` | | | -| `sparse_attention_config.sparse_local_blocks` | `` | `value` | | | -| `sparse_attention_config.sparse_num_index_heads` | `` | `value` | | | -| `sparse_attention_config.sparse_score_type` | `Literal['max']` | `categorical` | | `max` | -| `sparse_attention_config.sparse_topk_blocks` | `` | `value` | | | -| `sparse_attention_config.topk` | `Optional[int]` | `value` | | | -| `sparse_attention_config.topr` | `Union[int, float, NoneType]` | `value` | | | -| `sparse_attention_config.use_cute_dsl_paged_mqa_logits` | `` | `value` | | | -| `sparse_attention_config.use_cute_dsl_topk` | `` | `value` | | | -| `sparse_attention_config.window_size` | `` | `value` | | | -| `speculative_config.acceptance_rate_threshold` | `Optional[float]` | `value` | | | -| `speculative_config.acceptance_rate_window_size` | `Optional[Annotated[int, Ge(ge=0)]]` | `value` | | | -| `speculative_config.allow_advanced_sampling` | `` | `value` | | | -| `speculative_config.begin_thinking_phase_token` | `` | `value` | | | -| `speculative_config.decoding_type` | `Literal['AUTO']` | `categorical` | | `AUTO`, `DFlash`, `Draft_Target`, `Eagle3`, `Eagle`, `Lookahead`, `MTP`, `Medusa`, `NGram`, `PARD`, `SA`, `SaveState`, `User_Provided` | -| `speculative_config.dynamic_tree_max_topK` | `Optional[int]` | `value` | | | -| `speculative_config.eagle3_layers_to_capture` | `Optional[Set[int]]` | `value` | | | -| `speculative_config.eagle3_model_arch` | `Literal['llama3', 'mistral_large3']` | `categorical` | | `llama3`, `mistral_large3` | -| `speculative_config.eagle3_one_model` | `Optional[bool]` | `value` | | | -| `speculative_config.eagle_choices` | `Optional[List[List[int]]]` | `value` | | | -| `speculative_config.enable_global_pool` | `` | `value` | | | -| `speculative_config.end_thinking_phase_token` | `` | `value` | | | -| `speculative_config.global_pool_size` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` | | | -| `speculative_config.greedy_sampling` | `Optional[bool]` | `value` | | | -| `speculative_config.is_keep_all` | `` | `value` | | | -| `speculative_config.is_public_pool` | `` | `value` | | | -| `speculative_config.is_use_oldest` | `` | `value` | | | -| `speculative_config.mask_token_id` | `Optional[int]` | `value` | | | -| `speculative_config.max_concurrency` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` | | | -| `speculative_config.max_draft_len` | `Optional[Annotated[int, Ge(ge=0)]]` | `value` | | | -| `speculative_config.max_matching_ngram_size` | `` | `value` | | | -| `speculative_config.max_ngram_size` | `` | `value` | | | -| `speculative_config.max_non_leaves_per_layer` | `Optional[int]` | `value` | | | -| `speculative_config.max_total_draft_tokens` | `Optional[int]` | `value` | | | -| `speculative_config.max_verification_set_size` | `` | `value` | | | -| `speculative_config.max_window_size` | `` | `value` | | | -| `speculative_config.medusa_choices` | `Optional[List[List[int]]]` | `value` | | | -| `speculative_config.mtp_eagle_one_model` | `` | `value` | | | -| `speculative_config.num_eagle_layers` | `Optional[int]` | `value` | | | -| `speculative_config.num_medusa_heads` | `Optional[int]` | `value` | | | -| `speculative_config.num_nextn_predict_layers` | `Optional[int]` | `value` | | | -| `speculative_config.posterior_threshold` | `Optional[float]` | `value` | | | -| `speculative_config.relaxed_delta` | `` | `value` | | | -| `speculative_config.relaxed_topk` | `` | `value` | | | -| `speculative_config.sa_config.enable_global_pool` | `` | `value` | | | -| `speculative_config.sa_config.threshold` | `` | `value` | | | -| `speculative_config.target_layer_ids` | `Optional[List[int]]` | `value` | | | -| `speculative_config.use_dynamic_tree` | `Optional[bool]` | `value` | | | -| `speculative_config.use_mtp_vanilla` | `` | `value` | | | -| `speculative_config.use_rejection_sampling` | `` | `value` | | | -| `speculative_config.use_relaxed_acceptance_for_thinking` | `` | `value` | | | -| `speculative_config.write_interval` | `` | `value` | | | -| `telemetry_config.disabled` | `` | `value` | | | -| `telemetry_config.usage_context` | `` | `categorical` | | `unknown`, `llm_class`, `cli_serve`, `cli_bench`, `cli_eval` | -| `tensor_parallel_size` | `` | `value` | | | -| `tokenizer_mode` | `Literal['auto', 'slow']` | `categorical` | | `auto`, `slow` | -| `trust_remote_code` | `` | `value` | | | diff --git a/docs/source/examples/customization.md b/docs/source/examples/customization.md index 4c357554f504..3a7aa5b3f9e8 100644 --- a/docs/source/examples/customization.md +++ b/docs/source/examples/customization.md @@ -2,70 +2,51 @@ ## Quantization -TensorRT LLM can quantize the Hugging Face model automatically. By setting the appropriate flags in the `LLM` instance. For example, to perform an Int4 AWQ quantization, the following code triggers the model quantization. Please refer to complete list of [supported flags](https://nvidia.github.io/TensorRT-LLM/_modules/tensorrt_llm/quantization/mode.html#QuantAlgo) and acceptable values. +TensorRT LLM runs quantized models from pre-quantized checkpoints. Use a checkpoint quantized with [NVIDIA TensorRT Model Optimizer](https://github.com/NVIDIA/TensorRT-Model-Optimizer) (for example, the ready-made FP8/NVFP4 checkpoints published on the [NVIDIA Hugging Face hub](https://huggingface.co/nvidia)), and the quantization configuration is detected automatically when the model loads: ``` python -from tensorrt_llm.llmapi import QuantConfig, QuantAlgo +from tensorrt_llm import LLM -quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_AWQ) - -llm = LLM(, quant_config=quant_config) +llm = LLM("nvidia/Llama-3.1-8B-Instruct-FP8") ``` +Refer to the [quantization feature documentation](../features/quantization.md) for the supported formats per GPU architecture and instructions on quantizing your own model. + ## Sampling -SamplingParams can customize the sampling strategy to control LLM generated responses, such as beam search, temperature, and [others](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/llmapi/utils.py#L55-L76). +SamplingParams can customize the sampling strategy to control LLM generated responses, such as beam search, temperature, and many others. -As an example, to enable beam search with a beam size of 4, set the `sampling_params` as follows: +As an example, to enable beam search with a beam width of 4, configure the engine limit with `max_beam_width` and request beam search through `SamplingParams`: ```python -from tensorrt_llm.llmapi import LLM, SamplingParams, BuildConfig +from tensorrt_llm import LLM, SamplingParams -build_config = BuildConfig() -build_config.max_beam_width = 4 - -llm = LLM(, build_config=build_config) +llm = LLM(, max_beam_width=4) # Let the LLM object generate text with the default sampling strategy, or # you can create a SamplingParams object as well with several fields set manually -sampling_params = SamplingParams(beam_width=4) # current limitation: beam_width should be equal to max_beam_width +sampling_params = SamplingParams(n=4, use_beam_search=True) for output in llm.generate(, sampling_params=sampling_params): print(output) ``` -`SamplingParams` manages and dispatches fields to C++ classes including: - -* [SamplingConfig](https://nvidia.github.io/TensorRT-LLM/_cpp_gen/runtime.html#_CPPv4N12tensorrt_llm7runtime14SamplingConfigE) -* [OutputConfig](https://nvidia.github.io/TensorRT-LLM/_cpp_gen/executor.html#_CPPv4N12tensorrt_llm8executor12OutputConfigE) - -Refer to the [class documentation](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html#tensorrt_llm.llmapi.SamplingParams) for more details. - -## Build Configuration - -Apart from the arguments mentioned above, you can also customize the build configuration with the `build_config` class and other arguments borrowed from the trtllm-build CLI. These build configuration options provide flexibility in building engines for the target hardware and use cases. Refer to the following example: - -```python -llm = LLM(, - build_config=BuildConfig( - max_num_tokens=4096, - max_batch_size=128, - max_beam_width=4)) -``` -Refer to the [buildconfig documentation](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/builder.py#L470-L501) for more details. +Refer to the [class documentation](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html#tensorrt_llm.llmapi.SamplingParams) for the complete list of fields. ## Runtime Customization -Similar to `build_config`, you can also customize the runtime configuration with the `runtime_config`, `peft_cache_config` or other [arguments](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/llmapi/llm_utils.py#L186-L223) borrowed from the Executor APIs. These runtime configuration options provide additional flexibility with respect to KV cache management, GPU memory allocation and so on. Refer to the following example: - +Runtime behavior such as KV cache management and GPU memory allocation can be customized with dedicated configuration classes like `kv_cache_config` and `peft_cache_config` passed to the `LLM` constructor. Refer to the following example: ```python -from tensorrt_llm.llmapi import LLM, KvCacheConfig +from tensorrt_llm import LLM +from tensorrt_llm.llmapi import KvCacheConfig llm = LLM(, kv_cache_config=KvCacheConfig( free_gpu_memory_fraction=0.8)) ``` +Refer to the [LLM API reference](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) for all available configuration classes. + ## Tokenizer Customization By default, the LLM API uses transformers’ `AutoTokenizer`. You can override it with your own tokenizer by passing it when creating the LLM object. Refer to the following example: @@ -90,8 +71,8 @@ for output in llm.generate([32, 12]): For performance considerations, you can disable the tokenizer by passing `skip_tokenizer_init=True` when creating `LLM`. In this case, `LLM.generate` and `LLM.generate_async` will expect prompt token ids as input. Refer to the following example: ```python -llm = LLM() -for output in llm.generate([[32, 12]], skip_tokenizer_init=True): +llm = LLM(, skip_tokenizer_init=True) +for output in llm.generate([[32, 12]]): print(output) ``` diff --git a/examples/auto_deploy/paragraf/create_standalone_package.py b/examples/auto_deploy/paragraf/create_standalone_package.py index 45f12c85dc57..d2e75e5cbf39 100644 --- a/examples/auto_deploy/paragraf/create_standalone_package.py +++ b/examples/auto_deploy/paragraf/create_standalone_package.py @@ -555,14 +555,6 @@ def insert_optional_trtllm_guard() -> None: ensure_imports(trtllm_import.start(), "os", "pytest") insert_optional_trtllm_guard() - if optional_trtllm_guards: - # The standalone package can rely on the installed trtllm-bench entrypoint, - # but it does not ship TensorRT-LLM's source-tree benchmarks/ directory. - content = content.replace( - ' script_dir = Path(root_dir, "benchmarks")\n', - " script_dir = Path(temp_dir)\n", - ) - replacements = sum(1 for a, b in zip(original, content) if a != b) # rough count if content != original: with open(filepath, "w") as f: diff --git a/examples/layer_wise_benchmarks/sample_performance_alignment.sh b/examples/layer_wise_benchmarks/sample_performance_alignment.sh index 4b9bf883c3c9..30140c280f9b 100755 --- a/examples/layer_wise_benchmarks/sample_performance_alignment.sh +++ b/examples/layer_wise_benchmarks/sample_performance_alignment.sh @@ -14,8 +14,9 @@ export TLLM_AUTOTUNER_CACHE_PATH="$PROFILE_DIR/sample_performance_alignment_cach mkdir -p -- "$PROFILE_DIR" mkdir -p -- "$(dirname -- "$TLLM_AUTOTUNER_CACHE_PATH")" -python3 ../../benchmarks/prepare_dataset.py \ - --tokenizer "$MODEL" \ +trtllm-bench \ + --model "$MODEL" \ + prepare-dataset \ --stdout \ --random-seed 42 \ token-norm-dist \ diff --git a/examples/quantization/README.md b/examples/quantization/README.md index b3b2e35b20ff..23bfa5d40c8e 100644 --- a/examples/quantization/README.md +++ b/examples/quantization/README.md @@ -1,249 +1,19 @@ -# TensorRT LLM Quantization Toolkit Installation Guide +# Model Quantization -## Introduction +To run quantized models with TensorRT LLM: -This document introduces: +- Use a pre-quantized Hugging Face checkpoint (for example the FP8/NVFP4 + checkpoints published on the [NVIDIA Hugging Face hub](https://huggingface.co/nvidia)). + Quantization settings are detected automatically when the model loads. +- To quantize your own model, use the + [NVIDIA TensorRT Model Optimizer](https://github.com/NVIDIA/TensorRT-Model-Optimizer) + Hugging Face export flow (`examples/llm_ptq` in that repository). -- The steps to install the TensorRT LLM quantization toolkit. -- The Python APIs to quantize the models. +See the [quantization feature documentation](https://nvidia.github.io/TensorRT-LLM/features/quantization.html) +for supported formats per GPU architecture. -The detailed LLM quantization recipe is distributed to the README.md of the corresponding model examples. +## Mixed-precision MoE checkpoints -## Installation - -The NVIDIA Model Optimizer quantization toolkit is installed automatically as a dependency of TensorRT-LLM. - -```bash -# Install the additional requirements -cd examples/quantization -pip install -r requirements.txt -``` - -## Usage - -```bash -# FP8 quantization. -python quantize.py --model_dir $MODEL_PATH --qformat fp8 --kv_cache_dtype fp8 --output_dir $OUTPUT_PATH - -# INT4_AWQ tp4 quantization. -python quantize.py --model_dir $MODEL_PATH --qformat int4_awq --awq_block_size 64 --tp_size 4 --output_dir $OUTPUT_PATH - -# INT8 SQ with INT8 kv cache. -python quantize.py --model_dir $MODEL_PATH --qformat int8_sq --kv_cache_dtype int8 --output_dir $OUTPUT_PATH - -# Auto quantization(e.g. fp8 + int4_awq + w4a8_awq) using average weights bits 5 -python quantize.py --model_dir $MODEL_PATH --autoq_format fp8,int4_awq,w4a8_awq --output_dir $OUTPUT_PATH --auto_quantize_bits 5 --tp_size 2 - -# FP8 quantization for NeMo model. -python quantize.py --nemo_ckpt_path nemotron-3-8b-base-4k/Nemotron-3-8B-Base-4k.nemo \ - --dtype bfloat16 \ - --batch_size 64 \ - --qformat fp8 \ - --output_dir nemotron-3-8b/trt_ckpt/fp8/1-gpu - -# FP8 quantization for Medusa model. -python quantize.py --model_dir $MODEL_PATH\ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir $OUTPUT_PATH \ - --calib_size 512 \ - --tp_size 1 \ - --medusa_model_dir /path/to/medusa_head/ \ - --num_medusa_heads 4 -``` -Checkpoint saved in `output_dir` can be directly passed to `trtllm-build`. - -### Quantization Arguments: - -- model_dir: Hugging Face model path. -- qformat: Specify the quantization algorithm applied to the checkpoint. - - nvfp4: Weights are quantized to NVFP4 block-wise with size 16. Activation global scale are calibrated. - - fp8: Weights are quantized to FP8 tensor wise. Activation ranges are calibrated tensor wise. - - fp8_pc_pt: Weights are quantized to FP8 per-channel. Activation ranges are calibrated and quantized per-token. - - int8_sq: Weights are smoothed and quantized to INT8 channel wise. Activation ranges are calibrated tensor wise. - - int4_awq: Weights are re-scaled and block-wise quantized to INT4. Block size is specified by `awq_block_size`. - - w4a8_awq: Weights are re-scaled and block-wise quantized to INT4. Block size is specified by `awq_block_size`. Activation ranges are calibrated tensor wise. - - int8_wo: Actually nothing is applied to weights. Weights are quantized to INT8 channel wise when TRTLLM building the engine. - - int4_wo: Same as int8_wo but in INT4. - - full_prec: No quantization. -- autoq_format: Specific quantization algorithms are searched in auto quantization. The algorithm must in ['fp8', 'int4_awq', 'w4a8_awq', 'int8_sq'] and you can use ',' to separate more than one quantization algorithms, such as `--autoq_format fp8,int4_awq,w4a8_awq`. Please attention that using int8_sq and fp8 together is not supported. -- auto_quantize_bits: Effective bits constraint for auto quantization. If not set, regular quantization without auto quantization search is applied. Note: it must be set within correct range otherwise it will be set by lowest value if possible. For example, the weights of LLMs have 16 bits defaultly and it results in a weight compression rate of 40% if we set `auto_quantize_bits` to 9.6 (9.6 / 16 = 0.6), which means the average bits of the weights are 9.6 but not 16. However, which format to choose is determined by solving an optimization problem, so you need to generate the according checkpoint manually if you want to customize your checkpoint formats. The format of mixed precision checkpoint is described in detail below. -- output_dir: Path to save the quantized checkpoint. -- dtype: Specify data type of model when loading from Hugging Face. -- kv_cache_dtype: Specify kv cache data type. - - int8: Use int8 kv cache. - - fp8: Use FP8 kv cache. - - None (default): Use kv cache as model dtype. -- batch_size: Batch size for calibration. Default is 1. -- calib_size: Number of samples. Default is 512. -- calib_max_seq_length: Max sequence length of calibration samples. Default is 512. -- tp_size: Checkpoint is tensor paralleled by tp_size. Default is 1. -- pp_size: Checkpoint is pipeline paralleled by pp_size. Default is 1. -- awq_block_size: AWQ algorithm specific parameter. Indicate the block size when quantizing weights. 64 and 128 are supported by TRTLLM. -- quantize_lm_head: Enable quantization of lm_head layer. This is only supported for FP8 quantization. Default is false. - -#### NeMo model specific arguments: - -- nemo_ckpt_path: NeMo checkpoint path. -- calib_tp_size: TP size for NeMo checkpoint calibration. -- calib_pp_size: PP size for NeMo checkpoint calibration. - -#### Medusa specific arguments: - -- medusa_model_dir: Model path of medusa. -- quant_medusa_head: Whether to quantize the weights of medusa heads. -- num_medusa_heads: Number of medusa heads. -- num_medusa_layers: Number of medusa layers. -- max_draft_len: Max length of draft. -- medusa_hidden_act: Activation function of medusa. - -### Building Arguments: - -There are several arguments for the building stage which relate to quantization. -- use_fp8_context_fmha: This is Hopper-only feature. Use FP8 Gemm to calculate the attention operation. - -```python -qkv scale = 1.0 -FP_O = quantize(softmax(FP8_Q * FP8_K), scale=1.0) * FP8_V -FP_O * output_scale = FP8_O -``` - -### Checkpoint Conversion Arguments (not supported by all models) - -- FP8 - - use_fp8_rowwise: Enable FP8 per-token per-channel quantization for linear layer. (FP8 from `quantize.py` is per-tensor). -- INT8 - - smoothquant: Enable INT8 quantization for linear layer. Set the α parameter (see https://arxiv.org/pdf/2211.10438.pdf) to Smoothquant the model, and output int8 weights. A good first try is 0.5. Must be in [0, 1]. - - per_channel: Using per-channel quantization for weight when `smoothquant` is enabled. - - per_token: Using per-token quantization for activation when `smoothquant` is enabled. -- Weight-Only - - use_weight_only: Weights are quantized to INT4 or INT8 channel wise. - - weight_only_precision: Indicate `int4` or `int8` when `use_weight_only` is enabled. Or `int4_gptq` when `quant_ckpt_path` is provided which means checkpoint is for GPTQ. - - quant_ckpt_path: Path of a GPTQ quantized model checkpoint in `.safetensors` format. - - group_size: Group size used in GPTQ quantization. - - per_group: Should be enabled when load from GPTQ. -- KV Cache - - int8_kv_cache: By default, we use dtype for KV cache. int8_kv_cache chooses int8 quantization for KV cache. - - fp8_kv_cache: By default, we use dtype for KV cache. fp8_kv_cache chooses fp8 quantization for KV cache. - -### Format of Mixed Precision Checkpoints - -ModelOpt can produce a mixed precision TensorRT LLM checkpoint. After producing the quantized checkpoint, you can build engine directly by `trtllm-build` command: -```bash -trtllm-build --checkpoint_dir --output_dir $OUTPUT_PATH -``` -If you have some special needs about the model weights, such as int4 for MLP and int8 for the rest, you need to generate the checkpoint and config files by yourself. - -The `trtllm-build` command consumes the same format of weights, which is presented in [TensorRT LLM checkpoint formats](https://nvidia.github.io/TensorRT-LLM/architecture/checkpoint.html), but has different quantization method for every linear. Therefore, each layer, such as layer30.mlp.fc, layer30.attention.dense, and so on, keeps the same model weights according to the quantization formats in TensorRT LLM checkpoint. What's more, the `quantization` field in `config.json` will be like this: -``` - "quantization": { - "quant_algo": "MIXED_PRECISION", - "kv_cache_quant_algo": "FP8" // The quant_algo of KV cache may change - }, -``` -There will be another file about per-layer quantization information named `quant_cfg.json` in the same directory, the format of it is like: -``` -{ - "quant_algo": "MIXED_PRECISION", - "kv_cache_quant_algo": "FP8", - "quantized_layers": { // one more filed presents per-layer's information - "transformer.layers.0.attention.qkv": { - "quant_algo": "FP8" // specific algorithm for each linear - }, - "transformer.layers.0.attention.dense": { - "quant_algo": "FP8" - }, - "transformer.layers.0.mlp.fc": { - "quant_algo": "W4A16_AWQ", - "group_size": 128, - "has_zero_point": false, - "pre_quant_scale": true - }, - "transformer.layers.0.mlp.proj": { - "quant_algo": "W8A8_SQ_PER_CHANNEL" - }, - ... - "transformer.layers.31.mlp.proj": { - "quant_algo": "FP8" - } - } -} -``` - -TensorRT LLM will automatically read `quant_cfg.json` after recogniziong the `MIXED_PRECISION` quantization method in `config.json`. All the specific algorithm keeps the same as what in `quantization` field before. If some layers are not listed, they'll be treated as no quantization. - -## APIs - -[`quantize.py`](./quantize.py) uses the quantization toolkit to calibrate the PyTorch models and export TensorRT LLM checkpoints. Each TensorRT LLM checkpoint contains a config file (in .json format) and one or several rank weight files (in .safetensors format). It will produce one another quantization config for per-layer's information when setting auto quantization. The checkpoints can be directly used by `trtllm-build` command to build TensorRT LLM engines. See this [`doc`](../../docs/source/architecture/checkpoint.md) for more details on the TensorRT LLM checkpoint format. - -> *This quantization step may take a long time to finish and requires large GPU memory. Please use a server grade GPU if a GPU out-of-memory error occurs* - -> *If the model is trained with multi-GPU with tensor parallelism, the PTQ calibration process requires the same amount of GPUs as the training time too.* - - -### PTQ (Post Training Quantization) - -PTQ can be achieved with simple calibration on a small set of training or evaluation data (typically 128-512 samples) after converting a regular PyTorch model to a quantized model. - -```python -import torch -from torch.utils.data import DataLoader -from transformers import AutoModelForCausalLM -import modelopt.torch.quantization as mtq -import modelopt.torch.utils.dataset_utils as dataset_utils - -model = AutoModelForCausalLM.from_pretrained(...) - -# Select the quantization config, for example, FP8 -config = mtq.FP8_DEFAULT_CFG - -# Prepare the calibration set and define a forward loop -calib_dataloader = DataLoader(...) -calibrate_loop = dataset_utils.create_forward_loop( - calib_dataloader, dataloader=calib_dataloader -) - -# PTQ with in-place replacement to quantized modules -with torch.no_grad(): - mtq.quantize(model, config, forward_loop=calibrate_loop) - -# or PTQ with auto quantization -with torch.no_grad(): - model, search_history = mtq.auto_quantize( - model, - data_loader=calib_dataloader, - loss_func=lambda output, batch: output.loss, - constraints={"effective_bits": auto_quantize_bits}, # The average bits of quantized weights - forward_step=lambda model, batch: model(**batch), - quantization_formats=[quant_algo1, quant_algo2,...] + [None], - num_score_steps=min( - num_calib_steps=len(calib_dataloader), - len(calib_dataloader), 128 // batch_size - ), # Limit the number of score steps to avoid long calibration time - verbose=True, - ) -``` - -### Export Quantized Model - -After the model is quantized, it can be exported to a TensorRT LLM checkpoint, which includes - -- One json file recording the model structure and metadata, and -- One or several rank weight files storing quantized model weights and scaling factors. - -The export API is - -```python -from modelopt.torch.export import export_tensorrt_llm_checkpoint - -with torch.inference_mode(): - export_tensorrt_llm_checkpoint( - model, # The quantized model. - decoder_type, # The type of the model as str, e.g gptj, llama or gptnext. - dtype, # The exported weights data type as torch.dtype. - export_dir, # The directory where the exported files will be stored. - inference_tensor_parallel=tp_size, # The tensor parallelism size for inference. - inference_pipeline_parallel=pp_size, # The pipeline parallelism size for inference. - ) -``` +[`quantize_mixed_precision_moe.py`](quantize_mixed_precision_moe.py) builds a +mixed-precision MoE checkpoint from separately quantized checkpoints; see the +script's argparse help for usage. diff --git a/examples/quantization/quantize.py b/examples/quantization/quantize.py deleted file mode 100644 index 29c2fc5ca179..000000000000 --- a/examples/quantization/quantize.py +++ /dev/null @@ -1,208 +0,0 @@ -import argparse - -import torch.multiprocessing as mp - -from tensorrt_llm.quantization import (quantize_and_export, - quantize_nemo_and_export) - -if __name__ == "__main__": - mp.set_start_method("spawn", force=True) - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--model_dir", - help="Specify where the HuggingFace model is", - default=None) - parser.add_argument('--nemo_ckpt_path', - help="Specify where the NeMo checkpoint is", - default=None) - parser.add_argument( - '--decoder_type', - type=str, - default='gptnext', - choices=['gptnext', 'llama'], - help="Decoder type; effective for NeMo checkpoint only.") - parser.add_argument( - '--device', - help= - "The device to run calibration; effective for HuggingFace model only.", - default='cuda', - choices=['cuda', 'cpu']) - parser.add_argument( - "--device_map", - help="How to map the model on the devices", - default="auto", - choices=["auto", "sequential", "cpu", "gpu"], - ) - parser.add_argument( - '--calib_dataset', - type=str, - default='cnn_dailymail', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - parser.add_argument( - '--calib_tp_size', - type=int, - default=1, - help= - "Tensor parallel size for calibration; effective for NeMo checkpoint only." - ) - parser.add_argument( - '--calib_pp_size', - type=int, - default=1, - help= - "Pipeline parallel size for calibration; effective for NeMo checkpoint only." - ) - - parser.add_argument( - '--dtype', - type=str, - default='auto', - choices=['auto', 'float16', 'bfloat16', 'float32'], - help= - "The data type for the model weights and activations of the non-quantized part, e.g., embedding and lm_head. " - "If 'auto', the data type is automatically inferred from the source model; " - "however, if the source dtype is float32, it is converted to float16.") - parser.add_argument( - "--qformat", - help="Quantization format.", - default="full_prec", - choices=[ - "nvfp4", - "fp8", - "fp8_pc_pt", - "int8_sq", - "int4_awq", - "w4a8_awq", - "int8_wo", - "int4_wo", - "full_prec", - ], - ) - parser.add_argument( - "--seed", - help="Seed the generate random numbers, the value will be used to call" - "random.seed(value) and numpy.random.seed(value)", - type=int, - default=1234) - parser.add_argument("--tokenizer_max_seq_length", - help="Max sequence length to init the tokenizers", - type=int, - default=2048) - - parser.add_argument("--batch_size", - help="Batch size for calibration.", - type=int, - default=1) - parser.add_argument("--calib_size", - help="Number of samples for calibration.", - type=int, - default=512) - parser.add_argument("--calib_max_seq_length", - help="Max sequence length for calibration", - type=int, - default=512) - parser.add_argument("--output_dir", default="exported_model") - parser.add_argument("--tp_size", type=int, default=1) - parser.add_argument("--pp_size", type=int, default=1) - parser.add_argument("--cp_size", type=int, default=1) - parser.add_argument("--awq_block_size", type=int, default=128) - parser.add_argument("--kv_cache_dtype", - help="KV Cache dtype.", - default=None, - choices=["int8", "fp8", None]) - parser.add_argument("--quantize_lm_head", - action='store_true', - default=False) - # Medusa - parser.add_argument('--num_medusa_heads', type=int, default=4) - parser.add_argument('--num_medusa_layers', type=int, default=1) - parser.add_argument('--max_draft_len', type=int, default=63) - parser.add_argument('--medusa_hidden_act', type=str, default="silu") - parser.add_argument('--medusa_model_dir', type=str, default=None) - parser.add_argument('--quant_medusa_head', - default=False, - action='store_true', - help="whether to quantize the weights of medusa heads") - - # auto quantization - parser.add_argument( - '--autoq_format', - default=None, - type=str, - help= - "Specific quantization algorithms will be searched in auto quantization." - "The algorithm must in ['fp8', 'int4_awq', 'w4a8_awq', 'int8_sq']." - "You can use ',' to separate more than one quantization algorithms(e.g. --autoq_format fp8,int4_awq,w4a8_awq)." - "Notice: fp8 and int8_sq can't be used at the same time.") - parser.add_argument( - '--auto_quantize_bits', - type=float, - default=None, - help="Effective bits constraint for auto quantization. If not set, " - "regular quantization without auto quantization search will be applied." - "You can't set it lower than the num_bits of most aggressive quantization format." - "For example, if 'int4_awq' is in autoq_format, it can't be lower than 4.0." - ) - - args = parser.parse_args() - - # auto_quantize_bits check - if args.autoq_format: - lower_bound, upper_bound = 4 if '4' in args.autoq_format else 8, 16 - if args.auto_quantize_bits is None or args.auto_quantize_bits < lower_bound or args.auto_quantize_bits > upper_bound: - print( - f"invalid auto_quantize_bits value, will be set to {lower_bound}" - ) - args.auto_quantize_bits = lower_bound - - if args.model_dir is not None: - quantize_and_export( - model_dir=args.model_dir, - device=args.device, - calib_dataset=args.calib_dataset, - dtype=args.dtype, - qformat=args.qformat - if args.auto_quantize_bits is None else args.autoq_format, - kv_cache_dtype=args.kv_cache_dtype, - calib_size=args.calib_size, - batch_size=args.batch_size, - calib_max_seq_length=args.calib_max_seq_length, - awq_block_size=args.awq_block_size, - output_dir=args.output_dir, - tp_size=args.tp_size, - pp_size=args.pp_size, - cp_size=args.cp_size, - seed=args.seed, - tokenizer_max_seq_length=args.tokenizer_max_seq_length, - num_medusa_heads=args.num_medusa_heads, - num_medusa_layers=args.num_medusa_layers, - max_draft_len=args.max_draft_len, - medusa_hidden_act=args.medusa_hidden_act, - medusa_model_dir=args.medusa_model_dir, - quant_medusa_head=args.quant_medusa_head, - auto_quantize_bits=args.auto_quantize_bits, - device_map=args.device_map, - quantize_lm_head=args.quantize_lm_head) - elif args.nemo_ckpt_path is not None: - quantize_nemo_and_export(nemo_ckpt_path=args.nemo_ckpt_path, - decoder_type=args.decoder_type, - calib_dataset=args.calib_dataset, - calib_tp_size=args.calib_tp_size, - calib_pp_size=args.calib_pp_size, - dtype=args.dtype, - qformat=args.qformat, - kv_cache_dtype=args.kv_cache_dtype, - calib_size=args.calib_size, - batch_size=args.batch_size, - calib_max_seq_length=args.calib_max_seq_length, - awq_block_size=args.awq_block_size, - output_dir=args.output_dir, - tp_size=args.tp_size, - pp_size=args.pp_size, - cp_size=args.cp_size, - seed=args.seed) - else: - raise ValueError( - "One of source checkpoint (model_dir, nemo_ckpt_path) must be specified" - ) diff --git a/legacy-files.txt b/legacy-files.txt index de8f7f433ec0..021a4fb301a8 100644 --- a/legacy-files.txt +++ b/legacy-files.txt @@ -1,14 +1,6 @@ .devcontainer/make_env.py .github/scripts/label_community_user.py .github/scripts/pr_checklist_check.py -benchmarks/__init__.py -benchmarks/prepare_dataset.py -benchmarks/utils/__init__.py -benchmarks/utils/convert_nemo_dataset.py -benchmarks/utils/generate_rand_loras.py -benchmarks/utils/prepare_real_data.py -benchmarks/utils/prepare_synthetic_data.py -benchmarks/utils/utils.py cpp/conanfile.py cpp/kernels/fmha_v2/conftest.py cpp/kernels/fmha_v2/fmha_test.py @@ -40,24 +32,10 @@ examples/apps/chat.py examples/apps/fastapi_server.py examples/disaggregated/clients/disagg_client.py examples/disaggregated/slurm/benchmark/submit.py -examples/dora/normalize_weights.py -examples/eagle/convert_checkpoint.py -examples/eval_long_context.py -examples/generate_checkpoint_config.py -examples/generate_xgrammar_tokenizer_info.py -examples/hf_lora_convert.py examples/infinitebench/args.py examples/infinitebench/compute_scores.py examples/infinitebench/construct_synthetic_dataset.py examples/infinitebench/eval_utils.py -examples/llm-api/_tensorrt_engine/llm_eagle2_decoding.py -examples/llm-api/_tensorrt_engine/llm_eagle_decoding.py -examples/llm-api/_tensorrt_engine/llm_inference_customize.py -examples/llm-api/_tensorrt_engine/llm_inference_kv_events.py -examples/llm-api/_tensorrt_engine/llm_lookahead_decoding.py -examples/llm-api/_tensorrt_engine/llm_medusa_decoding.py -examples/llm-api/_tensorrt_engine/llm_quantization.py -examples/llm-api/_tensorrt_engine/quickstart_example.py examples/llm-api/llm_guided_decoding.py examples/llm-api/llm_inference.py examples/llm-api/llm_inference_async.py @@ -77,122 +55,17 @@ examples/llm-api/quickstart_advanced.py examples/llm-api/quickstart_example.py examples/llm-api/quickstart_multimodal.py examples/llm-api/star_attention.py -examples/llm-eval/lm-eval-harness/lm_eval_tensorrt_llm.py examples/longbench/eval_longbench_v1.py -examples/medusa/convert_checkpoint.py -examples/mmlu.py -examples/models/contrib/baichuan/convert_checkpoint.py -examples/models/contrib/bloom/convert_checkpoint.py -examples/models/contrib/chatglm-6b/tokenization_chatglm.py -examples/models/contrib/chatglm2-6b/tokenization_chatglm.py -examples/models/contrib/chatglm3-6b-32k/tokenization_chatglm.py -examples/models/contrib/cogvlm/convert_checkpoint.py -examples/models/contrib/dbrx/convert_checkpoint.py -examples/models/contrib/deepseek_v1/__init__.py -examples/models/contrib/deepseek_v1/convert_checkpoint.py -examples/models/contrib/deepseek_v2/convert_checkpoint.py -examples/models/contrib/dit/convert_checkpoint.py -examples/models/contrib/dit/diffusion.py -examples/models/contrib/dit/sample.py -examples/models/contrib/dit/utils_modelopt.py -examples/models/contrib/dit/vae_decoder_trt.py -examples/models/contrib/falcon/convert_checkpoint.py -examples/models/contrib/gptj/convert_checkpoint.py -examples/models/contrib/gptneox/convert_checkpoint.py -examples/models/contrib/grok/convert_checkpoint.py -examples/models/contrib/mmdit/convert_checkpoint.py -examples/models/contrib/mmdit/sample.py -examples/models/contrib/mpt/convert_checkpoint.py -examples/models/contrib/opt/convert_checkpoint.py -examples/models/contrib/sdxl/build_sdxl_unet.py -examples/models/contrib/sdxl/pipeline_stable_diffusion_xl.py -examples/models/contrib/sdxl/run_sdxl.py -examples/models/contrib/stdit/aspect.py -examples/models/contrib/stdit/convert_checkpoint.py -examples/models/contrib/stdit/pipeline_tllm.py -examples/models/contrib/stdit/sample.py -examples/models/contrib/stdit/scheduler.py -examples/models/contrib/stdit/text_encoder.py -examples/models/contrib/stdit/utils.py -examples/models/contrib/stdit/vae.py -examples/models/contrib/stdit/video_transforms.py -examples/models/core/bert/__init__.py -examples/models/core/bert/convert_checkpoint.py -examples/models/core/bert/run.py -examples/models/core/bert/utils.py -examples/models/core/commandr/convert_checkpoint.py -examples/models/core/enc_dec/__init__.py -examples/models/core/enc_dec/convert_checkpoint.py -examples/models/core/enc_dec/helper.py -examples/models/core/enc_dec/run.py -examples/models/core/gemma/convert_checkpoint.py -examples/models/core/glm-4-9b/convert_checkpoint.py -examples/models/core/glm-4-9b/tokenization_chatglm.py -examples/models/core/gpt/convert_checkpoint.py -examples/models/core/gpt/merge_ptuning_tables.py -examples/models/core/gpt/nemo_lora_convert.py -examples/models/core/gpt/nemo_prompt_convert.py -examples/models/core/gpt/run_hf.py examples/models/core/gpt_oss/openai_chat_client_function_calling.py -examples/models/core/internlm2/convert_checkpoint.py examples/models/core/kimi_k2/kimi_k2_tool_calling_example.py -examples/models/core/llama/convert_checkpoint.py -examples/models/core/llama/summarize_long.py -examples/models/core/mamba/convert_checkpoint.py -examples/models/core/mllama/convert_checkpoint.py -examples/models/core/multimodal/__init__.py -examples/models/core/multimodal/build_multimodal_engine.py -examples/models/core/multimodal/eval.py -examples/models/core/multimodal/run.py -examples/models/core/multimodal/utils.py -examples/models/core/nemotron_nas/calibration_utils.py -examples/models/core/nemotron_nas/convert_checkpoint.py -examples/models/core/phi/convert_checkpoint.py -examples/models/core/qwen/convert_checkpoint.py -examples/models/core/qwen2audio/run.py -examples/models/core/qwen2audio/run_chat.py -examples/models/core/qwen2audio/utils.py -examples/models/core/qwenvl/run.py -examples/models/core/qwenvl/run_chat.py -examples/models/core/qwenvl/show_pic.py -examples/models/core/qwenvl/vit_onnx_trt.py -examples/models/core/recurrentgemma/convert_checkpoint.py -examples/models/core/vit/convert_checkpoint.py -examples/models/core/whisper/convert_checkpoint.py -examples/models/core/whisper/distil_whisper/convert_from_distil_whisper.py -examples/models/core/whisper/run.py -examples/models/core/whisper/tokenizer.py -examples/models/core/whisper/whisper_utils.py -examples/ngram/run_dtm_ngram.py -examples/openai_triton/manual_plugin/build.py -examples/openai_triton/manual_plugin/fmha_triton.py -examples/openai_triton/manual_plugin/plugin.py -examples/openai_triton/manual_plugin/run.py -examples/openai_triton/plugin_autogen/build_engine.py -examples/openai_triton/plugin_autogen/kernel_config.py -examples/openai_triton/plugin_autogen/run_engine.py -examples/python_plugin/build_lookup.py -examples/python_plugin/plugin_lib/__init__.py -examples/python_plugin/plugin_lib/lookup_kernel.py -examples/python_plugin/plugin_lib/lookup_plugin.py -examples/python_plugin/run_lookup.py -examples/quantization/quantize.py examples/quantization/quantize_mixed_precision_moe.py examples/ray_orchestrator/llm_inference_async_ray.py examples/ray_orchestrator/llm_inference_distributed_ray.py -examples/redrafter/convert_checkpoint.py -examples/run.py examples/scaffolding/contrib/AsyncGeneration/stream_generation_controller.py examples/scaffolding/contrib/DeepConf/run_generation.py examples/scaffolding/contrib/Dynasor/scaffolding_dynasor_run.py examples/scaffolding/contrib/TreeInference/run_mcts_example.py examples/scaffolding/contrib/TreeInference/run_tot_example.py -examples/scaffolding/contrib/mcp/e2b/e2bserver.py -examples/scaffolding/contrib/mcp/e2b/main.py -examples/scaffolding/contrib/mcp/mcptest.py -examples/scaffolding/contrib/mcp/weather/weather.py -examples/scaffolding/contrib/mcp/websearch/main.py -examples/scaffolding/contrib/mcp/websearch/websearch.py examples/scaffolding/run_basic_generation.py examples/scaffolding/run_best_of_n_with_reward.py examples/scaffolding/run_majority_vote_aime24.py @@ -202,8 +75,6 @@ examples/serve/openai_chat_client_for_multimodal.py examples/serve/openai_completion_client.py examples/serve/openai_completion_client_for_lora.py examples/serve/openai_completion_client_json_schema.py -examples/summarize.py -examples/utils.py examples/wide_ep/ep_load_balancer/generate_eplb_config.py examples/wide_ep/ep_load_balancer/report_load_statistics.py examples/wide_ep/ep_load_balancer/utils.py @@ -215,7 +86,6 @@ scripts/build_wheel.py scripts/check_test_list.py scripts/dco_check.py scripts/format_test_list.py -scripts/generate_duration.py scripts/generate_lock_file.py scripts/get_wheel_from_package.py scripts/git_replace.py @@ -226,7 +96,6 @@ scripts/test_to_stage_mapping.py setup.py tensorrt_llm/__init__.py tensorrt_llm/_ray_utils.py -tensorrt_llm/_tensorrt_engine/__init__.py tensorrt_llm/_torch/__init__.py tensorrt_llm/_torch/attention_backend/__init__.py tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -430,7 +299,6 @@ tensorrt_llm/_torch/pyexecutor/grammar_matcher.py tensorrt_llm/_torch/pyexecutor/guided_decoder.py tensorrt_llm/_torch/pyexecutor/handle_additional_outputs.py tensorrt_llm/_torch/pyexecutor/handle_logits.py -tensorrt_llm/_torch/pyexecutor/kv_cache_connector.py tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py tensorrt_llm/_torch/pyexecutor/layerwise_nvtx_marker.py tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -468,11 +336,6 @@ tensorrt_llm/bench/benchmark/utils/__init__.py tensorrt_llm/bench/benchmark/utils/asynchronous.py tensorrt_llm/bench/benchmark/utils/general.py tensorrt_llm/bench/benchmark/utils/processes.py -tensorrt_llm/bench/build/__init__.py -tensorrt_llm/bench/build/build.py -tensorrt_llm/bench/build/dataclasses.py -tensorrt_llm/bench/build/tuning.py -tensorrt_llm/bench/build/utils.py tensorrt_llm/bench/dataclasses/__init__.py tensorrt_llm/bench/dataclasses/configuration.py tensorrt_llm/bench/dataclasses/engine.py @@ -482,13 +345,9 @@ tensorrt_llm/bench/dataclasses/reporting.py tensorrt_llm/bench/dataclasses/statistics.py tensorrt_llm/bench/utils/__init__.py tensorrt_llm/bench/utils/data.py -tensorrt_llm/builder.py tensorrt_llm/commands/__init__.py tensorrt_llm/commands/bench.py -tensorrt_llm/commands/build.py tensorrt_llm/commands/eval.py -tensorrt_llm/commands/prune.py -tensorrt_llm/commands/refit.py tensorrt_llm/commands/serve.py tensorrt_llm/evaluate/__init__.py tensorrt_llm/evaluate/cnn_dailymail.py @@ -524,23 +383,7 @@ tensorrt_llm/inputs/evs.py tensorrt_llm/inputs/multimodal.py tensorrt_llm/inputs/registry.py tensorrt_llm/inputs/utils.py -tensorrt_llm/layers/__init__.py -tensorrt_llm/layers/activation.py -tensorrt_llm/layers/attention.py -tensorrt_llm/layers/cast.py -tensorrt_llm/layers/conv.py -tensorrt_llm/layers/embedding.py -tensorrt_llm/layers/language_adapter.py -tensorrt_llm/layers/linear.py -tensorrt_llm/layers/lora.py -tensorrt_llm/layers/mlp.py -tensorrt_llm/layers/moe.py -tensorrt_llm/layers/normalization.py -tensorrt_llm/layers/pooling.py -tensorrt_llm/layers/recurrent.py -tensorrt_llm/layers/ssm.py tensorrt_llm/llmapi/__init__.py -tensorrt_llm/llmapi/build_cache.py tensorrt_llm/llmapi/disagg_utils.py tensorrt_llm/llmapi/kv_cache_type.py tensorrt_llm/llmapi/llm.py @@ -563,179 +406,17 @@ tensorrt_llm/metrics/collector.py tensorrt_llm/metrics/enums.py tensorrt_llm/models/__init__.py tensorrt_llm/models/automodel.py -tensorrt_llm/models/baichuan/__init__.py -tensorrt_llm/models/baichuan/config.py -tensorrt_llm/models/baichuan/convert.py -tensorrt_llm/models/baichuan/model.py -tensorrt_llm/models/bert/__init__.py -tensorrt_llm/models/bert/config.py -tensorrt_llm/models/bert/convert.py -tensorrt_llm/models/bert/model.py -tensorrt_llm/models/bloom/__init__.py -tensorrt_llm/models/bloom/model.py -tensorrt_llm/models/chatglm/__init__.py -tensorrt_llm/models/chatglm/config.py -tensorrt_llm/models/chatglm/convert.py -tensorrt_llm/models/chatglm/model.py -tensorrt_llm/models/clip/__init__.py -tensorrt_llm/models/clip/model.py -tensorrt_llm/models/cogvlm/__init__.py -tensorrt_llm/models/cogvlm/config.py -tensorrt_llm/models/cogvlm/convert.py -tensorrt_llm/models/cogvlm/model.py -tensorrt_llm/models/commandr/__init__.py -tensorrt_llm/models/commandr/config.py -tensorrt_llm/models/commandr/model.py tensorrt_llm/models/convert_utils.py -tensorrt_llm/models/dbrx/__init__.py -tensorrt_llm/models/dbrx/config.py -tensorrt_llm/models/dbrx/model.py -tensorrt_llm/models/deepseek_v1/__init__.py -tensorrt_llm/models/deepseek_v1/config.py -tensorrt_llm/models/deepseek_v1/convert.py -tensorrt_llm/models/deepseek_v1/model.py -tensorrt_llm/models/deepseek_v2/__init__.py -tensorrt_llm/models/deepseek_v2/config.py -tensorrt_llm/models/deepseek_v2/convert.py -tensorrt_llm/models/deepseek_v2/model.py -tensorrt_llm/models/dit/__init__.py -tensorrt_llm/models/dit/model.py -tensorrt_llm/models/eagle/__init__.py -tensorrt_llm/models/eagle/config.py -tensorrt_llm/models/eagle/model.py -tensorrt_llm/models/enc_dec/__init__.py -tensorrt_llm/models/enc_dec/model.py -tensorrt_llm/models/falcon/__init__.py -tensorrt_llm/models/falcon/config.py -tensorrt_llm/models/falcon/convert.py -tensorrt_llm/models/falcon/model.py -tensorrt_llm/models/gemma/__init__.py -tensorrt_llm/models/gemma/config.py -tensorrt_llm/models/gemma/convert.py -tensorrt_llm/models/gemma/model.py -tensorrt_llm/models/gemma/smoothquant.py -tensorrt_llm/models/gemma/utils/__init__.py -tensorrt_llm/models/gemma/utils/layers.py -tensorrt_llm/models/gemma/utils/modules.py -tensorrt_llm/models/gemma/utils/params.py -tensorrt_llm/models/gemma/utils/positional_embeddings.py -tensorrt_llm/models/gemma/utils/sampler.py -tensorrt_llm/models/gemma/utils/transformer.py -tensorrt_llm/models/gemma/weight.py -tensorrt_llm/models/generation_mixin.py -tensorrt_llm/models/gpt/__init__.py -tensorrt_llm/models/gpt/config.py -tensorrt_llm/models/gpt/convert.py -tensorrt_llm/models/gpt/model.py -tensorrt_llm/models/gptj/__init__.py -tensorrt_llm/models/gptj/config.py -tensorrt_llm/models/gptj/convert.py -tensorrt_llm/models/gptj/model.py -tensorrt_llm/models/gptneox/__init__.py -tensorrt_llm/models/gptneox/model.py -tensorrt_llm/models/grok/__init__.py -tensorrt_llm/models/grok/convert.py -tensorrt_llm/models/grok/model.py -tensorrt_llm/models/grok/weight.py -tensorrt_llm/models/llama/__init__.py -tensorrt_llm/models/llama/config.py -tensorrt_llm/models/llama/convert.py -tensorrt_llm/models/llama/model.py -tensorrt_llm/models/mamba/__init__.py -tensorrt_llm/models/mamba/config.py -tensorrt_llm/models/mamba/convert.py -tensorrt_llm/models/mamba/model.py -tensorrt_llm/models/medusa/__init__.py -tensorrt_llm/models/medusa/config.py -tensorrt_llm/models/medusa/model.py -tensorrt_llm/models/medusa/weight.py -tensorrt_llm/models/mllama/__init__.py -tensorrt_llm/models/mllama/config.py -tensorrt_llm/models/mllama/model.py -tensorrt_llm/models/mmdit_sd3/__init__.py -tensorrt_llm/models/mmdit_sd3/config.py -tensorrt_llm/models/mmdit_sd3/model.py -tensorrt_llm/models/model_weights_loader.py tensorrt_llm/models/modeling_utils.py -tensorrt_llm/models/mpt/__init__.py -tensorrt_llm/models/mpt/model.py -tensorrt_llm/models/multimodal_encoders/__init__.py -tensorrt_llm/models/multimodal_encoders/config.py -tensorrt_llm/models/multimodal_encoders/model.py -tensorrt_llm/models/nemotron_nas/__init__.py -tensorrt_llm/models/nemotron_nas/config.py -tensorrt_llm/models/nemotron_nas/convert.py -tensorrt_llm/models/nemotron_nas/layer_config.py -tensorrt_llm/models/nemotron_nas/model.py -tensorrt_llm/models/opt/__init__.py -tensorrt_llm/models/opt/model.py -tensorrt_llm/models/phi/__init__.py -tensorrt_llm/models/phi/config.py -tensorrt_llm/models/phi/convert.py -tensorrt_llm/models/phi/model.py -tensorrt_llm/models/phi3/__init__.py -tensorrt_llm/models/phi3/config.py -tensorrt_llm/models/phi3/convert.py -tensorrt_llm/models/phi3/model.py -tensorrt_llm/models/phi3/split_weights.py -tensorrt_llm/models/qwen/__init__.py -tensorrt_llm/models/qwen/config.py -tensorrt_llm/models/qwen/convert.py -tensorrt_llm/models/qwen/model.py -tensorrt_llm/models/qwen/utils.py -tensorrt_llm/models/recurrentgemma/__init__.py -tensorrt_llm/models/recurrentgemma/model.py -tensorrt_llm/models/redrafter/__init__.py -tensorrt_llm/models/redrafter/drafter.py -tensorrt_llm/models/redrafter/model.py -tensorrt_llm/models/redrafter/redrafter_helper.py -tensorrt_llm/models/stdit/__init__.py -tensorrt_llm/models/stdit/config.py -tensorrt_llm/models/stdit/model.py -tensorrt_llm/models/unet/__init__.py -tensorrt_llm/models/unet/attention.py -tensorrt_llm/models/unet/embeddings.py -tensorrt_llm/models/unet/pp/__init__.py -tensorrt_llm/models/unet/pp/attention.py -tensorrt_llm/models/unet/pp/conv2d.py -tensorrt_llm/models/unet/pp/groupnorm.py -tensorrt_llm/models/unet/pp/unet_pp.py -tensorrt_llm/models/unet/resnet.py -tensorrt_llm/models/unet/unet_2d_blocks.py -tensorrt_llm/models/unet/unet_2d_condition.py -tensorrt_llm/models/unet/weights.py -tensorrt_llm/network.py -tensorrt_llm/parameter.py -tensorrt_llm/plugin/__init__.py -tensorrt_llm/plugin/plugin.py tensorrt_llm/quantization/__init__.py tensorrt_llm/quantization/functional.py -tensorrt_llm/quantization/image_processing.py -tensorrt_llm/quantization/layers.py tensorrt_llm/quantization/mode.py -tensorrt_llm/quantization/quantize.py -tensorrt_llm/quantization/quantize_by_modelopt.py tensorrt_llm/quantization/utils/__init__.py tensorrt_llm/quantization/utils/fp4_utils.py tensorrt_llm/quantization/utils/fp8_utils.py tensorrt_llm/ray_stub.py tensorrt_llm/runtime/__init__.py -tensorrt_llm/runtime/enc_dec_model_runner.py -tensorrt_llm/runtime/generation.py -tensorrt_llm/runtime/kv_cache_manager.py -tensorrt_llm/runtime/medusa_utils.py tensorrt_llm/runtime/memory_pools/__init__.py -tensorrt_llm/runtime/memory_pools/memory_pools_allocator.py -tensorrt_llm/runtime/memory_pools/pool.py -tensorrt_llm/runtime/memory_pools/pools_kv_cache_manager.py -tensorrt_llm/runtime/model_runner.py -tensorrt_llm/runtime/model_runner_cpp.py -tensorrt_llm/runtime/multimodal_model_runner.py -tensorrt_llm/runtime/processor_wrapper/__init__.py -tensorrt_llm/runtime/processor_wrapper/mllama_processor_wrapper.py -tensorrt_llm/runtime/processor_wrapper/processor_wrapper.py -tensorrt_llm/runtime/redrafter_utils.py -tensorrt_llm/runtime/session.py tensorrt_llm/scaffolding/__init__.py tensorrt_llm/scaffolding/benchmark.py tensorrt_llm/scaffolding/contrib/AsyncGeneration/stream_generation.py @@ -790,12 +471,7 @@ tensorrt_llm/serve/tool_parser/utils.py tensorrt_llm/tokenizer/tokenizer.py tensorrt_llm/tools/__init__.py tensorrt_llm/tools/importlib_utils.py -tensorrt_llm/tools/multimodal_builder.py -tensorrt_llm/tools/onnx_utils.py tensorrt_llm/tools/plugin_gen/__init__.py -tensorrt_llm/tools/plugin_gen/core.py -tensorrt_llm/tools/plugin_gen/plugin_gen.py -tensorrt_llm/tools/plugin_gen/shape_infer.py tensorrt_llm/tools/ppl.py tensorrt_llm/tools/profiler/nsys_profile_tools/gputrc2graph.py tensorrt_llm/version.py @@ -804,9 +480,7 @@ tests/integration/defs/accuracy/__init__.py tests/integration/defs/accuracy/accuracy_core.py tests/integration/defs/accuracy/scripts/collect_evaluated_accuracies.py tests/integration/defs/accuracy/scripts/compute_theta_and_thresholds.py -tests/integration/defs/accuracy/test_cli_flow.py tests/integration/defs/accuracy/test_disaggregated_serving.py -tests/integration/defs/accuracy/test_llm_api.py tests/integration/defs/accuracy/test_llm_api_autodeploy.py tests/integration/defs/accuracy/test_llm_api_pytorch.py tests/integration/defs/accuracy/test_llm_api_pytorch_ray.py @@ -822,53 +496,22 @@ tests/integration/defs/disaggregated/test_disaggregated.py tests/integration/defs/disaggregated/test_disaggregated_etcd.py tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py tests/integration/defs/disaggregated/test_workers.py -tests/integration/defs/examples/run_llm_fp8_quant_llama_70b.py tests/integration/defs/examples/run_llm_quickstart_atexit.py tests/integration/defs/examples/serve/test_serve.py tests/integration/defs/examples/serve/test_serve_negative.py tests/integration/defs/examples/test_ad_guided_decoding.py -tests/integration/defs/examples/test_bert.py -tests/integration/defs/examples/test_bindings.py -tests/integration/defs/examples/test_chatglm.py -tests/integration/defs/examples/test_commandr.py -tests/integration/defs/examples/test_draft_target_model.py -tests/integration/defs/examples/test_eagle.py -tests/integration/defs/examples/test_enc_dec.py -tests/integration/defs/examples/test_exaone.py -tests/integration/defs/examples/test_gemma.py tests/integration/defs/examples/test_gpt.py -tests/integration/defs/examples/test_gptj.py -tests/integration/defs/examples/test_granite.py -tests/integration/defs/examples/test_internlm.py -tests/integration/defs/examples/test_llama.py tests/integration/defs/examples/test_llm_api_with_mpi.py -tests/integration/defs/examples/test_mamba.py -tests/integration/defs/examples/test_medusa.py -tests/integration/defs/examples/test_mistral.py -tests/integration/defs/examples/test_mixtral.py -tests/integration/defs/examples/test_multimodal.py -tests/integration/defs/examples/test_nemotron.py -tests/integration/defs/examples/test_nemotron_nas.py -tests/integration/defs/examples/test_ngram.py -tests/integration/defs/examples/test_openai.py tests/integration/defs/examples/test_phi.py -tests/integration/defs/examples/test_qwen.py -tests/integration/defs/examples/test_qwen2audio.py -tests/integration/defs/examples/test_qwenvl.py tests/integration/defs/examples/test_ray.py -tests/integration/defs/examples/test_recurrentgemma.py -tests/integration/defs/examples/test_redrafter.py -tests/integration/defs/examples/test_whisper.py tests/integration/defs/llmapi/__init__.py tests/integration/defs/llmapi/_run_llmapi_llm.py tests/integration/defs/llmapi/test_llm_api_connector.py tests/integration/defs/llmapi/test_llm_api_qa.py -tests/integration/defs/llmapi/test_llm_e2e.py tests/integration/defs/llmapi/test_llm_examples.py tests/integration/defs/local_venv.py tests/integration/defs/perf/__init__.py tests/integration/defs/perf/allowed_configs.py -tests/integration/defs/perf/build.py tests/integration/defs/perf/create_perf_comparison_report.py tests/integration/defs/perf/data.py tests/integration/defs/perf/data_export.py @@ -892,34 +535,18 @@ tests/integration/defs/test_list_validation.py tests/integration/defs/test_sanity.py tests/integration/defs/test_unittests.py tests/integration/defs/triton_server/__init__.py -tests/integration/defs/triton_server/build_engines.py tests/integration/defs/triton_server/common.py tests/integration/defs/triton_server/conftest.py -tests/integration/defs/triton_server/local_venv.py -tests/integration/defs/triton_server/rcca/bug_4323566/inflight_batcher_llm_client_with_end_id.py -tests/integration/defs/triton_server/runner_interface.py tests/integration/defs/triton_server/test_list_parser.py -tests/integration/defs/triton_server/test_triton.py -tests/integration/defs/triton_server/test_triton_llm.py -tests/integration/defs/triton_server/test_triton_memleak.py -tests/integration/defs/triton_server/test_triton_multi_node.py -tests/integration/defs/triton_server/test_triton_rcca.py tests/integration/defs/triton_server/trt_test_alternative.py tests/integration/defs/trt_test_alternative.py tests/integration/defs/utils/__init__.py tests/integration/defs/utils/periodic_junit.py tests/integration/defs/utils/timeout_manager.py tests/microbenchmarks/all_reduce.py -tests/microbenchmarks/build_time_benchmark.py -tests/microbenchmarks/build_time_dashboard.py tests/scripts/allreduce_perf/allreduce_heuristic_code_gen.py tests/scripts/allreduce_perf/allreduce_perf_viz.py tests/scripts/iteration_log_parser.py -tests/scripts/perf-sanity/parse_benchmark_results.py -tests/scripts/perf-sanity/run_benchmark_serve.py -tests/unittest/_torch/attention/sparse/test_dsa_indexer.py -tests/unittest/_torch/attention/sparse/test_flash_mla.py -tests/unittest/_torch/attention/sparse/test_rocketkv.py tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py tests/unittest/_torch/attention/test_attention.py tests/unittest/_torch/attention/test_attention_mla.py @@ -940,7 +567,6 @@ tests/unittest/_torch/misc/test_share_tensor.py tests/unittest/_torch/misc/test_virtual_memory.py tests/unittest/_torch/modeling/test_modeling_bert.py tests/unittest/_torch/modeling/test_modeling_clip.py -tests/unittest/_torch/modeling/test_modeling_exaone4.py tests/unittest/_torch/modeling/test_modeling_gemma3.py tests/unittest/_torch/modeling/test_modeling_gpt_oss.py tests/unittest/_torch/modeling/test_modeling_llama.py @@ -964,8 +590,6 @@ tests/unittest/_torch/modules/test_moe_load_balancer.py tests/unittest/_torch/modules/test_moe_routing.py tests/unittest/_torch/modules/test_rotary_embedding.py tests/unittest/_torch/modules/test_triton_linear.py -tests/unittest/_torch/modules/tests_lora_modules/test_lora_attention_pytorch_flow_vs_trt.py -tests/unittest/_torch/modules/tests_lora_modules/test_lora_plugin_vs_lora_op.py tests/unittest/_torch/multi_gpu/test_allreduce.py tests/unittest/_torch/multi_gpu/test_alltoall.py tests/unittest/_torch/multi_gpu/test_ar_residual_norm.py @@ -992,24 +616,10 @@ tests/unittest/_torch/ray_orchestrator/single_gpu/test_cache_transceiver_comm.py tests/unittest/_torch/sampler/test_beam_search.py tests/unittest/_torch/sampler/test_best_of_n.py tests/unittest/_torch/sampler/test_trtllm_sampler.py -tests/unittest/_torch/speculative/test_draft_target.py -tests/unittest/_torch/speculative/test_draft_token_tree_sampling.py -tests/unittest/_torch/speculative/test_draft_token_tree_verification.py -tests/unittest/_torch/speculative/test_dynamic_spec_decode.py tests/unittest/_torch/speculative/test_eagle3.py -tests/unittest/_torch/speculative/test_kv_cache_reuse.py -tests/unittest/_torch/speculative/test_mtp.py -tests/unittest/_torch/speculative/test_ngram.py -tests/unittest/_torch/speculative/test_save_state.py -tests/unittest/_torch/speculative/test_spec_gate.py -tests/unittest/_torch/speculative/test_torch_rejection_sampling.py -tests/unittest/_torch/speculative/test_user_provided.py tests/unittest/_torch/test_connector.py tests/unittest/_torch/test_torch_multi_arange.py tests/unittest/_torch/thop/parallel/deep_gemm_tests.py -tests/unittest/_torch/thop/parallel/test_causal_conv1d_op.py -tests/unittest/_torch/thop/parallel/test_cublas_mm.py -tests/unittest/_torch/thop/parallel/test_custom_ops.py tests/unittest/_torch/thop/parallel/test_dsv3_fused_a_gemm.py tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py tests/unittest/_torch/thop/parallel/test_finegrained_mixed_dtype_gemm.py @@ -1023,11 +633,6 @@ tests/unittest/_torch/thop/parallel/test_fp8_linear.py tests/unittest/_torch/thop/parallel/test_fp8_per_tensor_scale_tllmg_gemm.py tests/unittest/_torch/thop/parallel/test_fp8_quantize.py tests/unittest/_torch/thop/parallel/test_fp8_rowwise_linear.py -tests/unittest/_torch/thop/parallel/test_fused_qk_norm_rope.py -tests/unittest/_torch/thop/parallel/test_logits_bitmask_op.py -tests/unittest/_torch/thop/parallel/test_mamba2_chunk_ss_update.py -tests/unittest/_torch/thop/parallel/test_mamba_conv1d_op.py -tests/unittest/_torch/thop/parallel/test_noaux_tc.py tests/unittest/_torch/thop/parallel/test_scaled_mm.py tests/unittest/_torch/thop/parallel/test_selective_scan_op.py tests/unittest/_torch/thop/parallel/test_tinygemm2.py @@ -1071,12 +676,10 @@ tests/unittest/llmapi/apps/_test_openai_chat_guided_decoding.py tests/unittest/llmapi/apps/_test_openai_chat_harmony.py tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py tests/unittest/llmapi/apps/_test_openai_completions.py -tests/unittest/llmapi/apps/_test_openai_consistent_chat.py tests/unittest/llmapi/apps/_test_openai_lora.py tests/unittest/llmapi/apps/_test_openai_metrics.py tests/unittest/llmapi/apps/_test_openai_misc.py tests/unittest/llmapi/apps/_test_openai_mmencoder.py -tests/unittest/llmapi/apps/_test_openai_multi_chat.py tests/unittest/llmapi/apps/_test_openai_multi_gpu.py tests/unittest/llmapi/apps/_test_openai_multi_nodes.py tests/unittest/llmapi/apps/_test_openai_perf_metrics.py @@ -1099,15 +702,12 @@ tests/unittest/llmapi/run_llm.py tests/unittest/llmapi/run_llm_exit.py tests/unittest/llmapi/run_llm_with_postproc.py tests/unittest/llmapi/test_additional_model_outputs.py -tests/unittest/llmapi/test_build_cache.py tests/unittest/llmapi/test_executor.py tests/unittest/llmapi/test_gc_utils.py tests/unittest/llmapi/test_llm.py tests/unittest/llmapi/test_llm_args.py tests/unittest/llmapi/test_llm_download.py tests/unittest/llmapi/test_llm_kv_cache_events.py -tests/unittest/llmapi/test_llm_models.py -tests/unittest/llmapi/test_llm_multi_gpu.py tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py tests/unittest/llmapi/test_llm_pytorch.py tests/unittest/llmapi/test_llm_quant.py @@ -1118,24 +718,14 @@ tests/unittest/llmapi/test_reasoning_parser.py tests/unittest/llmapi/test_serialization.py tests/unittest/llmapi/test_utils.py tests/unittest/others/__init__.py -tests/unittest/others/test_builder.py tests/unittest/others/test_convert_spec_decoding_mask_to_packed_mask.py -tests/unittest/others/test_debugging_api.py tests/unittest/others/test_exception.py tests/unittest/others/test_export.py -tests/unittest/others/test_graph_rewriter.py -tests/unittest/others/test_kv_cache_manager.py tests/unittest/others/test_kv_cache_transceiver.py tests/unittest/others/test_kv_cache_update.py -tests/unittest/others/test_layer.py tests/unittest/others/test_mapping.py -tests/unittest/others/test_model_dtype.py -tests/unittest/others/test_module.py tests/unittest/others/test_multimodal_registry.py -tests/unittest/others/test_plugins.py -tests/unittest/others/test_precision_control.py tests/unittest/others/test_pretrained_config.py -tests/unittest/others/test_session.py tests/unittest/others/test_time_breakdown.py tests/unittest/profile_utils.py tests/unittest/scaffolding/__init__.py @@ -1144,141 +734,14 @@ tests/unittest/scaffolding/test_parallel_process.py tests/unittest/scaffolding/test_scaffolding.py tests/unittest/scaffolding/test_task_collection.py tests/unittest/scaffolding/test_worker.py -tests/unittest/test_model_runner_cpp.py tests/unittest/test_pip_install.py tests/unittest/tools/__init__.py -tests/unittest/tools/plugin_gen/__init__.py -tests/unittest/tools/plugin_gen/kernel_config.py -tests/unittest/tools/plugin_gen/test_core.py -tests/unittest/tools/plugin_gen/test_plugin_gen.py -tests/unittest/tools/plugin_gen/test_shape_infer.py tests/unittest/tools/test_prepare_dataset.py tests/unittest/tools/test_test_to_stage_mapping.py -tests/unittest/trt/__init__.py -tests/unittest/trt/attention/test_bert_attention.py -tests/unittest/trt/attention/test_gpt_attention.py -tests/unittest/trt/attention/test_gpt_attention_IFB.py -tests/unittest/trt/attention/test_gpt_attention_no_cache.py -tests/unittest/trt/attention/test_sage_attention.py -tests/unittest/trt/functional/__init__.py -tests/unittest/trt/functional/test_alibi.py -tests/unittest/trt/functional/test_allreduce_norm.py -tests/unittest/trt/functional/test_allreduce_prepost_residual_norm.py -tests/unittest/trt/functional/test_arange.py -tests/unittest/trt/functional/test_argmax.py -tests/unittest/trt/functional/test_assertion.py -tests/unittest/trt/functional/test_avg_pool2d.py -tests/unittest/trt/functional/test_cast.py -tests/unittest/trt/functional/test_conv2d.py -tests/unittest/trt/functional/test_conv3d.py -tests/unittest/trt/functional/test_cos.py -tests/unittest/trt/functional/test_cumsum.py -tests/unittest/trt/functional/test_dora.py -tests/unittest/trt/functional/test_einsum.py -tests/unittest/trt/functional/test_embedding_single_gpu.py -tests/unittest/trt/functional/test_exp.py -tests/unittest/trt/functional/test_expand.py -tests/unittest/trt/functional/test_flatten.py -tests/unittest/trt/functional/test_flip.py -tests/unittest/trt/functional/test_fp4_gemm.py -tests/unittest/trt/functional/test_fp4_gemm_ootb.py -tests/unittest/trt/functional/test_gather.py -tests/unittest/trt/functional/test_gather_nd.py -tests/unittest/trt/functional/test_geglu.py -tests/unittest/trt/functional/test_gelu.py -tests/unittest/trt/functional/test_gemm_swiglu.py -tests/unittest/trt/functional/test_group_norm.py -tests/unittest/trt/functional/test_identity.py -tests/unittest/trt/functional/test_index_select.py -tests/unittest/trt/functional/test_interpolate.py -tests/unittest/trt/functional/test_logsoftmax.py -tests/unittest/trt/functional/test_lora.py -tests/unittest/trt/functional/test_low_latency_gemm.py -tests/unittest/trt/functional/test_mamba_conv1d.py -tests/unittest/trt/functional/test_masked_scatter.py -tests/unittest/trt/functional/test_masked_select.py -tests/unittest/trt/functional/test_matmul.py -tests/unittest/trt/functional/test_meshgrid2d.py -tests/unittest/trt/functional/test_moe.py -tests/unittest/trt/functional/test_nccl.py -tests/unittest/trt/functional/test_nonzero.py -tests/unittest/trt/functional/test_outer.py -tests/unittest/trt/functional/test_pad.py -tests/unittest/trt/functional/test_permute.py -tests/unittest/trt/functional/test_pp_reduce_scatter.py -tests/unittest/trt/functional/test_quant.py -tests/unittest/trt/functional/test_rearrange.py -tests/unittest/trt/functional/test_repeat.py -tests/unittest/trt/functional/test_repeat_interleave.py -tests/unittest/trt/functional/test_rg_lru.py -tests/unittest/trt/functional/test_sample.py -tests/unittest/trt/functional/test_scatter.py -tests/unittest/trt/functional/test_scatter_nd.py -tests/unittest/trt/functional/test_select.py -tests/unittest/trt/functional/test_selective_scan.py -tests/unittest/trt/functional/test_sigmoid.py -tests/unittest/trt/functional/test_silu.py -tests/unittest/trt/functional/test_sin.py -tests/unittest/trt/functional/test_slice.py -tests/unittest/trt/functional/test_softplus.py -tests/unittest/trt/functional/test_split.py -tests/unittest/trt/functional/test_squeeze.py -tests/unittest/trt/functional/test_swiglu.py -tests/unittest/trt/functional/test_topk.py -tests/unittest/trt/functional/test_transpose.py -tests/unittest/trt/functional/test_unbind.py -tests/unittest/trt/functional/test_unsqueeze.py -tests/unittest/trt/functional/test_view.py -tests/unittest/trt/functional/test_where.py -tests/unittest/trt/model/__init__.py -tests/unittest/trt/model/eagle/test_decode_draft_tokens_plugin.py -tests/unittest/trt/model/eagle/test_prepare_drafter_inputs_plugin.py -tests/unittest/trt/model/eagle/test_sample_accept_draft_tokens_plugin.py -tests/unittest/trt/model/redrafter/test_beams2tree.py -tests/unittest/trt/model/redrafter/test_draft_token.py -tests/unittest/trt/model/redrafter/test_draft_token_indices.py -tests/unittest/trt/model/redrafter/test_gather_beams.py -tests/unittest/trt/model/redrafter/test_mask.py -tests/unittest/trt/model/redrafter/test_packed_position_ids.py -tests/unittest/trt/model/redrafter/test_prefix_match_indices.py -tests/unittest/trt/model/redrafter/test_prepare_input.py -tests/unittest/trt/model/redrafter/test_process_logits.py -tests/unittest/trt/model/redrafter/test_top1.py -tests/unittest/trt/model/redrafter/test_unpack_gen_data.py -tests/unittest/trt/model/redrafter/test_validate.py -tests/unittest/trt/model/test_gpt.py -tests/unittest/trt/model/test_gpt_e2e.py -tests/unittest/trt/model/test_llama.py -tests/unittest/trt/model/test_mamba.py -tests/unittest/trt/model/test_mistral.py -tests/unittest/trt/model/test_nemotron_nas.py -tests/unittest/trt/model/test_phi.py -tests/unittest/trt/model/test_unet.py -tests/unittest/trt/model_api/test_model_api_multi_gpu.py -tests/unittest/trt/model_api/test_model_level_api.py -tests/unittest/trt/model_api/test_model_quantization.py -tests/unittest/trt/python_plugin/plugin_wrapper_utils.py -tests/unittest/trt/python_plugin/test_plugin_wrapper.py -tests/unittest/trt/quantization/__init__.py -tests/unittest/trt/quantization/_utils.py -tests/unittest/trt/quantization/test_fp8_quantization.py -tests/unittest/trt/quantization/test_fp8_rowwise_gemm.py -tests/unittest/trt/quantization/test_functional.py -tests/unittest/trt/quantization/test_mode.py -tests/unittest/trt/quantization/test_moe_weight_only_quant_matmul.py -tests/unittest/trt/quantization/test_qserve_gemm.py -tests/unittest/trt/quantization/test_quant.py -tests/unittest/trt/quantization/test_quant_layer.py -tests/unittest/trt/quantization/test_smooth_quant_gemm.py -tests/unittest/trt/quantization/test_smooth_quant_layer_norm.py -tests/unittest/trt/quantization/test_smooth_quant_rms_norm.py -tests/unittest/trt/quantization/test_weight_only_groupwise_quant_matmul.py -tests/unittest/trt/quantization/test_weight_only_quant_matmul.py tests/unittest/utils/__init__.py tests/unittest/utils/cpp_paths.py tests/unittest/utils/llm_data.py tests/unittest/utils/runtime_defaults.py -tests/unittest/utils/test_medusa_utils.py tests/unittest/utils/test_prebuilt_whl_cpp_extensions.py tests/unittest/utils/test_util.py tests/unittest/utils/torch_ref.py diff --git a/pyproject.toml b/pyproject.toml index 49c5584d5680..2efd23c84d48 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,14 +58,6 @@ exclude = [ ".devcontainer/make_env.py", ".github/scripts/label_community_user.py", ".github/scripts/pr_checklist_check.py", - "benchmarks/__init__.py", - "benchmarks/prepare_dataset.py", - "benchmarks/utils/__init__.py", - "benchmarks/utils/convert_nemo_dataset.py", - "benchmarks/utils/generate_rand_loras.py", - "benchmarks/utils/prepare_real_data.py", - "benchmarks/utils/prepare_synthetic_data.py", - "benchmarks/utils/utils.py", "cpp/conanfile.py", "cpp/kernels/fmha_v2/conftest.py", "cpp/kernels/fmha_v2/fmha_test.py", @@ -97,24 +89,10 @@ exclude = [ "examples/apps/fastapi_server.py", "examples/disaggregated/clients/disagg_client.py", "examples/disaggregated/slurm/benchmark/submit.py", - "examples/dora/normalize_weights.py", - "examples/eagle/convert_checkpoint.py", - "examples/eval_long_context.py", - "examples/generate_checkpoint_config.py", - "examples/generate_xgrammar_tokenizer_info.py", - "examples/hf_lora_convert.py", "examples/infinitebench/args.py", "examples/infinitebench/compute_scores.py", "examples/infinitebench/construct_synthetic_dataset.py", "examples/infinitebench/eval_utils.py", - "examples/llm-api/_tensorrt_engine/llm_eagle2_decoding.py", - "examples/llm-api/_tensorrt_engine/llm_eagle_decoding.py", - "examples/llm-api/_tensorrt_engine/llm_inference_customize.py", - "examples/llm-api/_tensorrt_engine/llm_inference_kv_events.py", - "examples/llm-api/_tensorrt_engine/llm_lookahead_decoding.py", - "examples/llm-api/_tensorrt_engine/llm_medusa_decoding.py", - "examples/llm-api/_tensorrt_engine/llm_quantization.py", - "examples/llm-api/_tensorrt_engine/quickstart_example.py", "examples/llm-api/llm_guided_decoding.py", "examples/llm-api/llm_inference.py", "examples/llm-api/llm_inference_async.py", @@ -134,122 +112,17 @@ exclude = [ "examples/llm-api/quickstart_example.py", "examples/llm-api/quickstart_multimodal.py", "examples/llm-api/star_attention.py", - "examples/llm-eval/lm-eval-harness/lm_eval_tensorrt_llm.py", "examples/longbench/eval_longbench_v1.py", - "examples/medusa/convert_checkpoint.py", - "examples/mmlu.py", - "examples/models/contrib/baichuan/convert_checkpoint.py", - "examples/models/contrib/bloom/convert_checkpoint.py", - "examples/models/contrib/chatglm-6b/tokenization_chatglm.py", - "examples/models/contrib/chatglm2-6b/tokenization_chatglm.py", - "examples/models/contrib/chatglm3-6b-32k/tokenization_chatglm.py", - "examples/models/contrib/cogvlm/convert_checkpoint.py", - "examples/models/contrib/dbrx/convert_checkpoint.py", - "examples/models/contrib/deepseek_v1/__init__.py", - "examples/models/contrib/deepseek_v1/convert_checkpoint.py", - "examples/models/contrib/deepseek_v2/convert_checkpoint.py", - "examples/models/contrib/dit/convert_checkpoint.py", - "examples/models/contrib/dit/diffusion.py", - "examples/models/contrib/dit/sample.py", - "examples/models/contrib/dit/utils_modelopt.py", - "examples/models/contrib/dit/vae_decoder_trt.py", - "examples/models/contrib/falcon/convert_checkpoint.py", - "examples/models/contrib/gptj/convert_checkpoint.py", - "examples/models/contrib/gptneox/convert_checkpoint.py", - "examples/models/contrib/grok/convert_checkpoint.py", - "examples/models/contrib/mmdit/convert_checkpoint.py", - "examples/models/contrib/mmdit/sample.py", - "examples/models/contrib/mpt/convert_checkpoint.py", - "examples/models/contrib/opt/convert_checkpoint.py", - "examples/models/contrib/sdxl/build_sdxl_unet.py", - "examples/models/contrib/sdxl/pipeline_stable_diffusion_xl.py", - "examples/models/contrib/sdxl/run_sdxl.py", - "examples/models/contrib/stdit/aspect.py", - "examples/models/contrib/stdit/convert_checkpoint.py", - "examples/models/contrib/stdit/pipeline_tllm.py", - "examples/models/contrib/stdit/sample.py", - "examples/models/contrib/stdit/scheduler.py", - "examples/models/contrib/stdit/text_encoder.py", - "examples/models/contrib/stdit/utils.py", - "examples/models/contrib/stdit/vae.py", - "examples/models/contrib/stdit/video_transforms.py", - "examples/models/core/bert/__init__.py", - "examples/models/core/bert/convert_checkpoint.py", - "examples/models/core/bert/run.py", - "examples/models/core/bert/utils.py", - "examples/models/core/commandr/convert_checkpoint.py", - "examples/models/core/enc_dec/__init__.py", - "examples/models/core/enc_dec/convert_checkpoint.py", - "examples/models/core/enc_dec/helper.py", - "examples/models/core/enc_dec/run.py", - "examples/models/core/gemma/convert_checkpoint.py", - "examples/models/core/glm-4-9b/convert_checkpoint.py", - "examples/models/core/glm-4-9b/tokenization_chatglm.py", - "examples/models/core/gpt/convert_checkpoint.py", - "examples/models/core/gpt/merge_ptuning_tables.py", - "examples/models/core/gpt/nemo_lora_convert.py", - "examples/models/core/gpt/nemo_prompt_convert.py", - "examples/models/core/gpt/run_hf.py", "examples/models/core/gpt_oss/openai_chat_client_function_calling.py", - "examples/models/core/internlm2/convert_checkpoint.py", "examples/models/core/kimi_k2/kimi_k2_tool_calling_example.py", - "examples/models/core/llama/convert_checkpoint.py", - "examples/models/core/llama/summarize_long.py", - "examples/models/core/mamba/convert_checkpoint.py", - "examples/models/core/mllama/convert_checkpoint.py", - "examples/models/core/multimodal/__init__.py", - "examples/models/core/multimodal/build_multimodal_engine.py", - "examples/models/core/multimodal/eval.py", - "examples/models/core/multimodal/run.py", - "examples/models/core/multimodal/utils.py", - "examples/models/core/nemotron_nas/calibration_utils.py", - "examples/models/core/nemotron_nas/convert_checkpoint.py", - "examples/models/core/phi/convert_checkpoint.py", - "examples/models/core/qwen/convert_checkpoint.py", - "examples/models/core/qwen2audio/run.py", - "examples/models/core/qwen2audio/run_chat.py", - "examples/models/core/qwen2audio/utils.py", - "examples/models/core/qwenvl/run.py", - "examples/models/core/qwenvl/run_chat.py", - "examples/models/core/qwenvl/show_pic.py", - "examples/models/core/qwenvl/vit_onnx_trt.py", - "examples/models/core/recurrentgemma/convert_checkpoint.py", - "examples/models/core/vit/convert_checkpoint.py", - "examples/models/core/whisper/convert_checkpoint.py", - "examples/models/core/whisper/distil_whisper/convert_from_distil_whisper.py", - "examples/models/core/whisper/run.py", - "examples/models/core/whisper/tokenizer.py", - "examples/models/core/whisper/whisper_utils.py", - "examples/ngram/run_dtm_ngram.py", - "examples/openai_triton/manual_plugin/build.py", - "examples/openai_triton/manual_plugin/fmha_triton.py", - "examples/openai_triton/manual_plugin/plugin.py", - "examples/openai_triton/manual_plugin/run.py", - "examples/openai_triton/plugin_autogen/build_engine.py", - "examples/openai_triton/plugin_autogen/kernel_config.py", - "examples/openai_triton/plugin_autogen/run_engine.py", - "examples/python_plugin/build_lookup.py", - "examples/python_plugin/plugin_lib/__init__.py", - "examples/python_plugin/plugin_lib/lookup_kernel.py", - "examples/python_plugin/plugin_lib/lookup_plugin.py", - "examples/python_plugin/run_lookup.py", - "examples/quantization/quantize.py", "examples/quantization/quantize_mixed_precision_moe.py", "examples/ray_orchestrator/llm_inference_async_ray.py", "examples/ray_orchestrator/llm_inference_distributed_ray.py", - "examples/redrafter/convert_checkpoint.py", - "examples/run.py", "examples/scaffolding/contrib/AsyncGeneration/stream_generation_controller.py", "examples/scaffolding/contrib/DeepConf/run_generation.py", "examples/scaffolding/contrib/Dynasor/scaffolding_dynasor_run.py", "examples/scaffolding/contrib/TreeInference/run_mcts_example.py", "examples/scaffolding/contrib/TreeInference/run_tot_example.py", - "examples/scaffolding/contrib/mcp/e2b/e2bserver.py", - "examples/scaffolding/contrib/mcp/e2b/main.py", - "examples/scaffolding/contrib/mcp/mcptest.py", - "examples/scaffolding/contrib/mcp/weather/weather.py", - "examples/scaffolding/contrib/mcp/websearch/main.py", - "examples/scaffolding/contrib/mcp/websearch/websearch.py", "examples/scaffolding/run_basic_generation.py", "examples/scaffolding/run_best_of_n_with_reward.py", "examples/scaffolding/run_majority_vote_aime24.py", @@ -259,8 +132,6 @@ exclude = [ "examples/serve/openai_completion_client.py", "examples/serve/openai_completion_client_for_lora.py", "examples/serve/openai_completion_client_json_schema.py", - "examples/summarize.py", - "examples/utils.py", "examples/wide_ep/ep_load_balancer/generate_eplb_config.py", "examples/wide_ep/ep_load_balancer/report_load_statistics.py", "examples/wide_ep/ep_load_balancer/utils.py", @@ -272,7 +143,6 @@ exclude = [ "scripts/check_test_list.py", "scripts/dco_check.py", "scripts/format_test_list.py", - "scripts/generate_duration.py", "scripts/generate_lock_file.py", "scripts/get_wheel_from_package.py", "scripts/git_replace.py", @@ -283,7 +153,6 @@ exclude = [ "setup.py", "tensorrt_llm/__init__.py", "tensorrt_llm/_ray_utils.py", - "tensorrt_llm/_tensorrt_engine/__init__.py", "tensorrt_llm/_torch/__init__.py", "tensorrt_llm/_torch/attention_backend/__init__.py", "tensorrt_llm/_torch/attention_backend/flashinfer.py", @@ -487,7 +356,6 @@ exclude = [ "tensorrt_llm/_torch/pyexecutor/guided_decoder.py", "tensorrt_llm/_torch/pyexecutor/handle_additional_outputs.py", "tensorrt_llm/_torch/pyexecutor/handle_logits.py", - "tensorrt_llm/_torch/pyexecutor/kv_cache_connector.py", "tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py", "tensorrt_llm/_torch/pyexecutor/layerwise_nvtx_marker.py", "tensorrt_llm/_torch/pyexecutor/llm_request.py", @@ -525,11 +393,6 @@ exclude = [ "tensorrt_llm/bench/benchmark/utils/asynchronous.py", "tensorrt_llm/bench/benchmark/utils/general.py", "tensorrt_llm/bench/benchmark/utils/processes.py", - "tensorrt_llm/bench/build/__init__.py", - "tensorrt_llm/bench/build/build.py", - "tensorrt_llm/bench/build/dataclasses.py", - "tensorrt_llm/bench/build/tuning.py", - "tensorrt_llm/bench/build/utils.py", "tensorrt_llm/bench/dataclasses/__init__.py", "tensorrt_llm/bench/dataclasses/configuration.py", "tensorrt_llm/bench/dataclasses/engine.py", @@ -539,13 +402,9 @@ exclude = [ "tensorrt_llm/bench/dataclasses/statistics.py", "tensorrt_llm/bench/utils/__init__.py", "tensorrt_llm/bench/utils/data.py", - "tensorrt_llm/builder.py", "tensorrt_llm/commands/__init__.py", "tensorrt_llm/commands/bench.py", - "tensorrt_llm/commands/build.py", "tensorrt_llm/commands/eval.py", - "tensorrt_llm/commands/prune.py", - "tensorrt_llm/commands/refit.py", "tensorrt_llm/commands/serve.py", "tensorrt_llm/evaluate/__init__.py", "tensorrt_llm/evaluate/cnn_dailymail.py", @@ -581,23 +440,7 @@ exclude = [ "tensorrt_llm/inputs/multimodal.py", "tensorrt_llm/inputs/registry.py", "tensorrt_llm/inputs/utils.py", - "tensorrt_llm/layers/__init__.py", - "tensorrt_llm/layers/activation.py", - "tensorrt_llm/layers/attention.py", - "tensorrt_llm/layers/cast.py", - "tensorrt_llm/layers/conv.py", - "tensorrt_llm/layers/embedding.py", - "tensorrt_llm/layers/language_adapter.py", - "tensorrt_llm/layers/linear.py", - "tensorrt_llm/layers/lora.py", - "tensorrt_llm/layers/mlp.py", - "tensorrt_llm/layers/moe.py", - "tensorrt_llm/layers/normalization.py", - "tensorrt_llm/layers/pooling.py", - "tensorrt_llm/layers/recurrent.py", - "tensorrt_llm/layers/ssm.py", "tensorrt_llm/llmapi/__init__.py", - "tensorrt_llm/llmapi/build_cache.py", "tensorrt_llm/llmapi/disagg_utils.py", "tensorrt_llm/llmapi/kv_cache_type.py", "tensorrt_llm/llmapi/llm.py", @@ -620,179 +463,17 @@ exclude = [ "tensorrt_llm/metrics/enums.py", "tensorrt_llm/models/__init__.py", "tensorrt_llm/models/automodel.py", - "tensorrt_llm/models/baichuan/__init__.py", - "tensorrt_llm/models/baichuan/config.py", - "tensorrt_llm/models/baichuan/convert.py", - "tensorrt_llm/models/baichuan/model.py", - "tensorrt_llm/models/bert/__init__.py", - "tensorrt_llm/models/bert/config.py", - "tensorrt_llm/models/bert/convert.py", - "tensorrt_llm/models/bert/model.py", - "tensorrt_llm/models/bloom/__init__.py", - "tensorrt_llm/models/bloom/model.py", - "tensorrt_llm/models/chatglm/__init__.py", - "tensorrt_llm/models/chatglm/config.py", - "tensorrt_llm/models/chatglm/convert.py", - "tensorrt_llm/models/chatglm/model.py", - "tensorrt_llm/models/clip/__init__.py", - "tensorrt_llm/models/clip/model.py", - "tensorrt_llm/models/cogvlm/__init__.py", - "tensorrt_llm/models/cogvlm/config.py", - "tensorrt_llm/models/cogvlm/convert.py", - "tensorrt_llm/models/cogvlm/model.py", - "tensorrt_llm/models/commandr/__init__.py", - "tensorrt_llm/models/commandr/config.py", - "tensorrt_llm/models/commandr/model.py", "tensorrt_llm/models/convert_utils.py", - "tensorrt_llm/models/dbrx/__init__.py", - "tensorrt_llm/models/dbrx/config.py", - "tensorrt_llm/models/dbrx/model.py", - "tensorrt_llm/models/deepseek_v1/__init__.py", - "tensorrt_llm/models/deepseek_v1/config.py", - "tensorrt_llm/models/deepseek_v1/convert.py", - "tensorrt_llm/models/deepseek_v1/model.py", - "tensorrt_llm/models/deepseek_v2/__init__.py", - "tensorrt_llm/models/deepseek_v2/config.py", - "tensorrt_llm/models/deepseek_v2/convert.py", - "tensorrt_llm/models/deepseek_v2/model.py", - "tensorrt_llm/models/dit/__init__.py", - "tensorrt_llm/models/dit/model.py", - "tensorrt_llm/models/eagle/__init__.py", - "tensorrt_llm/models/eagle/config.py", - "tensorrt_llm/models/eagle/model.py", - "tensorrt_llm/models/enc_dec/__init__.py", - "tensorrt_llm/models/enc_dec/model.py", - "tensorrt_llm/models/falcon/__init__.py", - "tensorrt_llm/models/falcon/config.py", - "tensorrt_llm/models/falcon/convert.py", - "tensorrt_llm/models/falcon/model.py", - "tensorrt_llm/models/gemma/__init__.py", - "tensorrt_llm/models/gemma/config.py", - "tensorrt_llm/models/gemma/convert.py", - "tensorrt_llm/models/gemma/model.py", - "tensorrt_llm/models/gemma/smoothquant.py", - "tensorrt_llm/models/gemma/utils/__init__.py", - "tensorrt_llm/models/gemma/utils/layers.py", - "tensorrt_llm/models/gemma/utils/modules.py", - "tensorrt_llm/models/gemma/utils/params.py", - "tensorrt_llm/models/gemma/utils/positional_embeddings.py", - "tensorrt_llm/models/gemma/utils/sampler.py", - "tensorrt_llm/models/gemma/utils/transformer.py", - "tensorrt_llm/models/gemma/weight.py", - "tensorrt_llm/models/generation_mixin.py", - "tensorrt_llm/models/gpt/__init__.py", - "tensorrt_llm/models/gpt/config.py", - "tensorrt_llm/models/gpt/convert.py", - "tensorrt_llm/models/gpt/model.py", - "tensorrt_llm/models/gptj/__init__.py", - "tensorrt_llm/models/gptj/config.py", - "tensorrt_llm/models/gptj/convert.py", - "tensorrt_llm/models/gptj/model.py", - "tensorrt_llm/models/gptneox/__init__.py", - "tensorrt_llm/models/gptneox/model.py", - "tensorrt_llm/models/grok/__init__.py", - "tensorrt_llm/models/grok/convert.py", - "tensorrt_llm/models/grok/model.py", - "tensorrt_llm/models/grok/weight.py", - "tensorrt_llm/models/llama/__init__.py", - "tensorrt_llm/models/llama/config.py", - "tensorrt_llm/models/llama/convert.py", - "tensorrt_llm/models/llama/model.py", - "tensorrt_llm/models/mamba/__init__.py", - "tensorrt_llm/models/mamba/config.py", - "tensorrt_llm/models/mamba/convert.py", - "tensorrt_llm/models/mamba/model.py", - "tensorrt_llm/models/medusa/__init__.py", - "tensorrt_llm/models/medusa/config.py", - "tensorrt_llm/models/medusa/model.py", - "tensorrt_llm/models/medusa/weight.py", - "tensorrt_llm/models/mllama/__init__.py", - "tensorrt_llm/models/mllama/config.py", - "tensorrt_llm/models/mllama/model.py", - "tensorrt_llm/models/mmdit_sd3/__init__.py", - "tensorrt_llm/models/mmdit_sd3/config.py", - "tensorrt_llm/models/mmdit_sd3/model.py", - "tensorrt_llm/models/model_weights_loader.py", "tensorrt_llm/models/modeling_utils.py", - "tensorrt_llm/models/mpt/__init__.py", - "tensorrt_llm/models/mpt/model.py", - "tensorrt_llm/models/multimodal_encoders/__init__.py", - "tensorrt_llm/models/multimodal_encoders/config.py", - "tensorrt_llm/models/multimodal_encoders/model.py", - "tensorrt_llm/models/nemotron_nas/__init__.py", - "tensorrt_llm/models/nemotron_nas/config.py", - "tensorrt_llm/models/nemotron_nas/convert.py", - "tensorrt_llm/models/nemotron_nas/layer_config.py", - "tensorrt_llm/models/nemotron_nas/model.py", - "tensorrt_llm/models/opt/__init__.py", - "tensorrt_llm/models/opt/model.py", - "tensorrt_llm/models/phi/__init__.py", - "tensorrt_llm/models/phi/config.py", - "tensorrt_llm/models/phi/convert.py", - "tensorrt_llm/models/phi/model.py", - "tensorrt_llm/models/phi3/__init__.py", - "tensorrt_llm/models/phi3/config.py", - "tensorrt_llm/models/phi3/convert.py", - "tensorrt_llm/models/phi3/model.py", - "tensorrt_llm/models/phi3/split_weights.py", - "tensorrt_llm/models/qwen/__init__.py", - "tensorrt_llm/models/qwen/config.py", - "tensorrt_llm/models/qwen/convert.py", - "tensorrt_llm/models/qwen/model.py", - "tensorrt_llm/models/qwen/utils.py", - "tensorrt_llm/models/recurrentgemma/__init__.py", - "tensorrt_llm/models/recurrentgemma/model.py", - "tensorrt_llm/models/redrafter/__init__.py", - "tensorrt_llm/models/redrafter/drafter.py", - "tensorrt_llm/models/redrafter/model.py", - "tensorrt_llm/models/redrafter/redrafter_helper.py", - "tensorrt_llm/models/stdit/__init__.py", - "tensorrt_llm/models/stdit/config.py", - "tensorrt_llm/models/stdit/model.py", - "tensorrt_llm/models/unet/__init__.py", - "tensorrt_llm/models/unet/attention.py", - "tensorrt_llm/models/unet/embeddings.py", - "tensorrt_llm/models/unet/pp/__init__.py", - "tensorrt_llm/models/unet/pp/attention.py", - "tensorrt_llm/models/unet/pp/conv2d.py", - "tensorrt_llm/models/unet/pp/groupnorm.py", - "tensorrt_llm/models/unet/pp/unet_pp.py", - "tensorrt_llm/models/unet/resnet.py", - "tensorrt_llm/models/unet/unet_2d_blocks.py", - "tensorrt_llm/models/unet/unet_2d_condition.py", - "tensorrt_llm/models/unet/weights.py", - "tensorrt_llm/network.py", - "tensorrt_llm/parameter.py", - "tensorrt_llm/plugin/__init__.py", - "tensorrt_llm/plugin/plugin.py", "tensorrt_llm/quantization/__init__.py", "tensorrt_llm/quantization/functional.py", - "tensorrt_llm/quantization/image_processing.py", - "tensorrt_llm/quantization/layers.py", "tensorrt_llm/quantization/mode.py", - "tensorrt_llm/quantization/quantize.py", - "tensorrt_llm/quantization/quantize_by_modelopt.py", "tensorrt_llm/quantization/utils/__init__.py", "tensorrt_llm/quantization/utils/fp4_utils.py", "tensorrt_llm/quantization/utils/fp8_utils.py", "tensorrt_llm/ray_stub.py", "tensorrt_llm/runtime/__init__.py", - "tensorrt_llm/runtime/enc_dec_model_runner.py", - "tensorrt_llm/runtime/generation.py", - "tensorrt_llm/runtime/kv_cache_manager.py", - "tensorrt_llm/runtime/medusa_utils.py", "tensorrt_llm/runtime/memory_pools/__init__.py", - "tensorrt_llm/runtime/memory_pools/memory_pools_allocator.py", - "tensorrt_llm/runtime/memory_pools/pool.py", - "tensorrt_llm/runtime/memory_pools/pools_kv_cache_manager.py", - "tensorrt_llm/runtime/model_runner.py", - "tensorrt_llm/runtime/model_runner_cpp.py", - "tensorrt_llm/runtime/multimodal_model_runner.py", - "tensorrt_llm/runtime/processor_wrapper/__init__.py", - "tensorrt_llm/runtime/processor_wrapper/mllama_processor_wrapper.py", - "tensorrt_llm/runtime/processor_wrapper/processor_wrapper.py", - "tensorrt_llm/runtime/redrafter_utils.py", - "tensorrt_llm/runtime/session.py", "tensorrt_llm/scaffolding/__init__.py", "tensorrt_llm/scaffolding/benchmark.py", "tensorrt_llm/scaffolding/contrib/AsyncGeneration/stream_generation.py", @@ -847,12 +528,7 @@ exclude = [ "tensorrt_llm/tokenizer/tokenizer.py", "tensorrt_llm/tools/__init__.py", "tensorrt_llm/tools/importlib_utils.py", - "tensorrt_llm/tools/multimodal_builder.py", - "tensorrt_llm/tools/onnx_utils.py", "tensorrt_llm/tools/plugin_gen/__init__.py", - "tensorrt_llm/tools/plugin_gen/core.py", - "tensorrt_llm/tools/plugin_gen/plugin_gen.py", - "tensorrt_llm/tools/plugin_gen/shape_infer.py", "tensorrt_llm/tools/ppl.py", "tensorrt_llm/tools/profiler/nsys_profile_tools/gputrc2graph.py", "tensorrt_llm/version.py", @@ -861,9 +537,7 @@ exclude = [ "tests/integration/defs/accuracy/accuracy_core.py", "tests/integration/defs/accuracy/scripts/collect_evaluated_accuracies.py", "tests/integration/defs/accuracy/scripts/compute_theta_and_thresholds.py", - "tests/integration/defs/accuracy/test_cli_flow.py", "tests/integration/defs/accuracy/test_disaggregated_serving.py", - "tests/integration/defs/accuracy/test_llm_api.py", "tests/integration/defs/accuracy/test_llm_api_autodeploy.py", "tests/integration/defs/accuracy/test_llm_api_pytorch.py", "tests/integration/defs/accuracy/test_llm_api_pytorch_ray.py", @@ -879,53 +553,22 @@ exclude = [ "tests/integration/defs/disaggregated/test_disaggregated_etcd.py", "tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py", "tests/integration/defs/disaggregated/test_workers.py", - "tests/integration/defs/examples/run_llm_fp8_quant_llama_70b.py", "tests/integration/defs/examples/run_llm_quickstart_atexit.py", "tests/integration/defs/examples/serve/test_serve.py", "tests/integration/defs/examples/serve/test_serve_negative.py", "tests/integration/defs/examples/test_ad_guided_decoding.py", - "tests/integration/defs/examples/test_bert.py", - "tests/integration/defs/examples/test_bindings.py", - "tests/integration/defs/examples/test_chatglm.py", - "tests/integration/defs/examples/test_commandr.py", - "tests/integration/defs/examples/test_draft_target_model.py", - "tests/integration/defs/examples/test_eagle.py", - "tests/integration/defs/examples/test_enc_dec.py", - "tests/integration/defs/examples/test_exaone.py", - "tests/integration/defs/examples/test_gemma.py", "tests/integration/defs/examples/test_gpt.py", - "tests/integration/defs/examples/test_gptj.py", - "tests/integration/defs/examples/test_granite.py", - "tests/integration/defs/examples/test_internlm.py", - "tests/integration/defs/examples/test_llama.py", "tests/integration/defs/examples/test_llm_api_with_mpi.py", - "tests/integration/defs/examples/test_mamba.py", - "tests/integration/defs/examples/test_medusa.py", - "tests/integration/defs/examples/test_mistral.py", - "tests/integration/defs/examples/test_mixtral.py", - "tests/integration/defs/examples/test_multimodal.py", - "tests/integration/defs/examples/test_nemotron.py", - "tests/integration/defs/examples/test_nemotron_nas.py", - "tests/integration/defs/examples/test_ngram.py", - "tests/integration/defs/examples/test_openai.py", "tests/integration/defs/examples/test_phi.py", - "tests/integration/defs/examples/test_qwen.py", - "tests/integration/defs/examples/test_qwen2audio.py", - "tests/integration/defs/examples/test_qwenvl.py", "tests/integration/defs/examples/test_ray.py", - "tests/integration/defs/examples/test_recurrentgemma.py", - "tests/integration/defs/examples/test_redrafter.py", - "tests/integration/defs/examples/test_whisper.py", "tests/integration/defs/llmapi/__init__.py", "tests/integration/defs/llmapi/_run_llmapi_llm.py", "tests/integration/defs/llmapi/test_llm_api_connector.py", "tests/integration/defs/llmapi/test_llm_api_qa.py", - "tests/integration/defs/llmapi/test_llm_e2e.py", "tests/integration/defs/llmapi/test_llm_examples.py", "tests/integration/defs/local_venv.py", "tests/integration/defs/perf/__init__.py", "tests/integration/defs/perf/allowed_configs.py", - "tests/integration/defs/perf/build.py", "tests/integration/defs/perf/create_perf_comparison_report.py", "tests/integration/defs/perf/data.py", "tests/integration/defs/perf/data_export.py", @@ -949,34 +592,18 @@ exclude = [ "tests/integration/defs/test_sanity.py", "tests/integration/defs/test_unittests.py", "tests/integration/defs/triton_server/__init__.py", - "tests/integration/defs/triton_server/build_engines.py", "tests/integration/defs/triton_server/common.py", "tests/integration/defs/triton_server/conftest.py", - "tests/integration/defs/triton_server/local_venv.py", - "tests/integration/defs/triton_server/rcca/bug_4323566/inflight_batcher_llm_client_with_end_id.py", - "tests/integration/defs/triton_server/runner_interface.py", "tests/integration/defs/triton_server/test_list_parser.py", - "tests/integration/defs/triton_server/test_triton.py", - "tests/integration/defs/triton_server/test_triton_llm.py", - "tests/integration/defs/triton_server/test_triton_memleak.py", - "tests/integration/defs/triton_server/test_triton_multi_node.py", - "tests/integration/defs/triton_server/test_triton_rcca.py", "tests/integration/defs/triton_server/trt_test_alternative.py", "tests/integration/defs/trt_test_alternative.py", "tests/integration/defs/utils/__init__.py", "tests/integration/defs/utils/periodic_junit.py", "tests/integration/defs/utils/timeout_manager.py", "tests/microbenchmarks/all_reduce.py", - "tests/microbenchmarks/build_time_benchmark.py", - "tests/microbenchmarks/build_time_dashboard.py", "tests/scripts/allreduce_perf/allreduce_heuristic_code_gen.py", "tests/scripts/allreduce_perf/allreduce_perf_viz.py", "tests/scripts/iteration_log_parser.py", - "tests/scripts/perf-sanity/parse_benchmark_results.py", - "tests/scripts/perf-sanity/run_benchmark_serve.py", - "tests/unittest/_torch/attention/sparse/test_dsa_indexer.py", - "tests/unittest/_torch/attention/sparse/test_flash_mla.py", - "tests/unittest/_torch/attention/sparse/test_rocketkv.py", "tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py", "tests/unittest/_torch/attention/test_attention.py", "tests/unittest/_torch/attention/test_attention_mla.py", @@ -997,7 +624,6 @@ exclude = [ "tests/unittest/_torch/misc/test_virtual_memory.py", "tests/unittest/_torch/modeling/test_modeling_bert.py", "tests/unittest/_torch/modeling/test_modeling_clip.py", - "tests/unittest/_torch/modeling/test_modeling_exaone4.py", "tests/unittest/_torch/modeling/test_modeling_gemma3.py", "tests/unittest/_torch/modeling/test_modeling_gpt_oss.py", "tests/unittest/_torch/modeling/test_modeling_llama.py", @@ -1021,8 +647,6 @@ exclude = [ "tests/unittest/_torch/modules/test_moe_routing.py", "tests/unittest/_torch/modules/test_rotary_embedding.py", "tests/unittest/_torch/modules/test_triton_linear.py", - "tests/unittest/_torch/modules/tests_lora_modules/test_lora_attention_pytorch_flow_vs_trt.py", - "tests/unittest/_torch/modules/tests_lora_modules/test_lora_plugin_vs_lora_op.py", "tests/unittest/_torch/multi_gpu/test_allreduce.py", "tests/unittest/_torch/multi_gpu/test_alltoall.py", "tests/unittest/_torch/multi_gpu/test_ar_residual_norm.py", @@ -1049,24 +673,10 @@ exclude = [ "tests/unittest/_torch/sampler/test_beam_search.py", "tests/unittest/_torch/sampler/test_best_of_n.py", "tests/unittest/_torch/sampler/test_trtllm_sampler.py", - "tests/unittest/_torch/speculative/test_draft_target.py", - "tests/unittest/_torch/speculative/test_draft_token_tree_sampling.py", - "tests/unittest/_torch/speculative/test_draft_token_tree_verification.py", - "tests/unittest/_torch/speculative/test_dynamic_spec_decode.py", "tests/unittest/_torch/speculative/test_eagle3.py", - "tests/unittest/_torch/speculative/test_kv_cache_reuse.py", - "tests/unittest/_torch/speculative/test_mtp.py", - "tests/unittest/_torch/speculative/test_ngram.py", - "tests/unittest/_torch/speculative/test_save_state.py", - "tests/unittest/_torch/speculative/test_spec_gate.py", - "tests/unittest/_torch/speculative/test_torch_rejection_sampling.py", - "tests/unittest/_torch/speculative/test_user_provided.py", "tests/unittest/_torch/test_connector.py", "tests/unittest/_torch/test_torch_multi_arange.py", "tests/unittest/_torch/thop/parallel/deep_gemm_tests.py", - "tests/unittest/_torch/thop/parallel/test_causal_conv1d_op.py", - "tests/unittest/_torch/thop/parallel/test_cublas_mm.py", - "tests/unittest/_torch/thop/parallel/test_custom_ops.py", "tests/unittest/_torch/thop/parallel/test_dsv3_fused_a_gemm.py", "tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py", "tests/unittest/_torch/thop/parallel/test_finegrained_mixed_dtype_gemm.py", @@ -1080,11 +690,6 @@ exclude = [ "tests/unittest/_torch/thop/parallel/test_fp8_per_tensor_scale_tllmg_gemm.py", "tests/unittest/_torch/thop/parallel/test_fp8_quantize.py", "tests/unittest/_torch/thop/parallel/test_fp8_rowwise_linear.py", - "tests/unittest/_torch/thop/parallel/test_fused_qk_norm_rope.py", - "tests/unittest/_torch/thop/parallel/test_logits_bitmask_op.py", - "tests/unittest/_torch/thop/parallel/test_mamba2_chunk_ss_update.py", - "tests/unittest/_torch/thop/parallel/test_mamba_conv1d_op.py", - "tests/unittest/_torch/thop/parallel/test_noaux_tc.py", "tests/unittest/_torch/thop/parallel/test_scaled_mm.py", "tests/unittest/_torch/thop/parallel/test_selective_scan_op.py", "tests/unittest/_torch/thop/parallel/test_tinygemm2.py", @@ -1128,12 +733,10 @@ exclude = [ "tests/unittest/llmapi/apps/_test_openai_chat_harmony.py", "tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py", "tests/unittest/llmapi/apps/_test_openai_completions.py", - "tests/unittest/llmapi/apps/_test_openai_consistent_chat.py", "tests/unittest/llmapi/apps/_test_openai_lora.py", "tests/unittest/llmapi/apps/_test_openai_metrics.py", "tests/unittest/llmapi/apps/_test_openai_misc.py", "tests/unittest/llmapi/apps/_test_openai_mmencoder.py", - "tests/unittest/llmapi/apps/_test_openai_multi_chat.py", "tests/unittest/llmapi/apps/_test_openai_multi_gpu.py", "tests/unittest/llmapi/apps/_test_openai_multi_nodes.py", "tests/unittest/llmapi/apps/_test_openai_perf_metrics.py", @@ -1156,15 +759,12 @@ exclude = [ "tests/unittest/llmapi/run_llm_exit.py", "tests/unittest/llmapi/run_llm_with_postproc.py", "tests/unittest/llmapi/test_additional_model_outputs.py", - "tests/unittest/llmapi/test_build_cache.py", "tests/unittest/llmapi/test_executor.py", "tests/unittest/llmapi/test_gc_utils.py", "tests/unittest/llmapi/test_llm.py", "tests/unittest/llmapi/test_llm_args.py", "tests/unittest/llmapi/test_llm_download.py", "tests/unittest/llmapi/test_llm_kv_cache_events.py", - "tests/unittest/llmapi/test_llm_models.py", - "tests/unittest/llmapi/test_llm_multi_gpu.py", "tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py", "tests/unittest/llmapi/test_llm_pytorch.py", "tests/unittest/llmapi/test_llm_quant.py", @@ -1175,24 +775,14 @@ exclude = [ "tests/unittest/llmapi/test_serialization.py", "tests/unittest/llmapi/test_utils.py", "tests/unittest/others/__init__.py", - "tests/unittest/others/test_builder.py", "tests/unittest/others/test_convert_spec_decoding_mask_to_packed_mask.py", - "tests/unittest/others/test_debugging_api.py", "tests/unittest/others/test_exception.py", "tests/unittest/others/test_export.py", - "tests/unittest/others/test_graph_rewriter.py", - "tests/unittest/others/test_kv_cache_manager.py", "tests/unittest/others/test_kv_cache_transceiver.py", "tests/unittest/others/test_kv_cache_update.py", - "tests/unittest/others/test_layer.py", "tests/unittest/others/test_mapping.py", - "tests/unittest/others/test_model_dtype.py", - "tests/unittest/others/test_module.py", "tests/unittest/others/test_multimodal_registry.py", - "tests/unittest/others/test_plugins.py", - "tests/unittest/others/test_precision_control.py", "tests/unittest/others/test_pretrained_config.py", - "tests/unittest/others/test_session.py", "tests/unittest/others/test_time_breakdown.py", "tests/unittest/profile_utils.py", "tests/unittest/scaffolding/__init__.py", @@ -1201,141 +791,14 @@ exclude = [ "tests/unittest/scaffolding/test_scaffolding.py", "tests/unittest/scaffolding/test_task_collection.py", "tests/unittest/scaffolding/test_worker.py", - "tests/unittest/test_model_runner_cpp.py", "tests/unittest/test_pip_install.py", "tests/unittest/tools/__init__.py", - "tests/unittest/tools/plugin_gen/__init__.py", - "tests/unittest/tools/plugin_gen/kernel_config.py", - "tests/unittest/tools/plugin_gen/test_core.py", - "tests/unittest/tools/plugin_gen/test_plugin_gen.py", - "tests/unittest/tools/plugin_gen/test_shape_infer.py", "tests/unittest/tools/test_prepare_dataset.py", "tests/unittest/tools/test_test_to_stage_mapping.py", - "tests/unittest/trt/__init__.py", - "tests/unittest/trt/attention/test_bert_attention.py", - "tests/unittest/trt/attention/test_gpt_attention.py", - "tests/unittest/trt/attention/test_gpt_attention_IFB.py", - "tests/unittest/trt/attention/test_gpt_attention_no_cache.py", - "tests/unittest/trt/attention/test_sage_attention.py", - "tests/unittest/trt/functional/__init__.py", - "tests/unittest/trt/functional/test_alibi.py", - "tests/unittest/trt/functional/test_allreduce_norm.py", - "tests/unittest/trt/functional/test_allreduce_prepost_residual_norm.py", - "tests/unittest/trt/functional/test_arange.py", - "tests/unittest/trt/functional/test_argmax.py", - "tests/unittest/trt/functional/test_assertion.py", - "tests/unittest/trt/functional/test_avg_pool2d.py", - "tests/unittest/trt/functional/test_cast.py", - "tests/unittest/trt/functional/test_conv2d.py", - "tests/unittest/trt/functional/test_conv3d.py", - "tests/unittest/trt/functional/test_cos.py", - "tests/unittest/trt/functional/test_cumsum.py", - "tests/unittest/trt/functional/test_dora.py", - "tests/unittest/trt/functional/test_einsum.py", - "tests/unittest/trt/functional/test_embedding_single_gpu.py", - "tests/unittest/trt/functional/test_exp.py", - "tests/unittest/trt/functional/test_expand.py", - "tests/unittest/trt/functional/test_flatten.py", - "tests/unittest/trt/functional/test_flip.py", - "tests/unittest/trt/functional/test_fp4_gemm.py", - "tests/unittest/trt/functional/test_fp4_gemm_ootb.py", - "tests/unittest/trt/functional/test_gather.py", - "tests/unittest/trt/functional/test_gather_nd.py", - "tests/unittest/trt/functional/test_geglu.py", - "tests/unittest/trt/functional/test_gelu.py", - "tests/unittest/trt/functional/test_gemm_swiglu.py", - "tests/unittest/trt/functional/test_group_norm.py", - "tests/unittest/trt/functional/test_identity.py", - "tests/unittest/trt/functional/test_index_select.py", - "tests/unittest/trt/functional/test_interpolate.py", - "tests/unittest/trt/functional/test_logsoftmax.py", - "tests/unittest/trt/functional/test_lora.py", - "tests/unittest/trt/functional/test_low_latency_gemm.py", - "tests/unittest/trt/functional/test_mamba_conv1d.py", - "tests/unittest/trt/functional/test_masked_scatter.py", - "tests/unittest/trt/functional/test_masked_select.py", - "tests/unittest/trt/functional/test_matmul.py", - "tests/unittest/trt/functional/test_meshgrid2d.py", - "tests/unittest/trt/functional/test_moe.py", - "tests/unittest/trt/functional/test_nccl.py", - "tests/unittest/trt/functional/test_nonzero.py", - "tests/unittest/trt/functional/test_outer.py", - "tests/unittest/trt/functional/test_pad.py", - "tests/unittest/trt/functional/test_permute.py", - "tests/unittest/trt/functional/test_pp_reduce_scatter.py", - "tests/unittest/trt/functional/test_quant.py", - "tests/unittest/trt/functional/test_rearrange.py", - "tests/unittest/trt/functional/test_repeat.py", - "tests/unittest/trt/functional/test_repeat_interleave.py", - "tests/unittest/trt/functional/test_rg_lru.py", - "tests/unittest/trt/functional/test_sample.py", - "tests/unittest/trt/functional/test_scatter.py", - "tests/unittest/trt/functional/test_scatter_nd.py", - "tests/unittest/trt/functional/test_select.py", - "tests/unittest/trt/functional/test_selective_scan.py", - "tests/unittest/trt/functional/test_sigmoid.py", - "tests/unittest/trt/functional/test_silu.py", - "tests/unittest/trt/functional/test_sin.py", - "tests/unittest/trt/functional/test_slice.py", - "tests/unittest/trt/functional/test_softplus.py", - "tests/unittest/trt/functional/test_split.py", - "tests/unittest/trt/functional/test_squeeze.py", - "tests/unittest/trt/functional/test_swiglu.py", - "tests/unittest/trt/functional/test_topk.py", - "tests/unittest/trt/functional/test_transpose.py", - "tests/unittest/trt/functional/test_unbind.py", - "tests/unittest/trt/functional/test_unsqueeze.py", - "tests/unittest/trt/functional/test_view.py", - "tests/unittest/trt/functional/test_where.py", - "tests/unittest/trt/model/__init__.py", - "tests/unittest/trt/model/eagle/test_decode_draft_tokens_plugin.py", - "tests/unittest/trt/model/eagle/test_prepare_drafter_inputs_plugin.py", - "tests/unittest/trt/model/eagle/test_sample_accept_draft_tokens_plugin.py", - "tests/unittest/trt/model/redrafter/test_beams2tree.py", - "tests/unittest/trt/model/redrafter/test_draft_token.py", - "tests/unittest/trt/model/redrafter/test_draft_token_indices.py", - "tests/unittest/trt/model/redrafter/test_gather_beams.py", - "tests/unittest/trt/model/redrafter/test_mask.py", - "tests/unittest/trt/model/redrafter/test_packed_position_ids.py", - "tests/unittest/trt/model/redrafter/test_prefix_match_indices.py", - "tests/unittest/trt/model/redrafter/test_prepare_input.py", - "tests/unittest/trt/model/redrafter/test_process_logits.py", - "tests/unittest/trt/model/redrafter/test_top1.py", - "tests/unittest/trt/model/redrafter/test_unpack_gen_data.py", - "tests/unittest/trt/model/redrafter/test_validate.py", - "tests/unittest/trt/model/test_gpt.py", - "tests/unittest/trt/model/test_gpt_e2e.py", - "tests/unittest/trt/model/test_llama.py", - "tests/unittest/trt/model/test_mamba.py", - "tests/unittest/trt/model/test_mistral.py", - "tests/unittest/trt/model/test_nemotron_nas.py", - "tests/unittest/trt/model/test_phi.py", - "tests/unittest/trt/model/test_unet.py", - "tests/unittest/trt/model_api/test_model_api_multi_gpu.py", - "tests/unittest/trt/model_api/test_model_level_api.py", - "tests/unittest/trt/model_api/test_model_quantization.py", - "tests/unittest/trt/python_plugin/plugin_wrapper_utils.py", - "tests/unittest/trt/python_plugin/test_plugin_wrapper.py", - "tests/unittest/trt/quantization/__init__.py", - "tests/unittest/trt/quantization/_utils.py", - "tests/unittest/trt/quantization/test_fp8_quantization.py", - "tests/unittest/trt/quantization/test_fp8_rowwise_gemm.py", - "tests/unittest/trt/quantization/test_functional.py", - "tests/unittest/trt/quantization/test_mode.py", - "tests/unittest/trt/quantization/test_moe_weight_only_quant_matmul.py", - "tests/unittest/trt/quantization/test_qserve_gemm.py", - "tests/unittest/trt/quantization/test_quant.py", - "tests/unittest/trt/quantization/test_quant_layer.py", - "tests/unittest/trt/quantization/test_smooth_quant_gemm.py", - "tests/unittest/trt/quantization/test_smooth_quant_layer_norm.py", - "tests/unittest/trt/quantization/test_smooth_quant_rms_norm.py", - "tests/unittest/trt/quantization/test_weight_only_groupwise_quant_matmul.py", - "tests/unittest/trt/quantization/test_weight_only_quant_matmul.py", "tests/unittest/utils/__init__.py", "tests/unittest/utils/cpp_paths.py", "tests/unittest/utils/llm_data.py", "tests/unittest/utils/runtime_defaults.py", - "tests/unittest/utils/test_medusa_utils.py", "tests/unittest/utils/test_prebuilt_whl_cpp_extensions.py", "tests/unittest/utils/test_util.py", "tests/unittest/utils/torch_ref.py", diff --git a/ruff-legacy-baseline.json b/ruff-legacy-baseline.json index 0eb256fe39f6..eae44689bafc 100644 --- a/ruff-legacy-baseline.json +++ b/ruff-legacy-baseline.json @@ -1,22 +1,12 @@ { "_meta": { "generated_by": "scripts/legacy_utils.py lint-update-violations", - "total_violations": 2327, - "total_files": 266 + "total_violations": 2096, + "total_files": 224 }, ".github/scripts/label_community_user.py": { "D212": 1 }, - "benchmarks/utils/convert_nemo_dataset.py": { - "E741": 1 - }, - "benchmarks/utils/prepare_real_data.py": { - "D202": 1, - "D205": 1, - "D410": 7, - "D411": 8, - "D415": 1 - }, "cpp/kernels/fmha_v2/setup.py": { "E712": 27, "E731": 1 @@ -80,15 +70,6 @@ "D212": 1, "D300": 1 }, - "examples/dora/normalize_weights.py": { - "D200": 3, - "D202": 1, - "D205": 1, - "D212": 4 - }, - "examples/eval_long_context.py": { - "E711": 2 - }, "examples/infinitebench/compute_scores.py": { "D200": 1, "D205": 1, @@ -136,102 +117,10 @@ "D212": 1, "F821": 1 }, - "examples/mmlu.py": { - "D205": 1, - "D415": 1, - "E741": 1 - }, - "examples/models/contrib/chatglm-6b/tokenization_chatglm.py": { - "D205": 4, - "D210": 4, - "D212": 7, - "D301": 2, - "D415": 4 - }, - "examples/models/contrib/chatglm2-6b/tokenization_chatglm.py": { - "D205": 1, - "D210": 3, - "D212": 3, - "D415": 3 - }, - "examples/models/contrib/chatglm3-6b-32k/tokenization_chatglm.py": { - "D205": 1, - "D210": 3, - "D212": 3, - "D415": 3 - }, - "examples/models/contrib/sdxl/pipeline_stable_diffusion_xl.py": { - "D202": 1, - "D205": 6, - "D212": 10, - "D414": 1, - "D415": 2 - }, - "examples/models/core/multimodal/run.py": { - "E731": 1 - }, - "examples/models/core/qwen2audio/run.py": { - "D200": 1, - "D212": 1, - "D415": 1, - "E722": 1 - }, - "examples/models/core/qwen2audio/run_chat.py": { - "E722": 1 - }, - "examples/models/core/qwenvl/run.py": { - "E722": 1 - }, - "examples/models/core/qwenvl/run_chat.py": { - "E722": 1 - }, - "examples/ngram/run_dtm_ngram.py": { - "D200": 1, - "D205": 1, - "D212": 2, - "D415": 2, - "E741": 3 - }, - "examples/openai_triton/manual_plugin/build.py": { - "D202": 1, - "D205": 1, - "D212": 1, - "D300": 1 - }, - "examples/openai_triton/manual_plugin/fmha_triton.py": { - "D205": 1, - "D212": 1, - "D415": 1 - }, - "examples/openai_triton/manual_plugin/plugin.py": { - "D202": 1, - "D205": 1, - "D212": 1, - "D300": 1 - }, - "examples/openai_triton/plugin_autogen/build_engine.py": { - "D202": 2, - "D205": 2, - "D212": 2, - "D300": 2 - }, - "examples/openai_triton/plugin_autogen/kernel_config.py": { - "F403": 1, - "F405": 22 - }, - "examples/python_plugin/plugin_lib/lookup_plugin.py": { - "E731": 1 - }, "examples/quantization/quantize_mixed_precision_moe.py": { "E741": 1, "F601": 2 }, - "examples/run.py": { - "E711": 4 - }, - "examples/summarize.py": { - "E741": 1 - }, "jenkins/scripts/open_search_db.py": { "D205": 1, "D212": 7 @@ -392,18 +281,6 @@ "D411": 1, "D415": 1 }, - "tensorrt_llm/bench/build/build.py": { - "D205": 1, - "D210": 2, - "D411": 1 - }, - "tensorrt_llm/bench/build/dataclasses.py": { - "D210": 3 - }, - "tensorrt_llm/bench/build/tuning.py": { - "D205": 2, - "D210": 2 - }, "tensorrt_llm/bench/dataclasses/reporting.py": { "D200": 1, "D212": 1 @@ -605,17 +482,6 @@ "E731": 1, "F821": 1 }, - "tensorrt_llm/quantization/quantize_by_modelopt.py": { - "D200": 1, - "D205": 1, - "D208": 2, - "D212": 2, - "D300": 1, - "D415": 2, - "E722": 1, - "F601": 1, - "F821": 2 - }, "tensorrt_llm/quantization/utils/fp4_utils.py": { "D200": 2, "D202": 1, @@ -825,79 +691,17 @@ "D415": 15, "E722": 1 }, - "tests/integration/defs/examples/test_bert.py": { - "D300": 1, - "D415": 1 - }, - "tests/integration/defs/examples/test_bindings.py": { - "D300": 1, - "D415": 1 - }, - "tests/integration/defs/examples/test_chatglm.py": { - "D300": 1 - }, - "tests/integration/defs/examples/test_commandr.py": { - "D300": 1 - }, "tests/integration/defs/examples/test_gpt.py": { "D202": 3, "D300": 5, "D403": 1, "D415": 4 }, - "tests/integration/defs/examples/test_granite.py": { - "D202": 1, - "D300": 1 - }, - "tests/integration/defs/examples/test_internlm.py": { - "D300": 1, - "D415": 1 - }, - "tests/integration/defs/examples/test_llama.py": { - "D202": 2, - "D300": 4, - "D403": 2, - "D415": 2 - }, - "tests/integration/defs/examples/test_mamba.py": { - "D300": 2, - "D415": 2 - }, - "tests/integration/defs/examples/test_mistral.py": { - "D202": 1 - }, - "tests/integration/defs/examples/test_mixtral.py": { - "D300": 1, - "D403": 1 - }, - "tests/integration/defs/examples/test_multimodal.py": { - "D202": 1, - "D300": 2, - "D415": 2 - }, - "tests/integration/defs/examples/test_openai.py": { - "D300": 3, - "D403": 1, - "D415": 3 - }, "tests/integration/defs/examples/test_phi.py": { "D202": 1, "D300": 2, "D415": 2 }, - "tests/integration/defs/examples/test_qwen.py": { - "D202": 1, - "D300": 2, - "D403": 1 - }, - "tests/integration/defs/examples/test_qwen2audio.py": { - "D300": 2, - "D415": 1 - }, - "tests/integration/defs/examples/test_qwenvl.py": { - "D300": 1, - "D415": 1 - }, "tests/integration/defs/llmapi/test_llm_api_qa.py": { "D200": 1, "D212": 1, diff --git a/ruff-legacy.toml b/ruff-legacy.toml index 3aadfee4149c..39a722295b41 100644 --- a/ruff-legacy.toml +++ b/ruff-legacy.toml @@ -18,14 +18,6 @@ include = [ ".devcontainer/make_env.py", ".github/scripts/label_community_user.py", ".github/scripts/pr_checklist_check.py", - "benchmarks/__init__.py", - "benchmarks/prepare_dataset.py", - "benchmarks/utils/__init__.py", - "benchmarks/utils/convert_nemo_dataset.py", - "benchmarks/utils/generate_rand_loras.py", - "benchmarks/utils/prepare_real_data.py", - "benchmarks/utils/prepare_synthetic_data.py", - "benchmarks/utils/utils.py", "cpp/conanfile.py", "cpp/kernels/fmha_v2/conftest.py", "cpp/kernels/fmha_v2/fmha_test.py", @@ -57,24 +49,10 @@ include = [ "examples/apps/fastapi_server.py", "examples/disaggregated/clients/disagg_client.py", "examples/disaggregated/slurm/benchmark/submit.py", - "examples/dora/normalize_weights.py", - "examples/eagle/convert_checkpoint.py", - "examples/eval_long_context.py", - "examples/generate_checkpoint_config.py", - "examples/generate_xgrammar_tokenizer_info.py", - "examples/hf_lora_convert.py", "examples/infinitebench/args.py", "examples/infinitebench/compute_scores.py", "examples/infinitebench/construct_synthetic_dataset.py", "examples/infinitebench/eval_utils.py", - "examples/llm-api/_tensorrt_engine/llm_eagle2_decoding.py", - "examples/llm-api/_tensorrt_engine/llm_eagle_decoding.py", - "examples/llm-api/_tensorrt_engine/llm_inference_customize.py", - "examples/llm-api/_tensorrt_engine/llm_inference_kv_events.py", - "examples/llm-api/_tensorrt_engine/llm_lookahead_decoding.py", - "examples/llm-api/_tensorrt_engine/llm_medusa_decoding.py", - "examples/llm-api/_tensorrt_engine/llm_quantization.py", - "examples/llm-api/_tensorrt_engine/quickstart_example.py", "examples/llm-api/llm_guided_decoding.py", "examples/llm-api/llm_inference.py", "examples/llm-api/llm_inference_async.py", @@ -94,122 +72,17 @@ include = [ "examples/llm-api/quickstart_example.py", "examples/llm-api/quickstart_multimodal.py", "examples/llm-api/star_attention.py", - "examples/llm-eval/lm-eval-harness/lm_eval_tensorrt_llm.py", "examples/longbench/eval_longbench_v1.py", - "examples/medusa/convert_checkpoint.py", - "examples/mmlu.py", - "examples/models/contrib/baichuan/convert_checkpoint.py", - "examples/models/contrib/bloom/convert_checkpoint.py", - "examples/models/contrib/chatglm-6b/tokenization_chatglm.py", - "examples/models/contrib/chatglm2-6b/tokenization_chatglm.py", - "examples/models/contrib/chatglm3-6b-32k/tokenization_chatglm.py", - "examples/models/contrib/cogvlm/convert_checkpoint.py", - "examples/models/contrib/dbrx/convert_checkpoint.py", - "examples/models/contrib/deepseek_v1/__init__.py", - "examples/models/contrib/deepseek_v1/convert_checkpoint.py", - "examples/models/contrib/deepseek_v2/convert_checkpoint.py", - "examples/models/contrib/dit/convert_checkpoint.py", - "examples/models/contrib/dit/diffusion.py", - "examples/models/contrib/dit/sample.py", - "examples/models/contrib/dit/utils_modelopt.py", - "examples/models/contrib/dit/vae_decoder_trt.py", - "examples/models/contrib/falcon/convert_checkpoint.py", - "examples/models/contrib/gptj/convert_checkpoint.py", - "examples/models/contrib/gptneox/convert_checkpoint.py", - "examples/models/contrib/grok/convert_checkpoint.py", - "examples/models/contrib/mmdit/convert_checkpoint.py", - "examples/models/contrib/mmdit/sample.py", - "examples/models/contrib/mpt/convert_checkpoint.py", - "examples/models/contrib/opt/convert_checkpoint.py", - "examples/models/contrib/sdxl/build_sdxl_unet.py", - "examples/models/contrib/sdxl/pipeline_stable_diffusion_xl.py", - "examples/models/contrib/sdxl/run_sdxl.py", - "examples/models/contrib/stdit/aspect.py", - "examples/models/contrib/stdit/convert_checkpoint.py", - "examples/models/contrib/stdit/pipeline_tllm.py", - "examples/models/contrib/stdit/sample.py", - "examples/models/contrib/stdit/scheduler.py", - "examples/models/contrib/stdit/text_encoder.py", - "examples/models/contrib/stdit/utils.py", - "examples/models/contrib/stdit/vae.py", - "examples/models/contrib/stdit/video_transforms.py", - "examples/models/core/bert/__init__.py", - "examples/models/core/bert/convert_checkpoint.py", - "examples/models/core/bert/run.py", - "examples/models/core/bert/utils.py", - "examples/models/core/commandr/convert_checkpoint.py", - "examples/models/core/enc_dec/__init__.py", - "examples/models/core/enc_dec/convert_checkpoint.py", - "examples/models/core/enc_dec/helper.py", - "examples/models/core/enc_dec/run.py", - "examples/models/core/gemma/convert_checkpoint.py", - "examples/models/core/glm-4-9b/convert_checkpoint.py", - "examples/models/core/glm-4-9b/tokenization_chatglm.py", - "examples/models/core/gpt/convert_checkpoint.py", - "examples/models/core/gpt/merge_ptuning_tables.py", - "examples/models/core/gpt/nemo_lora_convert.py", - "examples/models/core/gpt/nemo_prompt_convert.py", - "examples/models/core/gpt/run_hf.py", "examples/models/core/gpt_oss/openai_chat_client_function_calling.py", - "examples/models/core/internlm2/convert_checkpoint.py", "examples/models/core/kimi_k2/kimi_k2_tool_calling_example.py", - "examples/models/core/llama/convert_checkpoint.py", - "examples/models/core/llama/summarize_long.py", - "examples/models/core/mamba/convert_checkpoint.py", - "examples/models/core/mllama/convert_checkpoint.py", - "examples/models/core/multimodal/__init__.py", - "examples/models/core/multimodal/build_multimodal_engine.py", - "examples/models/core/multimodal/eval.py", - "examples/models/core/multimodal/run.py", - "examples/models/core/multimodal/utils.py", - "examples/models/core/nemotron_nas/calibration_utils.py", - "examples/models/core/nemotron_nas/convert_checkpoint.py", - "examples/models/core/phi/convert_checkpoint.py", - "examples/models/core/qwen/convert_checkpoint.py", - "examples/models/core/qwen2audio/run.py", - "examples/models/core/qwen2audio/run_chat.py", - "examples/models/core/qwen2audio/utils.py", - "examples/models/core/qwenvl/run.py", - "examples/models/core/qwenvl/run_chat.py", - "examples/models/core/qwenvl/show_pic.py", - "examples/models/core/qwenvl/vit_onnx_trt.py", - "examples/models/core/recurrentgemma/convert_checkpoint.py", - "examples/models/core/vit/convert_checkpoint.py", - "examples/models/core/whisper/convert_checkpoint.py", - "examples/models/core/whisper/distil_whisper/convert_from_distil_whisper.py", - "examples/models/core/whisper/run.py", - "examples/models/core/whisper/tokenizer.py", - "examples/models/core/whisper/whisper_utils.py", - "examples/ngram/run_dtm_ngram.py", - "examples/openai_triton/manual_plugin/build.py", - "examples/openai_triton/manual_plugin/fmha_triton.py", - "examples/openai_triton/manual_plugin/plugin.py", - "examples/openai_triton/manual_plugin/run.py", - "examples/openai_triton/plugin_autogen/build_engine.py", - "examples/openai_triton/plugin_autogen/kernel_config.py", - "examples/openai_triton/plugin_autogen/run_engine.py", - "examples/python_plugin/build_lookup.py", - "examples/python_plugin/plugin_lib/__init__.py", - "examples/python_plugin/plugin_lib/lookup_kernel.py", - "examples/python_plugin/plugin_lib/lookup_plugin.py", - "examples/python_plugin/run_lookup.py", - "examples/quantization/quantize.py", "examples/quantization/quantize_mixed_precision_moe.py", "examples/ray_orchestrator/llm_inference_async_ray.py", "examples/ray_orchestrator/llm_inference_distributed_ray.py", - "examples/redrafter/convert_checkpoint.py", - "examples/run.py", "examples/scaffolding/contrib/AsyncGeneration/stream_generation_controller.py", "examples/scaffolding/contrib/DeepConf/run_generation.py", "examples/scaffolding/contrib/Dynasor/scaffolding_dynasor_run.py", "examples/scaffolding/contrib/TreeInference/run_mcts_example.py", "examples/scaffolding/contrib/TreeInference/run_tot_example.py", - "examples/scaffolding/contrib/mcp/e2b/e2bserver.py", - "examples/scaffolding/contrib/mcp/e2b/main.py", - "examples/scaffolding/contrib/mcp/mcptest.py", - "examples/scaffolding/contrib/mcp/weather/weather.py", - "examples/scaffolding/contrib/mcp/websearch/main.py", - "examples/scaffolding/contrib/mcp/websearch/websearch.py", "examples/scaffolding/run_basic_generation.py", "examples/scaffolding/run_best_of_n_with_reward.py", "examples/scaffolding/run_majority_vote_aime24.py", @@ -219,8 +92,6 @@ include = [ "examples/serve/openai_completion_client.py", "examples/serve/openai_completion_client_for_lora.py", "examples/serve/openai_completion_client_json_schema.py", - "examples/summarize.py", - "examples/utils.py", "examples/wide_ep/ep_load_balancer/generate_eplb_config.py", "examples/wide_ep/ep_load_balancer/report_load_statistics.py", "examples/wide_ep/ep_load_balancer/utils.py", @@ -232,7 +103,6 @@ include = [ "scripts/check_test_list.py", "scripts/dco_check.py", "scripts/format_test_list.py", - "scripts/generate_duration.py", "scripts/generate_lock_file.py", "scripts/get_wheel_from_package.py", "scripts/git_replace.py", @@ -243,7 +113,6 @@ include = [ "setup.py", "tensorrt_llm/__init__.py", "tensorrt_llm/_ray_utils.py", - "tensorrt_llm/_tensorrt_engine/__init__.py", "tensorrt_llm/_torch/__init__.py", "tensorrt_llm/_torch/attention_backend/__init__.py", "tensorrt_llm/_torch/attention_backend/flashinfer.py", @@ -447,7 +316,6 @@ include = [ "tensorrt_llm/_torch/pyexecutor/guided_decoder.py", "tensorrt_llm/_torch/pyexecutor/handle_additional_outputs.py", "tensorrt_llm/_torch/pyexecutor/handle_logits.py", - "tensorrt_llm/_torch/pyexecutor/kv_cache_connector.py", "tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py", "tensorrt_llm/_torch/pyexecutor/layerwise_nvtx_marker.py", "tensorrt_llm/_torch/pyexecutor/llm_request.py", @@ -485,11 +353,6 @@ include = [ "tensorrt_llm/bench/benchmark/utils/asynchronous.py", "tensorrt_llm/bench/benchmark/utils/general.py", "tensorrt_llm/bench/benchmark/utils/processes.py", - "tensorrt_llm/bench/build/__init__.py", - "tensorrt_llm/bench/build/build.py", - "tensorrt_llm/bench/build/dataclasses.py", - "tensorrt_llm/bench/build/tuning.py", - "tensorrt_llm/bench/build/utils.py", "tensorrt_llm/bench/dataclasses/__init__.py", "tensorrt_llm/bench/dataclasses/configuration.py", "tensorrt_llm/bench/dataclasses/engine.py", @@ -499,13 +362,9 @@ include = [ "tensorrt_llm/bench/dataclasses/statistics.py", "tensorrt_llm/bench/utils/__init__.py", "tensorrt_llm/bench/utils/data.py", - "tensorrt_llm/builder.py", "tensorrt_llm/commands/__init__.py", "tensorrt_llm/commands/bench.py", - "tensorrt_llm/commands/build.py", "tensorrt_llm/commands/eval.py", - "tensorrt_llm/commands/prune.py", - "tensorrt_llm/commands/refit.py", "tensorrt_llm/commands/serve.py", "tensorrt_llm/evaluate/__init__.py", "tensorrt_llm/evaluate/cnn_dailymail.py", @@ -541,23 +400,7 @@ include = [ "tensorrt_llm/inputs/multimodal.py", "tensorrt_llm/inputs/registry.py", "tensorrt_llm/inputs/utils.py", - "tensorrt_llm/layers/__init__.py", - "tensorrt_llm/layers/activation.py", - "tensorrt_llm/layers/attention.py", - "tensorrt_llm/layers/cast.py", - "tensorrt_llm/layers/conv.py", - "tensorrt_llm/layers/embedding.py", - "tensorrt_llm/layers/language_adapter.py", - "tensorrt_llm/layers/linear.py", - "tensorrt_llm/layers/lora.py", - "tensorrt_llm/layers/mlp.py", - "tensorrt_llm/layers/moe.py", - "tensorrt_llm/layers/normalization.py", - "tensorrt_llm/layers/pooling.py", - "tensorrt_llm/layers/recurrent.py", - "tensorrt_llm/layers/ssm.py", "tensorrt_llm/llmapi/__init__.py", - "tensorrt_llm/llmapi/build_cache.py", "tensorrt_llm/llmapi/disagg_utils.py", "tensorrt_llm/llmapi/kv_cache_type.py", "tensorrt_llm/llmapi/llm.py", @@ -580,179 +423,17 @@ include = [ "tensorrt_llm/metrics/enums.py", "tensorrt_llm/models/__init__.py", "tensorrt_llm/models/automodel.py", - "tensorrt_llm/models/baichuan/__init__.py", - "tensorrt_llm/models/baichuan/config.py", - "tensorrt_llm/models/baichuan/convert.py", - "tensorrt_llm/models/baichuan/model.py", - "tensorrt_llm/models/bert/__init__.py", - "tensorrt_llm/models/bert/config.py", - "tensorrt_llm/models/bert/convert.py", - "tensorrt_llm/models/bert/model.py", - "tensorrt_llm/models/bloom/__init__.py", - "tensorrt_llm/models/bloom/model.py", - "tensorrt_llm/models/chatglm/__init__.py", - "tensorrt_llm/models/chatglm/config.py", - "tensorrt_llm/models/chatglm/convert.py", - "tensorrt_llm/models/chatglm/model.py", - "tensorrt_llm/models/clip/__init__.py", - "tensorrt_llm/models/clip/model.py", - "tensorrt_llm/models/cogvlm/__init__.py", - "tensorrt_llm/models/cogvlm/config.py", - "tensorrt_llm/models/cogvlm/convert.py", - "tensorrt_llm/models/cogvlm/model.py", - "tensorrt_llm/models/commandr/__init__.py", - "tensorrt_llm/models/commandr/config.py", - "tensorrt_llm/models/commandr/model.py", "tensorrt_llm/models/convert_utils.py", - "tensorrt_llm/models/dbrx/__init__.py", - "tensorrt_llm/models/dbrx/config.py", - "tensorrt_llm/models/dbrx/model.py", - "tensorrt_llm/models/deepseek_v1/__init__.py", - "tensorrt_llm/models/deepseek_v1/config.py", - "tensorrt_llm/models/deepseek_v1/convert.py", - "tensorrt_llm/models/deepseek_v1/model.py", - "tensorrt_llm/models/deepseek_v2/__init__.py", - "tensorrt_llm/models/deepseek_v2/config.py", - "tensorrt_llm/models/deepseek_v2/convert.py", - "tensorrt_llm/models/deepseek_v2/model.py", - "tensorrt_llm/models/dit/__init__.py", - "tensorrt_llm/models/dit/model.py", - "tensorrt_llm/models/eagle/__init__.py", - "tensorrt_llm/models/eagle/config.py", - "tensorrt_llm/models/eagle/model.py", - "tensorrt_llm/models/enc_dec/__init__.py", - "tensorrt_llm/models/enc_dec/model.py", - "tensorrt_llm/models/falcon/__init__.py", - "tensorrt_llm/models/falcon/config.py", - "tensorrt_llm/models/falcon/convert.py", - "tensorrt_llm/models/falcon/model.py", - "tensorrt_llm/models/gemma/__init__.py", - "tensorrt_llm/models/gemma/config.py", - "tensorrt_llm/models/gemma/convert.py", - "tensorrt_llm/models/gemma/model.py", - "tensorrt_llm/models/gemma/smoothquant.py", - "tensorrt_llm/models/gemma/utils/__init__.py", - "tensorrt_llm/models/gemma/utils/layers.py", - "tensorrt_llm/models/gemma/utils/modules.py", - "tensorrt_llm/models/gemma/utils/params.py", - "tensorrt_llm/models/gemma/utils/positional_embeddings.py", - "tensorrt_llm/models/gemma/utils/sampler.py", - "tensorrt_llm/models/gemma/utils/transformer.py", - "tensorrt_llm/models/gemma/weight.py", - "tensorrt_llm/models/generation_mixin.py", - "tensorrt_llm/models/gpt/__init__.py", - "tensorrt_llm/models/gpt/config.py", - "tensorrt_llm/models/gpt/convert.py", - "tensorrt_llm/models/gpt/model.py", - "tensorrt_llm/models/gptj/__init__.py", - "tensorrt_llm/models/gptj/config.py", - "tensorrt_llm/models/gptj/convert.py", - "tensorrt_llm/models/gptj/model.py", - "tensorrt_llm/models/gptneox/__init__.py", - "tensorrt_llm/models/gptneox/model.py", - "tensorrt_llm/models/grok/__init__.py", - "tensorrt_llm/models/grok/convert.py", - "tensorrt_llm/models/grok/model.py", - "tensorrt_llm/models/grok/weight.py", - "tensorrt_llm/models/llama/__init__.py", - "tensorrt_llm/models/llama/config.py", - "tensorrt_llm/models/llama/convert.py", - "tensorrt_llm/models/llama/model.py", - "tensorrt_llm/models/mamba/__init__.py", - "tensorrt_llm/models/mamba/config.py", - "tensorrt_llm/models/mamba/convert.py", - "tensorrt_llm/models/mamba/model.py", - "tensorrt_llm/models/medusa/__init__.py", - "tensorrt_llm/models/medusa/config.py", - "tensorrt_llm/models/medusa/model.py", - "tensorrt_llm/models/medusa/weight.py", - "tensorrt_llm/models/mllama/__init__.py", - "tensorrt_llm/models/mllama/config.py", - "tensorrt_llm/models/mllama/model.py", - "tensorrt_llm/models/mmdit_sd3/__init__.py", - "tensorrt_llm/models/mmdit_sd3/config.py", - "tensorrt_llm/models/mmdit_sd3/model.py", - "tensorrt_llm/models/model_weights_loader.py", "tensorrt_llm/models/modeling_utils.py", - "tensorrt_llm/models/mpt/__init__.py", - "tensorrt_llm/models/mpt/model.py", - "tensorrt_llm/models/multimodal_encoders/__init__.py", - "tensorrt_llm/models/multimodal_encoders/config.py", - "tensorrt_llm/models/multimodal_encoders/model.py", - "tensorrt_llm/models/nemotron_nas/__init__.py", - "tensorrt_llm/models/nemotron_nas/config.py", - "tensorrt_llm/models/nemotron_nas/convert.py", - "tensorrt_llm/models/nemotron_nas/layer_config.py", - "tensorrt_llm/models/nemotron_nas/model.py", - "tensorrt_llm/models/opt/__init__.py", - "tensorrt_llm/models/opt/model.py", - "tensorrt_llm/models/phi/__init__.py", - "tensorrt_llm/models/phi/config.py", - "tensorrt_llm/models/phi/convert.py", - "tensorrt_llm/models/phi/model.py", - "tensorrt_llm/models/phi3/__init__.py", - "tensorrt_llm/models/phi3/config.py", - "tensorrt_llm/models/phi3/convert.py", - "tensorrt_llm/models/phi3/model.py", - "tensorrt_llm/models/phi3/split_weights.py", - "tensorrt_llm/models/qwen/__init__.py", - "tensorrt_llm/models/qwen/config.py", - "tensorrt_llm/models/qwen/convert.py", - "tensorrt_llm/models/qwen/model.py", - "tensorrt_llm/models/qwen/utils.py", - "tensorrt_llm/models/recurrentgemma/__init__.py", - "tensorrt_llm/models/recurrentgemma/model.py", - "tensorrt_llm/models/redrafter/__init__.py", - "tensorrt_llm/models/redrafter/drafter.py", - "tensorrt_llm/models/redrafter/model.py", - "tensorrt_llm/models/redrafter/redrafter_helper.py", - "tensorrt_llm/models/stdit/__init__.py", - "tensorrt_llm/models/stdit/config.py", - "tensorrt_llm/models/stdit/model.py", - "tensorrt_llm/models/unet/__init__.py", - "tensorrt_llm/models/unet/attention.py", - "tensorrt_llm/models/unet/embeddings.py", - "tensorrt_llm/models/unet/pp/__init__.py", - "tensorrt_llm/models/unet/pp/attention.py", - "tensorrt_llm/models/unet/pp/conv2d.py", - "tensorrt_llm/models/unet/pp/groupnorm.py", - "tensorrt_llm/models/unet/pp/unet_pp.py", - "tensorrt_llm/models/unet/resnet.py", - "tensorrt_llm/models/unet/unet_2d_blocks.py", - "tensorrt_llm/models/unet/unet_2d_condition.py", - "tensorrt_llm/models/unet/weights.py", - "tensorrt_llm/network.py", - "tensorrt_llm/parameter.py", - "tensorrt_llm/plugin/__init__.py", - "tensorrt_llm/plugin/plugin.py", "tensorrt_llm/quantization/__init__.py", "tensorrt_llm/quantization/functional.py", - "tensorrt_llm/quantization/image_processing.py", - "tensorrt_llm/quantization/layers.py", "tensorrt_llm/quantization/mode.py", - "tensorrt_llm/quantization/quantize.py", - "tensorrt_llm/quantization/quantize_by_modelopt.py", "tensorrt_llm/quantization/utils/__init__.py", "tensorrt_llm/quantization/utils/fp4_utils.py", "tensorrt_llm/quantization/utils/fp8_utils.py", "tensorrt_llm/ray_stub.py", "tensorrt_llm/runtime/__init__.py", - "tensorrt_llm/runtime/enc_dec_model_runner.py", - "tensorrt_llm/runtime/generation.py", - "tensorrt_llm/runtime/kv_cache_manager.py", - "tensorrt_llm/runtime/medusa_utils.py", "tensorrt_llm/runtime/memory_pools/__init__.py", - "tensorrt_llm/runtime/memory_pools/memory_pools_allocator.py", - "tensorrt_llm/runtime/memory_pools/pool.py", - "tensorrt_llm/runtime/memory_pools/pools_kv_cache_manager.py", - "tensorrt_llm/runtime/model_runner.py", - "tensorrt_llm/runtime/model_runner_cpp.py", - "tensorrt_llm/runtime/multimodal_model_runner.py", - "tensorrt_llm/runtime/processor_wrapper/__init__.py", - "tensorrt_llm/runtime/processor_wrapper/mllama_processor_wrapper.py", - "tensorrt_llm/runtime/processor_wrapper/processor_wrapper.py", - "tensorrt_llm/runtime/redrafter_utils.py", - "tensorrt_llm/runtime/session.py", "tensorrt_llm/scaffolding/__init__.py", "tensorrt_llm/scaffolding/benchmark.py", "tensorrt_llm/scaffolding/contrib/AsyncGeneration/stream_generation.py", @@ -807,12 +488,7 @@ include = [ "tensorrt_llm/tokenizer/tokenizer.py", "tensorrt_llm/tools/__init__.py", "tensorrt_llm/tools/importlib_utils.py", - "tensorrt_llm/tools/multimodal_builder.py", - "tensorrt_llm/tools/onnx_utils.py", "tensorrt_llm/tools/plugin_gen/__init__.py", - "tensorrt_llm/tools/plugin_gen/core.py", - "tensorrt_llm/tools/plugin_gen/plugin_gen.py", - "tensorrt_llm/tools/plugin_gen/shape_infer.py", "tensorrt_llm/tools/ppl.py", "tensorrt_llm/tools/profiler/nsys_profile_tools/gputrc2graph.py", "tensorrt_llm/version.py", @@ -821,9 +497,7 @@ include = [ "tests/integration/defs/accuracy/accuracy_core.py", "tests/integration/defs/accuracy/scripts/collect_evaluated_accuracies.py", "tests/integration/defs/accuracy/scripts/compute_theta_and_thresholds.py", - "tests/integration/defs/accuracy/test_cli_flow.py", "tests/integration/defs/accuracy/test_disaggregated_serving.py", - "tests/integration/defs/accuracy/test_llm_api.py", "tests/integration/defs/accuracy/test_llm_api_autodeploy.py", "tests/integration/defs/accuracy/test_llm_api_pytorch.py", "tests/integration/defs/accuracy/test_llm_api_pytorch_ray.py", @@ -839,53 +513,22 @@ include = [ "tests/integration/defs/disaggregated/test_disaggregated_etcd.py", "tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py", "tests/integration/defs/disaggregated/test_workers.py", - "tests/integration/defs/examples/run_llm_fp8_quant_llama_70b.py", "tests/integration/defs/examples/run_llm_quickstart_atexit.py", "tests/integration/defs/examples/serve/test_serve.py", "tests/integration/defs/examples/serve/test_serve_negative.py", "tests/integration/defs/examples/test_ad_guided_decoding.py", - "tests/integration/defs/examples/test_bert.py", - "tests/integration/defs/examples/test_bindings.py", - "tests/integration/defs/examples/test_chatglm.py", - "tests/integration/defs/examples/test_commandr.py", - "tests/integration/defs/examples/test_draft_target_model.py", - "tests/integration/defs/examples/test_eagle.py", - "tests/integration/defs/examples/test_enc_dec.py", - "tests/integration/defs/examples/test_exaone.py", - "tests/integration/defs/examples/test_gemma.py", "tests/integration/defs/examples/test_gpt.py", - "tests/integration/defs/examples/test_gptj.py", - "tests/integration/defs/examples/test_granite.py", - "tests/integration/defs/examples/test_internlm.py", - "tests/integration/defs/examples/test_llama.py", "tests/integration/defs/examples/test_llm_api_with_mpi.py", - "tests/integration/defs/examples/test_mamba.py", - "tests/integration/defs/examples/test_medusa.py", - "tests/integration/defs/examples/test_mistral.py", - "tests/integration/defs/examples/test_mixtral.py", - "tests/integration/defs/examples/test_multimodal.py", - "tests/integration/defs/examples/test_nemotron.py", - "tests/integration/defs/examples/test_nemotron_nas.py", - "tests/integration/defs/examples/test_ngram.py", - "tests/integration/defs/examples/test_openai.py", "tests/integration/defs/examples/test_phi.py", - "tests/integration/defs/examples/test_qwen.py", - "tests/integration/defs/examples/test_qwen2audio.py", - "tests/integration/defs/examples/test_qwenvl.py", "tests/integration/defs/examples/test_ray.py", - "tests/integration/defs/examples/test_recurrentgemma.py", - "tests/integration/defs/examples/test_redrafter.py", - "tests/integration/defs/examples/test_whisper.py", "tests/integration/defs/llmapi/__init__.py", "tests/integration/defs/llmapi/_run_llmapi_llm.py", "tests/integration/defs/llmapi/test_llm_api_connector.py", "tests/integration/defs/llmapi/test_llm_api_qa.py", - "tests/integration/defs/llmapi/test_llm_e2e.py", "tests/integration/defs/llmapi/test_llm_examples.py", "tests/integration/defs/local_venv.py", "tests/integration/defs/perf/__init__.py", "tests/integration/defs/perf/allowed_configs.py", - "tests/integration/defs/perf/build.py", "tests/integration/defs/perf/create_perf_comparison_report.py", "tests/integration/defs/perf/data.py", "tests/integration/defs/perf/data_export.py", @@ -909,34 +552,18 @@ include = [ "tests/integration/defs/test_sanity.py", "tests/integration/defs/test_unittests.py", "tests/integration/defs/triton_server/__init__.py", - "tests/integration/defs/triton_server/build_engines.py", "tests/integration/defs/triton_server/common.py", "tests/integration/defs/triton_server/conftest.py", - "tests/integration/defs/triton_server/local_venv.py", - "tests/integration/defs/triton_server/rcca/bug_4323566/inflight_batcher_llm_client_with_end_id.py", - "tests/integration/defs/triton_server/runner_interface.py", "tests/integration/defs/triton_server/test_list_parser.py", - "tests/integration/defs/triton_server/test_triton.py", - "tests/integration/defs/triton_server/test_triton_llm.py", - "tests/integration/defs/triton_server/test_triton_memleak.py", - "tests/integration/defs/triton_server/test_triton_multi_node.py", - "tests/integration/defs/triton_server/test_triton_rcca.py", "tests/integration/defs/triton_server/trt_test_alternative.py", "tests/integration/defs/trt_test_alternative.py", "tests/integration/defs/utils/__init__.py", "tests/integration/defs/utils/periodic_junit.py", "tests/integration/defs/utils/timeout_manager.py", "tests/microbenchmarks/all_reduce.py", - "tests/microbenchmarks/build_time_benchmark.py", - "tests/microbenchmarks/build_time_dashboard.py", "tests/scripts/allreduce_perf/allreduce_heuristic_code_gen.py", "tests/scripts/allreduce_perf/allreduce_perf_viz.py", "tests/scripts/iteration_log_parser.py", - "tests/scripts/perf-sanity/parse_benchmark_results.py", - "tests/scripts/perf-sanity/run_benchmark_serve.py", - "tests/unittest/_torch/attention/sparse/test_dsa_indexer.py", - "tests/unittest/_torch/attention/sparse/test_flash_mla.py", - "tests/unittest/_torch/attention/sparse/test_rocketkv.py", "tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py", "tests/unittest/_torch/attention/test_attention.py", "tests/unittest/_torch/attention/test_attention_mla.py", @@ -957,7 +584,6 @@ include = [ "tests/unittest/_torch/misc/test_virtual_memory.py", "tests/unittest/_torch/modeling/test_modeling_bert.py", "tests/unittest/_torch/modeling/test_modeling_clip.py", - "tests/unittest/_torch/modeling/test_modeling_exaone4.py", "tests/unittest/_torch/modeling/test_modeling_gemma3.py", "tests/unittest/_torch/modeling/test_modeling_gpt_oss.py", "tests/unittest/_torch/modeling/test_modeling_llama.py", @@ -981,8 +607,6 @@ include = [ "tests/unittest/_torch/modules/test_moe_routing.py", "tests/unittest/_torch/modules/test_rotary_embedding.py", "tests/unittest/_torch/modules/test_triton_linear.py", - "tests/unittest/_torch/modules/tests_lora_modules/test_lora_attention_pytorch_flow_vs_trt.py", - "tests/unittest/_torch/modules/tests_lora_modules/test_lora_plugin_vs_lora_op.py", "tests/unittest/_torch/multi_gpu/test_allreduce.py", "tests/unittest/_torch/multi_gpu/test_alltoall.py", "tests/unittest/_torch/multi_gpu/test_ar_residual_norm.py", @@ -1009,24 +633,10 @@ include = [ "tests/unittest/_torch/sampler/test_beam_search.py", "tests/unittest/_torch/sampler/test_best_of_n.py", "tests/unittest/_torch/sampler/test_trtllm_sampler.py", - "tests/unittest/_torch/speculative/test_draft_target.py", - "tests/unittest/_torch/speculative/test_draft_token_tree_sampling.py", - "tests/unittest/_torch/speculative/test_draft_token_tree_verification.py", - "tests/unittest/_torch/speculative/test_dynamic_spec_decode.py", "tests/unittest/_torch/speculative/test_eagle3.py", - "tests/unittest/_torch/speculative/test_kv_cache_reuse.py", - "tests/unittest/_torch/speculative/test_mtp.py", - "tests/unittest/_torch/speculative/test_ngram.py", - "tests/unittest/_torch/speculative/test_save_state.py", - "tests/unittest/_torch/speculative/test_spec_gate.py", - "tests/unittest/_torch/speculative/test_torch_rejection_sampling.py", - "tests/unittest/_torch/speculative/test_user_provided.py", "tests/unittest/_torch/test_connector.py", "tests/unittest/_torch/test_torch_multi_arange.py", "tests/unittest/_torch/thop/parallel/deep_gemm_tests.py", - "tests/unittest/_torch/thop/parallel/test_causal_conv1d_op.py", - "tests/unittest/_torch/thop/parallel/test_cublas_mm.py", - "tests/unittest/_torch/thop/parallel/test_custom_ops.py", "tests/unittest/_torch/thop/parallel/test_dsv3_fused_a_gemm.py", "tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py", "tests/unittest/_torch/thop/parallel/test_finegrained_mixed_dtype_gemm.py", @@ -1040,11 +650,6 @@ include = [ "tests/unittest/_torch/thop/parallel/test_fp8_per_tensor_scale_tllmg_gemm.py", "tests/unittest/_torch/thop/parallel/test_fp8_quantize.py", "tests/unittest/_torch/thop/parallel/test_fp8_rowwise_linear.py", - "tests/unittest/_torch/thop/parallel/test_fused_qk_norm_rope.py", - "tests/unittest/_torch/thop/parallel/test_logits_bitmask_op.py", - "tests/unittest/_torch/thop/parallel/test_mamba2_chunk_ss_update.py", - "tests/unittest/_torch/thop/parallel/test_mamba_conv1d_op.py", - "tests/unittest/_torch/thop/parallel/test_noaux_tc.py", "tests/unittest/_torch/thop/parallel/test_scaled_mm.py", "tests/unittest/_torch/thop/parallel/test_selective_scan_op.py", "tests/unittest/_torch/thop/parallel/test_tinygemm2.py", @@ -1088,12 +693,10 @@ include = [ "tests/unittest/llmapi/apps/_test_openai_chat_harmony.py", "tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py", "tests/unittest/llmapi/apps/_test_openai_completions.py", - "tests/unittest/llmapi/apps/_test_openai_consistent_chat.py", "tests/unittest/llmapi/apps/_test_openai_lora.py", "tests/unittest/llmapi/apps/_test_openai_metrics.py", "tests/unittest/llmapi/apps/_test_openai_misc.py", "tests/unittest/llmapi/apps/_test_openai_mmencoder.py", - "tests/unittest/llmapi/apps/_test_openai_multi_chat.py", "tests/unittest/llmapi/apps/_test_openai_multi_gpu.py", "tests/unittest/llmapi/apps/_test_openai_multi_nodes.py", "tests/unittest/llmapi/apps/_test_openai_perf_metrics.py", @@ -1116,15 +719,12 @@ include = [ "tests/unittest/llmapi/run_llm_exit.py", "tests/unittest/llmapi/run_llm_with_postproc.py", "tests/unittest/llmapi/test_additional_model_outputs.py", - "tests/unittest/llmapi/test_build_cache.py", "tests/unittest/llmapi/test_executor.py", "tests/unittest/llmapi/test_gc_utils.py", "tests/unittest/llmapi/test_llm.py", "tests/unittest/llmapi/test_llm_args.py", "tests/unittest/llmapi/test_llm_download.py", "tests/unittest/llmapi/test_llm_kv_cache_events.py", - "tests/unittest/llmapi/test_llm_models.py", - "tests/unittest/llmapi/test_llm_multi_gpu.py", "tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py", "tests/unittest/llmapi/test_llm_pytorch.py", "tests/unittest/llmapi/test_llm_quant.py", @@ -1135,24 +735,14 @@ include = [ "tests/unittest/llmapi/test_serialization.py", "tests/unittest/llmapi/test_utils.py", "tests/unittest/others/__init__.py", - "tests/unittest/others/test_builder.py", "tests/unittest/others/test_convert_spec_decoding_mask_to_packed_mask.py", - "tests/unittest/others/test_debugging_api.py", "tests/unittest/others/test_exception.py", "tests/unittest/others/test_export.py", - "tests/unittest/others/test_graph_rewriter.py", - "tests/unittest/others/test_kv_cache_manager.py", "tests/unittest/others/test_kv_cache_transceiver.py", "tests/unittest/others/test_kv_cache_update.py", - "tests/unittest/others/test_layer.py", "tests/unittest/others/test_mapping.py", - "tests/unittest/others/test_model_dtype.py", - "tests/unittest/others/test_module.py", "tests/unittest/others/test_multimodal_registry.py", - "tests/unittest/others/test_plugins.py", - "tests/unittest/others/test_precision_control.py", "tests/unittest/others/test_pretrained_config.py", - "tests/unittest/others/test_session.py", "tests/unittest/others/test_time_breakdown.py", "tests/unittest/profile_utils.py", "tests/unittest/scaffolding/__init__.py", @@ -1161,141 +751,14 @@ include = [ "tests/unittest/scaffolding/test_scaffolding.py", "tests/unittest/scaffolding/test_task_collection.py", "tests/unittest/scaffolding/test_worker.py", - "tests/unittest/test_model_runner_cpp.py", "tests/unittest/test_pip_install.py", "tests/unittest/tools/__init__.py", - "tests/unittest/tools/plugin_gen/__init__.py", - "tests/unittest/tools/plugin_gen/kernel_config.py", - "tests/unittest/tools/plugin_gen/test_core.py", - "tests/unittest/tools/plugin_gen/test_plugin_gen.py", - "tests/unittest/tools/plugin_gen/test_shape_infer.py", "tests/unittest/tools/test_prepare_dataset.py", "tests/unittest/tools/test_test_to_stage_mapping.py", - "tests/unittest/trt/__init__.py", - "tests/unittest/trt/attention/test_bert_attention.py", - "tests/unittest/trt/attention/test_gpt_attention.py", - "tests/unittest/trt/attention/test_gpt_attention_IFB.py", - "tests/unittest/trt/attention/test_gpt_attention_no_cache.py", - "tests/unittest/trt/attention/test_sage_attention.py", - "tests/unittest/trt/functional/__init__.py", - "tests/unittest/trt/functional/test_alibi.py", - "tests/unittest/trt/functional/test_allreduce_norm.py", - "tests/unittest/trt/functional/test_allreduce_prepost_residual_norm.py", - "tests/unittest/trt/functional/test_arange.py", - "tests/unittest/trt/functional/test_argmax.py", - "tests/unittest/trt/functional/test_assertion.py", - "tests/unittest/trt/functional/test_avg_pool2d.py", - "tests/unittest/trt/functional/test_cast.py", - "tests/unittest/trt/functional/test_conv2d.py", - "tests/unittest/trt/functional/test_conv3d.py", - "tests/unittest/trt/functional/test_cos.py", - "tests/unittest/trt/functional/test_cumsum.py", - "tests/unittest/trt/functional/test_dora.py", - "tests/unittest/trt/functional/test_einsum.py", - "tests/unittest/trt/functional/test_embedding_single_gpu.py", - "tests/unittest/trt/functional/test_exp.py", - "tests/unittest/trt/functional/test_expand.py", - "tests/unittest/trt/functional/test_flatten.py", - "tests/unittest/trt/functional/test_flip.py", - "tests/unittest/trt/functional/test_fp4_gemm.py", - "tests/unittest/trt/functional/test_fp4_gemm_ootb.py", - "tests/unittest/trt/functional/test_gather.py", - "tests/unittest/trt/functional/test_gather_nd.py", - "tests/unittest/trt/functional/test_geglu.py", - "tests/unittest/trt/functional/test_gelu.py", - "tests/unittest/trt/functional/test_gemm_swiglu.py", - "tests/unittest/trt/functional/test_group_norm.py", - "tests/unittest/trt/functional/test_identity.py", - "tests/unittest/trt/functional/test_index_select.py", - "tests/unittest/trt/functional/test_interpolate.py", - "tests/unittest/trt/functional/test_logsoftmax.py", - "tests/unittest/trt/functional/test_lora.py", - "tests/unittest/trt/functional/test_low_latency_gemm.py", - "tests/unittest/trt/functional/test_mamba_conv1d.py", - "tests/unittest/trt/functional/test_masked_scatter.py", - "tests/unittest/trt/functional/test_masked_select.py", - "tests/unittest/trt/functional/test_matmul.py", - "tests/unittest/trt/functional/test_meshgrid2d.py", - "tests/unittest/trt/functional/test_moe.py", - "tests/unittest/trt/functional/test_nccl.py", - "tests/unittest/trt/functional/test_nonzero.py", - "tests/unittest/trt/functional/test_outer.py", - "tests/unittest/trt/functional/test_pad.py", - "tests/unittest/trt/functional/test_permute.py", - "tests/unittest/trt/functional/test_pp_reduce_scatter.py", - "tests/unittest/trt/functional/test_quant.py", - "tests/unittest/trt/functional/test_rearrange.py", - "tests/unittest/trt/functional/test_repeat.py", - "tests/unittest/trt/functional/test_repeat_interleave.py", - "tests/unittest/trt/functional/test_rg_lru.py", - "tests/unittest/trt/functional/test_sample.py", - "tests/unittest/trt/functional/test_scatter.py", - "tests/unittest/trt/functional/test_scatter_nd.py", - "tests/unittest/trt/functional/test_select.py", - "tests/unittest/trt/functional/test_selective_scan.py", - "tests/unittest/trt/functional/test_sigmoid.py", - "tests/unittest/trt/functional/test_silu.py", - "tests/unittest/trt/functional/test_sin.py", - "tests/unittest/trt/functional/test_slice.py", - "tests/unittest/trt/functional/test_softplus.py", - "tests/unittest/trt/functional/test_split.py", - "tests/unittest/trt/functional/test_squeeze.py", - "tests/unittest/trt/functional/test_swiglu.py", - "tests/unittest/trt/functional/test_topk.py", - "tests/unittest/trt/functional/test_transpose.py", - "tests/unittest/trt/functional/test_unbind.py", - "tests/unittest/trt/functional/test_unsqueeze.py", - "tests/unittest/trt/functional/test_view.py", - "tests/unittest/trt/functional/test_where.py", - "tests/unittest/trt/model/__init__.py", - "tests/unittest/trt/model/eagle/test_decode_draft_tokens_plugin.py", - "tests/unittest/trt/model/eagle/test_prepare_drafter_inputs_plugin.py", - "tests/unittest/trt/model/eagle/test_sample_accept_draft_tokens_plugin.py", - "tests/unittest/trt/model/redrafter/test_beams2tree.py", - "tests/unittest/trt/model/redrafter/test_draft_token.py", - "tests/unittest/trt/model/redrafter/test_draft_token_indices.py", - "tests/unittest/trt/model/redrafter/test_gather_beams.py", - "tests/unittest/trt/model/redrafter/test_mask.py", - "tests/unittest/trt/model/redrafter/test_packed_position_ids.py", - "tests/unittest/trt/model/redrafter/test_prefix_match_indices.py", - "tests/unittest/trt/model/redrafter/test_prepare_input.py", - "tests/unittest/trt/model/redrafter/test_process_logits.py", - "tests/unittest/trt/model/redrafter/test_top1.py", - "tests/unittest/trt/model/redrafter/test_unpack_gen_data.py", - "tests/unittest/trt/model/redrafter/test_validate.py", - "tests/unittest/trt/model/test_gpt.py", - "tests/unittest/trt/model/test_gpt_e2e.py", - "tests/unittest/trt/model/test_llama.py", - "tests/unittest/trt/model/test_mamba.py", - "tests/unittest/trt/model/test_mistral.py", - "tests/unittest/trt/model/test_nemotron_nas.py", - "tests/unittest/trt/model/test_phi.py", - "tests/unittest/trt/model/test_unet.py", - "tests/unittest/trt/model_api/test_model_api_multi_gpu.py", - "tests/unittest/trt/model_api/test_model_level_api.py", - "tests/unittest/trt/model_api/test_model_quantization.py", - "tests/unittest/trt/python_plugin/plugin_wrapper_utils.py", - "tests/unittest/trt/python_plugin/test_plugin_wrapper.py", - "tests/unittest/trt/quantization/__init__.py", - "tests/unittest/trt/quantization/_utils.py", - "tests/unittest/trt/quantization/test_fp8_quantization.py", - "tests/unittest/trt/quantization/test_fp8_rowwise_gemm.py", - "tests/unittest/trt/quantization/test_functional.py", - "tests/unittest/trt/quantization/test_mode.py", - "tests/unittest/trt/quantization/test_moe_weight_only_quant_matmul.py", - "tests/unittest/trt/quantization/test_qserve_gemm.py", - "tests/unittest/trt/quantization/test_quant.py", - "tests/unittest/trt/quantization/test_quant_layer.py", - "tests/unittest/trt/quantization/test_smooth_quant_gemm.py", - "tests/unittest/trt/quantization/test_smooth_quant_layer_norm.py", - "tests/unittest/trt/quantization/test_smooth_quant_rms_norm.py", - "tests/unittest/trt/quantization/test_weight_only_groupwise_quant_matmul.py", - "tests/unittest/trt/quantization/test_weight_only_quant_matmul.py", "tests/unittest/utils/__init__.py", "tests/unittest/utils/cpp_paths.py", "tests/unittest/utils/llm_data.py", "tests/unittest/utils/runtime_defaults.py", - "tests/unittest/utils/test_medusa_utils.py", "tests/unittest/utils/test_prebuilt_whl_cpp_extensions.py", "tests/unittest/utils/test_util.py", "tests/unittest/utils/torch_ref.py", diff --git a/setup.py b/setup.py index 97ded73d1f8a..e71819c66fb9 100644 --- a/setup.py +++ b/setup.py @@ -181,8 +181,6 @@ def has_ext_modules(self): package_data += [ 'bindings/*.pyi', 'bindings/**/*.pyi', - 'tools/plugin_gen/templates/*', - 'bench/build/benchmark_config.yml', 'evaluate/lm_eval_tasks/**/*', "_torch/auto_deploy/config/*.yaml", # Include CUDA source for fused MoE align extension so runtime JIT can find it in wheels diff --git a/tensorrt_llm/_deprecation.py b/tensorrt_llm/_deprecation.py deleted file mode 100644 index 1d436040ef9a..000000000000 --- a/tensorrt_llm/_deprecation.py +++ /dev/null @@ -1,64 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Legacy-workflow warnings for the TensorRT engine-build path. - -The TensorRT engine-build workflow (convert_checkpoint.py -> trtllm-build -> -run.py) is a legacy path. The PyTorch backend (trtllm-serve / LLM API) is the -recommended approach for new projects. - -This module provides a shared warning function that can be called from legacy -scripts and internal chokepoints to inform users about the recommended -migration path. -""" - -import warnings - -_DEPRECATION_DOCS_URL = "https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html" - -_warned: set = set() - - -def emit_engine_arch_deprecation(caller_name: str) -> None: - """Emit a one-time FutureWarning for legacy engine-architecture usage. - - Each unique *caller_name* triggers the warning at most once per process, - so hot paths like ``builder.build()`` don't spam the console. - - Args: - caller_name: Human-readable identifier for the caller - (e.g., ``"convert_checkpoint.py"``, ``"trtllm-build"``, - ``"builder.build()"``). - """ - if caller_name in _warned: - return - _warned.add(caller_name) - - warnings.warn( - f"\n{'=' * 70}\n" - f"LEGACY WARNING: {caller_name}\n" - f"{'=' * 70}\n" - f"This is part of the legacy TensorRT engine-build workflow.\n" - f"New projects should use the PyTorch backend instead.\n\n" - f" # Serve a model (recommended):\n" - f" trtllm-serve \n\n" - f" # Python API:\n" - f" from tensorrt_llm import LLM\n" - f" llm = LLM(model='')\n" - f" output = llm.generate(['Hello, how are you?'])\n\n" - f"Documentation: {_DEPRECATION_DOCS_URL}\n" - f"{'=' * 70}\n", - FutureWarning, - stacklevel=2, - ) diff --git a/tensorrt_llm/bench/benchmark/__init__.py b/tensorrt_llm/bench/benchmark/__init__.py index 8386e115f505..7dcf3cc9bbed 100644 --- a/tensorrt_llm/bench/benchmark/__init__.py +++ b/tensorrt_llm/bench/benchmark/__init__.py @@ -6,9 +6,9 @@ from tensorrt_llm import LLM as PyTorchLLM from tensorrt_llm.bench.benchmark.utils.processes import IterationWriter -from tensorrt_llm.bench.build.build import get_model_config from tensorrt_llm.bench.dataclasses.configuration import RuntimeConfig from tensorrt_llm.bench.dataclasses.general import BenchmarkEnvironment +from tensorrt_llm.bench.tuning.settings import get_model_config from tensorrt_llm.commands.utils import \ collect_explicit_cli_keys as _collect_explicit_cli_keys from tensorrt_llm.logger import logger diff --git a/tensorrt_llm/bench/benchmark/utils/general.py b/tensorrt_llm/bench/benchmark/utils/general.py index e28121b18aba..ba3c312bd45a 100755 --- a/tensorrt_llm/bench/benchmark/utils/general.py +++ b/tensorrt_llm/bench/benchmark/utils/general.py @@ -9,12 +9,12 @@ from tensorrt_llm._torch.pyexecutor.model_loader import \ validate_and_set_kv_cache_quant -from tensorrt_llm.bench.build.build import (get_benchmark_engine_settings, - get_model_config) -from tensorrt_llm.bench.build.dataclasses import (NemotronHybridConfig, - Qwen3HybridConfig) from tensorrt_llm.bench.dataclasses.general import (DatasetMetadata, InferenceRequest) +from tensorrt_llm.bench.tuning.dataclasses import (NemotronHybridConfig, + Qwen3HybridConfig) +from tensorrt_llm.bench.tuning.settings import (get_benchmark_engine_settings, + get_model_config) from tensorrt_llm.logger import logger from tensorrt_llm.quantization.mode import QuantAlgo diff --git a/tensorrt_llm/bench/build/__init__.py b/tensorrt_llm/bench/build/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/tensorrt_llm/bench/build/utils.py b/tensorrt_llm/bench/build/utils.py deleted file mode 100644 index 18d0bac8f022..000000000000 --- a/tensorrt_llm/bench/build/utils.py +++ /dev/null @@ -1,34 +0,0 @@ -import pynvml - -DEFAULT_HF_MODEL_DIRS = { - 'BaichuanForCausalLM': 'baichuan-inc/Baichuan-13B-Chat', - 'BloomForCausalLM': 'bigscience/bloom-560m', - 'GLMModel': 'THUDM/glm-10b', - 'ChatGLMModel': 'THUDM/chatglm3-6b', - 'ChatGLMForCausalLM': 'THUDM/chatglm3-6b', - 'FalconForCausalLM': 'tiiuae/falcon-rw-1b', - 'GPTForCausalLM': 'gpt2-medium', - 'GPTJForCausalLM': 'EleutherAI/gpt-j-6b', - 'GPTNeoXForCausalLM': 'EleutherAI/gpt-neox-20b', - 'InternLMForCausalLM': 'internlm/internlm-chat-7b', - 'InternLM2ForCausalLM': 'internlm/internlm2-chat-7b', - 'LlamaForCausalLM': 'meta-llama/Llama-2-7b-hf', - 'MPTForCausalLM': 'mosaicml/mpt-7b', - 'PhiForCausalLM': 'microsoft/phi-2', - 'OPTForCausalLM': 'facebook/opt-350m', - 'QWenLMHeadModel': 'Qwen/Qwen-7B', - 'QWenForCausalLM': 'Qwen/Qwen-7B', - 'Qwen2ForCausalLM': 'Qwen/Qwen1.5-7B', - 'Qwen2MoeForCausalLM': 'Qwen/Qwen1.5-MoE-A2.7B', - 'RecurrentGemmaForCausalLM': 'google/recurrentgemma-2b', -} - - -def get_device_memory(): - pynvml.nvmlInit() - handle = pynvml.nvmlDeviceGetHandleByIndex(0) - mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle) - total_memory = mem_info.total / (1024**3) - pynvml.nvmlShutdown() - - return total_memory diff --git a/tensorrt_llm/bench/dataset/prepare_dataset.py b/tensorrt_llm/bench/dataset/prepare_dataset.py index aa7f4eb722e9..3f286a998553 100644 --- a/tensorrt_llm/bench/dataset/prepare_dataset.py +++ b/tensorrt_llm/bench/dataset/prepare_dataset.py @@ -25,7 +25,7 @@ class RootArgs(BaseModel): tokenizer: str - output: str + output: Optional[str] random_seed: int task_id: int trust_remote_code: bool = False @@ -54,6 +54,12 @@ def validate_tokenizer(self): @click.option( "--output", type=str, help="Output json filename.", default="preprocessed_dataset.json" ) +@click.option( + "--stdout", + is_flag=True, + default=False, + help="Print the dataset to stdout with a JSON entry on each line instead of writing a file.", +) @click.option( "--random-seed", required=False, type=int, help="random seed for token_ids", default=420 ) @@ -74,12 +80,15 @@ def validate_tokenizer(self): def prepare_dataset(ctx, **kwargs): """Prepare dataset for benchmarking with trtllm-bench.""" model = ctx.obj.model or ctx.obj.checkpoint_path - output_path = Path(kwargs["output"]) - output_path.parent.mkdir(parents=True, exist_ok=True) + # --stdout is encoded as a null output path (mutually exclusive with a file). + output = None if kwargs["stdout"] else kwargs["output"] + if output is not None: + output_path = Path(output) + output_path.parent.mkdir(parents=True, exist_ok=True) ctx.obj = RootArgs( tokenizer=model, - output=kwargs["output"], + output=output, random_seed=kwargs["random_seed"], task_id=kwargs["task_id"], rand_task_id=kwargs["rand_task_id"], diff --git a/tensorrt_llm/bench/dataset/utils.py b/tensorrt_llm/bench/dataset/utils.py index 2ecea0320a5b..39d507419207 100644 --- a/tensorrt_llm/bench/dataset/utils.py +++ b/tensorrt_llm/bench/dataset/utils.py @@ -103,6 +103,10 @@ def get_sample_from_population(population_range, sample_size): def write_dataset_to_file(dataset_generator, output_file): + if output_file is None: + for item in dataset_generator: + print(item) + return output_file = Path(output_file) os.makedirs(output_file.parent, exist_ok=True) with open(output_file, "w") as f: diff --git a/benchmarks/__init__.py b/tensorrt_llm/bench/tuning/__init__.py similarity index 100% rename from benchmarks/__init__.py rename to tensorrt_llm/bench/tuning/__init__.py diff --git a/tensorrt_llm/bench/build/dataclasses.py b/tensorrt_llm/bench/tuning/dataclasses.py similarity index 65% rename from tensorrt_llm/bench/build/dataclasses.py rename to tensorrt_llm/bench/tuning/dataclasses.py index 37d63062f23a..ca7e58a4fb51 100755 --- a/tensorrt_llm/bench/build/dataclasses.py +++ b/tensorrt_llm/bench/tuning/dataclasses.py @@ -1,40 +1,42 @@ -from typing import Optional, Literal -from pydantic import AliasPath, BaseModel, Field, AliasChoices, model_validator +import json +import os +import struct +from typing import Literal, Optional + import huggingface_hub from huggingface_hub.constants import ( SAFETENSORS_INDEX_FILE, SAFETENSORS_MAX_HEADER_LENGTH, SAFETENSORS_SINGLE_FILE, ) -from huggingface_hub.utils import SafetensorsRepoMetadata, SafetensorsFileMetadata, TensorInfo +from huggingface_hub.utils import SafetensorsFileMetadata, SafetensorsRepoMetadata, TensorInfo from huggingface_hub.utils import tqdm as hf_tqdm +from pydantic import AliasChoices, AliasPath, BaseModel, Field, model_validator from tqdm.contrib.concurrent import thread_map -import os -import json -import struct from tensorrt_llm._torch.pyexecutor.config_utils import ( - load_pretrained_config, get_qwen3_hybrid_layer_types) + get_qwen3_hybrid_layer_types, + load_pretrained_config, +) # Mapping from safetensors dtype strings to bytes per element. # Used to compute checkpoint size from per-dtype element counts. SAFETENSORS_DTYPE_BYTES = { - 'F64': 8, - 'F32': 4, - 'F16': 2, - 'BF16': 2, - 'I64': 8, - 'I32': 4, - 'I16': 2, - 'I8': 1, - 'U8': 1, - 'BOOL': 1, - 'F8_E4M3': 1, + "F64": 8, + "F32": 4, + "F16": 2, + "BF16": 2, + "I64": 8, + "I32": 4, + "I16": 2, + "I8": 1, + "U8": 1, + "BOOL": 1, + "F8_E4M3": 1, } def parse_safetensors_file_metadata(model_path, filename): - with open(os.path.join(model_path, filename), "rb") as f: metadata_size = f.read(8) metadata_size = struct.unpack(" None: files_metadata[filename] = parse_safetensors_file_metadata( - model_path=model_name_or_path, filename=filename) + model_path=model_name_or_path, filename=filename + ) thread_map( _parse, @@ -123,7 +124,8 @@ def _parse(filename: str) -> None: else: # Not a safetensors repo raise RuntimeError( - f"'{model_name_or_path}' is not a safetensors repo. Couldn't find '{SAFETENSORS_INDEX_FILE}' or '{SAFETENSORS_SINGLE_FILE}' files." + f"'{model_name_or_path}' is not a safetensors repo. Couldn't find " + f"'{SAFETENSORS_INDEX_FILE}' or '{SAFETENSORS_SINGLE_FILE}' files." ) else: return huggingface_hub.get_safetensors_metadata(model_name_or_path) @@ -134,23 +136,28 @@ class ModelConfig(BaseModel): The parameters are needed in engine setting calculation. """ + name: str model_type: str param_count: int checkpoint_size_in_gb: float = Field(default=0.0) - num_hidden_layers: int = Field(validation_alias=AliasChoices( - "num_hidden_layers", - "n_layer", - AliasPath("text_config", "num_hidden_layers"), - AliasPath("language_config", "num_hidden_layers"), - )) + num_hidden_layers: int = Field( + validation_alias=AliasChoices( + "num_hidden_layers", + "n_layer", + AliasPath("text_config", "num_hidden_layers"), + AliasPath("language_config", "num_hidden_layers"), + ) + ) num_attention_layers: Optional[int] = Field(default=None) - num_attention_heads: int = Field(validation_alias=AliasChoices( - "num_attention_heads", - "n_head", - AliasPath("text_config", "num_attention_heads"), - AliasPath("language_config", "num_attention_heads"), - )) + num_attention_heads: int = Field( + validation_alias=AliasChoices( + "num_attention_heads", + "n_head", + AliasPath("text_config", "num_attention_heads"), + AliasPath("language_config", "num_attention_heads"), + ) + ) num_key_value_heads: Optional[int] = Field( default=None, validation_alias=AliasChoices( @@ -160,33 +167,37 @@ class ModelConfig(BaseModel): AliasPath("language_config", "num_key_value_heads"), ), ) - hidden_size: int = Field(validation_alias=AliasChoices( - "hidden_size", - "n_embd", - AliasPath("text_config", "hidden_size"), - )) - head_size: Optional[int] = Field(default=None, - validation_alias=AliasChoices( - "head_size", - "head_dim", - "attention_head_dim", - AliasPath("text_config", "head_dim"), - )) + hidden_size: int = Field( + validation_alias=AliasChoices( + "hidden_size", + "n_embd", + AliasPath("text_config", "hidden_size"), + ) + ) + head_size: Optional[int] = Field( + default=None, + validation_alias=AliasChoices( + "head_size", + "head_dim", + "attention_head_dim", + AliasPath("text_config", "head_dim"), + ), + ) max_position_embeddings: Optional[int] = Field( default=None, validation_alias=AliasChoices( "max_position_embeddings", "n_positions", AliasPath("text_config", "max_position_embeddings"), - )) - dtype: Literal["float16", "bfloat16", "float32", - None] = Field(default="float16", - validation_alias=AliasChoices( - "dtype", "torch_dtype")) + ), + ) + dtype: Literal["float16", "bfloat16", "float32", None] = Field( + default="float16", validation_alias=AliasChoices("dtype", "torch_dtype") + ) @model_validator(mode="after") def set_values_if_none(self): - """ Set the values if cannot get values from HF config.json. """ + """Set the values if cannot get values from HF config.json.""" if not self.dtype: # for GPT-J self.dtype = "float16" if self.num_key_value_heads is None: @@ -222,28 +233,32 @@ def get_param_count_and_checkpoint_size(cls, model_hf_name, hf_model_path): # BF16 for non-quantized layers, F8_E4M3 for scales, etc.). checkpoint_size_in_bytes = sum( count * SAFETENSORS_DTYPE_BYTES.get(dtype, 1) - for dtype, count in metadata.parameter_count.items()) + for dtype, count in metadata.parameter_count.items() + ) checkpoint_size_in_gb = checkpoint_size_in_bytes / (1024**3) if not param_count: - raise ValueError(f"Can't get valid parameter count for model: " - f"{hf_model_path or model_hf_name}.") + raise ValueError( + f"Can't get valid parameter count for model: {hf_model_path or model_hf_name}." + ) return param_count, checkpoint_size_in_gb @classmethod def from_hf(cls, model_hf_name, hf_model_path): - pretrained_config = load_pretrained_config(hf_model_path - or model_hf_name, - trust_remote_code=True) + pretrained_config = load_pretrained_config( + hf_model_path or model_hf_name, trust_remote_code=True + ) hf_config = pretrained_config.to_dict() - param_count, checkpoint_size_in_gb = ( - cls.get_param_count_and_checkpoint_size(model_hf_name, - hf_model_path)) + param_count, checkpoint_size_in_gb = cls.get_param_count_and_checkpoint_size( + model_hf_name, hf_model_path + ) - return cls(name=model_hf_name, - param_count=param_count, - checkpoint_size_in_gb=checkpoint_size_in_gb, - **hf_config) + return cls( + name=model_hf_name, + param_count=param_count, + checkpoint_size_in_gb=checkpoint_size_in_gb, + **hf_config, + ) def extra_model_cache_in_gb(self, bytes_per_elem, target_seq_len=None): return 0 @@ -253,27 +268,35 @@ def cache_memory_fraction(self, cache_memory_fraction): class NemotronHybridConfig(ModelConfig): - hybrid_override_pattern: str = Field(validation_alias=AliasChoices( - "hybrid_override_pattern", - AliasPath("text_config", "hybrid_override_pattern"), - AliasPath("language_config", "hybrid_override_pattern"), - )) - num_hidden_layers: int = Field(validation_alias=AliasChoices( - "num_hidden_layers", - "n_layer", - AliasPath("text_config", "num_hidden_layers"), - AliasPath("language_config", "num_hidden_layers"), - )) - d_state: int = Field(validation_alias=AliasChoices( - "d_state", - "mamba_d_state", - "ssm_state_size", - )) - d_conv: int = Field(validation_alias=AliasChoices( - "d_conv", - "mamba_d_conv", - "conv_kernel", - )) + hybrid_override_pattern: str = Field( + validation_alias=AliasChoices( + "hybrid_override_pattern", + AliasPath("text_config", "hybrid_override_pattern"), + AliasPath("language_config", "hybrid_override_pattern"), + ) + ) + num_hidden_layers: int = Field( + validation_alias=AliasChoices( + "num_hidden_layers", + "n_layer", + AliasPath("text_config", "num_hidden_layers"), + AliasPath("language_config", "num_hidden_layers"), + ) + ) + d_state: int = Field( + validation_alias=AliasChoices( + "d_state", + "mamba_d_state", + "ssm_state_size", + ) + ) + d_conv: int = Field( + validation_alias=AliasChoices( + "d_conv", + "mamba_d_conv", + "conv_kernel", + ) + ) mamba_num_heads: int n_groups: int mamba_head_dim: int @@ -283,7 +306,7 @@ class NemotronHybridConfig(ModelConfig): @model_validator(mode="after") def set_values_if_none(self): - """ Set the values if cannot get values from HF config.json. """ + """Set the values if cannot get values from HF config.json.""" if not self.d_inner: self.d_inner = self.mamba_num_heads * self.mamba_head_dim if self.num_mamba_layers is None: @@ -298,12 +321,17 @@ def extra_model_cache_in_gb(self, bytes_per_elem, target_seq_len=None): conv_dim = self.d_inner + 2 * self.n_groups * self.d_state conv_state_elems = conv_dim * (self.d_conv - 1) ssm_state_elems = self.mamba_num_heads * self.mamba_head_dim * self.d_state - gb_per_mamba_cache = bytes_per_elem * self.num_mamba_layers * ( - conv_state_elems + ssm_state_elems) / (1024**3) + gb_per_mamba_cache = ( + bytes_per_elem + * self.num_mamba_layers + * (conv_state_elems + ssm_state_elems) + / (1024**3) + ) return gb_per_mamba_cache def cache_memory_fraction(self, cache_memory_fraction): - # Each mamba cache entry is pretty large (~50MB for 8B model), so we are more conservative when estimating the max batch size + # Each mamba cache entry is pretty large (~50MB for 8B model), so we are + # more conservative when estimating the max batch size return cache_memory_fraction**2 def set_mamba_ssm_cache_dtype(self, mamba_ssm_cache_dtype: str): @@ -311,13 +339,13 @@ def set_mamba_ssm_cache_dtype(self, mamba_ssm_cache_dtype: str): @classmethod def from_hf(cls, model_hf_name, hf_model_path): - pretrained_config = load_pretrained_config(hf_model_path - or model_hf_name, - trust_remote_code=True) + pretrained_config = load_pretrained_config( + hf_model_path or model_hf_name, trust_remote_code=True + ) hf_config = pretrained_config.to_dict() - param_count, checkpoint_size_in_gb = ( - cls.get_param_count_and_checkpoint_size(model_hf_name, - hf_model_path)) + param_count, checkpoint_size_in_gb = cls.get_param_count_and_checkpoint_size( + model_hf_name, hf_model_path + ) # HuggingFace PretrainedConfig.to_dict() only serializes attributes known to # the base class; custom configs (e.g. NemotronHConfig) have num_hidden_layers @@ -326,8 +354,8 @@ def from_hf(cls, model_hf_name, hf_model_path): text_config = getattr(pretrained_config, "text_config", None) language_config = getattr(pretrained_config, "language_config", None) for key in ( - "num_hidden_layers", - "hybrid_override_pattern", + "num_hidden_layers", + "hybrid_override_pattern", ): if hf_config.get(key) is None: value = None @@ -342,10 +370,12 @@ def from_hf(cls, model_hf_name, hf_model_path): if value is not None: hf_config[key] = value - return cls(name=model_hf_name, - param_count=param_count, - checkpoint_size_in_gb=checkpoint_size_in_gb, - **hf_config) + return cls( + name=model_hf_name, + param_count=param_count, + checkpoint_size_in_gb=checkpoint_size_in_gb, + **hf_config, + ) class Qwen3HybridConfig(ModelConfig): @@ -354,6 +384,7 @@ class Qwen3HybridConfig(ModelConfig): Maps Qwen3.5 linear-attention parameters to the same cache estimation formulas used by NemotronHybridConfig. """ + linear_key_head_dim: int # d_state linear_conv_kernel_dim: int # d_conv linear_num_value_heads: int # num_heads (mamba_num_heads) @@ -364,34 +395,38 @@ class Qwen3HybridConfig(ModelConfig): @classmethod def from_hf(cls, model_hf_name, hf_model_path): - pretrained_config = load_pretrained_config(hf_model_path - or model_hf_name, - trust_remote_code=True) + pretrained_config = load_pretrained_config( + hf_model_path or model_hf_name, trust_remote_code=True + ) hf_config = pretrained_config.to_dict() - param_count, checkpoint_size_in_gb = ( - cls.get_param_count_and_checkpoint_size(model_hf_name, - hf_model_path)) + param_count, checkpoint_size_in_gb = cls.get_param_count_and_checkpoint_size( + model_hf_name, hf_model_path + ) layer_types = get_qwen3_hybrid_layer_types(pretrained_config) - hf_config.setdefault("num_attention_layers", - layer_types.count("full_attention")) - hf_config.setdefault("num_linear_attention_layers", - layer_types.count("linear_attention")) - - return cls(name=model_hf_name, - param_count=param_count, - checkpoint_size_in_gb=checkpoint_size_in_gb, - **hf_config) + hf_config.setdefault("num_attention_layers", layer_types.count("full_attention")) + hf_config.setdefault("num_linear_attention_layers", layer_types.count("linear_attention")) + + return cls( + name=model_hf_name, + param_count=param_count, + checkpoint_size_in_gb=checkpoint_size_in_gb, + **hf_config, + ) def extra_model_cache_in_gb(self, bytes_per_elem, target_seq_len=None): d_inner = self.linear_value_head_dim * self.linear_num_value_heads conv_dim = d_inner + 2 * self.linear_num_key_heads * self.linear_key_head_dim conv_state_elems = conv_dim * (self.linear_conv_kernel_dim - 1) - ssm_state_elems = (self.linear_num_value_heads * - self.linear_value_head_dim * - self.linear_key_head_dim) - gb_per_cache = bytes_per_elem * self.num_linear_attention_layers * ( - conv_state_elems + ssm_state_elems) / (1024**3) + ssm_state_elems = ( + self.linear_num_value_heads * self.linear_value_head_dim * self.linear_key_head_dim + ) + gb_per_cache = ( + bytes_per_elem + * self.num_linear_attention_layers + * (conv_state_elems + ssm_state_elems) + / (1024**3) + ) return gb_per_cache def cache_memory_fraction(self, cache_memory_fraction): diff --git a/tensorrt_llm/bench/build/tuning.py b/tensorrt_llm/bench/tuning/heuristics.py similarity index 79% rename from tensorrt_llm/bench/build/tuning.py rename to tensorrt_llm/bench/tuning/heuristics.py index d77cf6591dbb..0b41fee9767f 100755 --- a/tensorrt_llm/bench/build/tuning.py +++ b/tensorrt_llm/bench/tuning/heuristics.py @@ -1,20 +1,25 @@ +import math from typing import Tuple import torch from tensorrt_llm._utils import str_dtype_to_torch +from tensorrt_llm.bench.tuning.dataclasses import ( + ModelConfig, + NemotronHybridConfig, + Qwen3HybridConfig, +) from tensorrt_llm.llmapi.llm_utils import QuantConfig from tensorrt_llm.logger import logger from tensorrt_llm.quantization.mode import QuantAlgo -from tensorrt_llm.bench.build.dataclasses import ModelConfig, NemotronHybridConfig, Qwen3HybridConfig + from .utils import get_device_memory -import math BYTES_PER_ELEM = { QuantAlgo.NO_QUANT: 2.0, QuantAlgo.FP8: 1.0, QuantAlgo.FP8_BLOCK_SCALES: 1.0, - QuantAlgo.NVFP4: .5, + QuantAlgo.NVFP4: 0.5, } @@ -28,13 +33,14 @@ def calc_engine_setting( kv_cache_gpu_mem_fraction: float = 0.95, enable_attention_dp: bool = False, ) -> Tuple[int, int]: - """ Calculate the engine build settings (max batch size and max num tokens) - for a specific model + parallelism mapping + dataset configuration. - trtllm-bench sets a slightly optimistic upper bound for max batch size - and max num tokens to avoid over-allocation of memory in activation, - runtime, and decoder buffers. In runtime, TRT-LLM relies on its runtime - tuning features to adjust the runtime max batch size according to - incoming traffic. + """Calculate the engine build settings (max batch size and max num tokens). + + For a specific model + parallelism mapping + dataset configuration, + trtllm-bench sets a slightly optimistic upper bound for max batch size + and max num tokens to avoid over-allocation of memory in activation, + runtime, and decoder buffers. In runtime, TRT-LLM relies on its runtime + tuning features to adjust the runtime max batch size according to + incoming traffic. Args: model_config (ModelConfig): Model specific configurations. @@ -63,11 +69,16 @@ def calc_engine_setting( # Each GPU in TP group has at least 1 kv head adjusted_num_kv_heads = max(tp_size, model_config.num_key_value_heads) - logger.info( - f"Number of attention layers: {model_config.num_attention_layers}") + logger.info(f"Number of attention layers: {model_config.num_attention_layers}") - gb_per_token = 2 * model_config.num_attention_layers * adjusted_num_kv_heads \ - * model_config.head_size * byte_per_kv_elem / (1024 ** 3) + gb_per_token = ( + 2 + * model_config.num_attention_layers + * adjusted_num_kv_heads + * model_config.head_size + * byte_per_kv_elem + / (1024**3) + ) # Number of GPU used for this run. n_gpus = tp_size * pp_size @@ -92,13 +103,11 @@ def calc_engine_setting( # Available memory to allocate KV cache. available_memory = total_gpu_memory - engine_size logger.info(f"Estimated engine size: {engine_size:.2f} GB") - logger.info("Estimated total available memory for KV cache: " - f"{available_memory:.2f} GB") + logger.info(f"Estimated total available memory for KV cache: {available_memory:.2f} GB") # Calculate max requests in KV cache based on target ISL and OSL. target_seq_len = target_input_len + target_output_len - cache_memory = available_memory * model_config.cache_memory_fraction( - kv_cache_gpu_mem_fraction) + cache_memory = available_memory * model_config.cache_memory_fraction(kv_cache_gpu_mem_fraction) bytes_per_elem = BYTES_PER_ELEM.get(QuantAlgo.NO_QUANT) if isinstance(model_config, (NemotronHybridConfig, Qwen3HybridConfig)): @@ -107,30 +116,31 @@ def calc_engine_setting( if str_dtype_to_torch(mamba_ssm_cache_dtype) == torch.float32: bytes_per_elem = 4.0 - gb_per_extra_cache = model_config.extra_model_cache_in_gb( - bytes_per_elem, target_seq_len) - kv_cache_max_requests = cache_memory / (gb_per_token * target_seq_len + - gb_per_extra_cache) + gb_per_extra_cache = model_config.extra_model_cache_in_gb(bytes_per_elem, target_seq_len) + kv_cache_max_requests = cache_memory / (gb_per_token * target_seq_len + gb_per_extra_cache) extra_cache_memory = gb_per_extra_cache * kv_cache_max_requests kv_cache_memory = cache_memory - extra_cache_memory kv_cache_max_tokens = kv_cache_memory / gb_per_token logger.info( - f"Estimated total cache memory: {cache_memory:.2f} GB. KV cache: {kv_cache_memory:.2f} GB, Extra cache: {extra_cache_memory:.2f} GB" + f"Estimated total cache memory: {cache_memory:.2f} GB. " + f"KV cache: {kv_cache_memory:.2f} GB, Extra cache: {extra_cache_memory:.2f} GB" ) logger.info(f"Estimated kv cache max tokens: {kv_cache_max_tokens:.2f}") - logger.info("Estimated max number of requests in KV cache memory: " - f"{kv_cache_max_requests:.2f}") + logger.info(f"Estimated max number of requests in KV cache memory: {kv_cache_max_requests:.2f}") # Fine-tune the max batch size and num token setting for performance. - # For mamba-attn hybrid models, we disable optimistic tuning because the mamba cache leaves less memory for the KV cache + # For mamba-attn hybrid models, we disable optimistic tuning because the + # mamba cache leaves less memory for the KV cache max_batch_size, max_num_tokens = finetune_setting( kv_cache_max_requests, target_input_len, target_output_len, pp_size, disable_optimistic_tuning=isinstance( - model_config, (NemotronHybridConfig, Qwen3HybridConfig))) + model_config, (NemotronHybridConfig, Qwen3HybridConfig) + ), + ) # Functional and performance if total_gpu_memory < engine_size: @@ -151,21 +161,26 @@ def calc_engine_setting( f"Number of Tensor Parallel Shards: {tp_size}\n" f"Number of Pipeline Parallel Stages: {pp_size}\n" f"KV Cache GPU Memory Fraction: {kv_cache_gpu_mem_fraction}\n" - "----------------------------------------------------------\n") + "----------------------------------------------------------\n" + ) if kv_cache_max_requests < 1: - raise RuntimeError("The amount of KV cache memory is insufficient to " - "run this model. Please try with more GPUs.") + raise RuntimeError( + "The amount of KV cache memory is insufficient to " + "run this model. Please try with more GPUs." + ) warning_gpu_count = pp_size if enable_attention_dp else n_gpus if cache_memory / warning_gpu_count < 10.0: logger.warning( - f"The KV cache memory per GPU is less than 10 GB. " + "The KV cache memory per GPU is less than 10 GB. " "Performance may be undesirable. Please consider using a different " - "mapping or more GPUs.") + "mapping or more GPUs." + ) if kv_cache_max_requests < 32: logger.warning( - f"The maximum number of requests in the KV cache is too " + "The maximum number of requests in the KV cache is too " "small. Performance may be undesirable. Please consider using more " - "GPUs or a different mapping to process more concurrent requests.") + "GPUs or a different mapping to process more concurrent requests." + ) return max_batch_size, max_num_tokens @@ -177,9 +192,10 @@ def finetune_setting( pp_size: int, disable_optimistic_tuning: bool = False, ) -> Tuple[int, int]: - """ Calculate and fine-tune the engine build settings (max batch size and - max num tokens). Both max batch size and max num tokens are fine-tuned - to be slightly optimistic. + """Calculate and fine-tune the engine build settings. + + Both max batch size and max num tokens are fine-tuned to be slightly + optimistic. Args: kv_cache_max_requests (float): Max number of requests that can fits in @@ -219,11 +235,12 @@ def finetune_setting( else: max_token = 1024 * math.ceil(raw_token / 1024) - logger.debug(f"Estimated max batch size (before fine-tune): " - f"{kv_cache_max_requests / pp_size:.2f}") + logger.debug( + f"Estimated max batch size (before fine-tune): {kv_cache_max_requests / pp_size:.2f}" + ) logger.debug( f"Estimated max num tokens (before fine-tune): " - f"{kv_cache_max_requests / pp_size * (1 + input_len / output_len) :.2f}" + f"{kv_cache_max_requests / pp_size * (1 + input_len / output_len):.2f}" ) logger.info(f"Estimated max batch size (after fine-tune): {max_bs}") logger.info(f"Estimated max num tokens (after fine-tune): {max_token}") diff --git a/tensorrt_llm/bench/build/build.py b/tensorrt_llm/bench/tuning/settings.py similarity index 77% rename from tensorrt_llm/bench/build/build.py rename to tensorrt_llm/bench/tuning/settings.py index 16bf6d23d9aa..9f47292975ac 100644 --- a/tensorrt_llm/bench/build/build.py +++ b/tensorrt_llm/bench/tuning/settings.py @@ -3,21 +3,28 @@ from pathlib import Path from typing import Tuple -from tensorrt_llm._torch.pyexecutor.config_utils import (is_nemotron_hybrid, - is_qwen3_hybrid, - load_pretrained_config) -from tensorrt_llm.bench.build.dataclasses import (ModelConfig, - NemotronHybridConfig, - Qwen3HybridConfig) -from tensorrt_llm.bench.build.tuning import calc_engine_setting +from tensorrt_llm._torch.pyexecutor.config_utils import ( + is_nemotron_hybrid, + is_qwen3_hybrid, + load_pretrained_config, +) +from tensorrt_llm.bench.tuning.dataclasses import ( + ModelConfig, + NemotronHybridConfig, + Qwen3HybridConfig, +) +from tensorrt_llm.bench.tuning.heuristics import calc_engine_setting from tensorrt_llm.llmapi.llm_args import TorchLlmArgs from tensorrt_llm.llmapi.llm_utils import QuantConfig from tensorrt_llm.logger import logger from tensorrt_llm.quantization.mode import QuantAlgo TUNED_QUANTS = { - QuantAlgo.NVFP4, QuantAlgo.FP8, QuantAlgo.FP8_BLOCK_SCALES, - QuantAlgo.NO_QUANT, None + QuantAlgo.NVFP4, + QuantAlgo.FP8, + QuantAlgo.FP8_BLOCK_SCALES, + QuantAlgo.NO_QUANT, + None, } # Sourced from TorchLlmArgs so the bench defaults track the args-class field # defaults and can't drift. @@ -35,7 +42,7 @@ def get_benchmark_engine_settings( kv_cache_gpu_mem_fraction: float = 0.95, enable_attention_dp: bool = False, ) -> Tuple[int, int]: - """ Retrieve benchmark settings for a specific model + configuration. + """Retrieve benchmark settings for a specific model + configuration. Args: model_config (ModelConfig): Model specific configurations. @@ -77,13 +84,14 @@ def get_benchmark_engine_settings( ) if max_batch_size <= 0 or max_num_tokens <= 0: - raise RuntimeError(f"Unable to obtain correct settings for benchmark.") + raise RuntimeError("Unable to obtain correct settings for benchmark.") return max_batch_size, max_num_tokens def get_model_config(model_name: str, model_path: Path = None) -> ModelConfig: - """ Obtain the model-related parameters from Hugging Face. + """Obtain the model-related parameters from Hugging Face. + Args: model_name (str): Huggingface model name. model_path (Path): Path to a local Huggingface checkpoint. @@ -91,8 +99,7 @@ def get_model_config(model_name: str, model_path: Path = None) -> ModelConfig: Raises: ValueError: When model is not supported. """ - pretrained_config = load_pretrained_config(model_path or model_name, - trust_remote_code=True) + pretrained_config = load_pretrained_config(model_path or model_name, trust_remote_code=True) if is_nemotron_hybrid(pretrained_config): return NemotronHybridConfig.from_hf(model_name, model_path) if is_qwen3_hybrid(pretrained_config): diff --git a/tensorrt_llm/bench/tuning/utils.py b/tensorrt_llm/bench/tuning/utils.py new file mode 100644 index 000000000000..37d92e197f9b --- /dev/null +++ b/tensorrt_llm/bench/tuning/utils.py @@ -0,0 +1,34 @@ +import pynvml + +DEFAULT_HF_MODEL_DIRS = { + "BaichuanForCausalLM": "baichuan-inc/Baichuan-13B-Chat", + "BloomForCausalLM": "bigscience/bloom-560m", + "GLMModel": "THUDM/glm-10b", + "ChatGLMModel": "THUDM/chatglm3-6b", + "ChatGLMForCausalLM": "THUDM/chatglm3-6b", + "FalconForCausalLM": "tiiuae/falcon-rw-1b", + "GPTForCausalLM": "gpt2-medium", + "GPTJForCausalLM": "EleutherAI/gpt-j-6b", + "GPTNeoXForCausalLM": "EleutherAI/gpt-neox-20b", + "InternLMForCausalLM": "internlm/internlm-chat-7b", + "InternLM2ForCausalLM": "internlm/internlm2-chat-7b", + "LlamaForCausalLM": "meta-llama/Llama-2-7b-hf", + "MPTForCausalLM": "mosaicml/mpt-7b", + "PhiForCausalLM": "microsoft/phi-2", + "OPTForCausalLM": "facebook/opt-350m", + "QWenLMHeadModel": "Qwen/Qwen-7B", + "QWenForCausalLM": "Qwen/Qwen-7B", + "Qwen2ForCausalLM": "Qwen/Qwen1.5-7B", + "Qwen2MoeForCausalLM": "Qwen/Qwen1.5-MoE-A2.7B", + "RecurrentGemmaForCausalLM": "google/recurrentgemma-2b", +} + + +def get_device_memory(): + pynvml.nvmlInit() + handle = pynvml.nvmlDeviceGetHandleByIndex(0) + mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle) + total_memory = mem_info.total / (1024**3) + pynvml.nvmlShutdown() + + return total_memory diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index e4ec7ac25137..8aa94fcb114a 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -90,5 +90,4 @@ 'PrometheusMetricsConfig', 'ThinkingBudgetLogitsProcessor', 'add_thinking_budget_logits_processor', - 'DeepSeekV4SparseAttentionConfig', ] diff --git a/tensorrt_llm/models/convert_utils.py b/tensorrt_llm/models/convert_utils.py index 58da3ec36349..5930532202de 100644 --- a/tensorrt_llm/models/convert_utils.py +++ b/tensorrt_llm/models/convert_utils.py @@ -4,7 +4,6 @@ from typing import Dict, List, Optional, Tuple, Union import torch -from datasets import load_dataset from .._utils import torch_dtype_to_str from ..logger import logger @@ -292,39 +291,6 @@ def has_safetensors(model_dir: str): return len(list(Path(model_dir).glob('*.safetensors'))) > 0 -DEFAULT_HF_DATASET_META = { - 'ccdv/cnn_dailymail': ('3.0.0', 'train', 'article'), - 'cnn_dailymail': ('3.0.0', 'train', 'article'), - 'lambada': (None, 'validation', 'text'), - '': (None, 'train', 'text'), # Default value in HF -} - - -def load_calib_dataset(dataset_name_or_dir: str, - config_name: Optional[str] = None, - split: Optional[str] = None, - key: Optional[str] = None, - trust_remote_code=True, - **kwargs): - if config_name is None: - for name, meta in DEFAULT_HF_DATASET_META.items(): - if name in dataset_name_or_dir: - if config_name is None: - config_name = meta[0] - if split is None: - split = meta[1] - if key is None: - key = meta[2] - break - - dataset = load_dataset(dataset_name_or_dir, - name=config_name, - split=split, - trust_remote_code=trust_remote_code, - **kwargs) - return dataset[key] - - @torch.no_grad() def apply_smoothing( scales: torch.Tensor, diff --git a/tensorrt_llm/models/unet/pp/__init__.py b/tensorrt_llm/models/unet/pp/__init__.py deleted file mode 100755 index e69de29bb2d1..000000000000 diff --git a/tensorrt_llm/quantization/__init__.py b/tensorrt_llm/quantization/__init__.py index 0ecdd664e60d..900ad46bd0e8 100644 --- a/tensorrt_llm/quantization/__init__.py +++ b/tensorrt_llm/quantization/__init__.py @@ -16,11 +16,9 @@ from .mode import (KV_CACHE_QUANT_ALGO_LIST, MODELOPT_FLOW_QUANTIZATIONS, QUANT_ALGO_LIST, W8A8_SQ_PLUGIN_LIST, GroupwiseQuantAlgo, QuantAlgo, QuantMode) -from .quantize_by_modelopt import quantize_and_export, quantize_nemo_and_export __all__ = [ 'QUANT_ALGO_LIST', 'KV_CACHE_QUANT_ALGO_LIST', 'W8A8_SQ_PLUGIN_LIST', 'MODELOPT_FLOW_QUANTIZATIONS', 'QuantAlgo', 'QuantMode', - 'GroupwiseQuantAlgo', 'quantize_and_export', 'quantize_nemo_and_export', - 'utils' + 'GroupwiseQuantAlgo', 'utils' ] diff --git a/tensorrt_llm/quantization/image_processing.py b/tensorrt_llm/quantization/image_processing.py deleted file mode 100644 index 68360ff32ee0..000000000000 --- a/tensorrt_llm/quantization/image_processing.py +++ /dev/null @@ -1,97 +0,0 @@ -import torch - - -class BaseImageProcessor: - - def __init__(self, tokenizer, device='auto'): - self.tokenizer = tokenizer - self.device = device - - def __call__(self, **kwargs): - return self.tokenizer(**kwargs) - - def preprocess_function(self, examples): - raise NotImplementedError( - "Each image processor must implement its own preprocess method") - - def collate_function(self, examples): - raise NotImplementedError( - "Each image processor must implement its own colloate method") - - -# A light Encapsulation for Huggingface MllamaImageProcessor -class MllamaImageProcessor(BaseImageProcessor): - - def preprocess_function(self, examples): - # Prepare prompts in a generic chat format - if 'question' in examples: - question = examples['question'] - else: - question = "Describe this image." - - if examples['image'] is not None: - if self.tokenizer.chat_template is not None: - prompt = self.tokenizer.apply_chat_template( - [{ - "role": - "user", - "content": [{ - "type": "image" - }, { - "type": "text", - "text": question - }], - }], - add_generation_prompt=True, - ) - else: - prompt = f"<|image|><|begin_of_text|>{question}" - - # Process images using the processor's image processor - values = self.tokenizer(text=prompt, - images=examples['image'], - return_tensors="pt").to(self.device) - else: - if self.tokenizer.chat_template is not None: - prompt = self.tokenizer.apply_chat_template( - [{ - "role": "user", - "content": [{ - "type": "text", - "text": question - }], - }], - add_generation_prompt=True, - ) - else: - prompt = question - - values = self.tokenizer(text=prompt, - images=None, - return_tensors="pt").to(self.device) - - values['pixel_values'] = None - values['aspect_ratio_ids'] = None - values['aspect_ratio_mask'] = None - values['cross_attention_mask'] = None - - return values - - # Define a collate function to process images during data loading - def collate_function(self, batch): - batch[0]['input_ids'] = torch.LongTensor(batch[0]['input_ids']).to( - self.device) - batch[0]['attention_mask'] = torch.LongTensor( - batch[0]['attention_mask']).to(self.device) - - if batch[0]['pixel_values'] is not None: - batch[0]['pixel_values'] = torch.Tensor( - batch[0]['pixel_values']).to(self.device) - batch[0]['aspect_ratio_ids'] = torch.LongTensor( - batch[0]['aspect_ratio_ids']).to(self.device) - batch[0]['aspect_ratio_mask'] = torch.LongTensor( - batch[0]['aspect_ratio_mask']).to(self.device) - batch[0]['cross_attention_mask'] = torch.LongTensor( - batch[0]['cross_attention_mask']).to(self.device) - - return batch[0] diff --git a/tensorrt_llm/quantization/quantize_by_modelopt.py b/tensorrt_llm/quantization/quantize_by_modelopt.py deleted file mode 100755 index be9236aa6cef..000000000000 --- a/tensorrt_llm/quantization/quantize_by_modelopt.py +++ /dev/null @@ -1,1332 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -""" -Adapted from examples/quantization/hf_ptq.py -""" - -import copy -import json -import os -import random -import sys -import time -from importlib.metadata import version - -import numpy as np -import torch -from accelerate.hooks import remove_hook_from_module -from datasets import load_dataset -from modelopt.torch.utils import print_rank_0 -from torch import nn -from torch.utils.data import DataLoader -from transformers import (AutoConfig, AutoModelForCausalLM, AutoProcessor, - AutoTokenizer) - -from .._utils import get_hf_rope_theta, release_gc, str_dtype_to_torch -from ..logger import logger -from .image_processing import MllamaImageProcessor -from .mode import QuantAlgo - -EMPTY_CFG = { - "quant_cfg": { - "*weight_quantizer": { - "enable": False, - }, - "*input_quantizer": { - "enable": False - }, - "*lm_head*": { - "enable": False - }, - "*output_layer*": { - "enable": False - }, - "default": { - "enable": False - }, - }, - "algorithm": "max", -} - -KV_CACHE_CFG = { - "*.query_key_value.output_quantizer": { - "num_bits": 8, - "axis": None, - "enable": True - }, - "*.Wqkv.output_quantizer": { - "num_bits": 8, - "axis": None, - "enable": True - }, - "*.W_pack.output_quantizer": { - "num_bits": 8, - "axis": None, - "enable": True - }, - "*.c_attn.output_quantizer": { - "num_bits": 8, - "axis": None, - "enable": True - }, - "*.k_proj.output_quantizer": { - "num_bits": 8, - "axis": None, - "enable": True - }, - "*.v_proj.output_quantizer": { - "num_bits": 8, - "axis": None, - "enable": True - }, - "*.k.output_quantizer": { - "num_bits": 8, - "axis": None, - "enable": True - }, - "*.v.output_quantizer": { - "num_bits": 8, - "axis": None, - "enable": True - }, -} - -KV_QUANT_CFG_CHOICES = { - "fp8": "FP8_KV_CFG", - "nvfp4": "NVFP4_KV_CFG", -} - - -def quant_cfg_choices(): - import modelopt.torch.quantization as mtq - QUANT_CFG_CHOICES = { - "int8_sq": mtq.INT8_SMOOTHQUANT_CFG, - "fp8": mtq.FP8_DEFAULT_CFG, - "fp8_pc_pt": mtq.FP8_PER_CHANNEL_PER_TOKEN_CFG, - "int4_awq": mtq.INT4_AWQ_CFG, - "w4a8_awq": mtq.W4A8_AWQ_BETA_CFG, - "int8_wo": EMPTY_CFG, - "int4_wo": EMPTY_CFG, - "full_prec": EMPTY_CFG, - } - if hasattr(mtq, "NVFP4_DEFAULT_CFG"): - QUANT_CFG_CHOICES["nvfp4"] = mtq.NVFP4_DEFAULT_CFG - return QUANT_CFG_CHOICES - - -def model_type_is_enc_dec(model_type): - return model_type in ["t5", "bart"] - - -MODEL_NAME_PATTERN_MAP = { - "GPT2": "gpt2", - "Xverse": "llama", - "MllamaForConditionalGeneration": "mllama", - "Llama": "llama", - "MllamaForCausalLM": "mllama", - "Mistral": "llama", - "GPTJ": "gptj", - "FalconForCausalLM": "falcon", - "RWForCausalLM": "falcon", - "baichuan": "baichuan", - "MPT": "mpt", - "Bloom": "bloom", - "ChatGLM": "chatglm", - "QWen": "qwen", - "Qwen2VLForConditionalGeneration": "qwen2_vl", - "RecurrentGemma": "recurrentgemma", - "Gemma3": "gemma3", - "Gemma2": "gemma2", - "Gemma": "gemma", - "MixtralForCausalLM": "llama", - "NemotronForCausalLM": "nemotron", - "GPTBigCodeForCausalLM": "gpt_bigcode", - "ArcticForCausalLM": "llama", - "PhiMoEForCausalLM": "phi3", - "Phi3SmallForCausalLM": "phi3small", - "Phi3ForCausalLM": "phi3", - "Phi3VForCausalLM": "phi3", - "Starcoder2ForCausalLM": "gptnext", - "GPTBigCodeForCausalLM": "gptnext", - "GLM": "glm", - "Exaone": "exaone", - "DeciLMForCausalLM": "deci", - "DeepseekForCausalLM": "deepseek", - "GraniteForCausalLM": "granite", - "GraniteMoeForCausalLM": "granitemoe", - "T5": "t5", - "Bart": "bart" -} - -MULTIMODAL_DATASETS = ['scienceqa', 'science_qa'] - - -class _CustomDataset(torch.utils.data.Dataset): - - def __init__(self, encodings): - self.encodings = encodings - - def __getitem__(self, idx): - item = { - key: val[idx].clone().detach().requires_grad_(False) - for key, val in self.encodings.items() - } - return item - - def __len__(self): - return len(self.encodings["input_ids"]) - - -class EncDecModelWrapper(torch.nn.Module): - - def __init__(self, hf_model=None): - super().__init__() - self.hf_model = hf_model - self.model_type = get_model_type(hf_model) - - def forward(self, **kwargs): - self.hf_model.generate(**kwargs) - - def __getattr__(self, name): - try: - return super().__getattr__(name) - except AttributeError: - return getattr(self.hf_model, name) - - -def get_tokenizer(ckpt_path, max_seq_length=2048, model_type=None): - logger.info(f"Initializing tokenizer from {ckpt_path}") - tokenizer = AutoTokenizer.from_pretrained( - ckpt_path, - model_max_length=max_seq_length, - padding_side="left", - trust_remote_code=True, - ) - - if tokenizer.pad_token is None: - if model_type and model_type == "qwen": - # qwen use token id 151643 as pad and eos tokens - tokenizer.eos_token = tokenizer.convert_ids_to_tokens(151643) - tokenizer.pad_token = tokenizer.convert_ids_to_tokens(151643) - elif model_type and model_type == "qwen2_vl": - # qwen use token id 151643 as pad and 151643 and 151645 as eos tokens - tokenizer.eos_token = [ - tokenizer.convert_ids_to_tokens(151643), - tokenizer.convert_ids_to_tokens(151645) - ] - tokenizer.pad_token = tokenizer.convert_ids_to_tokens(151643) - else: - tokenizer.pad_token = tokenizer.eos_token - assert tokenizer.pad_token is not None, f"Pad token for {model_type} cannot be set!" - - return tokenizer - - -def get_processor(ckpt_path, max_seq_length=2048, model_type=None, device=None): - logger.info(f"Initializing tokenizer from {ckpt_path}") - processor = AutoProcessor.from_pretrained( - ckpt_path, - model_max_length=max_seq_length, - padding_side="left", - trust_remote_code=True, - ) - - if processor.tokenizer.pad_token is None: - if model_type and model_type == "qwen": - # qwen use token id 151643 as pad and eos tokens - processor.tokenizer.eos_token = processor.tokenizer.convert_ids_to_tokens( - 151643) - processor.tokenizer.pad_token = processor.tokenizer.convert_ids_to_tokens( - 151643) - else: - processor.tokenizer.pad_token = processor.tokenizer.eos_token - assert processor.tokenizer.pad_token is not None, f"Pad token for {model_type} cannot be set!" - - if model_type == 'mllama': - processor = MllamaImageProcessor(processor, device) - return processor - - -def _get_vila_model(model_dir): - sys.path.append(model_dir + "/../VILA") - from llava.model import LlavaLlamaConfig, LlavaLlamaModel # noqa - from transformers import AutoModel - model = AutoModel.from_pretrained( - model_dir, - device_map='auto', - trust_remote_code=True, - ) - return model.llm - - -def get_hf_config(ckpt_path): - if "mpt" in ckpt_path: - # MPT-7B cannot get initialized from AutoConfig - from transformers import MptConfig - return MptConfig.from_pretrained(ckpt_path) - else: - return AutoConfig.from_pretrained(ckpt_path, trust_remote_code=True) - - -class _MixtralBlockSparseTop2MLPCompat(nn.Module): - # Per-expert 4.x-style MLP using zero-copy parameter views into the - # transformers-5.x fused MoE tensors. Uses real nn.Linear modules so - # modelopt's quantizer wrapping (which keys off type name "Linear") - # produces per-expert weight scaling factors during calibration. - def __init__(self, hidden_dim, intermediate_dim, w1_view, w2_view, w3_view, - act_fn): - super().__init__() - self.w1 = nn.Linear(hidden_dim, intermediate_dim, bias=False) - self.w2 = nn.Linear(intermediate_dim, hidden_dim, bias=False) - self.w3 = nn.Linear(hidden_dim, intermediate_dim, bias=False) - self.w1.weight = nn.Parameter(w1_view, requires_grad=False) - self.w2.weight = nn.Parameter(w2_view, requires_grad=False) - self.w3.weight = nn.Parameter(w3_view, requires_grad=False) - self.act_fn = act_fn - - def forward(self, hidden_states): - return self.w2( - self.act_fn(self.w1(hidden_states)) * self.w3(hidden_states)) - - -class _MixtralSparseMoeBlockCompat(nn.Module): - # Drop-in replacement for transformers-5.x MixtralSparseMoeBlock that - # restores the pre-5.x ModuleList(MixtralBlockSparseTop2MLP) layout that - # nvidia-modelopt 0.37 iterates (len(experts), experts[i].{w1,w2,w3}.weight), - # while sharing weight storage with the 5.x fused tensors via parameter - # views (zero-copy). The class name contains "MixtralSparseMoeBlock" so - # modelopt's is_moe substring check still matches. - def __init__(self, mlp): - super().__init__() - experts = mlp.experts # MixtralExperts - self.top_k = mlp.top_k - self.jitter_noise = mlp.jitter_noise - self.num_experts = experts.num_experts - self.hidden_dim = experts.hidden_dim - self.intermediate_dim = experts.intermediate_dim - self.gate = mlp.gate # MixtralTopKRouter, returns (logits, scores, indices) - - gate_up = experts.gate_up_proj # [N, 2*I, H] - down = experts.down_proj # [N, H, I] - act_fn = experts.act_fn - - new_experts = nn.ModuleList() - for i in range(self.num_experts): - new_experts.append( - _MixtralBlockSparseTop2MLPCompat( - self.hidden_dim, - self.intermediate_dim, - gate_up[i, :self.intermediate_dim, :], # w1: gate - down[i], # w2: down - gate_up[i, self.intermediate_dim:, :], # w3: up - act_fn, - )) - self.experts = new_experts - - def forward(self, hidden_states): - batch_size, sequence_length, hidden_dim = hidden_states.shape - if self.training and self.jitter_noise > 0: - hidden_states = hidden_states * torch.empty_like( - hidden_states).uniform_(1.0 - self.jitter_noise, - 1.0 + self.jitter_noise) - hidden_states_flat = hidden_states.view(-1, hidden_dim) - _, top_k_weights, top_k_index = self.gate(hidden_states_flat) - - final_hidden_states = torch.zeros_like(hidden_states_flat) - expert_mask = torch.nn.functional.one_hot( - top_k_index, num_classes=self.num_experts).permute(2, 1, 0) - expert_hit = (expert_mask.sum(dim=(-1, -2)) > 0).nonzero() - - for expert_idx in expert_hit: - expert_idx = expert_idx[0] - if expert_idx == self.num_experts: - continue - top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) - current_state = hidden_states_flat[token_idx] - current_hidden_states = self.experts[expert_idx](current_state) - current_hidden_states = current_hidden_states * top_k_weights[ - token_idx, top_k_pos, None] - final_hidden_states.index_add_( - 0, token_idx, - current_hidden_states.to(final_hidden_states.dtype)) - - return final_hidden_states.reshape(batch_size, sequence_length, - hidden_dim) - - -def _unfuse_mixtral_for_modelopt(model: nn.Module) -> None: - # transformers 5.x stores Mixtral experts as a single MixtralExperts module - # with 3D fused tensors (gate_up_proj, down_proj) under layer.mlp, replacing - # the per-expert ModuleList layout that nvidia-modelopt 0.37 iterates - # (len(experts), experts[i].w1.weight, ...). Without this swap, the export - # call fails with: TypeError: object of type 'MixtralExperts' has no len(). - # - # An earlier attempt added a sibling layer.block_sparse_moe with the legacy - # layout, but modelopt iterates decoder_layer.named_children() and hits the - # original layer.mlp (a MixtralSparseMoeBlock matching is_moe) first, so the - # exception still triggered. Replace layer.mlp in place so calibration uses - # per-expert nn.Linear modules (allowing modelopt to attach quantizers) and - # export reads the per-expert ModuleList layout. - try: - from transformers.models.mixtral.modeling_mixtral import MixtralExperts - except ImportError: - return - - if not (hasattr(model, "model") and hasattr(model.model, "layers")): - return - - for layer in model.model.layers: - mlp = getattr(layer, "mlp", None) - if mlp is None or not isinstance(getattr(mlp, "experts", None), - MixtralExperts): - continue - # nn.Module() defaults to training=True; mirror the original mlp's - # mode so a prior model.eval() is preserved (avoids re-enabling the - # router-input jitter during calibration/export). - compat_mlp = _MixtralSparseMoeBlockCompat(mlp) - compat_mlp.train(mlp.training) - layer.mlp = compat_mlp - - -def _get_llava_qwen_model(model_dir, dtype, device): - if "hf" in model_dir: - from transformers import LlavaOnevisionForConditionalGeneration - model = LlavaOnevisionForConditionalGeneration.from_pretrained( - model_dir, dtype=dtype, device_map=device) - model = model.language_model - else: - from llava.model.builder import load_pretrained_model - _, model, _, _ = load_pretrained_model(model_dir, - None, - 'llava_qwen', - torch_dtype=dtype, - device_map=device) - return model - - -def get_model(ckpt_path: str, - dtype: str = 'bfloat16', - device: str = 'cuda', - device_map: str = "auto"): - logger.info(f"Initializing model from {ckpt_path}") - # Note: VILA model is not in public HF model zoo yet. We need to explicitly import from the git repo - hf_config = get_hf_config(ckpt_path) - torch_dtype = str_dtype_to_torch(dtype) - - model_cls = AutoModelForCausalLM - if hf_config.model_type == "llava": - from transformers import LlavaForConditionalGeneration - model_cls = LlavaForConditionalGeneration - elif hf_config.model_type == "mpt": - from transformers import MptForCausalLM - model_cls = MptForCausalLM - elif hf_config.model_type == 'mllama': - from transformers import MllamaForConditionalGeneration - model_cls = MllamaForConditionalGeneration - elif hf_config.model_type == 'qwen2_vl': - from transformers import Qwen2VLForConditionalGeneration - model_cls = Qwen2VLForConditionalGeneration - - if "vila" in ckpt_path: - model = _get_vila_model(ckpt_path) - elif "llava-onevision-qwen2" in ckpt_path: - model = _get_llava_qwen_model(ckpt_path, dtype, device) - elif hf_config.model_type == "glm": - from transformers import AutoModelForSeq2SeqLM - model = AutoModelForSeq2SeqLM.from_pretrained(ckpt_path, - device_map="cuda", - dtype=torch_dtype, - trust_remote_code=True) - elif model_type_is_enc_dec(hf_config.model_type): - from transformers import AutoModelForSeq2SeqLM - model = AutoModelForSeq2SeqLM.from_pretrained(ckpt_path, - device_map=device, - dtype=torch_dtype, - trust_remote_code=True) - model = EncDecModelWrapper(hf_model=model) - else: - model = model_cls.from_pretrained( - ckpt_path, - device_map=device_map if device != "cpu" else "cpu", - dtype="auto", - trust_remote_code=True) - if hf_config.model_type in ["llava", "internvl_chat"]: - model = model.language_model - elif hf_config.model_type == "qwen2_vl": - #WAR for Qwen2-VL because its lm_head is outside of LLM - lm_head = model.lm_head - model = model.model - model.lm_head = lm_head - - model.eval() - - # transformers 5.x changed Mixtral MoE to a fused 3D layout that - # nvidia-modelopt 0.37 cannot iterate. Restore the per-expert layout - # so both calibration and export work; see _unfuse_mixtral_for_modelopt. - if hf_config.model_type == "mixtral": - _unfuse_mixtral_for_modelopt(model) - - model_dtype = next(model.parameters()).dtype - if torch_dtype != model_dtype: - logger.info( - f"[TensorRT-LLM][WARNING] The manually set model data type is {dtype}, " - f"but the data type of the HuggingFace model is {model_dtype}.") - - return model - - -def get_model_type(model): - if type(model).__name__ == "EncDecModelWrapper": - return model.model_type - if type(model).__name__ in MODEL_NAME_PATTERN_MAP: - return MODEL_NAME_PATTERN_MAP[type(model).__name__] - for k, v in MODEL_NAME_PATTERN_MAP.items(): - if k.lower() in type(model).__name__.lower(): - return v - return None - - -def _is_cnn_dailymail_local_repo(path: str) -> bool: - if not os.path.isdir(path): - return False - # The loader only uses the "3.0.0" config. - if os.path.isdir(os.path.join(path, "3.0.0")): - return True - if os.path.isfile(os.path.join(path, "cnn_dailymail.py")): - return True - return False - - -def get_calib_dataloader(dataset_name_or_dir="cnn_dailymail", - tokenizer=None, - batch_size=1, - calib_size=512, - block_size=512, - device=None, - include_labels=False): - logger.info("Loading calibration dataset") - if dataset_name_or_dir == "pileval": - dataset = load_dataset( - "json", - data_files="https://the-eye.eu/public/AI/pile/val.jsonl.zst", - split="train", - trust_remote_code=True) - dataset = dataset["text"][:calib_size] - elif "scienceqa" in dataset_name_or_dir.lower( - ) or "science_qa" in dataset_name_or_dir.lower(): - if os.path.isdir(dataset_name_or_dir): - dataset = load_dataset(dataset_name_or_dir, - split="train", - trust_remote_code=True) - else: - dataset = load_dataset("derek-thomas/ScienceQA", - split="train", - trust_remote_code=True) - dataset = dataset.select(range(calib_size)) - elif "cnn_dailymail" in dataset_name_or_dir or _is_cnn_dailymail_local_repo( - dataset_name_or_dir): - # Bare "cnn_dailymail" id is rejected by newer huggingface_hub; use the namespaced repo. - if dataset_name_or_dir == "cnn_dailymail": - dataset_name_or_dir = "abisee/cnn_dailymail" - dataset = load_dataset( - dataset_name_or_dir, - name="3.0.0", - split="train", - trust_remote_code=True, - ) - dataset = dataset["article"][:calib_size] - elif os.path.isdir(dataset_name_or_dir): - logger.info( - f"Recognized local dataset repo {dataset_name_or_dir} for calibration; " - "assuming the calibration data are in the train split and text column." - ) - dataset = load_dataset(dataset_name_or_dir, - split="train", - trust_remote_code=True) - dataset = dataset["text"][:calib_size] - else: - raise NotImplementedError( - f"Unsupported dataset name or local repo directory: {dataset_name_or_dir}." - ) - - is_multimodal = False - for dataset_name in MULTIMODAL_DATASETS: - if dataset_name in dataset_name_or_dir: - is_multimodal = True - if is_multimodal: - # Apply the preprocessing function to the dataset - processed_dataset = dataset.map(tokenizer.preprocess_function, - batched=False, - remove_columns=dataset.column_names) - - # Create DataLoader with the custom collate function - calib_dataloader = DataLoader(processed_dataset, - batch_size=batch_size, - shuffle=False, - collate_fn=tokenizer.collate_function) - else: - batch_encoded = tokenizer(dataset, - return_tensors="pt", - padding=True, - truncation=True, - max_length=block_size) - if device: - batch_encoded = batch_encoded.to(device) - - if include_labels: - # Labels are needed when backward is called in the model. - # The labels should be a shifted version of the input_ids. - # However, we should not shift the input_ids here since the labels are shifted by - # Huggingface models during loss calculation as shown here - - # https://github.com/huggingface/transformers/blob/7f79a97399bb52aad8460e1da2f36577d5dccfed/src/transformers/models/llama/modeling_llama.py#L1093-L1095 - batch_encoded["labels"] = torch.where( - batch_encoded["attention_mask"] > 0.5, - batch_encoded["input_ids"], -100) - batch_encoded = _CustomDataset(batch_encoded) - else: - # For backward compatibility, if labels are not needed, we only return input_ids. - batch_encoded = _CustomDataset( - {"input_ids": batch_encoded["input_ids"]}) - - calib_dataloader = DataLoader(batch_encoded, - batch_size=batch_size, - shuffle=False) - - return calib_dataloader - - -def quantize_model(model, quant_cfg, calib_dataloader, batch_size, qformat, - auto_quantize_bits): - import modelopt.torch.quantization as mtq - - # NOTE: for ModelOpt v0.19 release - # calibrate_loop = dataset_utils.create_forward_loop( - # calib_dataloader, dataloader=calib_dataloader) - - def calibrate_loop(): - if calib_dataloader is None: - return - with torch.no_grad(): - low_mem_mode = False - for idx, data in enumerate(calib_dataloader): - logger.debug(f"Calibrating batch {idx}") - batch_size = data[list(data.keys())[0]].shape[0] - if batch_size == 1: - model(**data) - elif not low_mem_mode: - # Try running the forward once. - # If output memory, we try running inference with split input tensors - try: - model(**data) - except torch.OutOfMemoryError: - print( - "Warning: torch.OutOfMemoryError detected, try reducing the batch size..." - ) - low_mem_mode = True - - if low_mem_mode: - split_data_1 = { - key: data[key][:batch_size // 2, ...] - for key in data - } - model(**split_data_1) - - split_data_2 = { - key: data[key][batch_size // 2:, ...] - for key in data - } - model(**split_data_2) - - QUANT_CFG_CHOICES = { - "int8": "INT8_DEFAULT_CFG", - "int8_sq": "INT8_SMOOTHQUANT_CFG", - "fp8": "FP8_DEFAULT_CFG", - "fp8_pc_pt": "FP8_PER_CHANNEL_PER_TOKEN_CFG", - "int4_awq": "INT4_AWQ_CFG", - "w4a8_awq": "W4A8_AWQ_BETA_CFG", - } - - logger.info("Starting quantization...") - start_time = time.time() - if auto_quantize_bits: - logger.info("Starting mixed precision quantization...") - - from packaging import version as v - opt_kwargs = {} - modelopt_version = version('nvidia-modelopt') - if v.parse(modelopt_version) > v.parse("0.21"): - opt_kwargs['disabled_layers'] = ["*lm_head*"] - - model, search_history = mtq.auto_quantize( - model, - data_loader=calib_dataloader, - loss_func=lambda output, batch: output.loss, - constraints={"effective_bits": auto_quantize_bits}, - forward_step=lambda model, batch: model(**batch), - quantization_formats=[ - QUANT_CFG_CHOICES[item] for item in qformat.split(",") - ] + [None], - num_calib_steps=len(calib_dataloader), - num_score_steps=min( - len(calib_dataloader), 128 // batch_size - ), # Limit the number of score steps to avoid long calibration time - verbose=True, - **opt_kwargs) - mtq.print_quant_summary(model) - - # We need to explicitly calibrate for kv cache quantization - enable_kv_cache_quantization = "int8" not in qformat - if enable_kv_cache_quantization: - mtq.set_quantizer_by_cfg( - model, - quant_cfg={ - "*output_quantizer": { - "num_bits": (4, 3), - "axis": None, - "enable": True - } - }, - ) - # Lets calibrate only the output quantizer this time. Let's disable all other quantizers. - with mtq.set_quantizer_by_cfg_context(model, { - "*": { - "enable": False - }, - "*output_quantizer": { - "enable": True - } - }): - mtq.calibrate(model, - algorithm="max", - forward_loop=calibrate_loop) - else: - mtq.quantize(model, quant_cfg, forward_loop=calibrate_loop) - end_time = time.time() - logger.info( - "Quantization done. Total time used: {:.2f} s.".format(end_time - - start_time)) - return model - - -def quantize_and_export(*, - model_dir, - device, - calib_dataset, - dtype, - qformat, - kv_cache_dtype, - calib_size, - batch_size, - calib_max_seq_length, - awq_block_size, - output_dir, - tp_size, - pp_size, - cp_size, - seed, - tokenizer_max_seq_length, - num_medusa_heads=None, - num_medusa_layers=None, - max_draft_len=None, - medusa_hidden_act=None, - medusa_model_dir=None, - quant_medusa_head=None, - auto_quantize_bits=None, - device_map="auto", - quantize_lm_head=False): - ''' - Load model from the model_dir, call Modelopt to quantize the model, and then export - the quantized model as TRT-LLM checkpoint - ''' - try: - import modelopt # noqa - except ImportError as e: - logger.error( - "Failed to import modelopt, pls check the Modelopt installation. Currently it is known to be unsupported on Windows OS" - ) - raise e - - import modelopt.torch.quantization as mtq - from modelopt.torch.export import export_tensorrt_llm_checkpoint - - from tensorrt_llm.models.convert_utils import infer_dtype - - if not torch.cuda.is_available(): - raise EnvironmentError("GPU is required for inference.") - - random.seed(seed) - np.random.seed(seed) - - # Check that only one quantization format is provided for non auto_quant case - if not auto_quantize_bits: - assert (len(qformat.split(",")) == 1 - ), "Quantization supports only one quantization format." - - hf_config = get_hf_config(model_dir) - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - model = get_model(model_dir, dtype, device=device, device_map=device_map) - model_type = get_model_type(model) - is_enc_dec = model_type_is_enc_dec(model_type) - if "vila" in model_dir: - tokenizer = get_tokenizer(model_dir + "/llm", - max_seq_length=tokenizer_max_seq_length, - model_type=model_type) - elif model_type == "mllama": - tokenizer = get_processor(model_dir, - max_seq_length=tokenizer_max_seq_length, - model_type=model_type, - device=device) - else: - tokenizer = get_tokenizer(model_dir, - max_seq_length=tokenizer_max_seq_length, - model_type=model_type) - - if qformat in ["full_prec", "int8_wo", "int4_wo" - ] and kv_cache_dtype is None: - logger.info(f"No quantization applied, export {dtype} model") - else: - if "awq" in qformat: - if calib_size > 32: - logger.info( - f"AWQ calibration could take longer with calib_size = {calib_size}, Using" - " calib_size=32 instead") - calib_size = 32 - logger.info( - "\nAWQ calibration could take longer than other calibration methods. Please" - " increase the batch size to speed up the calibration process. Batch size can be" - " set by adding the argument --batch_size to the command line.\n" - ) - - quant_cfg = None - if not auto_quantize_bits: - if qformat in quant_cfg_choices(): - quant_cfg = quant_cfg_choices()[qformat] - else: - raise ValueError(f"Unsupported quantization format: {qformat}") - - if "awq" in qformat: - quant_cfg = copy.deepcopy(quant_cfg_choices()[qformat]) - weight_quantizer = quant_cfg["quant_cfg"]["*weight_quantizer"] - if isinstance(weight_quantizer, list): - weight_quantizer = weight_quantizer[0] - if awq_block_size: - weight_quantizer["block_sizes"][-1] = awq_block_size - - # Coarser optimal scale search seems to resolve the overflow in TRT-LLM for some models - if "w4a8_awq" == qformat and model_type in ["gemma", "mpt"]: - quant_cfg["algorithm"] = { - "method": "awq_lite", - "alpha_step": 1 - } - - if kv_cache_dtype is not None: - if kv_cache_dtype == "fp8": - kv_cache_quant_cfg = getattr( - mtq, KV_QUANT_CFG_CHOICES[kv_cache_dtype])["quant_cfg"] - quant_cfg["quant_cfg"].update(kv_cache_quant_cfg) - else: - quant_cfg["quant_cfg"].update(KV_CACHE_CFG) # type: ignore - - # Gemma 7B has accuracy regression using alpha 1. We set 0.5 instead. - if model_type == "gemma" and "int8_sq" in qformat: - quant_cfg["algorithm"] = {"method": "smoothquant", "alpha": 0.5} - - if qformat == 'fp8' and quantize_lm_head: - print_rank_0("Quantizing lm_head layer") - del quant_cfg["quant_cfg"]["*lm_head*"] - - calib_dataloader = get_calib_dataloader( - dataset_name_or_dir=calib_dataset, - tokenizer=tokenizer, - batch_size=batch_size, - calib_size=calib_size, - block_size=calib_max_seq_length, - device=model.device, - include_labels=auto_quantize_bits is not None, - ) - - model = quantize_model(model, quant_cfg, calib_dataloader, batch_size, - qformat, auto_quantize_bits) - - with torch.inference_mode(): - if model_type is None: - logger.info( - f"Unknown model type {type(model).__name__}. Continue exporting..." - ) - model_type = f"unknown:{type(model).__name__}" - - architecture = type(model).__name__ - - export_path = output_dir - start_time = time.time() - - # Move meta tensor back to device before exporting. - remove_hook_from_module(model, recurse=True) - - QUANT_ALGO = { - "int8": "INT8", - "int8_sq": "W8A8_SQ_PER_CHANNEL", - "fp8": "FP8", - "int4_awq": "W4A16_AWQ", - "w4a8_awq": "W4A8_AWQ", - } - - if model_type == 'mllama': - model = model.language_model - - export_tensorrt_llm_checkpoint( - model.hf_model if is_enc_dec else model, - model_type, - getattr(torch, dtype), - export_dir=export_path, - inference_tensor_parallel=tp_size, - inference_pipeline_parallel=pp_size, - ) - - export_paths = [] - tensorrt_llm_configs = [] - if not is_enc_dec: - with open(f"{export_path}/config.json", "r") as f: - tensorrt_llm_config = json.load(f) - tensorrt_llm_configs.append(tensorrt_llm_config) - export_paths.append(export_path) - else: - for component in ["encoder", "decoder"]: - with open(f"{export_path}/{component}/config.json", "r") as f: - tensorrt_llm_config = json.load(f) - tensorrt_llm_configs.append(tensorrt_llm_config) - export_paths.append(f"{export_path}/{component}") - - for export_path, tensorrt_llm_config in zip(export_paths, - tensorrt_llm_configs): - - tensorrt_llm_config["model_type"] = model_type - if not is_enc_dec: - tensorrt_llm_config["architecture"] = architecture - - # Workaround for wo quantization - if qformat in ["int8_wo", "int4_wo", "full_prec"]: - if qformat == "int8_wo": - tensorrt_llm_config["quantization"][ - "quant_algo"] = QuantAlgo.W8A16 - elif qformat == "int4_wo": - tensorrt_llm_config["quantization"][ - "quant_algo"] = QuantAlgo.W4A16 - else: - tensorrt_llm_config["quantization"]["quant_algo"] = None - - # HF uses rope_scaling while tensorrt_llm uses rotary_scaling - if hasattr(model.config, "rope_scaling" - ) and "rotary_scaling" not in tensorrt_llm_config: - tensorrt_llm_config["rotary_scaling"] = getattr( - model.config, "rope_scaling") - with open(f"{export_path}/config.json", "w") as f: - json.dump(tensorrt_llm_config, f, indent=4) - - # Workaround for Modelopt 0.9.x fp8_kv_cache knob issue - if qformat in ['fp8', 'nvfp4'] and kv_cache_dtype is None: - with open(f"{export_path}/config.json", "r") as f: - tensorrt_llm_config = json.load(f) - tensorrt_llm_config["quantization"][ - "kv_cache_quant_algo"] = None - with open(f"{export_path}/config.json", "w") as f: - json.dump(tensorrt_llm_config, f, indent=4) - - # Workaround for qwen version - if model_type == 'qwen' or model_type == 'qwen2_vl': - with open(f"{export_path}/config.json", "r") as f: - tensorrt_llm_config = json.load(f) - qwen_config = AutoConfig.from_pretrained(model_dir, - trust_remote_code=True) - try: - from transformers import LlavaOnevisionConfig - if isinstance(qwen_config, LlavaOnevisionConfig): - qwen_config = qwen_config.text_config - except: - pass - tensorrt_llm_config["qwen_type"] = qwen_config.model_type - if qwen_config.model_type == "qwen2": - tensorrt_llm_config[ - "norm_epsilon"] = qwen_config.rms_norm_eps - tensorrt_llm_config["rotary_base"] = get_hf_rope_theta( - qwen_config, 100000.0) - tensorrt_llm_config[ - "intermediate_size"] = qwen_config.intermediate_size - with open(f"{export_path}/config.json", "w") as f: - json.dump(tensorrt_llm_config, f, indent=4) - - # Set rotary parameters correctly for chatglm. - if model_type == 'chatglm': - rotary_base = 10000.0 - rotary_embedding_scaling = None - chatglm_config = AutoConfig.from_pretrained( - model_dir, trust_remote_code=True) - chatglm_version = tensorrt_llm_config['chatglm_version'] - rope_ratio = tensorrt_llm_config.get('rope_ratio', 1.0) - if chatglm_version == 'chatglm2': - if rope_ratio > 1: - rotary_embedding_scaling = { - 'type': 'linear', - 'factor': rope_ratio - } - elif chatglm_version == 'chatglm3': - rotary_base *= rope_ratio - - with open(f"{export_path}/config.json", "r") as f: - tensorrt_llm_config = json.load(f) - tensorrt_llm_config['rotary_base'] = rotary_base - tensorrt_llm_config['rotary_scaling'] = rotary_embedding_scaling - tensorrt_llm_config['rotary_pct'] = 0.5 - with open(f"{export_path}/config.json", "w") as f: - json.dump(tensorrt_llm_config, f, indent=4) - - # context parallel - if cp_size > 1: - with open(f"{export_path}/config.json", "r") as f: - tensorrt_llm_config = json.load(f) - tensorrt_llm_config["mapping"]["cp_size"] = cp_size - tensorrt_llm_config["mapping"]["attn_tp_size"] = -1 - tensorrt_llm_config["mapping"]["attn_cp_size"] = -1 - tensorrt_llm_config["mapping"]["world_size"] *= cp_size - with open(f"{export_path}/config.json", "w") as f: - json.dump(tensorrt_llm_config, f, indent=4) - - if model_type == 'gptnext': - with open(f"{export_path}/config.json", "r") as f: - tensorrt_llm_config = json.load(f) - if tensorrt_llm_config['max_position_embeddings'] is None: - tensorrt_llm_config['max_position_embeddings'] = getattr( - model.config, "n_positions", None) - with open(f"{export_path}/config.json", "w") as f: - json.dump(tensorrt_llm_config, f, indent=4) - - end_time = time.time() - logger.info( - "Quantized model exported to {} \nTotal time used {:.2f} s.".format( - export_path, end_time - start_time)) - - # Need to delete the model and release memory explicitly; - # otherwise torch may retain its GPU memory until a delayed GC running, - # which reduces the available GPU memory for subsequent stages. - del model - release_gc() - - -def unwrap_model(model, module_instances=None): - # Reference: https://github.com/NVIDIA/Megatron-LM/blob/core_r0.8.0/megatron/training/utils.py - from megatron.core import DistributedDataParallel as DDP - from megatron.core.transformer.module import Float16Module - - if module_instances is None: - module_instances = (DDP, Float16Module) - - return_list = True - if not isinstance(model, list): - model = [model] - return_list = False - unwrapped_model = [] - for model_module in model: - while isinstance(model_module, module_instances): - model_module = model_module.module - unwrapped_model.append(model_module) - if not return_list: - return unwrapped_model[0] - return unwrapped_model - - -def get_nemo_calib_dataloader(dataset_name_or_dir="cnn_dailymail", - batch_size=64, - calib_size=512, - max_sequence_length=512): - if dataset_name_or_dir == "pileval": - dataset = load_dataset( - "json", - data_files="https://the-eye.eu/public/AI/pile/val.jsonl.zst", - split="train", - trust_remote_code=True) - text_column = "text" - elif "wikitext" in dataset_name_or_dir: - dataset = load_dataset(dataset_name_or_dir, - "wikitext-103-v1", - split="train", - trust_remote_code=True) - text_column = "text" - elif "cnn_dailymail" in dataset_name_or_dir or _is_cnn_dailymail_local_repo( - dataset_name_or_dir): - # Bare "cnn_dailymail" id is rejected by newer huggingface_hub; use the namespaced repo. - if dataset_name_or_dir == "cnn_dailymail": - dataset_name_or_dir = "abisee/cnn_dailymail" - dataset = load_dataset(dataset_name_or_dir, - name="3.0.0", - split="train", - trust_remote_code=True) - text_column = "article" - elif os.path.isdir(dataset_name_or_dir): - logger.info( - f"Recognized local dataset repo {dataset_name_or_dir} for calibration; " - "assuming the calibration data are in the train split and text column." - ) - dataset = load_dataset(dataset_name_or_dir, - split="train", - trust_remote_code=True) - text_column = "text" - else: - raise NotImplementedError( - f"Unsupported dataset name or local repo directory: {dataset_name_or_dir}." - ) - calib_size = max(min(len(dataset), calib_size), batch_size) - for i in range(calib_size // batch_size): - batch = dataset[i * batch_size:(i + 1) * batch_size][text_column] - for j in range(len(batch)): - batch[j] = batch[j][:max_sequence_length] - yield batch - - -def quantize_nemo_and_export(*, nemo_ckpt_path, decoder_type, calib_dataset, - calib_tp_size, calib_pp_size, dtype, qformat, - kv_cache_dtype, calib_size, batch_size, - calib_max_seq_length, awq_block_size, output_dir, - tp_size, pp_size, cp_size, seed): - try: - import modelopt # noqa - except ImportError as e: - logger.error( - "Failed to import modelopt, pls check the modelopt installation. Currently it is known to be unsupported on Windows OS" - ) - raise e - - import modelopt.torch.quantization as mtq - from megatron.core import parallel_state - from megatron.core.transformer.module import Float16Module - from modelopt.torch.export import export_tensorrt_llm_checkpoint - from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import \ - MegatronGPTModel - from nemo.collections.nlp.modules.common.text_generation_strategy import \ - GPTModelTextGenerationStrategy - from nemo.collections.nlp.parts.nlp_overrides import ( - NLPDDPStrategy, NLPSaveRestoreConnector) - from nemo.utils.model_utils import load_config, save_artifacts - from omegaconf.omegaconf import open_dict - from pytorch_lightning.trainer.trainer import Trainer - - if not torch.cuda.is_available(): - raise EnvironmentError("GPU is required for the inference.") - - random.seed(seed) - np.random.seed(seed) - - model_cfg = load_config(nemo_ckpt_path) - - # dtype is used for non-quantized layers - supported_dtype = ["auto", "float16", "bfloat16"] - assert dtype in supported_dtype, f"{dtype} not supported. Supported dtypes are {supported_dtype}" - - if dtype == 'auto': - dtype = model_cfg.get('precision', None) - if dtype is None: - dtype = 'float16' - elif 'bf16' in dtype or 'bfloat16' in dtype: - dtype = 'bfloat16' - else: - dtype = 'float16' - logger.info(f"Specified dtype 'auto'; inferred dtype {dtype!r}.") - torch_dtype = getattr(torch, dtype) - - with open_dict(model_cfg): - model_cfg.activations_checkpoint_method = None - model_cfg.activations_checkpoint_granularity = None - model_cfg.tensor_model_parallel_size = calib_tp_size - model_cfg.pipeline_model_parallel_size = calib_pp_size - model_cfg.sequence_parallel = False - # Only custom modelopt spec is supported for PTQ: this custom spec is largely based on local Megatron-LM - # layer definitions to avoid Transformer Engine implementations that are currently not supported. - model_cfg.name = "modelopt" - - # trainer required for restoring model parallel models - trainer_config = { - 'devices': calib_tp_size * calib_pp_size, - 'num_nodes': 1, - 'accelerator': 'gpu', - 'logger': False, - 'precision': model_cfg.precision, - 'enable_checkpointing': False, - } - trainer = Trainer(strategy=NLPDDPStrategy(), **trainer_config) - connector = NLPSaveRestoreConnector() - - model = MegatronGPTModel.restore_from( - restore_path=nemo_ckpt_path, - trainer=trainer, - override_config_path=model_cfg, - save_restore_connector=connector, - ) - model.freeze() - - print_rank_0(model) - # Have to turn off activations_checkpoint_method for inference - try: - model.model.module.language_model.encoder.activations_checkpoint_method = None - except AttributeError: - pass - - # Check whether the DDP is initialized - if parallel_state.is_unitialized(): - - def dummy(): - return - - if model.trainer.strategy.launcher is not None: - model.trainer.strategy.launcher.launch(dummy, trainer=model.trainer) - model.trainer.strategy.setup_environment() - - inference_config = { - 'greedy': False, - 'top_k': 0, - 'top_p': 0.9, - 'temperature': 1.0, - 'add_BOS': True, - 'tokens_to_generate': 30, - 'all_probs': False, - 'repetition_penalty': 1.2, - 'min_tokens_to_generate': 0, - 'compute_logprob': False, - 'batch_size': batch_size, - 'max_context_length': calib_max_seq_length, - 'strategy': GPTModelTextGenerationStrategy(model), - } - model.set_inference_config(inference_config) - - if qformat in ["full_prec", "int8_wo", "int4_wo" - ] and kv_cache_dtype is None: - print_rank_0(f"No quantization applied, export {dtype} model") - else: - if "awq" in qformat: - if calib_size > 32: - print_rank_0( - "AWQ calibration could take longer with calib_size =" - f" {calib_size}, Using calib_size=32 instead") - calib_size = 32 - print_rank_0( - "\nAWQ calibration could take longer than other calibration methods. Please" - " increase the batch size to speed up the calibration process. Batch size can be" - " set by adding the argument inference.batch_size= to the command" - " line.\n") - - dataloader = get_nemo_calib_dataloader( - dataset_name_or_dir=calib_dataset, - batch_size=batch_size, - calib_size=calib_size, - max_sequence_length=calib_max_seq_length, - ) - - # =================== Start Quantization ==================== - if qformat in quant_cfg_choices(): - quant_cfg = quant_cfg_choices()[qformat] - else: - raise ValueError(f"Unsupported quantization format: {qformat}") - - if "awq" in qformat: - quant_cfg = copy.deepcopy(quant_cfg_choices()[qformat]) - weight_quantizer = quant_cfg["quant_cfg"][ - "*weight_quantizer"] # type: ignore - if isinstance(weight_quantizer, list): - weight_quantizer = weight_quantizer[0] - weight_quantizer["block_sizes"][-1] = awq_block_size - - if kv_cache_dtype is not None: - if kv_cache_dtype == "fp8": - for value in KV_CACHE_CFG.values(): - value.update({"num_bits": (4, 3)}) # type: ignore - quant_cfg["quant_cfg"].update(KV_CACHE_CFG) # type: ignore - - print_rank_0(quant_cfg) - - # Always turn on FP8 kv cache to save memory footprint. - # For int8_sq, we use int8 kv cache. - # TODO: Investigate why enabling FP8 kv cache will cause accuracy regressions for nemotron. - # quant_cfg["quant_cfg"]["*output_quantizer"] = { # type: ignore[index] - # "num_bits": 8 if args.qformat == "int8_sq" else (4, 3), - # "axis": None, - # "enable": args.decoder_type != "gptnext", - # } - - dataloader = [data for data in dataloader] - - def forward_loop(model): - for i, batch in enumerate(dataloader): - print_rank_0(f"Calibrating batch {i}") - model.predict_step(batch, i) - - start_time = time.time() - model = mtq.quantize(model, quant_cfg, - forward_loop) # type: ignore[arg-type] - end_time = time.time() - tot_time = end_time - start_time - tput = calib_size / tot_time - print_rank_0( - f"Quantization done. Total time used {tot_time}s. Throughput {tput} samples/s" - ) - # =================== End Quantization ====================== - - if decoder_type == "gptnext": - # We found squared_relu may have an under-calibration problem. - # Clamp the scaling_factor with a min threshold to avoid under-calibration. - maxbound = 0 - if qformat == "fp8": - maxbound = 448 - elif qformat == "int8_sq": - maxbound = 127 - model = mtq.postprocess_amax( - model, "*input_quantizer", - lambda amax: torch.clamp(amax, min=0.01 * maxbound)) - - if torch.distributed.get_rank() == 0: - mtq.print_quant_summary(model) - - if model_cfg.megatron_amp_O2: - model.model = unwrap_model(model.model, Float16Module) - - start_time = time.time() - export_tensorrt_llm_checkpoint( - model, - decoder_type, - torch_dtype, - export_dir=output_dir, - inference_tensor_parallel=tp_size, - inference_pipeline_parallel=pp_size, - ) - - # context parallel - if cp_size > 1: - with open(f"{export_path}/config.json", "r") as f: - tensorrt_llm_config = json.load(f) - tensorrt_llm_config["mapping"]["cp_size"] = cp_size - tensorrt_llm_config["mapping"]["world_size"] *= cp_size - with open(f"{export_path}/config.json", "w") as f: - json.dump(tensorrt_llm_config, f, indent=4) - - end_time = time.time() - print_rank_0( - f"Model config exported to: {output_dir}. Total time used {end_time - start_time}s" - ) - if torch.distributed.get_rank() == 0: - save_artifacts(model, output_dir, use_abspath=True) - - # Need to delete the model and release memory explicitly; - # otherwise torch may retain its GPU memory until a delayed GC running, - # which reduces the available GPU memory for subsequent stages. - del model - release_gc() diff --git a/tensorrt_llm/serialization.py b/tensorrt_llm/serialization.py index ba548cc426f7..2311cec58f55 100644 --- a/tensorrt_llm/serialization.py +++ b/tensorrt_llm/serialization.py @@ -66,7 +66,6 @@ "KvCacheRetentionConfig.TokenRangeRetentionConfig", "PeftCacheConfig", "SchedulerConfig" ], - "tensorrt_llm.builder": ["BuildConfig"], "tensorrt_llm.disaggregated_params": ["DisaggregatedParams"], "tensorrt_llm.inputs.multimodal": ["MultimodalInput"], "tensorrt_llm.executor.postproc_worker": [ diff --git a/tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py b/tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py index 6da6cd6e9c0a..fad24e46bc1c 100644 --- a/tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py +++ b/tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py @@ -146,16 +146,14 @@ def _prepare_dataset(root_dir: str, temp_dir: str, model_path_or_name: str, num_ """Prepare a synthetic dataset for benchmarking.""" _DATASET_NAME = "synthetic_128_128.txt" dataset_path = Path(temp_dir, _DATASET_NAME) - dataset_tool = Path(root_dir, "benchmarks", "prepare_dataset.py") - script_dir = Path(root_dir, "benchmarks") # Generate a small dataset to run a test - matching workload configuration command = [ - "python3", - f"{dataset_tool}", - "--stdout", - "--tokenizer", + "trtllm-bench", + "--model", model_path_or_name, + "prepare-dataset", + "--stdout", "token-norm-dist", "--input-mean", "128", @@ -169,9 +167,7 @@ def _prepare_dataset(root_dir: str, temp_dir: str, model_path_or_name: str, num_ str(num_requests), ] print(f"Running command: {' '.join(command)}") - result = subprocess.run( - command, cwd=str(script_dir), capture_output=True, text=True, timeout=300 - ) + result = subprocess.run(command, cwd=str(temp_dir), capture_output=True, text=True, timeout=300) if result.returncode != 0: raise RuntimeError(f"Failed to prepare dataset: {result.stderr}") # Grab the stdout and write it to a dataset file for passing to suite. diff --git a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_trtllm_bench.py b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_trtllm_bench.py index f766576cd454..6ba35a0aa44d 100644 --- a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_trtllm_bench.py +++ b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_trtllm_bench.py @@ -65,7 +65,6 @@ def run_benchmark( def prepare_dataset(root_dir: str, temp_dir: str, model_path_or_name: str): _DATASET_NAME = "synthetic_128_128.txt" dataset_path = Path(temp_dir, _DATASET_NAME) - script_dir = Path(root_dir, "benchmarks") # Generate a small dataset to run a test - matching workload configuration command = [ @@ -88,9 +87,7 @@ def prepare_dataset(root_dir: str, temp_dir: str, model_path_or_name: str): "10", ] print(f"Running command: {' '.join(command)}") - result = subprocess.run( - command, cwd=str(script_dir), capture_output=True, text=True, timeout=300 - ) + result = subprocess.run(command, cwd=str(temp_dir), capture_output=True, text=True, timeout=300) if result.returncode != 0: raise RuntimeError(f"Failed to prepare dataset: {result.stderr}") diff --git a/tests/unittest/others/test_quantize_calib_dataset.py b/tests/unittest/others/test_quantize_calib_dataset.py deleted file mode 100644 index 7a2e0549baf4..000000000000 --- a/tests/unittest/others/test_quantize_calib_dataset.py +++ /dev/null @@ -1,83 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import os -import tempfile - -import torch -from utils.util import skip_no_modelopt - - -@skip_no_modelopt -def test_is_cnn_dailymail_local_repo(): - from tensorrt_llm.quantization.quantize_by_modelopt import _is_cnn_dailymail_local_repo - - # Non-existent path - assert not _is_cnn_dailymail_local_repo("/does/not/exist") - - # Empty directory - with tempfile.TemporaryDirectory() as d: - assert not _is_cnn_dailymail_local_repo(d) - - # 3.0.0 subdir - with tempfile.TemporaryDirectory() as d: - os.makedirs(os.path.join(d, "3.0.0")) - assert _is_cnn_dailymail_local_repo(d) - - # Other versions not detected - for version in ("1.0.0", "2.0.0"): - with tempfile.TemporaryDirectory() as d: - os.makedirs(os.path.join(d, version)) - assert not _is_cnn_dailymail_local_repo(d), ( - f"version subdir {version} should not be detected" - ) - - # Directory with the cnn_dailymail.py builder script - with tempfile.TemporaryDirectory() as d: - open(os.path.join(d, "cnn_dailymail.py"), "w").close() - assert _is_cnn_dailymail_local_repo(d) - - # Directory with unrelated content - with tempfile.TemporaryDirectory() as d: - os.makedirs(os.path.join(d, "train")) - open(os.path.join(d, "data.parquet"), "w").close() - assert not _is_cnn_dailymail_local_repo(d) - - -@skip_no_modelopt -def test_get_calib_dataloader_local_cnn_dailymail(monkeypatch): - from tensorrt_llm.quantization import quantize_by_modelopt - - captured = {} - - def fake_load_dataset(path, **kwargs): - captured["kwargs"] = kwargs - return {"article": ["calibration article"] * 2} - - def fake_tokenizer(dataset, **kwargs): - return {"input_ids": torch.ones(len(dataset), 4, dtype=torch.long)} - - monkeypatch.setattr(quantize_by_modelopt, "load_dataset", fake_load_dataset) - - with tempfile.TemporaryDirectory() as d: - # Name lacks "cnn_dailymail". - os.makedirs(os.path.join(d, "3.0.0")) - dataloader = quantize_by_modelopt.get_calib_dataloader( - dataset_name_or_dir=d, - tokenizer=fake_tokenizer, - calib_size=2, - ) - - assert captured["kwargs"].get("name") == "3.0.0" - assert len(list(dataloader)) == 2 diff --git a/tests/unittest/tools/test_prepare_dataset.py b/tests/unittest/tools/test_prepare_dataset.py index 08bcb0ab1e38..af54d1be9e52 100644 --- a/tests/unittest/tools/test_prepare_dataset.py +++ b/tests/unittest/tools/test_prepare_dataset.py @@ -16,15 +16,14 @@ _DEFAULT_OUTPUT_STDEV = 10 _TEST_TASK_IDS = [0, 1, 2] _TOKENIZER_SUBPATH = "llama-models-v2/tinyllama-tarot-v1/" -_PREPARE_DATASET_SCRIPT_PATH = "benchmarks/prepare_dataset.py" class TestPrepareDatasetLora: """ - Test suite for prepare_dataset.py CLI tool LoRA metadata generation + Test suite for the trtllm-bench prepare-dataset LoRA metadata generation functionality. - This test class validates that the prepare_dataset.py script correctly + This test class validates that trtllm-bench prepare-dataset correctly generates LoRA request metadata when LoRA-specific parameters are provided. It covers both fixed task ID and random task ID scenarios. """ @@ -50,7 +49,7 @@ def temp_lora_dir(self) -> str: def _build_base_command(self, output_path: Path) -> List[str]: """ - Build the base command for running prepare_dataset.py. + Build the base command for running trtllm-bench prepare-dataset. Args: output_path: Path to the output dataset file @@ -110,7 +109,7 @@ def _add_synthetic_data_arguments(self, cmd: List[str]) -> None: def _run_prepare_dataset(self, **kwargs) -> str: """ - Execute prepare_dataset.py with specified parameters and capture + Execute trtllm-bench prepare-dataset with specified parameters and capture output. Args: @@ -140,7 +139,7 @@ def _run_prepare_dataset(self, **kwargs) -> str: def _parse_json_output(self, output: str) -> List[Dict[str, Any]]: """ - Parse JSON lines from prepare_dataset.py output. + Parse JSON lines from the prepare-dataset output. Args: output: Raw stdout output containing JSON lines From 7d3a4e9796c4dded3e12ebfac21e7c6cf8403028 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Fri, 24 Jul 2026 11:01:57 +0800 Subject: [PATCH 05/13] [https://nvbugs/6426834][fix] Deflake test_kv_transfer: cap NIXL progress threads, pin UCX env (#16708) Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 3 --- tests/unittest/disaggregated/test_agent.py | 6 ++++-- .../disaggregated/test_agent_multi_backends.py | 7 +++++-- .../test_cache_transceiver_single_process.py | 8 ++++---- tests/unittest/disaggregated/test_kv_transfer.py | 15 +++++++++++++-- .../unittest/disaggregated/test_kv_transfer_mp.py | 8 +++++++- .../unittest/disaggregated/test_mamba_transfer.py | 6 ++++++ .../disaggregated/test_py_cache_transceiver_mp.py | 6 ++++-- .../unittest/others/test_kv_cache_transceiver.py | 6 ++++++ 9 files changed, 49 insertions(+), 16 deletions(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 75133bae9fcb..6e48b1ebc11d 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -365,9 +365,6 @@ unittest/_torch/visual_gen/multi_gpu/test_ulysses_async.py::test_capture_smoke S unittest/_torch/visual_gen/test_attention_integration.py::test_sage_attention_self_attention[fp8-2-1560] SKIP (https://nvbugs/6198760) unittest/_torch/visual_gen/test_attention_integration.py::test_sage_attention_self_attention[int8-2-1560] SKIP (https://nvbugs/6198760) unittest/auto_deploy/multigpu/custom_ops SKIP (https://nvbugs/6403920) -unittest/disaggregated/test_kv_transfer.py SKIP (https://nvbugs/6403793) -unittest/disaggregated/test_kv_transfer.py::test_transfer_worker_v2[tp1_pp4_to_tp2_pp2] SKIP (https://nvbugs/6445316) -unittest/disaggregated/test_kv_transfer.py::test_transfer_worker_v2[tp4_pp1_to_tp2_pp2] SKIP (https://nvbugs/6426834) unittest/executor/test_rpc.py::TestRpcCorrectness::test_incremental_task_async SKIP (https://nvbugs/5741476) unittest/executor/test_rpc_proxy.py SKIP (https://nvbugs/5605741) unittest/executor/test_rpc_worker.py SKIP (https://nvbugs/5605741) diff --git a/tests/unittest/disaggregated/test_agent.py b/tests/unittest/disaggregated/test_agent.py index f2212a6c5f82..c8425ab49e18 100644 --- a/tests/unittest/disaggregated/test_agent.py +++ b/tests/unittest/disaggregated/test_agent.py @@ -6,8 +6,10 @@ import pytest import torch -# Exclude IB (no fabric) and gdr_copy (UCX rcache SIGABRT at teardown). -os.environ.setdefault("UCX_TLS", "^ib,gdr_copy") +# Force a deterministic UCX/NIXL config regardless of what the cluster/CI +# injects; see test_kv_transfer.py for the full rationale. +os.environ["UCX_TLS"] = "^ib,gdr_copy" +os.environ["TRTLLM_NIXL_NUM_THREADS"] = "1" from tensorrt_llm import logger from tensorrt_llm._torch.disaggregation.base.agent import ( diff --git a/tests/unittest/disaggregated/test_agent_multi_backends.py b/tests/unittest/disaggregated/test_agent_multi_backends.py index d983af152542..d82815c99a79 100644 --- a/tests/unittest/disaggregated/test_agent_multi_backends.py +++ b/tests/unittest/disaggregated/test_agent_multi_backends.py @@ -3,8 +3,11 @@ import pytest -# Exclude IB (no fabric) and gdr_copy (UCX rcache SIGABRT at teardown). -os.environ.setdefault("UCX_TLS", "^ib,gdr_copy") +# Force a deterministic UCX/NIXL config regardless of what the cluster/CI +# injects; see test_kv_transfer.py for the full rationale. The subprocesses +# spawned below inherit these via os.environ.copy(). +os.environ["UCX_TLS"] = "^ib,gdr_copy" +os.environ["TRTLLM_NIXL_NUM_THREADS"] = "1" def test_load_agent_missing_module(): diff --git a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py index 651a347b2f08..26b9e2e9199d 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py +++ b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py @@ -28,7 +28,10 @@ # Exclude UCX IB transport (avoid NIXL setup hangs without IB) and gdr_copy # (avoid SIGSEGV at process exit from UCX rcache cleanup; gdr_copy disabled # falls back to cuda_ipc / cuda_copy without affecting correctness). -os.environ.setdefault("UCX_TLS", "^ib,gdr_copy") +# Force a deterministic UCX/NIXL config regardless of what the cluster/CI +# injects; see test_kv_transfer.py for the full rationale. +os.environ["UCX_TLS"] = "^ib,gdr_copy" +os.environ["TRTLLM_NIXL_NUM_THREADS"] = "1" from dataclasses import dataclass from typing import Dict, List, Optional @@ -55,9 +58,6 @@ AttentionTypeCpp = tensorrt_llm.bindings.internal.batch_manager.AttentionType -# Reduce NIXL threads for unit test: default 8 threads per agent causes heavy -# contention when creating multiple agents on a single GPU in the same process. -os.environ.setdefault("TRTLLM_NIXL_NUM_THREADS", "0") # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index 76a0d8db25b3..c2cfc1382e37 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.py @@ -5,8 +5,19 @@ import time import uuid -# Exclude IB (no fabric) and gdr_copy (UCX rcache SIGABRT at teardown). -os.environ.setdefault("UCX_TLS", "^ib,gdr_copy") +# Force a deterministic UCX config regardless of what the cluster/CI injects +# (the CI agent bootstrap exports UCX_TLS=tcp,cuda_copy,cuda_ipc before pytest +# starts, which a setdefault would leave in place): exclude IB (no fabric +# assumed) and gdr_copy (UCX rcache SIGABRT at teardown). +os.environ["UCX_TLS"] = "^ib,gdr_copy" +# Each NIXL agent spawns TRTLLM_NIXL_NUM_THREADS (default 8) busy-polling +# progress threads, and a single case builds up to 8 TransferWorkers (one per +# rank). On CI nodes shared with other single-GPU jobs the resulting CPU +# oversubscription inflates agent construction from ~3s to ~30s each, blowing +# the 120s per-test timeout intermittently (https://nvbugs/6426834). One +# progress thread is enough here: these tests verify transfer logic, not +# transfer-engine threading. +os.environ["TRTLLM_NIXL_NUM_THREADS"] = "1" from dataclasses import dataclass from typing import List, Optional diff --git a/tests/unittest/disaggregated/test_kv_transfer_mp.py b/tests/unittest/disaggregated/test_kv_transfer_mp.py index 76e5f4aa7573..a52d95bde7a4 100644 --- a/tests/unittest/disaggregated/test_kv_transfer_mp.py +++ b/tests/unittest/disaggregated/test_kv_transfer_mp.py @@ -8,7 +8,13 @@ import torch.multiprocessing as mp # Exclude IB (no fabric) and gdr_copy (UCX rcache SIGABRT at teardown). -os.environ.setdefault("UCX_TLS", "^ib,gdr_copy") +# Force a deterministic UCX config regardless of what the cluster/CI injects; +# see test_kv_transfer.py for the full rationale. +os.environ["UCX_TLS"] = "^ib,gdr_copy" +# Limit NIXL busy-polling progress threads; see test_kv_transfer.py for the +# full rationale (intermittent 120s timeouts on shared CI nodes, +# https://nvbugs/6426834). +os.environ["TRTLLM_NIXL_NUM_THREADS"] = "1" import tensorrt_llm import tensorrt_llm.bindings diff --git a/tests/unittest/disaggregated/test_mamba_transfer.py b/tests/unittest/disaggregated/test_mamba_transfer.py index 0cf0b3209899..8688d75f9a53 100644 --- a/tests/unittest/disaggregated/test_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_mamba_transfer.py @@ -12,6 +12,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import os import threading import uuid from typing import Dict, List @@ -19,6 +20,11 @@ import pytest import torch +# Force a deterministic UCX/NIXL config regardless of what the cluster/CI +# injects; see test_kv_transfer.py for the full rationale. +os.environ["UCX_TLS"] = "^ib,gdr_copy" +os.environ["TRTLLM_NIXL_NUM_THREADS"] = "1" + import tensorrt_llm import tensorrt_llm.bindings import tensorrt_llm.tensorrt_llm_transfer_agent_binding # noqa: F401 diff --git a/tests/unittest/disaggregated/test_py_cache_transceiver_mp.py b/tests/unittest/disaggregated/test_py_cache_transceiver_mp.py index a81e05026cbb..d414d18fb48b 100644 --- a/tests/unittest/disaggregated/test_py_cache_transceiver_mp.py +++ b/tests/unittest/disaggregated/test_py_cache_transceiver_mp.py @@ -15,8 +15,10 @@ import torch.distributed as dist import torch.multiprocessing as mp -# Exclude IB (no fabric) and gdr_copy (UCX rcache SIGABRT at teardown). -os.environ.setdefault("UCX_TLS", "^ib,gdr_copy") +# Force a deterministic UCX/NIXL config regardless of what the cluster/CI +# injects; see test_kv_transfer.py for the full rationale. +os.environ["UCX_TLS"] = "^ib,gdr_copy" +os.environ["TRTLLM_NIXL_NUM_THREADS"] = "1" import tensorrt_llm import tensorrt_llm.bindings diff --git a/tests/unittest/others/test_kv_cache_transceiver.py b/tests/unittest/others/test_kv_cache_transceiver.py index 0d6b31e367d1..241a4aee0d09 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -3,6 +3,7 @@ import gc import multiprocessing +import os import sys import time import uuid @@ -11,6 +12,11 @@ import pytest import torch +# Force a deterministic UCX/NIXL config regardless of what the cluster/CI +# injects; see disaggregated/test_kv_transfer.py for the full rationale. +os.environ["UCX_TLS"] = "^ib,gdr_copy" +os.environ["TRTLLM_NIXL_NUM_THREADS"] = "1" + import tensorrt_llm import tensorrt_llm.bindings import tensorrt_llm.bindings.executor as trtllm From 2801e9478987da398c18b868cef58ebf5484e5da Mon Sep 17 00:00:00 2001 From: "Wang, Xiao" <24860335+xwang233@users.noreply.github.com> Date: Thu, 23 Jul 2026 20:22:03 -0700 Subject: [PATCH 06/13] [TRTLLM-14575][fix] MoE: fp32 accumulation in deferred MoEAllReduce finalize (#16778) Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../moeAllReduceFusionKernels.cu | 31 ++++++++++++++----- 1 file changed, 23 insertions(+), 8 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.cu index d1f50a2fb9c8..5b0da0d0e80c 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.cu @@ -557,11 +557,20 @@ __global__ void moefinalize_allreduce_fusion_kernel_oneshot_lamport(MoeFinalizeA } // * MoE finalize - ACC_TYPE accumulator; + // Accumulate the top-k weighted expert sum and the shared-expert add in + // fp32 (local `facc`), rounding to DType (bf16/fp16) only once when packing + // into `accumulator` for the 128-bit Lamport all-reduce store below. + // Accumulating directly in DType here rounds after every one of the top_k + // terms; across the many routed MoE layers this rounding bias is large + // enough to visibly degrade the routed output, and with attention-DP + // disabled + MTP speculative decoding it drifts the target hidden states + // enough to lower the acceptance length. The non-deferred in-kernel + // finalize (do_finalize=true) already accumulates in fp32; match it here. + float facc[kElemsPerAccess]; #pragma unroll for (int i = 0; i < kElemsPerAccess; ++i) { - accumulator.unpacked[i] = static_cast(0); + facc[i] = 0.f; } for (int k = 0; k < top_k; k++) @@ -583,17 +592,15 @@ __global__ void moefinalize_allreduce_fusion_kernel_oneshot_lamport(MoeFinalizeA permuted_data.packed = reinterpret_cast(params.allreduce_in)[thread_offset_across_token / kElemsPerAccess]; - // * acc += scale(data) + // * acc += scale(data) (fp32 accumulation) #pragma unroll for (int i = 0; i < kElemsPerAccess; ++i) { - // assume computation is done in ScaleType - accumulator.unpacked[i] - += static_cast((static_cast(permuted_data.unpacked[i]) * block_scale)); + facc[i] += static_cast(permuted_data.unpacked[i]) * block_scale; } } - // * Add shared expert output + // * Add shared expert output (fp32 accumulation) if (params.shared_expert_output) { // * Load shared expert output @@ -604,10 +611,18 @@ __global__ void moefinalize_allreduce_fusion_kernel_oneshot_lamport(MoeFinalizeA #pragma unroll for (int i = 0; i < kElemsPerAccess; ++i) { - accumulator.unpacked[i] += shared_expert_output.unpacked[i]; + facc[i] += static_cast(shared_expert_output.unpacked[i]); } } + // Round the fp32 accumulator to DType once, packed for the Lamport AR store. + ACC_TYPE accumulator; +#pragma unroll + for (int i = 0; i < kElemsPerAccess; ++i) + { + accumulator.unpacked[i] = static_cast(facc[i]); + } + // * AR Store int access_id = token_id * params.hidden_dim / kElemsPerAccess + access_id_in_token; int idx = access_id; From 16bdc6d875cee996bb3066c0ba103aa4610992b0 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Fri, 24 Jul 2026 11:25:14 +0800 Subject: [PATCH 07/13] [None][feat] Support MARLIN MoE with MTP and attention DP + EP (#16597) Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_utils.py | 62 ++++++++++++------- .../modules/fused_moe/MOE_DEVELOPER_GUIDE.md | 2 +- .../_torch/modules/fused_moe/create_moe.py | 11 +++- .../modules/fused_moe/fused_moe_marlin.py | 34 +++------- .../_torch/modules/fused_moe/moe_scheduler.py | 8 ++- .../_torch/pyexecutor/model_loader.py | 10 ++- .../defs/accuracy/test_llm_api_pytorch.py | 40 ++++++++++++ .../test_lists/qa/llm_function_core.txt | 1 + .../test_lists/test-db/l0_dgx_h100.yml | 4 +- .../test_lists/test-db/l0_dgx_h200.yml | 1 + .../_torch/modules/moe/test_moe_backend.py | 41 +++++++++++- .../_torch/modules/moe/test_moe_module.py | 54 +++++++++++----- tests/unittest/_torch/test_model_config.py | 30 +++++++++ 13 files changed, 226 insertions(+), 72 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 237ed00af403..94d290d4f192 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -387,11 +387,46 @@ def __init__(self, model: TModel, *, config: ModelConfig[TConfig], self.pp_size = config.mapping.pp_size self.has_custom_lm_head = False + # Per-layer quant entry for lm_head (e.g. ModelOpt MIXED_PRECISION + # checkpoints that quantize lm_head to NVFP4). Model-specific + # config normalizers opt in by keeping/synthesizing this entry; + # exclude_modules still wins. Applies to both the attention-DP + # (replicated) and TP lm_head below. + lm_head_quant_config = None + if config.quant_config_dict is not None: + lm_head_quant_config = config.quant_config_dict.get("lm_head") + if (lm_head_quant_config is not None + and config.quant_config is not None and config.quant_config. + is_module_excluded_from_quantization("lm_head")): + lm_head_quant_config = None + if (lm_head_quant_config is not None and getattr( + config.pretrained_config, 'tie_word_embeddings', False)): + # Tied embeddings replace lm_head.weight with the dense + # bf16 embedding weight below, which would silently clash + # with a quantized (packed) weight and its quant method. + logger.info( + "Ignoring lm_head quant entry: tie_word_embeddings " + "shares the dense embedding weight, so lm_head stays " + "unquantized") + lm_head_quant_config = None + if (lm_head_quant_config is not None + and config.mapping.enable_attention_dp + and config.mapping.enable_lm_head_tp_in_adp): + # lm_head TP in ADP slices the dense weight at forward time + # for the spec-decoding head (see LMHead.forward), which is + # incompatible with quantized (packed) weights — LMHead + # rejects that combination at construction. + logger.info( + "Ignoring lm_head quant entry: lm_head TP in ADP " + "slices the dense weight, so lm_head stays unquantized") + lm_head_quant_config = None + if config.mapping.enable_attention_dp and not config.mapping.enable_lm_head_tp_in_adp: self.lm_head = LMHead( vocab_size, hidden_size, dtype=config.pretrained_config.torch_dtype, + quant_config=lm_head_quant_config, ) else: if (hasattr(config, 'lora_config') @@ -405,29 +440,10 @@ def __init__(self, model: TModel, *, config: ModelConfig[TConfig], self.has_custom_lm_head = True vocab_size = lora_loader.vocab_size - # Per-layer quant entry for lm_head (e.g. ModelOpt MIXED_PRECISION - # checkpoints that quantize lm_head to NVFP4). Model-specific - # config normalizers opt in by keeping/synthesizing this entry; - # exclude_modules still wins. - lm_head_quant_config = None - if not self.has_custom_lm_head and config.quant_config_dict is not None: - lm_head_quant_config = config.quant_config_dict.get("lm_head") - if (lm_head_quant_config is not None - and config.quant_config is not None - and config.quant_config. - is_module_excluded_from_quantization("lm_head")): - lm_head_quant_config = None - if (lm_head_quant_config is not None - and getattr(config.pretrained_config, - 'tie_word_embeddings', False)): - # Tied embeddings replace lm_head.weight with the dense - # bf16 embedding weight below, which would silently clash - # with a quantized (packed) weight and its quant method. - logger.info( - "Ignoring lm_head quant entry: tie_word_embeddings " - "shares the dense embedding weight, so lm_head stays " - "unquantized") - lm_head_quant_config = None + # A custom LoRA lm_head replaces the checkpoint weight with a + # dense bf16 tensor, so the quant entry must not apply. + if self.has_custom_lm_head: + lm_head_quant_config = None self.lm_head = LMHead( vocab_size, diff --git a/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md index b84f0e5c4577..77895edb07ff 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -150,7 +150,7 @@ Still on old path (standalone, with embedded communication): | `fused_moe_cute_dsl_b12x.py` | `CuteDslB12xFusedMoE` | SM120/SM121 | NVFP4 hybrid CUTLASS-prefill / FlashInfer NVFP4 MoE decode — best perf on RTX PRO 6000 (SM120) and DGX Spark (SM121); select via the `CUTEDSL` backend path (auto-promoted when flashinfer is importable) | `EXTERNAL_COMM` | | `mega_moe/mega_moe_deepgemm.py` | `MegaMoEDeepGemm` | SM100/SM103 | W4A8_MXFP4_MXFP8 via DeepGEMM `fp8_fp4_mega_moe` fused dispatch+GEMM+act+GEMM+combine kernel; requires `hidden_size % 512 == 0` | `FUSED_COMM` | | `mega_moe/mega_moe_cute_dsl.py` | `MegaMoECuteDsl` | SM100/SM103 | NVFP4 via ported CuteDSL `Sm100MegaMoEKernel` fused dispatch+FC1+act+FC2+combine kernel; requires CUDA 13 Cutlass DSL runtime (PR #14354) and NVSHMEM provider (hard gate); threads per-expert `fc31_alpha`/`fc2_alpha`/`fc1_norm_const` through the kernel ABI and supports SwiGLU clamp via `swiglu_limit`; default deepgemm graph (topk score folded before fc1-out quant, host `combine_output.sum(dim=1)`) | `FUSED_COMM` | -| `fused_moe_marlin.py` | `MarlinFusedMoE` | SM90 only | W4A16 NVFP4 on Hopper (BF16 activations + FP4 weights, fused single-launch `marlin_nvfp4_moe_gemm` kernel); no dynamic EPLB | `EXTERNAL_COMM` | +| `fused_moe_marlin.py` | `MarlinFusedMoE` | SM90 only | W4A16 NVFP4 on Hopper (BF16 activations + FP4 weights, fused single-launch `marlin_nvfp4_moe_gemm` kernel); supports attention-DP + EP via external comm (scheduler precomputes routing; dispatch payload is plain BF16, no activation scales); non-NVFP4 layers (e.g. unquantized MTP draft layers) fall back to Cutlass in `get_moe_cls`; no dynamic EPLB | `EXTERNAL_COMM` | | `fused_moe_triton.py` | `TritonFusedMoE` | SM90 only | GPT-OSS on Hopper (requires `swiglu_gptoss_style=True`) | (legacy path) | | `fused_moe_wide_ep.py` | `WideEPMoE` | All GPUs | Deprecating — use ConfigurableMoE instead | (legacy path) | | `fused_moe_vanilla.py` | `VanillaMoE` | All devices | Reference / debugging only | (legacy path) | diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py index 1f3e9b112dd0..3761fab5a06e 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py @@ -64,10 +64,15 @@ def get_moe_cls( quant_config = override_quant_config layer_prefix = f"[layer_idx={layer_idx}] " if layer_idx is not None else "" if moe_backend.upper() == "MARLIN": - # Marlin MoE is a Hopper-specific NVFP4 W4A16 backend. Require nvfp4 - # quantization explicitly so a misconfigured model fails fast. + # Marlin MoE is a Hopper-specific NVFP4 W4A16 backend. Layers without + # NVFP4 quantization (e.g. deliberately-unquantized MTP draft layers in + # MIXED_PRECISION checkpoints) fall back to CutlassFusedMoE, matching + # the CUTEDSL / DENSEGEMM / MEGAMOE_* fallback behavior below. if quant_config is None or not quant_config.quant_mode.has_nvfp4(): - raise ValueError("MarlinFusedMoE only supports NVFP4 quantization.") + logger.warning(f"{layer_prefix}MarlinFusedMoE only supports NVFP4 " + "quantization. Check out details in quant_config: " + f"{quant_config}. Using CutlassFusedMoE instead.") + return CutlassFusedMoE return MarlinFusedMoE if moe_backend.upper() == "CUTLASS": return CutlassFusedMoE diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py index e6e9f3c7ea7c..c4728be2b140 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py @@ -104,31 +104,6 @@ def _get_quant_method(self): def _supports_load_balancer(self) -> bool: return False - def validate_configurable_moe(self, moe) -> None: - """Reject configs that require external-communication MoE dispatch. - - Marlin is W4A16 (``quantize_input`` produces no activation scale) and - routes internally inside ``run_moe``. The host-side all-to-all - dispatch/combine path that ConfigurableMoE uses for attention-DP - expert parallelism needs the routing decided *before* dispatch and a - per-token scale payload, so it is incompatible with Marlin and fails - inside ``moe_a2a_dispatch``. - - ConfigurableMoE only creates an external communication strategy when - ``enable_attention_dp and dp_size > 1`` (see CommunicationFactory); - ``moe.comm`` is not assigned yet when this hook runs, so check that - same condition via the mapping. Single-GPU, and TP/EP without - attention DP, are supported. - """ - if moe.use_dp and moe.mapping.dp_size > 1: - raise ValueError( - "MarlinFusedMoE does not support external-communication MoE " - "(attention data parallelism combined with expert " - "parallelism): its W4A16 layout and internal routing are " - "incompatible with the all-to-all dispatch path. Use Marlin " - "single-node, or with TP/EP without attention DP." - ) - def _apply_activation(self, gemm1_out: torch.Tensor) -> torch.Tensor: """Apply the activation function to the gemm1 output. @@ -202,6 +177,15 @@ def run_moe( local_n = self.expert_size_per_partition if local_n != self.num_experts: + # EP: non-local token-expert pairs are clamped to local expert 0 + # with a zero final scale, so they still run through both GEMMs + # and are discarded at the combine. + # TODO(perf): skip them instead — e.g. mark non-local pairs with a + # sentinel expert id that moe_align_block_size drops. Requires + # bounds guards in moeAlignKernels.cu (the histogram and + # count_and_sort kernels index shared/cumsum buffers with the raw + # id) and zero-initializing unscheduled GEMM output rows before + # the index_add_ combine. slot_start = self.slot_start is_local = (token_selected_experts >= slot_start) & ( token_selected_experts < slot_start + local_n diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py index fe0153386763..5c6ca72d176f 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py @@ -374,8 +374,14 @@ def _forward_chunk_impl( moe._load_balancer_start_wait_gpu_stage(is_first_call) # ========== Step 2: Apply routing ========== + # External dispatch (Step 5) sends per-token expert/scale payloads, so + # routing must be precomputed whenever a comm strategy is active — even + # for backends whose run_moe can otherwise route internally from + # router_logits (e.g. MarlinFusedMoE under attention-DP + EP). requires_separated_routing = ( - moe.backend._supports_load_balancer() or moe.routing_method.requires_separated_routing + moe.backend._supports_load_balancer() + or moe.routing_method.requires_separated_routing + or moe.comm is not None ) if requires_separated_routing: # Separated routing: ConfigurableMoE calls routing_method diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 9f3218f2e192..f2284e06bb46 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -71,7 +71,7 @@ def validate_and_set_mamba_ssm_cache_dtype( def validate_and_set_kv_cache_quant(model_config: ModelConfig, - pyt_kv_cache_dtype: str) -> QuantAlgo: + pyt_kv_cache_dtype: str) -> None: logger.info( f'Validating KV Cache config against kv_cache_dtype="{pyt_kv_cache_dtype}"' ) @@ -105,6 +105,14 @@ def validate_and_set_kv_cache_quant(model_config: ModelConfig, # Apply explicit override from kv_cache_config.dtype. model_config.quant_config.kv_cache_quant_algo = mapped_pyt_quant + # MIXED_PRECISION checkpoints carry per-layer QuantConfigs in + # quant_config_dict; modules built from them (e.g. attention) must agree + # with the global config on the KV element size, otherwise the KV pool is + # allocated with the overridden dtype while attention layers read/write + # with the checkpoint dtype -> out-of-bounds access. + if model_config.quant_config_dict is not None: + for layer_quant_config in model_config.quant_config_dict.values(): + layer_quant_config.kv_cache_quant_algo = mapped_pyt_quant def validate_encoder_decoder_kv_cache_config(model_config: ModelConfig, diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 5e0d6e4cc74d..378aa74751db 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -7034,6 +7034,46 @@ def test_nvfp4_marlin_multi_gpus(self, tp_size): task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + @skip_pre_hopper + @skip_post_hopper + @pytest.mark.skip_less_device_memory(80000) + @pytest.mark.skip_less_mpi_world_size(4) + @parametrize_with_ids("mtp_nextn", [3]) + def test_nvfp4_marlin_adp_4gpus(self, mtp_nextn): + """Accuracy guard for MARLIN under attention DP + EP with MTP. + + The NVFP4 checkpoint is MIXED_PRECISION with deliberately-unquantized + MTP draft layers, so this also guards the per-layer + MARLIN -> Cutlass fallback in ``create_moe.get_moe_cls`` while the + main-model expert layers stay on MARLIN. Attention DP exercises the + external-comm dispatch path with scheduler-precomputed routing. + """ + model_path = f"{llm_models_root()}/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4" + mtp_config = MTPDecodingConfig(num_nextn_predict_layers=mtp_nextn, + mtp_eagle_one_model=True) + with LLM(model_path, + tensor_parallel_size=4, + moe_expert_parallel_size=4, + enable_attention_dp=True, + moe_config=MoeConfig(backend="MARLIN"), + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, + mamba_ssm_cache_dtype="float16", + free_gpu_memory_fraction=0.8, + ), + max_batch_size=32, + cuda_graph_config=CudaGraphConfig(max_batch_size=32, + enable_padding=True), + speculative_config=mtp_config, + nvfp4_gemm_config={"allowed_backends": ["marlin"]}) as llm: + assert llm.args.quant_config.quant_algo == QuantAlgo.MIXED_PRECISION + task = MMLU(self.MODEL_NAME) + task.evaluate(llm, + extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm, + extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + @skip_pre_hopper @pytest.mark.skip_less_mpi_world_size(4) @pytest.mark.skip_less_device_memory(40000) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index d675ad632e6c..0e44fdeab6ca 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -685,6 +685,7 @@ accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp_custom_op accuracy/test_llm_api_pytorch.py::TestNemotronV3Nano::test_nvfp4_marlin_multi_gpus[tp_size=8] accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_marlin_multi_gpus[tp_size=8] +accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_marlin_adp_4gpus[mtp_nextn=3] accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_marlin_8gpus accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_parallelism[TEP4_PP2] accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_parallelism[TEP8_PP1] diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index 65d4829c1811..6d63a7782a53 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -168,12 +168,14 @@ l0_dgx_h100: orchestrator: mpi tests: # ------------- MoE: test_multi_gpu (by backend x quant) --------------- - # Only CUTLASS backend runs on H100 (SM90). TRTLLM/CUTEDSL/DEEPGEMM require SM100+. + # Only CUTLASS and MARLIN backends run on H100 (SM90). TRTLLM/CUTEDSL/DEEPGEMM require SM100+. # --- CUTLASS --- - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and FP8_BLOCK_SCALES" - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W8A16" - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W4A16_MXFP4" - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W4A8_AWQ" + # --- MARLIN (SM90-only; focused DEP + ALLGATHER x NVFP4 matrix) --- + - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "MARLIN" # ------------- MoE: test_multi_gpu_eplb --------------- - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb - condition: diff --git a/tests/integration/test_lists/test-db/l0_dgx_h200.yml b/tests/integration/test_lists/test-db/l0_dgx_h200.yml index f707f7e61036..3b40b905220b 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h200.yml @@ -25,6 +25,7 @@ l0_dgx_h200: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload_mtp3_no_adp] - accuracy/test_llm_api_pytorch.py::TestNemotronV3Nano::test_nvfp4_marlin_multi_gpus[tp_size=8] - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_marlin_multi_gpus[tp_size=8] + - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_marlin_adp_4gpus[mtp_nextn=3] - accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_marlin_8gpus - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=0-overlap_scheduler=True] - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=0-overlap_scheduler=False] diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index 060472b5b8be..9daeb099adcf 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py @@ -57,7 +57,9 @@ DeepSeekV3MoeRoutingMethod, RenormalizeMoeRoutingMethod, ) -from tensorrt_llm._torch.modules.fused_moe.create_moe import create_moe_backend +from tensorrt_llm._torch.modules.fused_moe.create_moe import create_moe_backend, get_moe_cls +from tensorrt_llm._torch.modules.fused_moe.fused_moe_cutlass import CutlassFusedMoE +from tensorrt_llm._torch.modules.fused_moe.fused_moe_marlin import MarlinFusedMoE from tensorrt_llm._torch.modules.fused_moe.interface import MoE, MoEWeightLoadingMode from tensorrt_llm._torch.modules.fused_moe.mega_moe import MegaMoECuteDsl, MegaMoEDeepGemm from tensorrt_llm._torch.modules.fused_moe.quantization import ( @@ -69,7 +71,7 @@ from tensorrt_llm._torch.utils import ActivationType, is_gated_activation from tensorrt_llm._utils import mpi_rank from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.modeling_utils import QuantAlgo +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig logger = logging.getLogger(__name__) @@ -329,6 +331,41 @@ def test_marlin_moe_repack_is_transform_stage(): assert NVFP4MarlinFusedMoEMethod.post_load_weights is FusedMoEMethodBase.post_load_weights +def _marlin_model_config(quant_algo=QuantAlgo.NVFP4): + cfg = ModelConfig() + cfg.moe_backend = "MARLIN" + cfg.quant_config = QuantConfig(quant_algo=quant_algo) if quant_algo else None + return cfg + + +def test_get_moe_cls_marlin_selects_marlin_for_nvfp4(): + assert get_moe_cls(_marlin_model_config()) is MarlinFusedMoE + + +@pytest.mark.parametrize( + "quant_algo", + [ + pytest.param(None, id="unquantized"), + pytest.param(QuantAlgo.FP8, id="fp8"), + ], +) +def test_get_moe_cls_marlin_falls_back_to_cutlass_on_non_nvfp4(quant_algo): + """MARLIN + non-NVFP4 layers (e.g. unquantized MTP draft layers in + MIXED_PRECISION checkpoints) fall back to CutlassFusedMoE instead of + raising, matching CUTEDSL/DENSEGEMM fallback behavior.""" + assert get_moe_cls(_marlin_model_config(quant_algo)) is CutlassFusedMoE + + +def test_get_moe_cls_marlin_override_quant_config_per_layer(): + """Per-layer override (the MTP draft-layer path): an unquantized per-layer + override falls back to Cutlass even though the global config is NVFP4.""" + cfg = _marlin_model_config() + assert ( + get_moe_cls(cfg, override_quant_config=QuantConfig(quant_algo=None), layer_idx=52) + is CutlassFusedMoE + ) + + def test_megamoe_cutedsl_post_load_weights_uses_staged_hooks(): moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) torch.nn.Module.__init__(moe) diff --git a/tests/unittest/_torch/modules/moe/test_moe_module.py b/tests/unittest/_torch/modules/moe/test_moe_module.py index f9fcaec58a9c..46f95187b45a 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_module.py +++ b/tests/unittest/_torch/modules/moe/test_moe_module.py @@ -1343,24 +1343,33 @@ def generate_multi_gpu_test_params( return params -def _generate_megamoe_multi_gpu_test_params( +def _generate_focused_multi_gpu_test_params( *, backend_type, quant_algo, - should_skip_fn, + parallel_modes, + comm_methods, + should_skip_fn=None, ) -> List: - """Generate focused MegaMoE module multi-GPU coverage for one backend. - - Both MegaMoE backends share the same multi-GPU matrix shape; only the - backend/quant enum and capability skip hook differ between DeepGemm - (W4A8_MXFP4_MXFP8) and CuteDsl (NVFP4). The comm method is hardcoded to - the ``IGNORE`` sentinel because the fused kernel owns dispatch/combine - (the worker pops ``TRTLLM_FORCE_COMM_METHOD`` and takes the fused path). + """Generate focused module multi-GPU coverage for one backend. + + Shared by the backends that opt out of the full ``COMM_METHODS`` matrix: + + - MegaMoE (DeepGemm / CuteDsl): the fused kernel owns dispatch/combine, + so the comm method is the ``IGNORE`` sentinel (the worker pops + ``TRTLLM_FORCE_COMM_METHOD`` and takes the fused path) and each backend + passes its capability skip hook via ``should_skip_fn``. + - Marlin: an EXTERNAL_COMM backend whose ``run_moe`` can route internally, + but under attention-DP the scheduler precomputes routing and dispatches + plain BF16 activations (W4A16 — no activation scales). Coverage is + pinned to ALLGATHER, which is available on every SM90 box; the NVLink + a2a strategies in ``COMM_METHODS`` require MNNVL fabric that Hopper CI + nodes lack. """ params: List = [] seq_lens = [8] if IS_CI_MODE else SEQ_LENS - for parallel_mode, comm_method in product(MEGAMOE_PARALLEL_MODES, [MEGAMOE_IGNORE_COMM_METHOD]): + for parallel_mode, comm_method in product(parallel_modes, comm_methods): for ( swiglu_alpha, swiglu_beta, @@ -1382,7 +1391,7 @@ def _generate_megamoe_multi_gpu_test_params( [quant_algo], MULTI_GPU_ROUTING_METHODS, ): - if not skip_reason: + if not skip_reason and should_skip_fn is not None: skip_reason = should_skip_fn( parallel_mode, comm_method, @@ -1676,19 +1685,34 @@ def test_trtllm_gen_fp32_routing_bias(routing_method_cls, moe_model_config, quan seq_lens=[8] if IS_CI_MODE else SEQ_LENS, dtypes=DTYPES, backend_types=[ - b for b in BACKEND_TYPES if b != MoeBackendType.MARLIN - ], # Marlin doesn't support fused routing + # Marlin gets focused ALLGATHER coverage below; the NVLink a2a + # strategies in COMM_METHODS require MNNVL fabric unavailable on + # the Hopper (SM90) nodes Marlin runs on. + b + for b in BACKEND_TYPES + if b != MoeBackendType.MARLIN + ], quant_algos=QUANT_ALGOS, routing_methods=MULTI_GPU_ROUTING_METHODS, ) -MULTI_GPU_TEST_PARAMS += _generate_megamoe_multi_gpu_test_params( +MULTI_GPU_TEST_PARAMS += _generate_focused_multi_gpu_test_params( + backend_type=MoeBackendType.MARLIN, + quant_algo=QuantAlgo.NVFP4, + parallel_modes=["DEP"] if IS_CI_MODE else ["DEP", "TEP"], + comm_methods=["ALLGATHER"], +) +MULTI_GPU_TEST_PARAMS += _generate_focused_multi_gpu_test_params( backend_type=MoeBackendType.MEGAMOE_DEEPGEMM, quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8, + parallel_modes=MEGAMOE_PARALLEL_MODES, + comm_methods=[MEGAMOE_IGNORE_COMM_METHOD], should_skip_fn=should_skip_MegaMoEDeepGemm, ) -MULTI_GPU_TEST_PARAMS += _generate_megamoe_multi_gpu_test_params( +MULTI_GPU_TEST_PARAMS += _generate_focused_multi_gpu_test_params( backend_type=MoeBackendType.MEGAMOE_CUTEDSL, quant_algo=QuantAlgo.NVFP4, + parallel_modes=MEGAMOE_PARALLEL_MODES, + comm_methods=[MEGAMOE_IGNORE_COMM_METHOD], should_skip_fn=should_skip_MegaMoECuteDsl, ) diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py index 775f3e7b6be1..28b2b0f9a96a 100644 --- a/tests/unittest/_torch/test_model_config.py +++ b/tests/unittest/_torch/test_model_config.py @@ -134,6 +134,36 @@ def test_validate_and_set_kv_cache_quant_rejects_invalid_dtype(): validate_and_set_kv_cache_quant(model_config, "invalid_dtype") +def _make_mixed_precision_model_config(): + """MIXED_PRECISION checkpoint shape: global config plus per-layer entries + whose kv_cache_quant_algo comes only from hf_quant_config.json (None here).""" + return ModelConfig( + quant_config=QuantConfig(quant_algo=QuantAlgo.MIXED_PRECISION, kv_cache_quant_algo=None), + quant_config_dict={ + "model.layers.0.attention": QuantConfig(quant_algo=QuantAlgo.FP8), + "model.layers.0.mixer.experts": QuantConfig(quant_algo=QuantAlgo.NVFP4), + }, + ) + + +def test_validate_and_set_kv_cache_quant_propagates_to_quant_config_dict(): + """Explicit kv_cache_config.dtype must override the per-layer QuantConfigs + too, otherwise the KV pool (sized from the global config) and attention + modules (built from per-layer configs) disagree on KV element size.""" + model_config = _make_mixed_precision_model_config() + validate_and_set_kv_cache_quant(model_config, "fp8") + assert model_config.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 + for layer_quant_config in model_config.quant_config_dict.values(): + assert layer_quant_config.kv_cache_quant_algo == QuantAlgo.FP8 + + +def test_validate_and_set_kv_cache_quant_auto_keeps_quant_config_dict(): + model_config = _make_mixed_precision_model_config() + validate_and_set_kv_cache_quant(model_config, "auto") + for layer_quant_config in model_config.quant_config_dict.values(): + assert layer_quant_config.kv_cache_quant_algo is None + + def _write_safetensors_header(checkpoint_dir, tensor_dtype, tensor_shape): shard_name = "model-00001-of-00001.safetensors" header = { From 37f59acbb59a84312320d8535bfc357001d2c203 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Fri, 24 Jul 2026 11:48:35 +0800 Subject: [PATCH 08/13] [https://nvbugs/6485885][fix] Stop thinking-budget processor re-forcing the reasoning end tag (#16785) Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- tensorrt_llm/llmapi/thinking_budget.py | 40 +++++++++++-- tests/unittest/llmapi/test_sampling_params.py | 60 +++++++++++++++++++ 2 files changed, 94 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/llmapi/thinking_budget.py b/tensorrt_llm/llmapi/thinking_budget.py index d7ca6a5a58a8..35e68d48d1c8 100644 --- a/tensorrt_llm/llmapi/thinking_budget.py +++ b/tensorrt_llm/llmapi/thinking_budget.py @@ -33,6 +33,19 @@ def __init__( self.thinking_token_budget = budget self.reasoning_start_token_ids = list(reasoning_start_token_ids) self.reasoning_end_token_ids = list(reasoning_end_token_ids) + # Progress through the reasoning end sequence, keyed by (req_id, beam_idx): + # 0 = not closing, 1..len-1 = mid-sequence, >= len = fully emitted. A + # missing key defaults to 0; the key is popped when the block closes. + # + # We keep this instead of re-deriving from `token_ids` because the overlap + # scheduler runs the processor one step behind the sampled tokens: a stale + # view would not show the just-forced end sequence, so re-derivation would + # force it again and leak a duplicate end tag (e.g. ``). + # + # Keys are reclaimed only when a block is seen to close, so entries for + # blocks that never close persist for the processor's lifetime -- bounded + # per request, but a slow drip if one SamplingParams is shared across many. + self._end_progress: dict = {} def __call__( self, @@ -42,28 +55,41 @@ def __call__( stream_ptr: Optional[int], client_id: Optional[int], ) -> None: - del req_id, client_id + del client_id if stream_ptr is None: - self._apply(token_ids, logits) + self._apply(req_id, token_ids, logits) return with torch.cuda.stream(torch.cuda.ExternalStream(stream_ptr)): - self._apply(token_ids, logits) + self._apply(req_id, token_ids, logits) - def _apply(self, token_ids: List[List[int]], logits: torch.Tensor) -> None: + def _apply(self, req_id: int, token_ids: List[List[int]], logits: torch.Tensor) -> None: for beam_idx, beam_token_ids in enumerate(token_ids): - forced_token = self._forced_token(beam_token_ids) + forced_token = self._forced_token((req_id, beam_idx), beam_token_ids) if forced_token is not None: self._force_token(logits, beam_idx, len(token_ids), forced_token) - def _forced_token(self, token_ids: List[int]) -> Optional[int]: + def _forced_token(self, key, token_ids: List[int]) -> Optional[int]: start_idx = _find_last_sequence_index(token_ids, self.reasoning_start_token_ids) if start_idx == -1: return None end_idx = _find_last_sequence_index(token_ids, self.reasoning_end_token_ids) if end_idx > start_idx: + # Reasoning block closed; reset so a later block is budgeted again. + self._end_progress.pop(key, None) return None + # Block still open. If we have already begun forcing the end sequence, + # trust recorded progress (see _end_progress) over the lagging token view. + progress = self._end_progress.get(key, 0) + if progress >= len(self.reasoning_end_token_ids): + # Whole end sequence already forced; stop so we don't leak a second tag. + return None + if progress > 0: + # Mid end-sequence: emit the next token and advance (don't restart at 0). + self._end_progress[key] = progress + 1 + return self.reasoning_end_token_ids[progress] + reasoning_start = start_idx + len(self.reasoning_start_token_ids) reasoning_token_count = len(token_ids) - reasoning_start partial_end_len = _longest_suffix_prefix_len(token_ids, self.reasoning_end_token_ids) @@ -71,9 +97,11 @@ def _forced_token(self, token_ids: List[int]) -> Optional[int]: partial_end_len > 0 and reasoning_token_count - partial_end_len >= self.thinking_token_budget ): + self._end_progress[key] = partial_end_len + 1 return self.reasoning_end_token_ids[partial_end_len] if reasoning_token_count >= self.thinking_token_budget: + self._end_progress[key] = 1 return self.reasoning_end_token_ids[0] return None diff --git a/tests/unittest/llmapi/test_sampling_params.py b/tests/unittest/llmapi/test_sampling_params.py index 8988cc137b93..f9e4ad99e5e6 100644 --- a/tests/unittest/llmapi/test_sampling_params.py +++ b/tests/unittest/llmapi/test_sampling_params.py @@ -201,6 +201,66 @@ def test_thinking_budget_logits_processor_ignores_closed_reasoning_block(): assert torch.equal(logits, torch.zeros(1, 1, 8)) +def _run(processor, token_ids): + """Call the processor on a fresh logits row and return it.""" + logits = torch.zeros(1, 1, 8) + processor(0, logits, token_ids, None, None) + return logits + + +def test_thinking_budget_logits_processor_does_not_force_end_twice_on_stale_view(): + """Regression: a stale token view must not re-force the end tag. + + With the overlap scheduler, logits processors run one step + behind the sampled tokens, so the call right after forcing the end token + sees token_ids WITHOUT it. A stateless processor forces the end token a + second time, producing e.g. `` — the reasoning parser + consumes the first tag and the second leaks into message.content. + """ + processor = ThinkingBudgetLogitsProcessor( + thinking_token_budget=2, + reasoning_start_token_ids=[1], + reasoning_end_token_ids=[2], + ) + stale = [[1, 5, 6]] # budget spent, forced end not visible yet + + logits = _run(processor, stale) + assert logits[0, 0, 2] == 0 and torch.isneginf(logits[0, 0, 3]) + + # Next step: the forced end token is still not visible (stale view). + assert torch.equal(_run(processor, stale), torch.zeros(1, 1, 8)) + + # Once the view catches up (end tag present), progress state resets so a + # later reasoning block is budgeted again. + assert torch.equal(_run(processor, [[1, 5, 6, 2, 7]]), torch.zeros(1, 1, 8)) + logits = _run(processor, [[1, 5, 6, 2, 7, 1, 8, 9]]) + assert logits[0, 0, 2] == 0 and torch.isneginf(logits[0, 0, 3]) + + +def test_thinking_budget_logits_processor_continues_end_sequence_on_stale_view(): + """Multi-token end sequence must continue, not restart, on a stale view. + + Under the overlap scheduler the processor must continue from its + recorded progress rather than re-derive from token_ids. + """ + processor = ThinkingBudgetLogitsProcessor( + thinking_token_budget=2, + reasoning_start_token_ids=[1], + reasoning_end_token_ids=[2, 3], + ) + stale = [[1, 5, 6]] + + assert _run(processor, stale)[0, 0, 2] == 0 + + # Stale view: forced first end token not visible; must force the SECOND + # end token, not the first again. + logits = _run(processor, stale) + assert logits[0, 0, 3] == 0 and torch.isneginf(logits[0, 0, 2]) + + # Sequence complete: no further forcing even while the view is stale. + assert torch.equal(_run(processor, stale), torch.zeros(1, 1, 8)) + + def test_add_thinking_budget_logits_processor_uses_reasoning_parser_tokens(): class FakeTokenizer: def encode(self, text, add_special_tokens=False): From 0960d59d221eeac2017b467d4424159716e78c0f Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Fri, 24 Jul 2026 11:52:00 +0800 Subject: [PATCH 09/13] [#16767][fix] Fix DSpark rolling-window slot collision in disaggregated serving (#16772) Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- tensorrt_llm/_torch/speculative/dspark.py | 59 +++++++++- .../hw_agnostic/test_dspark_worker.py | 104 +++++++++++++++++- 2 files changed, 156 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 80425833b0be..7f43f103ec4e 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -29,6 +29,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping +from ..pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID from .interface import SpecMetadata, SpecWorkerBase if TYPE_CHECKING: @@ -112,9 +113,31 @@ def prepare(self): worker._ctx_len[slot] = 0 worker._kv_windows[slot].zero_() worker._free_slots.append(slot) - # Unknown request IDs (e.g. synthetic warmup requests) default to slot 0. + # Assign a persistent rolling-window slot to every real generation + # request that never ran a context/seed forward on this worker. In + # disaggregated serving the prompt is prefilled (and the window + # seeded) on the *context* server, so ``_seed_context_windows`` never + # runs on the generation server and ``_req_to_slot`` stays empty; + # without this, all concurrent gen requests fall through to the shared + # scratch row below and corrupt each other's draft window at batch + # size > 1 (GitHub #16767). Context-prefix entries are left to + # ``_seed_context_windows``; the ADP-idle (id 0) and CUDA-graph + # padding dummies are kept on the scratch row. + num_contexts = max(0, len(self.request_ids) - self.num_generations) + for rid in self.request_ids[num_contexts:]: + if ( + rid != ATTENTION_DP_DUMMY_REQUEST_ID + and rid < worker._graph_dummy_id_floor + and rid not in worker._req_to_slot + ): + worker._assign_slot(rid, reset=False) + # Unknown request IDs (synthetic warmup / CUDA-graph padding, ADP idle + # requests, or disagg seed forwards without a real id) map to the + # dedicated throwaway scratch row so they cannot overwrite a live + # request's rolling window (they previously aliased to slot 0). + scratch = worker._scratch_slot mapping = torch.tensor( - [worker._req_to_slot.get(rid, 0) for rid in self.request_ids], + [worker._req_to_slot.get(rid, scratch) for rid in self.request_ids], dtype=torch.long, device="cpu", pin_memory=prefer_pinned(), @@ -212,6 +235,10 @@ def __init__( self._req_to_slot = {} # request_id -> slot index self._free_slots = deque() # available slot indices self._batch_to_slot: Optional[torch.Tensor] = None # [max_batch] long, cuda + # Index of the throwaway "scratch" window row that absorbs padded / + # unknown request IDs (set in ``_lazy_init`` to ``max_batch``); it is + # never handed out through ``_free_slots``. + self._scratch_slot = 0 # The generation draft path is the batched, host-sync-free # ``_draft_gen_block_batched`` + ``DSparkDraftModel.forward_batched`` + @@ -244,19 +271,41 @@ def _lazy_init(self, draft_model, spec_metadata) -> None: self._win = int(draft_model._attn_params["window_size"]) head_dim = int(draft_model._attn_params["head_dim"]) + # Real requests occupy slots ``[0, max_batch)``; one extra "scratch" row + # at index ``max_batch`` absorbs padded / unknown request IDs (CUDA-graph + # padding, ADP idle requests, or disagg seed forwards that arrive without + # a real request id) so they can never overwrite a live request's rolling + # window. Previously such IDs aliased to slot 0 and corrupted whichever + # real request occupied it. The scratch row is never handed out through + # ``_free_slots`` and its contents are throwaway. + self._scratch_slot = max_batch + num_rows = max_batch + 1 + + # CUDA-graph padding requests carry ids in + # ``[CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len, CUDA_GRAPH_DUMMY_REQUEST_ID]``, + # while real request ids start at ``max_batch_size`` and grow, so a simple + # floor cleanly separates them. Together with ``ATTENTION_DP_DUMMY_REQUEST_ID`` + # (0) these dummies must route to the scratch row (see ``prepare()``) and + # never consume a real slot. Imported lazily to break the + # dspark -> cuda_graph_runner -> speculative.utils -> dspark import cycle. + from ..pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID + + self._graph_dummy_id_floor = CUDA_GRAPH_DUMMY_REQUEST_ID - self.max_draft_len + self._kv_windows = torch.zeros( - (max_batch, num_stages, self._win, head_dim), + (num_rows, num_stages, self._win, head_dim), dtype=torch.bfloat16, device="cuda", ) - self._ctx_len = torch.zeros(max_batch, dtype=torch.long, device="cuda") + self._ctx_len = torch.zeros(num_rows, dtype=torch.long, device="cuda") self._batch_to_slot = torch.zeros(max_batch, dtype=torch.long, device="cuda") self._free_slots = deque(range(max_batch)) self._req_to_slot = {} self._win_inited = True logger.info( f"DSpark: allocated rolling KV windows " - f"[{max_batch}, {num_stages}, {self._win}, {head_dim}]" + f"[{num_rows}, {num_stages}, {self._win}, {head_dim}] " + f"({max_batch} request slots + 1 scratch row)" ) def _assign_slot(self, req_id: int, reset: bool) -> int: diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py index 34ee08006ee9..4625d1b0115a 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -117,8 +117,13 @@ def test_worker_lazy_init_window_buffers(): dm = _fake_draft_model(num_stages=3, window_size=128, head_dim=64) meta = _make_metadata(max_num_requests=8) worker._lazy_init(dm, meta) - assert worker._kv_windows.shape == (8, 3, 128, 64) - assert worker._ctx_len.shape == (8,) + # max_batch (8) request slots + 1 scratch row for padded / unknown IDs. + assert worker._kv_windows.shape == (9, 3, 128, 64) + assert worker._ctx_len.shape == (9,) + assert worker._scratch_slot == 8 + # Dummy-id floor separates real request ids from CUDA-graph padding ids. + assert worker._graph_dummy_id_floor == (1 << 64) - 1 - worker.max_draft_len + # The scratch row is never handed out through the free pool. assert list(worker._free_slots) == list(range(8)) assert worker._batch_to_slot is not None assert worker._batch_to_slot.shape == (8,) @@ -255,6 +260,101 @@ def test_prepare_frees_stale_slots_on_batched_path(): assert int(worker._ctx_len[sa]) == 0 +def test_prepare_maps_unknown_request_to_scratch_row_not_slot_zero(): + """Padded / unknown request IDs route to the scratch row, never a live slot. + + Regression for the rolling-window aliasing bug (GitHub #16767): an unknown + request id (CUDA-graph padding, ADP idle request, or a disagg seed forward + without a real id) must not overwrite the request that owns slot 0. + """ + worker = _make_worker() + meta = _make_metadata(max_num_requests=4) + worker._lazy_init(_fake_draft_model(), meta) + meta._dspark_worker = worker + + # A live request takes the first free slot (0) and populates its window. + s_real = worker._assign_slot(100, reset=True) + assert s_real == 0 + worker._ctx_len[s_real] = 17 + worker._kv_windows[s_real].fill_(1.0) + + # Batch contains the live request plus an unknown id (e.g. graph padding). + meta.request_ids = [100, 999] + meta.prepare() + + # The unknown id maps to the scratch row, not to slot 0. + assert worker._batch_to_slot[:2].tolist() == [s_real, worker._scratch_slot] + assert worker._scratch_slot != s_real + # It is not silently registered as a real request and did not consume a slot. + assert 999 not in worker._req_to_slot + assert list(worker._free_slots) == [1, 2, 3] + # The live request's rolling window and position are untouched. + assert int(worker._ctx_len[s_real]) == 17 + assert torch.all(worker._kv_windows[s_real] == 1.0) + + +def test_prepare_assigns_slots_to_disagg_generation_requests(): + """Disagg gen requests (never seeded here) get distinct slots, not one shared row. + + Regression for GitHub #16767: on the disaggregated generation server the + prompt is prefilled (and the DSpark window seeded) on the *context* server, + so ``_seed_context_windows`` never runs here and ``_req_to_slot`` stays empty. + ``prepare()`` must therefore assign each real generation request its own + rolling-window slot instead of collapsing them all onto the shared scratch + row (which corrupts drafts and collapses accept length at batch size > 1). + """ + worker = _make_worker() + meta = _make_metadata(max_num_requests=4) + worker._lazy_init(_fake_draft_model(), meta) + meta._dspark_worker = worker + + # All-generation batch (num_contexts == 0), no prior seeding. + meta.request_ids = [1000, 1001] + meta.num_generations = 2 + meta.prepare() + + s0 = worker._req_to_slot[1000] + s1 = worker._req_to_slot[1001] + assert s0 != s1 + assert s0 != worker._scratch_slot and s1 != worker._scratch_slot + assert worker._batch_to_slot[:2].tolist() == [s0, s1] + + # Stable across steps: the same ids keep their slots (no churn / reassignment). + meta.prepare() + assert worker._req_to_slot[1000] == s0 + assert worker._req_to_slot[1001] == s1 + assert worker._batch_to_slot[:2].tolist() == [s0, s1] + + +def test_prepare_keeps_dummy_generation_requests_on_scratch_row(): + """ADP-idle (id 0) and CUDA-graph padding dummies never consume a real slot.""" + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID + from tensorrt_llm._torch.pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID + + worker = _make_worker() + meta = _make_metadata(max_num_requests=4) + worker._lazy_init(_fake_draft_model(), meta) + meta._dspark_worker = worker + + graph_dummy = CUDA_GRAPH_DUMMY_REQUEST_ID - worker.max_draft_len + meta.request_ids = [1000, ATTENTION_DP_DUMMY_REQUEST_ID, graph_dummy] + meta.num_generations = 3 + meta.prepare() + + s_real = worker._req_to_slot[1000] + assert s_real != worker._scratch_slot + # Dummies are neither registered nor given a real slot; they map to scratch. + assert ATTENTION_DP_DUMMY_REQUEST_ID not in worker._req_to_slot + assert graph_dummy not in worker._req_to_slot + assert worker._batch_to_slot[:3].tolist() == [ + s_real, + worker._scratch_slot, + worker._scratch_slot, + ] + # Exactly one real slot consumed. + assert list(worker._free_slots) == [1, 2, 3] + + def test_forward_mixed_batch_routes_through_base_entries(monkeypatch): """Mixed (context + gen) batch: ``forward`` must route acceptance and production through the unified ``SpecWorkerBase`` entries, one-hot-fill the From cf492257a86b26c16ecc798ad246412d4cdcb25e Mon Sep 17 00:00:00 2001 From: Emma Qiao Date: Fri, 24 Jul 2026 11:54:02 +0800 Subject: [PATCH 10/13] [None][infra] Fix release check failure for .test_durations (#16784) Signed-off-by: EmmaQiaoCh --- jenkins/UpdateTestDurations.groovy | 1 + jenkins/scripts/generate_duration.py | 1 + tests/integration/defs/.test_durations | 2 +- 3 files changed, 3 insertions(+), 1 deletion(-) diff --git a/jenkins/UpdateTestDurations.groovy b/jenkins/UpdateTestDurations.groovy index 56b25b926617..78938d084e76 100644 --- a/jenkins/UpdateTestDurations.groovy +++ b/jenkins/UpdateTestDurations.groovy @@ -147,6 +147,7 @@ pipeline { echo "Generated file size: \$(wc -l < new_test_durations.json) lines" echo "Sample output (first 5 lines):" head -5 new_test_durations.json + """ // Always archive the freshly generated file so the user can download diff --git a/jenkins/scripts/generate_duration.py b/jenkins/scripts/generate_duration.py index a5221a565dc4..74f9e1687c7b 100644 --- a/jenkins/scripts/generate_duration.py +++ b/jenkins/scripts/generate_duration.py @@ -222,6 +222,7 @@ def main(): with open(NEW_TEST_DURATION, "w") as file: json.dump(test_durations, file, indent=3) + file.write("\n") print("\nSummary:") print(f" OpenSearch index : {OPENSEARCH_INDEX}") diff --git a/tests/integration/defs/.test_durations b/tests/integration/defs/.test_durations index 56f666656e9f..23baf4fc8495 100644 --- a/tests/integration/defs/.test_durations +++ b/tests/integration/defs/.test_durations @@ -1462,4 +1462,4 @@ "verl/test_verl_cases.py::test_wake_sleep_cycle": 118.8309, "visual_gen/test_visual_gen_benchmark.py::test_offline_benchmark": 68.15085714285715, "visual_gen/test_visual_gen_benchmark.py::test_online_benchmark[openai-videos]": 102.87023076923077 -} \ No newline at end of file +} From 58b9d1589d124cd668157c80ff58bd19d0799d52 Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+Tabrizian@users.noreply.github.com> Date: Thu, 23 Jul 2026 20:58:15 -0700 Subject: [PATCH 11/13] [None][perf] spec one-model sampling: greedy rows via top_k=1 instead of unconditional vocab argmax (#16686) Signed-off-by: Iman Tabrizian <10105175+Tabrizian@users.noreply.github.com> --- .../pyexecutor/sampler/sampling_utils.py | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py index 7ca58c7edc08..20bc59316844 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py @@ -950,21 +950,18 @@ def sampling_batch_spec_dec_one_model( ) -> torch.Tensor: """CUDA-graph compatible sampling; supports mixed sampling params. Returns sampled tokens.""" top_k = sanitize_top_k(top_k, logits.shape[-1]) - # Greedy rows (temperature <= threshold) must return the argmax token, not a - # sample from the temperature-scaled distribution. Capture the argmax from the - # *original* logits up front; safely_apply_temperature_inplace then guards the division - # against the greedy sentinel, and torch.where restores the greedy rows below. - # All ops are branch-free (no data-dependent control flow), so this stays - # CUDA-graph safe. + # Greedy rows (temperature <= threshold) reduce to top_k=1 sampling: with the + # divisor clamped to 1.0 by safely_apply_temperature_inplace (order-preserving + # for those rows), flashinfer deterministically returns the max-probability + # token, i.e. the argmax of the original logits. All ops remain branch-free + # (no data-dependent control flow), so this stays CUDA-graph safe. is_greedy = temperatures <= vanilla.GREEDY_TEMPERATURE_THRESHOLD - greedy_tokens = logits.argmax(dim=-1) + top_k = torch.where(is_greedy, torch.ones_like(top_k), top_k) + top_p = torch.where(is_greedy, torch.ones_like(top_p), top_p) logits = vanilla.safely_apply_temperature_inplace(logits, temperatures) - sampled = flashinfer.top_k_top_p_sampling_from_logits_op( + return flashinfer.top_k_top_p_sampling_from_logits_op( logits, top_k, top_p, seed=seed, offset=offset ) - # argmax yields int64; cast so torch.where preserves the sampler's dtype - # (flashinfer returns int32) instead of promoting the result to int64. - return torch.where(is_greedy, greedy_tokens.to(sampled.dtype), sampled) @torch.compile(options={"max-autotune": True}) From 29919947fcc5a108c7af7806401797b5fe6ea84d Mon Sep 17 00:00:00 2001 From: Xianjie Qiao <5410381+qiaoxj07@users.noreply.github.com> Date: Fri, 24 Jul 2026 12:35:01 +0800 Subject: [PATCH 12/13] [None][perf] Avoid implicit device-scalar syncs in DeepSeek-V4 ctx sparse metadata (#16734) Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com> --- .../sparse/deepseek_v4/deepseek_v4.py | 42 ++++++++++++++++-- .../deepseek_v4/test_compressor_module.py | 44 +++++++++++++++++++ 2 files changed, 82 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py index e003d44ff3c8..0dcd458ec943 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py @@ -275,6 +275,7 @@ def __post_init__(self): super().__post_init__() self.num_total_compressed_tokens = {} self.max_ctx_compressed_tokens = {} + self._ctx_output_sizes: Optional[Dict[int, int]] = None sparse_metadata_params = self.sparse_metadata_params if not isinstance(sparse_metadata_params, DeepSeekV4MetadataParams): raise ValueError("DeepSeek-V4 sparse attention metadata params are not set") @@ -737,12 +738,18 @@ def prepare(self): kv_lens_slice = kv_lens[:num_requests] cached_slice = cached_token_lens[:num_requests] + # Host-side per-ratio ctx compressed-token counts (Python ints), so + # _compute_ctx_compressed_position_ids never reads a device scalar + # (implicit D2H + stream sync) for its arange size / slice bound. + ctx_output_sizes: Optional[Dict[int, int]] = None if num_contexts > 0: # Prefill path: need per-request tensor ops for ctx scalar metadata. + ctx_output_sizes = {} for compress_ratio in self.compress_ratio_set: new_comp_kv_lens = kv_lens_slice // compress_ratio - cached_slice // compress_ratio cu_new = new_comp_kv_lens.cumsum(0) num_ctx_compressed_tokens = cu_new[num_contexts - 1].item() + ctx_output_sizes[compress_ratio] = num_ctx_compressed_tokens num_gen_compressed_tokens = num_generations * ( (num_gen_tokens_per_seq + compress_ratio - 1) // compress_ratio ) @@ -761,12 +768,15 @@ def prepare(self): ) self.max_ctx_compressed_tokens[compress_ratio] = 0 + # Cached for on_update_kv_lens(); see the reuse gate there. + self._ctx_output_sizes = ctx_output_sizes + # 2) CUDA-side: fill *_cuda buffers on device. kv_lens_cuda = ( self.cached_token_lens_cuda[:num_requests] + self._seq_lens_cuda[:num_requests] ) cached_tokens_cuda = self.cached_token_lens_cuda[:num_requests] - self.prepare_compressed_kv_metadata(kv_lens_cuda, cached_tokens_cuda) + self.prepare_compressed_kv_metadata(kv_lens_cuda, cached_tokens_cuda, ctx_output_sizes) self._compute_compressed_mask( self.new_comp_kv_lens_cuda, @@ -781,6 +791,7 @@ def prepare_compressed_kv_metadata( self, kv_lens: torch.Tensor, cached_tokens: torch.Tensor, + ctx_output_sizes: Optional[Dict[int, int]] = None, ): """Compute per-ratio compressed KV lens and position IDs on device. @@ -789,6 +800,12 @@ def prepare_compressed_kv_metadata( Args: kv_lens: Total KV lengths per request (device tensor, [batch_size]). cached_tokens: Cached token counts per request (device tensor, [batch_size]). + ctx_output_sizes: Optional per-ratio host-computed ctx + compressed-token counts (Python ints); avoids implicit + device-scalar reads (D2H + stream sync) in the ctx position-id + computation. prepare() always passes it; on_update_kv_lens() + reuses the cached copy unless the extend_ctx path may have + mutated ctx-row kv_lens on device. """ batch_size = kv_lens.shape[0] num_contexts = self.num_contexts @@ -812,6 +829,7 @@ def prepare_compressed_kv_metadata( self.compressed_position_ids_cuda, num_contexts, self._compress_ratios_sorted, + ctx_output_sizes, ) if self.num_gen_tokens_per_seq > 0 and num_generations > 0: @@ -848,7 +866,13 @@ def on_update_kv_lens(self): num_gen_tokens // self.num_generations if self.num_generations > 0 else 0 ) - self.prepare_compressed_kv_metadata(kv_lens, cached_tokens) + # Reuse prepare()'s host-computed ctx sizes unless the extend_ctx path + # (num_chunked_ctx_requests > 0) may have mutated ctx-row kv_lens on + # device; every other path only changes gen rows. + ctx_output_sizes = ( + self._ctx_output_sizes if getattr(self, "num_chunked_ctx_requests", 0) == 0 else None + ) + self.prepare_compressed_kv_metadata(kv_lens, cached_tokens, ctx_output_sizes) self._compute_compressed_mask( self.new_comp_kv_lens_cuda, @@ -1004,14 +1028,24 @@ def _compute_ctx_compressed_position_ids( compressed_position_ids_bufs: Dict[int, torch.Tensor], num_contexts: int, compress_ratios: list, + ctx_output_sizes: Optional[Dict[int, int]] = None, ): - """Context-only compressed position IDs (eager, data-dependent shapes).""" + """Context-only compressed position IDs (eager, data-dependent shapes). + + ctx_output_sizes (host ints) keeps the arange size and slice bound off + the device; the 0-dim-CUDA fallback costs two implicit D2H syncs per + ratio. + """ device = past_kv_lens_bufs[compress_ratios[0]].device for compress_ratio in compress_ratios: past_kv = past_kv_lens_bufs[compress_ratio] cu_new_comp = cu_new_comp_kv_bufs[compress_ratio] - total_ctx_comp = cu_new_comp[num_contexts] + total_ctx_comp = ( + ctx_output_sizes[compress_ratio] + if ctx_output_sizes is not None + else cu_new_comp[num_contexts] + ) ctx_idx = torch.arange(total_ctx_comp, dtype=torch.int32, device=device) ctx_cu = cu_new_comp[: num_contexts + 1].to(torch.int32) ctx_req = torch.searchsorted(ctx_cu[1:], ctx_idx, right=True) diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py index 0784d7256264..3ba703006935 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py @@ -449,6 +449,50 @@ def test_mixed_context_generation_position_ids_follow_compact_output(): assert actual_position_ids == [0, 4, 4, 4] +@pytest.mark.parametrize( + "compress_ratio,cached_tokens,kv_lens", + [ + pytest.param(1, [0], [4096], id="cr1_single"), + pytest.param(4, [0, 75, 4000], [4096, 4171, 8096], id="cr4_multi_boundary"), + pytest.param(128, [0], [64], id="cr128_empty_output"), + pytest.param(128, [973632, 8064], [990016, 8192], id="cr128_chunked_long_ctx"), + ], +) +def test_ctx_position_ids_host_sizes_match_device_scalar_fallback( + compress_ratio, cached_tokens, kv_lens +): + """Host-int ctx_output_sizes must reproduce the device-scalar fallback exactly. + + prepare() threads host-computed ctx compressed-token counts into + _compute_ctx_compressed_position_ids so the arange size and slice bound + are Python ints (no implicit D2H + stream sync). Both paths must produce + identical position IDs, including the untouched padding tail. + """ + num_contexts = len(kv_lens) + cached = torch.tensor(cached_tokens, dtype=torch.int32, device=DEVICE) + kv = torch.tensor(kv_lens, dtype=torch.int32, device=DEVICE) + past = (cached // compress_ratio).to(torch.int32) + new_comp = (kv // compress_ratio).to(torch.int32) - past + cu = F.pad(torch.cumsum(new_comp, dim=0), (1, 0)).to(torch.int32) + total = int(cu[num_contexts].item()) + + def _run(ctx_output_sizes): + out = torch.full((total + 8,), -1, dtype=torch.int32, device=DEVICE) + DeepseekV4TrtllmAttentionMetadata._compute_ctx_compressed_position_ids( + {compress_ratio: past}, + {compress_ratio: cu}, + {compress_ratio: out}, + num_contexts, + [compress_ratio], + ctx_output_sizes, + ) + return out + + golden = _run(None) + fast = _run({compress_ratio: total}) + assert torch.equal(golden, fast) + + def precompute_freqs_cis( dim, seqlen, original_seq_len, base, factor, beta_fast, beta_slow ) -> torch.Tensor: From 929f153ea66fdc7e1ee6cce74f79a2eb1788a2bd Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+Tabrizian@users.noreply.github.com> Date: Thu, 23 Jul 2026 21:42:13 -0700 Subject: [PATCH 13/13] [TRTLLM-9920][feat] Add support for arbitrary KVCache transfer (#13055) Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 8 + .../batch_manager/kvCacheManager.h | 36 +++ .../executor/dataTransceiverState.h | 16 +- .../batch_manager/cacheFormatter.cpp | 66 ++-- .../batch_manager/cacheFormatter.h | 2 +- .../batch_manager/cacheTransceiver.cpp | 12 + .../batch_manager/dataTransceiver.cpp | 204 ++++++++++-- .../batch_manager/dataTransceiver.h | 17 +- .../batch_manager/kvCacheManager.cpp | 108 ++++++- .../batch_manager/mlaCacheFormatter.cpp | 30 +- .../batch_manager/rnnCacheFormatter.cpp | 10 +- .../ucx_utils/ucxCacheCommunicator.cpp | 38 ++- cpp/tensorrt_llm/executor/serialization.cpp | 4 + .../batch_manager/cacheTransceiver.cpp | 6 + .../multi_gpu/cacheTransceiverTest.cpp | 91 +++++- .../_torch/pyexecutor/kv_cache_transceiver.py | 7 + tensorrt_llm/executor/base_worker.py | 11 + tensorrt_llm/executor/executor.py | 7 + tensorrt_llm/executor/proxy.py | 10 + tensorrt_llm/executor/rpc_proxy.py | 4 + tensorrt_llm/llmapi/llm.py | 11 + tensorrt_llm/serve/openai_server.py | 16 + .../test_disaggregated_single_gpu.py | 296 +++++++++++++++++- .../test_lists/qa/llm_function_core.txt | 2 + .../integration/test_lists/test-db/l0_a10.yml | 2 + .../api_stability/api_stability_core.py | 11 +- .../api_stability/references/llm.yaml | 4 + 27 files changed, 923 insertions(+), 106 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 735716e59e92..ae0478a190a2 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -249,6 +249,12 @@ class BaseCacheTransceiver virtual bool cancelRequest(std::shared_ptr llmRequest) = 0; + /// Get the serialized DataTransceiverState (CacheState + CommState) for this transceiver. + [[nodiscard]] virtual std::vector getSerializedDataTransceiverState() const + { + return {}; + } + [[nodiscard]] virtual bool hasPoisonedTransferBuffer() const { return false; @@ -298,6 +304,8 @@ class CacheTransceiver : public BaseCacheTransceiver virtual bool cancelRequest(std::shared_ptr llmRequest) override; + [[nodiscard]] std::vector getSerializedDataTransceiverState() const override; + [[nodiscard]] bool hasPoisonedTransferBuffer() const override; private: diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 01adf276f878..04b6c230e75a 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -1251,14 +1251,29 @@ class WindowBlockManager return mEnablePartialReuse; } + //! \brief Look up the block chain matching blockKey in the reuse tree. [[nodiscard]] std::shared_ptr findBlocksInReuseTreeByBlockKey(BlockKey const& blockKey); + //! \brief Same lookup, additionally pinning matched blocks; on a miss all pins are + //! rolled back and pinnedBlockIds is cleared. + [[nodiscard]] std::shared_ptr findBlocksInReuseTreeByBlockKey( + BlockKey const& blockKey, std::vector& pinnedBlockIds); + [[nodiscard]] std::shared_ptr findBlocksInReuseTreeByBlockKeys( std::vector const& blockKeys); //! \brief Unpin blocks by block ids directly void unpinBlocksById(std::vector const& blockIds); + //! \brief Pin a block: claim it from the eviction policy if free, then take a reference. + //! Safe to call from cache-transceiver threads: block bookkeeping is serialized by the + //! lookup-tree mutex, which every mutating entry point acquires. + void pinBlock(BlockPtr const& block); + + //! \brief Inverse of pinBlock: drop one reference and release the block back to the + //! eviction policy once no references remain. + void unpinBlock(BlockPtr const& block); + void truncateBlocks(LlmRequest::VecTokens const& targetTokens, SizeType32 numTokensToKeep); void resetReuseState() @@ -1274,6 +1289,10 @@ class WindowBlockManager } private: + //! \brief Shared implementation of the findBlocksInReuseTreeByBlockKey overloads. + [[nodiscard]] std::shared_ptr findBlocksInReuseTreeByBlockKeyImpl( + BlockKey const& blockKey, bool pinBlocks, std::vector& pinnedBlockIds); + //! \brief Walk the reuse tree with precomputed per-block keys (no lock; callers must hold mLookupTree->getMutex()). [[nodiscard]] std::shared_ptr searchReuseTree(std::vector const& blockKeys); @@ -1836,6 +1855,12 @@ class BlockManager return mWindowBlockManagers.at(windowSize).findBlocksInReuseTreeByBlockKey(blockKey); } + [[nodiscard]] std::shared_ptr findBlocksInReuseTreeByBlockKey( + BlockKey const& blockKey, SizeType32 windowSize, std::vector& pinnedBlockIds) + { + return mWindowBlockManagers.at(windowSize).findBlocksInReuseTreeByBlockKey(blockKey, pinnedBlockIds); + } + [[nodiscard]] std::shared_ptr findBlocksInReuseTreeByBlockKeys( std::vector const& blockKeys, SizeType32 windowSize) { @@ -2210,6 +2235,11 @@ class BaseKVCacheManager BlockKey const& blockKey, SizeType32 windowSize) = 0; + //! \brief Pinning lookup: pins matched blocks and records their ids for unpinBlocksById. + [[nodiscard]] virtual std::shared_ptr findBlocksInReuseTreeByBlockKey( + BlockKey const& blockKey, SizeType32 windowSize, std::vector& pinnedBlockIds) + = 0; + [[nodiscard]] virtual std::shared_ptr findBlocksInReuseTreeByBlockKeys( std::vector const& blockKeys, SizeType32 windowSize) = 0; @@ -2650,6 +2680,12 @@ class KVCacheManager : public BaseKVCacheManager return mBlockManager.findBlocksInReuseTreeByBlockKey(blockKey, windowSize); } + std::shared_ptr findBlocksInReuseTreeByBlockKey( + BlockKey const& blockKey, SizeType32 windowSize, std::vector& pinnedBlockIds) override + { + return mBlockManager.findBlocksInReuseTreeByBlockKey(blockKey, windowSize, pinnedBlockIds); + } + std::shared_ptr findBlocksInReuseTreeByBlockKeys( std::vector const& blockKeys, SizeType32 windowSize) override { diff --git a/cpp/include/tensorrt_llm/executor/dataTransceiverState.h b/cpp/include/tensorrt_llm/executor/dataTransceiverState.h index b00d44d129e7..578e53b81dbf 100644 --- a/cpp/include/tensorrt_llm/executor/dataTransceiverState.h +++ b/cpp/include/tensorrt_llm/executor/dataTransceiverState.h @@ -619,9 +619,22 @@ class DataTransceiverState final return mCacheState.has_value() && mCacheState->hasRnnConfig(); } + /// @brief Set only when exported via CacheTransceiver::getSerializedDataTransceiverState: + /// transfers driven by such a state have no LlmRequest on the sender. + [[nodiscard]] bool isArbitraryTransferState() const noexcept + { + return mIsArbitraryTransferState; + } + + void setIsArbitraryTransferState(bool isArbitraryTransferState) noexcept + { + mIsArbitraryTransferState = isArbitraryTransferState; + } + [[nodiscard]] bool operator==(DataTransceiverState const& other) const noexcept { - return mCacheState == other.mCacheState && mCommState == other.mCommState; + return mCacheState == other.mCacheState && mCommState == other.mCommState + && mIsArbitraryTransferState == other.mIsArbitraryTransferState; } [[nodiscard]] std::string toString() const @@ -642,6 +655,7 @@ class DataTransceiverState final friend class Serialization; std::optional mCacheState; std::optional mCommState; + bool mIsArbitraryTransferState{false}; }; } // namespace tensorrt_llm::executor diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp index 9dc44531d409..6165ec7386b4 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp @@ -183,7 +183,7 @@ void sendAllBuffers(TransferSession& session, int deviceId, namespace tensorrt_llm::batch_manager::kv_cache_manager { -BlockRange getBlockRangeForSending(BaseKVCacheManager* cacheManager, LlmRequest const& llmRequest, +BlockRange getBlockRangeForSending(BaseKVCacheManager* cacheManager, std::optional llmRequest, BlockKey const& lastBlockKey, int32_t indexFromEnd, bool recvSideHasCP, SizeType32 ppSize) { auto poolNum = cacheManager->getBlockManager().getNumPools( @@ -197,9 +197,10 @@ BlockRange getBlockRangeForSending(BaseKVCacheManager* cacheManager, LlmRequest || lastBlockKey.uniqueTokens.size() == 0 || recvSideHasCP || ppSize > 1) { // disable reuse path, and vwsa don't support reuse. + TLLM_CHECK_WITH_INFO(llmRequest.has_value(), "LlmRequest required for non-reuse-tree transfer path"); bool needSendAllForWindow = common::getEnvKVCacheTransferAllBlocksForWindow(); - auto blockRange = BlockRange::fromAllBlockIds(*cacheManager, llmRequest.mRequestId); + auto blockRange = BlockRange::fromAllBlockIds(*cacheManager, (*llmRequest)->mRequestId); auto const& windowsMetadata = cacheManager->getBlockManager().getWindowSizesMetadata(); @@ -235,16 +236,20 @@ BlockRange getBlockRangeForSending(BaseKVCacheManager* cacheManager, LlmRequest TLLM_CHECK_WITH_INFO(lastBlockKey.uniqueTokens.size() > 0, "lastBlockKey must be non-empty when reuse is enabled"); - auto multimodalHashes = llmRequest.getMultimodalHashes(); - bool isMultimodal = multimodalHashes.has_value() && *multimodalHashes && !(*multimodalHashes)->empty(); - if (isMultimodal) + // No request on the reuse-tree path: fall through to the plain lastBlockKey lookup. + if (llmRequest.has_value()) { - auto tokensPerBlock = cacheManager->getBlockManager().getTokensPerBlock(); - auto const usableSize = static_cast(lastBlockKey.uniqueTokens.size()); - auto blockedUniqueTokens = chopVectorIntoBlocks( - lastBlockKey.uniqueTokens, usableSize, tokensPerBlock, /*allowPartial=*/true); - auto blockKeys = buildBlockKeys(blockedUniqueTokens, llmRequest); - return BlockRange::fromReuseTree(*cacheManager, blockKeys, indexFromEnd); + auto multimodalHashes = (*llmRequest)->getMultimodalHashes(); + bool isMultimodal = multimodalHashes.has_value() && *multimodalHashes && !(*multimodalHashes)->empty(); + if (isMultimodal) + { + auto tokensPerBlock = cacheManager->getBlockManager().getTokensPerBlock(); + auto const usableSize = static_cast(lastBlockKey.uniqueTokens.size()); + auto blockedUniqueTokens = chopVectorIntoBlocks( + lastBlockKey.uniqueTokens, usableSize, tokensPerBlock, /*allowPartial=*/true); + auto blockKeys = buildBlockKeys(blockedUniqueTokens, **llmRequest); + return BlockRange::fromReuseTree(*cacheManager, blockKeys, indexFromEnd); + } } return BlockRange::fromReuseTree(*cacheManager, lastBlockKey, indexFromEnd); @@ -363,11 +368,15 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio { NVTX3_SCOPED_RANGE(CacheFormatter_format); session.setTime(TransferSession::kTimeFormatter); - auto const& llmRequest = session.getLlmRequest(); - TLLM_LOG_DEBUG( - mpi::MpiComm::world().getRank(), "Start sending KV cache for request ID: %ld.", llmRequest.mRequestId); + auto llmRequest = session.getLlmRequest(); + if (llmRequest.has_value()) + { + TLLM_LOG_DEBUG( + mpi::MpiComm::world().getRank(), "Start sending KV cache for request ID: %ld.", (*llmRequest)->mRequestId); + TLLM_CHECK_WITH_INFO( + (*llmRequest)->mSamplingConfig.beamWidth == 1, "Currently, only beam width 1 is supported."); + } - TLLM_CHECK_WITH_INFO(llmRequest.mSamplingConfig.beamWidth == 1, "Currently, only beam width 1 is supported."); auto const& connections = session.getConnections(); auto const& selfConfig = session.getSelfState().getCacheState().value(); auto const& destConfig = session.getOtherState().getCacheState().value(); @@ -385,7 +394,6 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio size_t targetNum = pickUpConnections.size(); if (targetNum == 0) { - TLLM_LOG_DEBUG("No targets to send KV cache to for request ID: %ld", llmRequest.mRequestId); return; } @@ -413,13 +421,15 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio SizeType32 const numKvPools = static_cast(kvWindowSizes.size()); - TLLM_LOG_DEBUG("CacheFormatter::format: allWindowSizes=%zu, kvWindowSizes=%d, numPools=%d, requestId=%lu", - allWindowSizes.size(), numKvPools, numPools, llmRequest.mRequestId); + TLLM_LOG_DEBUG("CacheFormatter::format: allWindowSizes=%zu, kvWindowSizes=%d, numPools=%d, requestId=%s", + allWindowSizes.size(), numKvPools, numPools, + llmRequest.has_value() ? std::to_string((*llmRequest)->mRequestId).c_str() : ""); bool layerWise = common::getEnvDisaggLayerwise() && numKvPools == 1; if (layerWise) { - auto& progress = llmRequest.getContextProgress(); + TLLM_CHECK_WITH_INFO(llmRequest.has_value(), "LlmRequest required for layer-wise transfer"); + auto& progress = (*llmRequest)->getContextProgress(); SizeType32 const numLayers = blockManager.getNumLayers(); runtime::ITensor::Shape offset = runtime::ITensor::makeShape({0, 0}); for (SizeType32 layerIdx = 0; layerIdx < numLayers; layerIdx++) @@ -515,8 +525,11 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio } } } - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "End the sending of KV cache for the request ID: %ld.", - llmRequest.mRequestId); + if (llmRequest.has_value()) + { + TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "End the sending of KV cache for the request ID: %ld.", + (*llmRequest)->mRequestId); + } return; } @@ -635,15 +648,20 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio sendHolder.release(); session.setTime(TransferSession::kTimePostprocess); } - TLLM_LOG_DEBUG( - mpi::MpiComm::world().getRank(), "End the sending of KV cache for the request ID:%ld ", llmRequest.mRequestId); + if (llmRequest.has_value()) + { + TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "End the sending of KV cache for the request ID:%ld ", + (*llmRequest)->mRequestId); + } } void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& session) { NVTX3_SCOPED_RANGE(CacheFormatter_unformat); session.setTime(TransferSession::kTimeFormatter); - auto const& llmRequest = session.getLlmRequest(); + auto llmRequestOpt = session.getLlmRequest(); + TLLM_CHECK_WITH_INFO(llmRequestOpt.has_value(), "LlmRequest required for receiving KV cache"); + auto const& llmRequest = **llmRequestOpt; auto const ctxReqId = llmRequest.getContextPhaseParams().value().getReqId(); TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "Start receiving KV cache for request ID: %ld, context request ID: %ld.", llmRequest.mRequestId, ctxReqId); diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.h b/cpp/tensorrt_llm/batch_manager/cacheFormatter.h index 92d1b9d58a27..21356f3dd6c0 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.h +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.h @@ -176,7 +176,7 @@ inline std::pair, std::vector> pickRecvConnections(s namespace tensorrt_llm::batch_manager::kv_cache_manager { -BlockRange getBlockRangeForSending(BaseKVCacheManager* cacheManager, LlmRequest const& llmRequest, +BlockRange getBlockRangeForSending(BaseKVCacheManager* cacheManager, std::optional llmRequest, BlockKey const& lastBlockKey, SizeType32 indexFromEnd, bool recvSideHasCP = false, SizeType32 ppSize = 1); using DataContext = tensorrt_llm::executor::kv_cache::DataContext; diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 69adf4246591..89e6bd5d1b97 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -51,6 +51,7 @@ #include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/cache_transmission/mpi_utils/connection.h" #include "tensorrt_llm/executor/dataTransceiverState.h" +#include "tensorrt_llm/executor/serialization.h" #include "tensorrt_llm/executor/serializeUtils.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/runtime/utils/pgUtils.h" @@ -621,6 +622,17 @@ void CacheTransceiver::initializeCommState() mCommState = std::addressof(mCacheSender->getCommState()); } +std::vector CacheTransceiver::getSerializedDataTransceiverState() const +{ + TLLM_CHECK(mCommState != nullptr && mCacheState != nullptr); + executor::DataTransceiverState state; + state.setCommState(*mCommState); + state.setCacheState(*mCacheState); + // Only this API marks the state; context responses leave it unset. + state.setIsArbitraryTransferState(true); + return executor::Serialization::serialize(state); +} + void CacheTransceiver::setContextState(LlmRequest* llmRequest) { TLLM_CHECK(llmRequest && llmRequest->isContextOnlyRequest()); diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 851f0abb04ef..e406d6620468 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp @@ -37,6 +37,7 @@ #include #include #include +#include namespace tensorrt_llm::batch_manager { @@ -81,6 +82,11 @@ void TransferSession::send(size_t idx, void const* data, size_t size) } catch (std::exception const& e) { + // Request-free (llmRequest-agnostic) transfer: there is no valid ID to attach. + if (mRequest == nullptr) + { + TLLM_THROW("%s", e.what()); + } throw common::RequestSpecificException( __FILE__, __LINE__, e.what(), mRequest->mRequestId, common::RequestErrorCode::kNETWORK_ERROR); } @@ -94,15 +100,23 @@ void TransferSession::recv(size_t idx, void* data, size_t size) } catch (std::exception const& e) { + // Request-free (llmRequest-agnostic) transfer: there is no valid ID to attach. + if (mRequest == nullptr) + { + TLLM_THROW("%s", e.what()); + } throw common::RequestSpecificException( __FILE__, __LINE__, e.what(), mRequest->mRequestId, common::RequestErrorCode::kNETWORK_ERROR); } } -LlmRequest const& TransferSession::getLlmRequest() const +std::optional TransferSession::getLlmRequest() const { - TLLM_CHECK(mRequest != nullptr); - return *mRequest; + if (mRequest == nullptr) + { + return std::nullopt; + } + return mRequest; } void TransferSession::setLlmRequest(LlmRequest const& llmRequest) @@ -171,7 +185,8 @@ void TransferSession::poisonReservedRecvBuffers() noexcept void TransferSession::exportMeasure(std::ofstream& outFile, bool isContext) const { - if (!mTimes || mTimes->measures.empty()) + // Request-free transfers are excluded: the exported row is keyed by the LlmRequest. + if (!mTimes || mTimes->measures.empty() || mRequest == nullptr) { return; } @@ -275,7 +290,7 @@ RequestInfo::RequestInfo(LlmRequest::RequestIdType requestId, executor::DataTran bool RequestInfo::operator==(RequestInfo const& rhs) const { return mRequestId == rhs.mRequestId && mIndexFromEnd == rhs.mIndexFromEnd && mLastBlockKey == rhs.mLastBlockKey - && mTransState == rhs.mTransState; + && mIsArbitraryTransfer == rhs.mIsArbitraryTransfer && mTransState == rhs.mTransState; } LlmRequest::RequestIdType RequestInfo::getRequestId() const noexcept @@ -294,6 +309,7 @@ void RequestInfo::serialize(RequestInfo const& requestInfo, std::ostream& os) su::serialize(requestInfo.mRequestId, os); su::serialize(requestInfo.mIndexFromEnd, os); su::serialize(requestInfo.mLastBlockKey, os); + su::serialize(requestInfo.mIsArbitraryTransfer, os); su::serialize(requestInfo.mTransState, os); } @@ -303,8 +319,11 @@ RequestInfo RequestInfo::deserialize(std::istream& is) auto requestId = su::deserialize(is); auto indexFromEnd = su::deserialize(is); auto lastBlockKey = su::deserialize(is); + auto isArbitraryTransfer = su::deserialize(is); auto transState = su::deserialize(is); - return RequestInfo{requestId, std::move(transState), indexFromEnd, lastBlockKey}; + auto requestInfo = RequestInfo{requestId, std::move(transState), indexFromEnd, lastBlockKey}; + requestInfo.setIsArbitraryTransfer(isArbitraryTransfer); + return requestInfo; } std::size_t RequestInfo::serializedSize(RequestInfo const& requestInfo) @@ -314,6 +333,7 @@ std::size_t RequestInfo::serializedSize(RequestInfo const& requestInfo) totalSize += su::serializedSize(requestInfo.mRequestId); totalSize += su::serializedSize(requestInfo.mIndexFromEnd); totalSize += su::serializedSize(requestInfo.mLastBlockKey); + totalSize += su::serializedSize(requestInfo.mIsArbitraryTransfer); totalSize += su::serializedSize(requestInfo.mTransState); return totalSize; } @@ -467,7 +487,6 @@ class CacheSender::Impl : mManager->recvConnect(DataContext{TransceiverTag::kID_TAG, mTerminate}, &id, sizeof(id)); if (connection == nullptr) { - TLLM_LOG_WARNING("recvRequestInfo connection is nullptr, maybe the server is terminating"); return std::nullopt; } @@ -625,10 +644,23 @@ class CacheSender::Impl private: struct Response { - // shared_ptr so this struct co-owns the request until the promise resolves; - // protects worker-side dereferences and the promise itself from premature destruction. - std::shared_ptr mRequest; + // An LlmRequest (co-owned until the promise resolves) for normal transfers, or + // just the request id for llmRequest-agnostic reuse-tree transfers. + std::variant, RequestIdType> mRequestOrId; std::promise mPromise; + std::vector mPinnedBlockIds; + + [[nodiscard]] LlmRequest* getRequest() const + { + auto const* request = std::get_if>(&mRequestOrId); + return request != nullptr ? request->get() : nullptr; + } + + [[nodiscard]] RequestIdType getRequestId() const + { + auto const* request = getRequest(); + return request != nullptr ? request->mRequestId : std::get(mRequestOrId); + } }; struct AsyncSendResource @@ -661,21 +693,44 @@ class CacheSender::Impl resp = std::move(resource.mSendQueue.front()); resource.mSendQueue.pop_front(); } - // Sequence the read before the move: argument initializations - // are indeterminately sequenced, so inlining resp.mRequest->... - // alongside std::move(resp) is UB once mRequest is a shared_ptr. - TLLM_CHECK(resp.mRequest != nullptr); - auto const reqId = resp.mRequest->mRequestId; - sendAndRemoveResponse(reqId, std::move(resp)); + // Read before std::move(resp): argument evaluations are indeterminately sequenced. + auto const requestId = resp.getRequestId(); + sendAndRemoveResponse(requestId, std::move(resp)); } } + //! Must not throw: called from noexcept send/failure paths. + void releasePinnedBlocks(Response& response) noexcept + { + if (response.mPinnedBlockIds.empty()) + { + return; + } + try + { + mCacheTransferLayer.getCacheManager()->unpinBlocksById(response.mPinnedBlockIds); + } + catch (std::exception const& err) + { + TLLM_LOG_ERROR("Failed to unpin reuse-tree blocks: %s", err.what()); + } + response.mPinnedBlockIds.clear(); + } + void sendAndRemoveResponse(RequestIdType id, Response resp) noexcept { try { TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); - sendSync(*resp.mRequest); + if (auto* llmRequest = resp.getRequest(); llmRequest != nullptr) + { + sendSync(*llmRequest); + } + else + { + // Reuse tree path — no LlmRequest + sendSyncFromReuseTree(id); + } release(id); resp.mPromise.set_value(); } @@ -700,6 +755,7 @@ class CacheSender::Impl discardTransferState(id); failResponse(resp, exception); } + releasePinnedBlocks(resp); } void asyncSendAndRemoveResponse(RequestIdType id, Response resp) noexcept @@ -806,6 +862,37 @@ class CacheSender::Impl } } + void sendSyncFromReuseTree(RequestIdType requestId) + { + TransferSession* session = nullptr; + { + std::unique_lock lk(mMtxForMap); + auto it = mRequestToSession.find(requestId); + TLLM_CHECK(it != mRequestToSession.end()); + session = std::addressof(it->second); + } + // READY was already sent by response(); the receiver consumes exactly one per transfer. + mCacheTransferLayer.format(*session); + } + + // Pin the requested chain in the reuse tree; an empty result means no full match. + // The caller must unpin once the transfer settles. + std::vector pinReuseTreeBlocks(RequestIdType requestId) + { + std::unique_lock lk(mMtxForMap); + auto it = mRequestToSession.find(requestId); + auto const& lastBlockKey = it->second.getLastBlockKey(); + auto* cacheManager = mCacheTransferLayer.getCacheManager(); + auto windowSize = cacheManager->getBlockManager().getWindowSizesMetadata().begin()->first; + std::vector pinnedIds; + auto lastBlock = cacheManager->findBlocksInReuseTreeByBlockKey(lastBlockKey, windowSize, pinnedIds); + if (lastBlock == nullptr) + { + return {}; + } + return pinnedIds; + } + void response() noexcept { std::exception_ptr responseException; @@ -815,16 +902,13 @@ class CacheSender::Impl TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); while (true) { + if (mTerminate) { - std::unique_lock lock(mSenderMutex); - mSenderCv.wait(lock, - [this]() { return mTerminate || !mReadyResponses.empty() || !mCancelledRequests.empty(); }); - if (mTerminate) - { - break; - } + break; } + // Arbitrary transfers arrive without a pre-registered response; do not gate on + // mReadyResponses. auto requestInfo = recvRequestInfo(); if (!requestInfo.has_value() || mTerminate || !mManager->isRunning()) { @@ -837,22 +921,70 @@ class CacheSender::Impl mRemainSendCount[reqId] = getCounterpartsCount(reqId); } + if (requestInfo->isArbitraryTransfer()) { - std::unique_lock lock(mSenderMutex); - mCurrentRequest = reqId; - mSenderCv.wait(lock, - [this, reqId]() + // No LlmRequest will ever be registered; serve from the reuse tree off-thread. + { + std::scoped_lock lock(mSenderMutex); + mCurrentRequest = reqId; + } + auto countIt = mRemainSendCount.find(reqId); + auto const count = --countIt->second; + TLLM_CHECK(count >= 0); + if (count == 0) + { + mRemainSendCount.erase(countIt); + auto pinnedIds = pinReuseTreeBlocks(reqId); + if (pinnedIds.empty()) + { + TLLM_LOG_ERROR( + "Requested blocks do not exist in the source's reuse tree (request id: %lu). Notifying " + "receiver.", + reqId); + sendReadySignal(reqId, false); + discardTransferState(reqId); + } + else { - return mTerminate || mReadyResponses.find(reqId) != mReadyResponses.end() - || mCancelledRequests.find(reqId) != mCancelledRequests.end(); - }); - if (mTerminate) + sendReadySignal(reqId, true); + std::promise promise; + // Id-only response: the reuse-tree path has no LlmRequest. + Response resp{reqId, std::move(promise), std::move(pinnedIds)}; + if (dynamic_cast(mManager) != nullptr) + { + sendAndRemoveResponse(reqId, std::move(resp)); + } + else + { + asyncSendAndRemoveResponse(reqId, std::move(resp)); + } + } + } { + std::scoped_lock lock(mSenderMutex); mCurrentRequest = std::nullopt; - break; } } - sendResponse(reqId); + else + { + // The RequestInfo may race ahead of sendAsync; wait for the specific response. + { + std::unique_lock lock(mSenderMutex); + mCurrentRequest = reqId; + mSenderCv.wait(lock, + [this, reqId]() + { + return mTerminate || mReadyResponses.find(reqId) != mReadyResponses.end() + || mCancelledRequests.find(reqId) != mCancelledRequests.end(); + }); + if (mTerminate) + { + mCurrentRequest = std::nullopt; + break; + } + } + sendResponse(reqId); + } } } catch (std::exception const& err) @@ -930,6 +1062,7 @@ class CacheSender::Impl { TLLM_LOG_ERROR("Failed to set CacheSender response exception: %s", err.what()); } + releasePinnedBlocks(response); } void failPendingResponses(std::exception_ptr const& exception) noexcept @@ -1115,6 +1248,9 @@ class CacheReceiver::Impl requestInfo = RequestInfo(requestId, mSelfState, indexFromEnd, lastBlockKey); } } + // The state's provenance marks llmRequest-agnostic transfers: only + // getSerializedDataTransceiverState sets it; context responses leave it unset. + requestInfo.setIsArbitraryTransfer(contextState.isArbitraryTransferState()); auto* agentConnectionManager = dynamic_cast(mManager); std::vector recvHolders; diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h index c1da646916dc..fffe55dceab5 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h @@ -19,6 +19,7 @@ #include #include #include +#include #include #include @@ -121,7 +122,7 @@ class TransferSession void recv(size_t idx, void* data, size_t size); - [[nodiscard]] LlmRequest const& getLlmRequest() const; + [[nodiscard]] std::optional getLlmRequest() const; // in CacheSender, the LlmRequest is not available until the sendSync is called void setLlmRequest(LlmRequest const& llmRequest); @@ -234,6 +235,17 @@ class RequestInfo return mLastBlockKey; } + /// @brief Arbitrary (llmRequest-agnostic) transfer served from the sender's reuse tree. + [[nodiscard]] bool isArbitraryTransfer() const noexcept + { + return mIsArbitraryTransfer; + } + + void setIsArbitraryTransfer(bool isArbitraryTransfer) noexcept + { + mIsArbitraryTransfer = isArbitraryTransfer; + } + /// @brief Serialization. /// @param requestInfo Request information to be serialized. /// @param os The output stream to which the serialization result points. @@ -257,6 +269,9 @@ class RequestInfo // Last block key, used to derive other block keys on receiver BlockKey mLastBlockKey{}; + // True for arbitrary (llmRequest-agnostic) transfers served from the sender's reuse tree. + bool mIsArbitraryTransfer{false}; + // The state of the data transceiver. executor::DataTransceiverState mTransState; }; diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 266cdafc8d06..8c1ffb70e372 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -937,6 +937,7 @@ bool BlockManager::verifyQueueIntegrity(SizeType32 windowSize) const bool WindowBlockManager::verifyQueueIntegrity() const { + std::lock_guard lock(mLookupTree->getMutex()); return mEvictionPolicy->verifyQueueIntegrity(); } @@ -1219,6 +1220,7 @@ void BlockManager::startScheduling() void WindowBlockManager::startScheduling() { + std::lock_guard lock(mLookupTree->getMutex()); mSchedulingNumFreeBlocks = mEvictionPolicy->getNumFreeBlocks(kPrimaryLevel); for (auto& [requestId, slotAllocatedBlocks] : mAllocatedBlocksPerSeq) { @@ -1238,6 +1240,7 @@ void WindowBlockManager::freeLeafBlock(BlockPtr const& block) void WindowBlockManager::releaseSubtree(BlockPtr const& block) { + std::lock_guard lock(mLookupTree->getMutex()); // Iterative pre-order DFS over `block` and all its descendants. Collect // first, then detach in reverse order: cascade-prune in freeLeafBlock() // removes empty parent nodes from the trie, so leaves must be detached @@ -1275,6 +1278,7 @@ BlockPtr WindowBlockManager::getFreeBlock(GenerationRequest& sequence, executor: std::optional durationMs, executor::KvCacheTransferMode mode, std::string const& directory, bool wantPlaceholder) { + std::lock_guard lock(mLookupTree->getMutex()); // eviction policy get free primary block auto [block, canOffload] = mEvictionPolicy->getFreeBlock(kPrimaryLevel, wantPlaceholder); if (block->getUniqueTokens().empty()) @@ -1398,6 +1402,7 @@ void BlockManager::onboardBlock(GenerationRequest& sequence, BlockPtr const& off void WindowBlockManager::onboardBlock(GenerationRequest& sequence, BlockPtr const& offloadBlock, executor::KvCacheTransferMode mode, std::string const& directory) { + std::lock_guard lock(mLookupTree->getMutex()); if (!offloadBlock->isPlaceholder() && !offloadBlock->isPrimary()) { auto block = getFreeBlock( @@ -1426,6 +1431,7 @@ void BlockManager::offloadBlock( void WindowBlockManager::offloadBlock( BlockPtr const& block, executor::KvCacheTransferMode mode, std::string const& directory) { + std::lock_guard lock(mLookupTree->getMutex()); // The current default behavior is to offload the out-of-window block // to secondary block pool to allow more free primary blocks for reuse. // However, such behavior does not take account whether the offloaded @@ -2131,6 +2137,19 @@ bool WindowBlockManager::blockInRadixTree(BlockPtr const& block) } std::shared_ptr WindowBlockManager::findBlocksInReuseTreeByBlockKey(BlockKey const& blockKey) +{ + std::vector unusedPinnedBlockIds; + return findBlocksInReuseTreeByBlockKeyImpl(blockKey, /*pinBlocks=*/false, unusedPinnedBlockIds); +} + +std::shared_ptr WindowBlockManager::findBlocksInReuseTreeByBlockKey( + BlockKey const& blockKey, std::vector& pinnedBlockIds) +{ + return findBlocksInReuseTreeByBlockKeyImpl(blockKey, /*pinBlocks=*/true, pinnedBlockIds); +} + +std::shared_ptr WindowBlockManager::findBlocksInReuseTreeByBlockKeyImpl( + BlockKey const& blockKey, bool pinBlocks, std::vector& pinnedBlockIds) { std::lock_guard lock(mLookupTree->getMutex()); auto blockedUniqueTokens @@ -2143,7 +2162,42 @@ std::shared_ptr WindowBlockManager::findBlocksInReuseTreeByBlockKe blockKeys.emplace_back(blockKey.usesExtraIds, blockKey.loraTaskId, blockedUniqueTokensList, blockKey.extraKeys, blockKey.cacheSalt); } - return searchReuseTree(blockKeys); + auto searchRoot = mCachedBlocksRoot; + std::vector pinnedInScope; + for (auto const& blockKey : blockKeys) + { + auto [partialMatch, numMatched, matchingBlock] = searchRoot != nullptr + ? searchRoot->findMatchingBlock(blockKey, true, true) + : std::make_tuple(false, 0, nullptr); + + // A prefix-only match lacks KV for the requested tail, and pinning (transfer) + // lookups read primary-pool buffers directly, so offloaded and placeholder blocks + // are misses too (isPlaceholder first: isPrimary asserts on placeholders). + bool const fullyMatched + = matchingBlock != nullptr && numMatched == static_cast(blockKey.uniqueTokens.size()); + bool const transferable + = fullyMatched && (!pinBlocks || (!matchingBlock->isPlaceholder() && matchingBlock->isPrimary())); + if (!transferable) + { + // Roll back pins taken during the partial walk. + for (auto const& block : pinnedInScope) + { + unpinBlock(block); + } + pinnedBlockIds.clear(); + return nullptr; + } + + if (pinBlocks) + { + pinBlock(matchingBlock); + pinnedInScope.push_back(matchingBlock); + pinnedBlockIds.push_back(matchingBlock->getBlockId()); + } + + searchRoot = std::move(matchingBlock); + } + return searchRoot; } std::shared_ptr WindowBlockManager::findBlocksInReuseTreeByBlockKeys( @@ -2192,6 +2246,7 @@ void BlockManager::refreshBlocks() void WindowBlockManager::refreshBlocks() { + std::lock_guard lock(mLookupTree->getMutex()); mEvictionPolicy->refresh(); mTransferManager->syncTransfers(); } @@ -2322,6 +2377,7 @@ bool BlockManager::copyLinearAttentionBlock(GenerationRequest& sequence, LlmRequ bool WindowBlockManager::tryAllocatePlaceholderForLinearAttention(GenerationRequest& sequence, bool shareAmongBeams) { + std::lock_guard lock(mLookupTree->getMutex()); auto const beamWidth = sequence.getBeamWidth(); auto const newBlockIdx = sequence.getCacheBlockIds(mWindowSize).at(0).size(); // The first block is not a placeholder. @@ -2411,6 +2467,7 @@ bool WindowBlockManager::tryAllocatePlaceholderForLinearAttention(GenerationRequ void WindowBlockManager::allocateBlock(GenerationRequest& sequence, bool shareAmongBeams) { + std::lock_guard lock(mLookupTree->getMutex()); auto const beamWidth = sequence.getBeamWidth(); auto const requiredBlocks = shareAmongBeams ? 1 : beamWidth; @@ -2692,16 +2749,7 @@ std::pair> WindowBlockManager::sto if (pinBlocks) { - // If the block has no refs it sits in the eviction policy's free - // queue. Claim it first so that the later unpinBlocksById / - // releaseBlock cycle does not create a duplicate queue entry. - // Pass the block's existing priority and duration so that - // claimBlock does not clear its retention/expiry metadata. - if (!prevBlock->hasRefs()) - { - mEvictionPolicy->claimBlock(prevBlock, prevBlock->getPriority(), prevBlock->getDurationMs()); - } - prevBlock->incRefCount(); + pinBlock(prevBlock); pinnedBlockIds.push_back(prevBlock->getBlockId()); } } @@ -2736,6 +2784,7 @@ void BlockManager::replaceSharedBlock(GenerationRequest& sequence, SizeType32 wi void WindowBlockManager::replaceSharedBlock(GenerationRequest& sequence, SizeType32 blockIdx) { + std::lock_guard lock(mLookupTree->getMutex()); auto const requestId = sequence.getRequestId(); auto const beamWidth = sequence.getBeamWidth(); auto& allocatedBlocks = mAllocatedBlocksPerSeq.at(requestId); @@ -2787,6 +2836,7 @@ void BlockManager::releaseLastBlock(GenerationRequest& sequence, SizeType32 wind void WindowBlockManager::releaseLastBlock(GenerationRequest& sequence) { + std::lock_guard lock(mLookupTree->getMutex()); if (isRecurrentState()) { // In recurrent state, the last block always contains the current state and should not be released. @@ -2835,6 +2885,7 @@ void WindowBlockManager::releaseLastBlock(GenerationRequest& sequence) [[nodiscard]] SizeType32 WindowBlockManager::getNumFreeBlocks() const { + std::lock_guard lock(mLookupTree->getMutex()); auto const numFree = mEvictionPolicy->getNumFreeBlocks(kPrimaryLevel); TLLM_CHECK_WITH_INFO(numFree <= getMaxNumBlocks(), "%s::getNumFreeBlocks - primary free block count (%d) exceeds total block count (%d). " @@ -2846,6 +2897,7 @@ void WindowBlockManager::releaseLastBlock(GenerationRequest& sequence) [[nodiscard]] SizeType32 WindowBlockManager::getNumFreeSecondaryBlocks() const noexcept { + std::lock_guard lock(mLookupTree->getMutex()); return mEvictionPolicy->getNumFreeBlocks(kSecondaryLevel); } @@ -2976,18 +3028,41 @@ void BlockManager::unpinBlocksById(std::vector const& bloc firstManager.unpinBlocksById(blockIds); } +void WindowBlockManager::pinBlock(BlockPtr const& block) +{ + std::lock_guard lock(mLookupTree->getMutex()); + // Claim free blocks out of the eviction queue first, keeping their retention metadata. + if (!block->hasRefs()) + { + mEvictionPolicy->claimBlock(block, block->getPriority(), block->getDurationMs()); + } + block->incRefCount(); +} + +void WindowBlockManager::unpinBlock(BlockPtr const& block) +{ + std::lock_guard lock(mLookupTree->getMutex()); + block->decRefCount(); + if (!block->hasRefs()) + { + mEvictionPolicy->releaseBlock(block); + } +} + void WindowBlockManager::pinBlocks(GenerationRequest& sequence) { + std::lock_guard lock(mLookupTree->getMutex()); auto const requestId = sequence.getRequestId(); auto& allocatedBlocks = mAllocatedBlocksPerSeq.at(requestId); for (auto& block : allocatedBlocks) { - block->incRefCount(); + pinBlock(block); } } void WindowBlockManager::unpinBlocksById(std::vector const& blockIds) { + std::lock_guard lock(mLookupTree->getMutex()); if (blockIds.empty()) { return; @@ -3000,11 +3075,7 @@ void WindowBlockManager::unpinBlocksById(std::vector const auto block = mAllBlocksById[blockId]; if (block && block->getBlockId() != KVCacheBlock::kCachedBlocksRootId) { - block->decRefCount(); - if (!block->hasRefs()) - { - mEvictionPolicy->releaseBlock(block); - } + unpinBlock(block); } } } @@ -3020,6 +3091,7 @@ void BlockManager::storeNewBlock(GenerationRequest& sequence, OptionalRef llmRequest) { + std::lock_guard lock(mLookupTree->getMutex()); auto constexpr beamIdx = 0; auto const& uniqueTokens = llmRequest->getUniqueTokens(beamIdx); @@ -3126,6 +3198,7 @@ std::vector WindowBlockManager::storeBlocksForReuse( std::optional WindowBlockManager::releaseBlocks( GenerationRequest& sequence, OptionalRef llmRequest) { + std::lock_guard lock(mLookupTree->getMutex()); auto const requestId = sequence.getRequestId(); TLLM_LOG_DEBUG("%s::releaseBlocks - requestId=%lu, llmRequest.id=%s", mLogPrefix.c_str(), requestId, llmRequest.has_value() ? std::to_string(llmRequest->mRequestId).c_str() : "null"); @@ -3712,6 +3785,7 @@ bool KVCacheManager::copyLinearAttentionBlockBatch(std::vector lock(mLookupTree->getMutex()); // streamLLM is not supported at the moment. The out of window block will // always be the 0th block. TLLM_CHECK_WITH_INFO( diff --git a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp index 24b7182fe628..2631df9aa9a2 100644 --- a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp @@ -124,9 +124,13 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses { NVTX3_SCOPED_RANGE(MLACacheFormatter_format); session.setTime(TransferSession::kTimeFormatter); - auto const& llmRequest = session.getLlmRequest(); - TLLM_LOG_DEBUG( - mpi::MpiComm::world().getRank(), "Start sending KV cache for request ID: %ld.", llmRequest.mRequestId); + auto llmRequest = session.getLlmRequest(); + if (llmRequest.has_value()) + { + TLLM_LOG_DEBUG( + mpi::MpiComm::world().getRank(), "Start sending KV cache for request ID: %ld.", (*llmRequest)->mRequestId); + TLLM_CHECK_WITH_INFO((*llmRequest)->mSamplingConfig.beamWidth == 1, "Currently only supports beam width 1."); + } auto const& selfConfig = session.getSelfState().getCacheState().value(); auto const& destConfig = session.getOtherState().getCacheState().value(); auto const selfIdx = session.getSelfState().getCommState().value().getSelfIdx(); @@ -134,7 +138,6 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses auto const& lastBlockKey = session.getLastBlockKey(); auto const& connections = session.getConnections(); auto& bufferManager = session.getBufferManager(); - TLLM_CHECK_WITH_INFO(llmRequest.mSamplingConfig.beamWidth == 1, "Currently only supports beam width 1."); TLLM_CHECK(!connections.empty()); if (!needSendCache(selfConfig, destConfig, selfIdx)) { @@ -145,7 +148,6 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses auto targetNum = pickUpConnections.size(); if (targetNum == 0) { - TLLM_LOG_DEBUG("No targets to send KV cache to for request ID: %ld", llmRequest.mRequestId); return; } @@ -219,8 +221,11 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses } } - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "End the sending of KV cache for the request ID: %ld.", - llmRequest.mRequestId); + if (llmRequest.has_value()) + { + TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "End the sending of KV cache for the request ID: %ld.", + (*llmRequest)->mRequestId); + } return; } @@ -412,15 +417,20 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses session.setTime(TransferSession::kTimeTransmissions); session.setTime(TransferSession::kTimePostprocess); - TLLM_LOG_DEBUG( - mpi::MpiComm::world().getRank(), "End the sending of KV cache for the request ID: %ld.", llmRequest.mRequestId); + if (llmRequest.has_value()) + { + TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "End the sending of KV cache for the request ID: %ld.", + (*llmRequest)->mRequestId); + } } void MLACacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& session) { NVTX3_SCOPED_RANGE(MLACacheFormatter_unformat); session.setTime(TransferSession::kTimeFormatter); - auto const& llmRequest = session.getLlmRequest(); + auto llmRequestOpt = session.getLlmRequest(); + TLLM_CHECK_WITH_INFO(llmRequestOpt.has_value(), "LlmRequest required for receiving KV cache"); + auto const& llmRequest = **llmRequestOpt; TLLM_CHECK_WITH_INFO(llmRequest.mSamplingConfig.beamWidth == 1, "Currently only supports beam width 1."); auto const ctxReqId = llmRequest.getContextPhaseParams().value().getReqId(); TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), diff --git a/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp index c58733f25885..b7ceb727efa3 100644 --- a/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp @@ -47,7 +47,9 @@ void RnnCacheFormatter::format(TransferSession& session) NVTX3_SCOPED_RANGE(RnnCacheFormatter_formatUnifiedPool); session.setTime(TransferSession::kTimeFormatter); - auto const& llmRequest = session.getLlmRequest(); + auto llmRequestOpt = session.getLlmRequest(); + TLLM_CHECK_WITH_INFO(llmRequestOpt.has_value(), "LlmRequest required for RNN state transfer"); + auto const& llmRequest = **llmRequestOpt; TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "Start sending unified pool RNN state for request ID: %ld.", llmRequest.mRequestId); TLLM_CHECK_WITH_INFO(llmRequest.mSamplingConfig.beamWidth == 1, "Currently, only beam width 1 is supported."); @@ -79,7 +81,7 @@ void RnnCacheFormatter::format(TransferSession& session) bool const recvSideHasCP = destConfig.getParallelConfig().mContextParallelism > 1; auto const indexFromEnd = session.getIndexFromEnd(); auto blockRange = kv_cache_manager::getBlockRangeForSending( - mKvCacheManager, llmRequest, lastBlockKey, indexFromEnd, recvSideHasCP, ppSize); + mKvCacheManager, llmRequestOpt, lastBlockKey, indexFromEnd, recvSideHasCP, ppSize); auto const& blockIdsPerWindow = blockRange.getBlockIdsPerWindow(); auto const allWindowSizes = blockRange.getWindowSizes(); @@ -251,7 +253,9 @@ void RnnCacheFormatter::unformat(TransferSession& session) NVTX3_SCOPED_RANGE(RnnCacheFormatter_unformatUnifiedPool); session.setTime(TransferSession::kTimeFormatter); - auto const& llmRequest = session.getLlmRequest(); + auto llmRequestOpt = session.getLlmRequest(); + TLLM_CHECK_WITH_INFO(llmRequestOpt.has_value(), "LlmRequest required for RNN state transfer"); + auto const& llmRequest = **llmRequestOpt; TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "Start receiving unified pool RNN state for request ID: %ld.", llmRequest.mRequestId); TLLM_CHECK_WITH_INFO(llmRequest.mSamplingConfig.beamWidth == 1, "Currently, only beam width 1 is supported."); diff --git a/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/ucxCacheCommunicator.cpp b/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/ucxCacheCommunicator.cpp index 4ad1e7bffc86..3cd388625640 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/ucxCacheCommunicator.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/ucxCacheCommunicator.cpp @@ -611,23 +611,41 @@ UcxConnection::ConnectionIdType UcxConnectionManager::getNewConnectionId(std::sh Connection const* UcxConnectionManager::recvConnect(DataContext const& ctx, void* data, size_t size) { - std::vector buffer(size + sizeof(UcxConnection::ConnectionIdType)); - std::promise promise; - std::future future = promise.get_future(); - auto completionCallback = [&](ucs_status_t, ucxx::RequestCallbackUserData) -> void { promise.set_value(); }; + // Co-owned by the completion callback: the terminate path returns while the + // cancelled request may still write the buffer and fire the callback. + struct RecvState + { + std::vector mBuffer; + std::promise mPromise; + }; - std::shared_ptr req = mWorkersPool.front()->tagRecv( - buffer.data(), buffer.size(), ucxx::Tag(ctx.getTag()), ucxx::TagMask(0xFFFFFFFF), false, completionCallback); + auto state = std::make_shared(); + state->mBuffer.resize(size + sizeof(UcxConnection::ConnectionIdType)); + std::future future = state->mPromise.get_future(); + auto completionCallback + = [state](ucs_status_t, ucxx::RequestCallbackUserData) -> void { state->mPromise.set_value(); }; + + std::shared_ptr req = mWorkersPool.front()->tagRecv(state->mBuffer.data(), state->mBuffer.size(), + ucxx::Tag(ctx.getTag()), ucxx::TagMask(0xFFFFFFFF), false, completionCallback); if (!req->isCompleted()) { - future.get(); + // Poll with timeout to allow checking the terminate flag + auto const& terminate = ctx.getTransferTerminate(); + while (future.wait_for(std::chrono::milliseconds(100)) != std::future_status::ready) + { + if (terminate.load()) + { + req->cancel(); + return nullptr; + } + } } TLLM_CHECK_WITH_INFO(req->isCompleted(), "recv SendConnectionId should be completed"); req->checkError(); - memcpy(data, buffer.data(), size); + memcpy(data, state->mBuffer.data(), size); UcxConnection::ConnectionIdType connectionId - = *reinterpret_cast(buffer.data() + size); + = *reinterpret_cast(state->mBuffer.data() + size); std::scoped_lock lock(mConnectionsMutex, mConnectionFuturesMutex); TLLM_CHECK_WITH_INFO(mConnectionFutures.find(connectionId) != mConnectionFutures.end(), "connectionFuture not found In recvConnect connectionId : %lu , worldRank: %d", connectionId, mRank); @@ -642,7 +660,7 @@ Connection const* UcxConnectionManager::recvConnect(DataContext const& ctx, void TLLM_CHECK(!mConnections[connectionId]->isFromRequester()); TLLM_LOG_DEBUG(mRank, "recvConnect connectionId: %lu , sendIDData:%lu", connectionId, - *reinterpret_cast(buffer.data())); + *reinterpret_cast(state->mBuffer.data())); return mConnections[connectionId].get(); } diff --git a/cpp/tensorrt_llm/executor/serialization.cpp b/cpp/tensorrt_llm/executor/serialization.cpp index 4c2998c25d26..02077b6857ad 100644 --- a/cpp/tensorrt_llm/executor/serialization.cpp +++ b/cpp/tensorrt_llm/executor/serialization.cpp @@ -765,6 +765,8 @@ DataTransceiverState Serialization::deserializeDataTransceiverState(std::istream { state.setCacheState(std::move(cacheState).value()); } + auto isArbitraryTransferState = su::deserialize(is); + state.setIsArbitraryTransferState(isArbitraryTransferState); return state; } @@ -772,6 +774,7 @@ void Serialization::serialize(DataTransceiverState const& state, std::ostream& o { su::serialize(state.mCommState, os); su::serialize(state.mCacheState, os); + su::serialize(state.mIsArbitraryTransferState, os); } std::vector Serialization::serialize(DataTransceiverState const& state) @@ -790,6 +793,7 @@ size_t Serialization::serializedSize(DataTransceiverState const& state) size_t totalSize = 0; totalSize += su::serializedSize(state.mCommState); totalSize += su::serializedSize(state.mCacheState); + totalSize += su::serializedSize(state.mIsArbitraryTransferState); return totalSize; } diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp index d11838f68af5..2268e6f8ed03 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp @@ -92,6 +92,12 @@ void tb::CacheTransceiverBindings::initBindings(nb::module_& m) .def("request_and_receive_sync", &BaseCacheTransceiver::requestAndReceiveSync, nb::call_guard()) .def("request_and_receive_async", &BaseCacheTransceiver::requestAndReceiveAsync) + .def("get_serialized_data_transceiver_state", + [](tb::BaseCacheTransceiver& self) + { + auto serialized = self.getSerializedDataTransceiverState(); + return nb::bytes(serialized.data(), serialized.size()); + }) .def( "check_context_transfer_status", [](tb::BaseCacheTransceiver& self, std::optional const& atLeastRequestNum, bool markComplete = false) diff --git a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp index 5d1a958d37ce..814e60ee8262 100644 --- a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp @@ -286,7 +286,7 @@ class SymmetricalCacheTest : public ::testing::Test // NOLINT(cppcoreguidelines- return mWorldSize; } - void setUpCacheManager() + void setUpCacheManager(bool enableBlockReuse = false) { auto constexpr numLayers = 4; auto constexpr numHeads = 2; @@ -308,7 +308,6 @@ class SymmetricalCacheTest : public ::testing::Test // NOLINT(cppcoreguidelines- auto totalNumBlocks = mMaxNumSequences * numBlocksPerSeq; auto constexpr blocksInSecondaryPool = 0; - auto constexpr enableBlockReuse = false; auto constexpr dataType = tensorrt_llm::DataType::kFLOAT; using BlocksPerWindow = std::map>; @@ -422,6 +421,21 @@ class SymmetricalCacheTest : public ::testing::Test // NOLINT(cppcoreguidelines- return std::make_unique(mRequestId++, std::move(request)); } + // Generation-only request whose DataTransceiverState carries the arbitrary-transfer + // provenance marker, as getSerializedDataTransceiverState would produce. + auto makeArbitraryLlmRequest(SizeType32 length, LlmRequest::RequestIdType arbitraryId) + { + constexpr SizeType32 maxNewTokens{1}; + texec::Request request{VecTokens(length, length), maxNewTokens}; + auto state = std::make_unique(); + state->setCommState(*mContextCommState); + state->setCacheState(*mCacheState); + state->setIsArbitraryTransferState(true); + auto stats = texec::ContextPhaseParams({}, arbitraryId, state.release(), std::nullopt); + request.setContextPhaseParams(std::move(stats)); + return std::make_unique(arbitraryId, std::move(request)); + } + void addRequestAndTransportCache(std::shared_ptr const& llmRequest) { auto constexpr beamIdx{0}; @@ -519,6 +533,79 @@ TEST_F(SymmetricalCacheTest, SimpleTest) } } +TEST_F(SymmetricalCacheTest, ArbitraryTransferTest) +{ + auto worldSize = setUpCommunicator(); + if (worldSize != 2) + { + GTEST_SKIP() << "mpirun 2 processes is required to run this test."; + } + setUpCacheManager(/*enableBlockReuse=*/true); + setUpCacheTransceiver(); + + // 3 full blocks (tokensPerBlock = 8) so every requested chunk fully matches a stored block. + constexpr SizeType32 promptLen = 24; + constexpr SizeType32 missPromptLen = 40; + constexpr LlmRequest::RequestIdType arbitraryId = 4242; + auto constexpr beamIdx{0}; + auto constexpr beamWidth{1}; + + if (isSender) + { + // Store a patterned request in the reuse tree; no LlmRequest exists on the + // sender for the transfers below. + auto request = makeLlmRequest(promptLen); + mManager->addSequenceBatch( + {{{request->mRequestId, request->getNumTokens(beamIdx), beamWidth}}}, {std::ref(*request)}); + auto blockRange = BlockRange::fromAllBlockIds(*mManager, request->mRequestId); + for (auto const& windowSize : blockRange.getWindowSizes()) + { + auto blockRangeForWindow = blockRange.getBlockRangeForWindow(windowSize); + for (auto it = blockRangeForWindow.begin(); it != blockRangeForWindow.end(); ++it) + { + TLLM_CUDA_CHECK(cudaMemset(it->data(), request->getPromptLen(), it->getSizeInBytes())); + } + } + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*request); + // A completed context request carries one generated token beyond the prompt; + // without it, storeBlocksForReuse drops the trailing prompt token and the + // final block never enters the reuse tree. + request->addNewToken(0, beamIdx); + mManager->removeSequence(request->mRequestId, request); + } + else + { + // Hit: the requested tokens are stored in the sender's reuse tree. + std::shared_ptr request = makeArbitraryLlmRequest(promptLen, arbitraryId); + mManager->addSequenceBatch( + {{{request->mRequestId, request->getNumTokens(beamIdx), beamWidth}}}, {std::ref(*request)}); + auto future = mRequester->receiveAsync(request); + future.get(); + TLLM_CUDA_CHECK(cudaDeviceSynchronize()); + auto blockRange = BlockRange::fromAllBlockIds(*mManager, request->mRequestId); + for (auto const& windowSize : blockRange.getWindowSizes()) + { + auto blockRangeForWindow = blockRange.getBlockRangeForWindow(windowSize); + for (auto it = blockRangeForWindow.begin(); it != blockRangeForWindow.end(); ++it) + { + std::vector bytes(it->getSizeInBytes()); + TLLM_CUDA_CHECK(cudaMemcpy(bytes.data(), it->data(), it->getSizeInBytes(), cudaMemcpyDeviceToHost)); + EXPECT_TRUE(std::all_of(bytes.begin(), bytes.end(), [](uint8_t i) { return i == (promptLen & 0xff); })); + } + } + + // Miss: tokens never stored on the sender are rejected; the rejection + // surfaces as an exception on the receive future. + std::shared_ptr missRequest = makeArbitraryLlmRequest(missPromptLen, arbitraryId + 1); + mManager->addSequenceBatch( + {{{missRequest->mRequestId, missRequest->getNumTokens(beamIdx), beamWidth}}}, {std::ref(*missRequest)}); + auto missFuture = mRequester->receiveAsync(missRequest); + EXPECT_THROW(missFuture.get(), tensorrt_llm::common::TllmException); + } + // The sender must not tear down while the receiver's transfers are in flight. + tensorrt_llm::mpi::MpiComm::world().barrier(); +} + #if ENABLE_MULTI_DEVICE using AsymmetricTestParam = std::tuple Dict[str, Any]: def commit_blocks_for_reuse(self, req: LlmRequest) -> None: """Commit received KV blocks to the radix tree for prefix reuse. No-op by default.""" + def get_data_transceiver_state(self) -> bytes: + """Get the serialized DataTransceiverState (CacheState + CommState).""" + return b"" + def shutdown(self): """Shut down the transceiver and release registered resources.""" @@ -335,6 +339,9 @@ def prepare_context_requests(self, requests: List[LlmRequest]): # not implemented, an empty placeholder to allow being invoked unconditionally ... + def get_data_transceiver_state(self) -> bytes: + return self.impl.get_serialized_data_transceiver_state() + def get_disaggregated_params(self): # Cpp kv cache transceiver will set the disaggregated params to context response # Only new py cache transceiver will support gen-first disagg diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index a245a3cfe2e6..5e1413a4f0da 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -53,6 +53,7 @@ is_llm_response) if TYPE_CHECKING: + from .._torch.pyexecutor.kv_cache_transceiver import KvCacheTransceiver from ..disaggregated_params import DisaggregatedParams __all__ = [ @@ -972,6 +973,16 @@ def get_disaggregated_params(self) -> dict: return {} return self.engine.kv_cache_transceiver.get_disaggregated_params() + def get_cache_transceiver(self) -> Optional["KvCacheTransceiver"]: + if self.engine is None: + return None + return self.engine.kv_cache_transceiver + + def get_data_transceiver_state(self) -> bytes: + if self.engine is None or self.engine.kv_cache_transceiver is None: + return b"" + return self.engine.kv_cache_transceiver.get_data_transceiver_state() + @staticmethod def _stats_serializer(stats) -> str: # Per-rank path: stats is ("per_rank_dict", {..., "rank": N}). diff --git a/tensorrt_llm/executor/executor.py b/tensorrt_llm/executor/executor.py index 43adcbba93e6..98afe8daa44f 100644 --- a/tensorrt_llm/executor/executor.py +++ b/tensorrt_llm/executor/executor.py @@ -43,6 +43,7 @@ from .utils import IntraProcessQueue, ProcessPoolExecutorSession, RequestError if TYPE_CHECKING: + from .._torch.pyexecutor.kv_cache_transceiver import KvCacheTransceiver from .proxy import GenerationExecutorProxy from .worker import GenerationExecutorWorker @@ -464,6 +465,12 @@ def aget_kv_events(self, timeout=None) -> IterationResult: def get_disaggregated_params(self) -> dict: return {} + def get_cache_transceiver(self) -> Optional["KvCacheTransceiver"]: + return None + + def get_data_transceiver_state(self) -> bytes: + return b"" + @staticmethod def _create_ray_executor( worker_kwargs: Dict, diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index a7751cc6fe95..d1baa0792785 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -948,6 +948,16 @@ def get_disaggregated_params(self) -> dict: logger.warning(f"Error fetching disaggregated params via RPC: {e}") return {} + def get_data_transceiver_state(self) -> bytes: + """Get serialized DataTransceiverState from worker runtime via RPC.""" + if self.rpc_client is None: + return b"" + try: + return self.rpc_client.get_data_transceiver_state().remote() + except RPCError as e: + logger.error(f"Error fetching data transceiver state via RPC: {e}") + raise + def aget_stats(self, timeout: float) -> IterationResult: """Get iteration statistics from the runtime via RPC (async). diff --git a/tensorrt_llm/executor/rpc_proxy.py b/tensorrt_llm/executor/rpc_proxy.py index 680b8d5c84e4..0763e09ad032 100644 --- a/tensorrt_llm/executor/rpc_proxy.py +++ b/tensorrt_llm/executor/rpc_proxy.py @@ -254,6 +254,10 @@ def collective_rpc( return [remote_call.remote_future()] return [remote_call.remote()] + def get_data_transceiver_state(self) -> bytes: + """Get serialized DataTransceiverState from worker runtime via RPC.""" + return self.rpc_client.get_data_transceiver_state().remote() + def setup_engine_remote(self): return self.rpc_client.setup_engine().remote(need_response=True) diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index d3992365f651..577717586748 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -328,6 +328,17 @@ def llm_id(self) -> str: return self._llm_id + @set_api_status("prototype") + def get_data_transceiver_state(self) -> bytes: + """Get the serialized DataTransceiverState for arbitrary KV cache transfer. + + Returns: + bytes: Serialized DataTransceiverState, or empty bytes if no transceiver is configured. + """ + if self._executor is None: + return b"" + return self._executor.get_data_transceiver_state() + @property @set_api_status("beta") def disaggregated_params(self) -> dict: diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 153b35b5f6cb..eedd0be39285 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -795,6 +795,9 @@ def register_routes(self): methods=["GET"]) self.app.add_api_route("/version", self.version, methods=["GET"]) self.app.add_api_route("/v1/models", self.get_model, methods=["GET"]) + self.app.add_api_route("/v1/data_transceiver_state", + self.data_transceiver_state, + methods=["GET"]) # TODO: the metrics endpoint only reports iteration stats, not the runtime stats for now self.app.add_api_route("/metrics", self.get_iteration_stats, @@ -1101,6 +1104,19 @@ def register_visual_gen_routes(self): self.delete_video, methods=["DELETE"]) + async def data_transceiver_state(self) -> JSONResponse: + """Return the serialized DataTransceiverState as base64-encoded JSON.""" + state = self.generator.get_data_transceiver_state() + if not state: + return JSONResponse( + status_code=404, + content={"error": "No transceiver state available"}) + import base64 + return JSONResponse(content={ + "data_transceiver_state": + base64.b64encode(state).decode("utf-8") + }) + async def health(self) -> Response: if self._check_health(): return Response(status_code=200) diff --git a/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py b/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py index 2d5b5000005a..97634b686ea5 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py +++ b/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py @@ -138,6 +138,7 @@ async def run_worker(kv_cache_config, print(f"Sending ready signal to main process") intercomm.send(intercomm.Get_rank(), dest=0, tag=MPI_READY) + last_cached_tokens = 0 print(f"Waiting for requests") while True: try: @@ -146,6 +147,15 @@ async def run_worker(kv_cache_config, if requests is None: break + # Handle special commands + if requests == "GET_STATE": + state = llm.get_data_transceiver_state() + intercomm.send(state, dest=0, tag=MPI_RESULT) + continue + if requests == "GET_LAST_CACHED_TOKENS": + intercomm.send(last_cached_tokens, dest=0, tag=MPI_RESULT) + continue + request_metas = [] futures = [] for i, request in enumerate(requests): @@ -214,8 +224,12 @@ async def run_worker(kv_cache_config, f"Worker {rank}: awaiting future {i}/{len(futures)}", flush=True) result = await future - print(f"Worker {rank}: got result {i}, sending", - flush=True) + last_cached_tokens = getattr( + result, 'cached_tokens', 0) + print( + f"Worker {rank}: got result {i}, " + f"cached_tokens={last_cached_tokens}, sending", + flush=True) intercomm.send(result.outputs, dest=0, tag=MPI_RESULT) @@ -612,6 +626,7 @@ def test_disaggregated_spec_dec_batch_slot_limit(model, spec_dec_model_path, @pytest.mark.parametrize("generation_overlap", [False, True]) def test_disaggregated_logprobs(model, generation_overlap): """Verify that logprobs propagate correctly from prefill to decode. + Ensures first_gen_log_probs is carried in DisaggregatedParams so the generation_only worker receives one logprob per token. """ @@ -956,5 +971,282 @@ def test_disaggregated_logits(model, generation_overlap): print("All workers terminated.") +@pytest.mark.parametrize("model", ["TinyLlama-1.1B-Chat-v1.0"]) +@pytest.mark.parametrize("generation_overlap", [False]) +def test_arbitrary_kv_cache_transfer(model, generation_overlap): + """Test KV cache transfer from the reuse tree. + + Flow: + 1. Worker 0 runs a normal generate to fill the KV cache reuse tree. + 2. Retrieve the data_transceiver_state from worker 0. + 3. Worker 1 sends a generation_only request using the state from step 2. + The sender (worker 0) serves blocks directly from its reuse tree. + """ + worker_pytorch_configs = [] + + # Worker 0 (sender) + worker_pytorch_configs.append( + dict(disable_overlap_scheduler=True, + cuda_graph_config=CudaGraphConfig())) + + # Worker 1 (receiver) + worker_pytorch_configs.append( + dict(disable_overlap_scheduler=not generation_overlap, + cuda_graph_config=CudaGraphConfig())) + + kv_cache_configs = [ + KvCacheConfig(max_tokens=2048 * 8, enable_block_reuse=True) + for _ in range(2) + ] + cache_transceiver_configs = [ + CacheTransceiverConfig(backend="DEFAULT") for _ in range(2) + ] + model_names = [model_path(model) for _ in range(2)] + ranks = [0, 1] + worker_args = list( + zip(kv_cache_configs, cache_transceiver_configs, worker_pytorch_configs, + model_names, ranks)) + + port_name = mpi_publish_name() + + prompt = "What is the capital of Germany?" + + with MPIPoolExecutor(max_workers=2, + env={ + "UCX_TLS": "^ib,gdr_copy", + "UCX_MM_ERROR_HANDLING": "y" + }) as executor: + futures = [] + try: + for worker_arg in worker_args: + future = executor.submit(worker_entry_point, *worker_arg) + futures.append(future) + except Exception as e: + print(f"Error in worker {worker_arg}: {e}") + raise e + + intercomm = None + try: + print("Launched all the workers.", flush=True) + intercomm = mpi_initialize_intercomm(port_name) + + for _ in range(2): + intercomm.recv(tag=MPI_READY) + print("Received ready signal.") + + # Normal generate on worker 0 to fill KV cache reuse tree + print("Filling KV cache reuse tree on worker 0", flush=True) + requests = [(prompt, SamplingParams(max_tokens=10, + ignore_eos=True), None)] + responses = send_requests_to_worker(requests, 0, intercomm) + assert len(responses) == 1 + print(f"Worker 0 output: {responses[0][0].text}", flush=True) + + # Get data_transceiver_state from worker 0 + print("Getting data_transceiver_state from worker 0", flush=True) + intercomm.send("GET_STATE", dest=0, tag=MPI_REQUEST) + state = intercomm.recv(source=0, tag=MPI_RESULT) + assert isinstance(state, bytes) + assert len(state) > 0 + print(f"Got data_transceiver_state: {len(state)} bytes", flush=True) + + # generation_only on worker 1 using state from worker 0. + # max_tokens=1 ensures only one decode step runs, so the test + # relies on the KV cache transfer delivering all prompt tokens. + print("Sending generation_only to worker 1", flush=True) + DISAGG_REQ_ID = 42 + disagg_params = DisaggregatedParams( + request_type="generation_only", + opaque_state=state, + first_gen_tokens=[0], + disagg_request_id=DISAGG_REQ_ID, + ) + requests = [(prompt, SamplingParams(max_tokens=1, ignore_eos=True), + disagg_params)] + responses = send_requests_to_worker(requests, 1, intercomm) + assert len(responses) == 1 + output = responses[0][0] + print(f"Worker 1 output: {output.text}", flush=True) + print(f"Worker 1 token_ids: {output.token_ids}", flush=True) + assert len(output.token_ids) > 0, \ + "generation_only request should produce output tokens" + + # Send a normal (non-disagg) request to worker 1 to verify + # it still works properly after the KV cache transfer. + print("Normal request on worker 1 after transfer", flush=True) + requests = [(prompt, SamplingParams(max_tokens=10, + ignore_eos=True), None)] + responses = send_requests_to_worker(requests, 1, intercomm) + assert len(responses) == 1 + output2 = responses[0][0] + print(f"Worker 1 normal output: {output2.text}", flush=True) + assert len(output2.token_ids) > 0, \ + "Normal request should produce output tokens after transfer" + + # Query cached_tokens from worker 1's normal request + intercomm.send("GET_LAST_CACHED_TOKENS", dest=1, tag=MPI_REQUEST) + cached1 = intercomm.recv(source=1, tag=MPI_RESULT) + print(f"Worker 1 normal cached_tokens: {cached1}", flush=True) + + # Send a completely different prompt to worker 1 to confirm + # cached_tokens is 0 when there is no reuse match. + different_prompt = "Bonjour le monde, comment allez-vous aujourd'hui" + print("Different prompt on worker 1 (no match expected)", + flush=True) + requests = [(different_prompt, + SamplingParams(max_tokens=5, ignore_eos=True), None)] + responses = send_requests_to_worker(requests, 1, intercomm) + assert len(responses) == 1 + intercomm.send("GET_LAST_CACHED_TOKENS", dest=1, tag=MPI_REQUEST) + cached2 = intercomm.recv(source=1, tag=MPI_RESULT) + print(f"Worker 1 different prompt cached_tokens: {cached2}", + flush=True) + # BOS token may be shared, so at most 1 cached token is expected + assert cached2 <= 1, \ + f"Expected at most 1 cached token for unrelated prompt, got {cached2}" + assert cached2 < cached1, \ + "Unrelated prompt should have fewer cached tokens than transferred prompt" + + except Exception as e: + print(f"Exception encountered: {e}", flush=True) + raise + finally: + print("Sending termination request", flush=True) + mpi_send_termination_request(intercomm) + + print("Waiting for all workers to terminate. ", flush=True) + for future in futures: + future.result() + print("All workers terminated.") + + +@pytest.mark.parametrize("model", ["TinyLlama-1.1B-Chat-v1.0"]) +@pytest.mark.parametrize("generation_overlap", [False]) +def test_arbitrary_kv_cache_transfer_missing_blocks(model, generation_overlap): + """Test that missing-block transfers fail. + + When the receiver asks for blocks that don't exist on the sender, + the sender must notify the receiver so the request surfaces a + clear error instead of waiting for the hang detector. + + Flow: + 1. Worker 0 (sender) keeps an EMPTY KV cache reuse tree (no prior + generate call). Its reuse tree therefore has no blocks matching + any prompt the receiver might request. + 2. Retrieve the data_transceiver_state from worker 0. + 3. Worker 1 (receiver) sends a generation_only request using that + state. The sender attempts to look up the receiver's prompt + blocks in its empty reuse tree and fails with + "Couldn't find the requested block in the reuse tree". + 4. The receiver surfaces the failure as an error string. + """ + worker_pytorch_configs = [] + + # Worker 0 (sender) + worker_pytorch_configs.append( + dict(disable_overlap_scheduler=True, + cuda_graph_config=CudaGraphConfig())) + + # Worker 1 (receiver) + worker_pytorch_configs.append( + dict(disable_overlap_scheduler=not generation_overlap, + cuda_graph_config=CudaGraphConfig())) + + kv_cache_configs = [ + KvCacheConfig(max_tokens=2048 * 8, enable_block_reuse=True) + for _ in range(2) + ] + cache_transceiver_configs = [ + CacheTransceiverConfig(backend="DEFAULT") for _ in range(2) + ] + model_names = [model_path(model) for _ in range(2)] + ranks = [0, 1] + worker_args = list( + zip(kv_cache_configs, cache_transceiver_configs, worker_pytorch_configs, + model_names, ranks)) + + port_name = mpi_publish_name() + + prompt = "What is the capital of Germany?" + + with MPIPoolExecutor(max_workers=2, + env={ + "UCX_TLS": "^ib,gdr_copy", + "UCX_MM_ERROR_HANDLING": "y" + }) as executor: + futures = [] + try: + for worker_arg in worker_args: + future = executor.submit(worker_entry_point, *worker_arg) + futures.append(future) + except Exception as e: + print(f"Error in worker {worker_arg}: {e}") + raise e + + intercomm = None + try: + print("Launched all the workers.", flush=True) + intercomm = mpi_initialize_intercomm(port_name) + + for _ in range(2): + intercomm.recv(tag=MPI_READY) + print("Received ready signal.") + + # Get state from worker 0 + print("Getting data_transceiver_state from empty worker 0", + flush=True) + intercomm.send("GET_STATE", dest=0, tag=MPI_REQUEST) + state = intercomm.recv(source=0, tag=MPI_RESULT) + assert isinstance(state, bytes) + assert len(state) > 0 + print(f"Got data_transceiver_state: {len(state)} bytes", flush=True) + + # generation_only on worker 1 — sender has no blocks to serve. + print("Sending generation_only to worker 1 (sender has no blocks)", + flush=True) + DISAGG_REQ_ID = 43 + disagg_params = DisaggregatedParams( + request_type="generation_only", + opaque_state=state, + first_gen_tokens=[0], + disagg_request_id=DISAGG_REQ_ID, + ) + requests = [(prompt, SamplingParams(max_tokens=10, ignore_eos=True), + disagg_params)] + responses = send_requests_to_worker(requests, 1, intercomm) + assert len(responses) == 1 + response = responses[0] + print(f"Worker 1 response: {response}", flush=True) + assert isinstance(response, str), \ + f"Expected error string when sender has no blocks, got: {response!r}" + assert "cache transfer" in response.lower(), \ + f"Expected cache-transfer error, got: {response!r}" + + # Verify worker 1 still works for normal (non-disagg) requests + # after the failed transfer attempt. + print("Normal request on worker 1 after failed transfer", + flush=True) + requests = [(prompt, SamplingParams(max_tokens=10, + ignore_eos=True), None)] + responses = send_requests_to_worker(requests, 1, intercomm) + assert len(responses) == 1 + output = responses[0][0] + print(f"Worker 1 normal output: {output.text}", flush=True) + assert len(output.token_ids) > 0, \ + "Worker 1 should still serve normal requests after failed transfer" + + except Exception as e: + print(f"Exception encountered: {e}", flush=True) + raise + finally: + print("Sending termination request", flush=True) + mpi_send_termination_request(intercomm) + + print("Waiting for all workers to terminate. ", flush=True) + for future in futures: + future.result() + print("All workers terminated.") + + if __name__ == "__main__": pytest.main() diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 0e44fdeab6ca..8443df1caf14 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -883,6 +883,8 @@ disaggregated/test_disaggregated.py::test_disaggregated_overlap_gen_first[ctx_pp disaggregated/test_disaggregated.py::test_disaggregated_overlap_transceiver_runtime_python[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated.py::test_disaggregated_perf_metrics[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated.py::test_disaggregated_single_gpu[TinyLlama-1.1B-Chat-v1.0] +disaggregated/test_disaggregated_single_gpu.py::test_arbitrary_kv_cache_transfer[False-TinyLlama-1.1B-Chat-v1.0] +disaggregated/test_disaggregated_single_gpu.py::test_arbitrary_kv_cache_transfer_missing_blocks[False-TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_cancel_gen_requests[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_llama_context_capacity[False-False-DeepSeek-V3-Lite-fp8/fp8] disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_logits[False-TinyLlama-1.1B-Chat-v1.0] diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index f142378a7687..1a24dc44c789 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -105,6 +105,8 @@ l0_a10: - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_llama[False-True-TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_llama[True-False-TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_llama[True-True-TinyLlama-1.1B-Chat-v1.0] + - disaggregated/test_disaggregated_single_gpu.py::test_arbitrary_kv_cache_transfer[False-TinyLlama-1.1B-Chat-v1.0] + - disaggregated/test_disaggregated_single_gpu.py::test_arbitrary_kv_cache_transfer_missing_blocks[False-TinyLlama-1.1B-Chat-v1.0] - test_e2e.py::test_get_ci_container_port - test_e2e.py::test_openai_chat_multimodal_example ISOLATION - test_e2e.py::test_openai_mmencoder_example diff --git a/tests/unittest/api_stability/api_stability_core.py b/tests/unittest/api_stability/api_stability_core.py index 2a82d11fe3ee..c9b9a42388a7 100644 --- a/tests/unittest/api_stability/api_stability_core.py +++ b/tests/unittest/api_stability/api_stability_core.py @@ -157,9 +157,18 @@ def from_inspect(cls, method: MethodType): return_annotation = eval(return_annotation) return cls(parameters, return_annotation) + @classmethod + def _strip_api_status_tag(cls, docstring: str) -> str: + """Strip the :tag:`...` prefix added by @set_api_status decorator.""" + if docstring and docstring.startswith(":tag:"): + import re + docstring = re.sub(r'^:tag:`[^`]*`\s*', '', docstring) + return docstring + @classmethod def from_docstring(cls, method: MethodType): - doc = docstring_parser.parse(method.__doc__) + docstring = cls._strip_api_status_tag(method.__doc__) + doc = docstring_parser.parse(docstring) parameters = {} for param in doc.params: if param.args[0] == 'param': diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 959f249183e8..1f5e97d1b74f 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -441,6 +441,10 @@ methods: default: 2 return_annotation: tensorrt_llm.executor.result.IterationResult status: beta + get_data_transceiver_state: + parameters: {} + return_annotation: bytes + status: prototype shutdown: parameters: {} return_annotation: None