From 9843df94dec02c63952b72368db5e1e396ed8109 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Mon, 20 Jul 2026 07:22:57 +0000 Subject: [PATCH] [TRTLLM-14474][chore] Remove legacy python relics and refresh docs after the backend removal Trailing python-side cleanup after the legacy TensorRT backend removal (TRTLLM-14026, PR #16369): - Remove legacy quantization-export and plugin-gen relics, the engine-workflow serialization/deprecation shims, and other dead python modules (with their setup.py package_data entries). - Rename tensorrt_llm/bench/build to bench/tuning (build.py -> settings.py, tuning.py -> heuristics.py); the module no longer builds anything. The renamed files are also cleaned up to pass the standard ruff rules (docstrings, imports, quotes, line length) and removed from the legacy-lint allowlist. - Remove the top-level benchmarks/ folder entirely. benchmarks/prepare_dataset.py (already marked deprecated upstream) is deduplicated into the packaged tensorrt_llm/bench/dataset implementation behind trtllm-bench prepare-dataset, which gains --tokenizer and --stdout root options as a full replacement; all in-repo callers and docs are repointed. The remaining utils were orphaned legacy-engine helpers (generate_rand_loras.py wrote the removed .npy LoRA engine format; convert_nemo_dataset.py had no consumers). - Refresh docs and repo guidance to the current PyTorch/AutoDeploy state (customization.md rewrite, AGENTS.md backend section, examples/quantization/README.md, telemetry doc regeneration). - Prune stale legacy-lint entries (removed files) from legacy-files.txt, the generated ruff configs, and ruff-legacy-baseline.json. 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